feat(async): release rollout lanes during remote reward scoring

[zzhuoxin1508]

Summary

Two model-agnostic additions to the async RL rollout path:

1. Per-lane released-capacity state in RolloutPool — the main change.
Before, a rollout lane's capacity was freed only when its call fully resolved. Now a pending call may expose is_capacity_released(): the instant generation finishes, the lane is released — moved from _running into a new _released state — so it can admit the next prompt while the call keeps completing in the background. The unit only enters _completed once fully done. This splits "generation finished" (frees the GPU lane) from "result ready" (joins the batch), so downstream work can overlap the next generation instead of holding the lane.

2. A reward client that runs in the driver (DriverRewardClient).
A reward client hosted on the driver process (a thread pool), not on a GPU worker/slab. It sends generated media to a remote RewardService over HTTP and exposes the same launch_nowait / ready / result surface as a GPU Handle, so the per-lane manager drives it uniformly. chain_reward hooks it onto the released-capacity release point: generation done → stream to the remote reward → lane freed → the scored group joins the batch in reward-completion order. Opt-in via reward_client_on_driver: true; every training-node GPU then stays train+rollout while the reward model lives on its own node.

Both are generic and default-off — they work for any diffusion/AR rollout and any remote reward model. Also adds RemoteRewardBackend (HTTP client to the RewardService /score gateway, with per-request identity/idempotency for safe retries under async concurrency) and an opt-in example recipe.

Rebased onto current main after #289 and #304 merged. The diff now contains only the waiting-queue and driver-local remote-reward layer.

Related Issue

Builds on #289 and #304 (both merged).

Test Plan

Validated end-to-end with BAGEL-7B image-edit (it2i) as the policy and EditReward (MiMo-VL-7B) as the remote reward — the mechanism itself is model/reward agnostic. 2-node bring-up: a 1-GPU EditReward HTTP RewardService (unirl-reward-service) plus 8-GPU async training via the bagel_it2i_vllmomni_async_remote recipe (num_devices=8, reward_client_on_driver=true) wired by REWARD_SERVICE_URL. The streaming run advances at ~90 s/round with per-round reward scores logged and varying; lanes free at generation-done (_released) and reward scoring overlaps the next generation. See Performance below for the overlap ablation. Focused smoke tests for released-capacity dispatch and chained reward completion were run locally; the standalone BAGEL remote recipe also passes Hydra compose (not committed, per the repo tests/ policy #99/#267).

Performance

Speed ablation on the same setup (BAGEL-7B it2i policy, batch_size=8 × samples_per_prompt=8 = 64 images/round, remote EditReward over HTTP, single 8×H20 node; steady-state s/round after warmup):

#	Layout	Rollout engine	max_inflight	staleness	Overlap	Active GPUs/phase	s/round	vs #3
1	async, separate 4+4	vllm_omni	1	0	❌ none (forced serial)	4	~153	1.70×
2	async, separate 4+4	vllm_omni	1	1	✅ 1-deep pipeline	4 + 4	~89.9	~1.00×
3	async, separate 4+4	vllm_omni	1	2	✅ 1-deep pipeline	4 + 4	~90	1.00× (ref)
4	colocate, shared 8	vllm_omni	— (serial)	—	❌ none (single pool)	8	~106	1.18×
5	colocate, shared 8	trainside (bs=1)	— (serial)	—	❌ none (single pool)	8	~128.5	1.43×

Takeaways. The released-capacity overlap this PR enables is the dominant lever: turning it on (staleness 0 → 1) cuts −63 s/round (153 → 90) — the rollout slab generates round N+1 while the train slab trains round N. 1 → 2 adds nothing, since max_inflight=1 already caps in-flight work at one round. With overlap on, the separate 4+4 async layout (90 s) beats colocate-8 run serially (106 / 128.5 s) even though each phase uses only half the GPUs — generation is fully hidden behind training. An independent second axis is the rollout engine: trainside bs=1 eager generation vs vllm_omni costs +22.5 s/round on colocate (128.5 → 106).

Compatibility / Risk

Opt-in: reward_client_on_driver defaults off; existing reward placement is unchanged when off. Requires a reachable RewardService (REWARD_SERVICE_URL).

Reviewer Notes

Rebased onto current main after #289/#304 merged; the diff is now limited to the waiting-queue and driver-local remote-reward changes.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.

[zzhuoxin1508]

Test result. Under async training（max_inflight=1 staleness=2）, evaluated on BAGEL it2i (image editing) as the policy with EditReward as the reward model. The green curve scores each rollout the moment a single prompt finishes — its 8 images are streamed to the reward service immediately (per-prompt streaming); the pink curve waits for a whole batch to complete and scores it all at once (batched). The pink (batched) run averages 125 s/round, while the green (per-prompt streaming) run averages 88 s/round — a ~30% reduction in per-round time (125 s → 88 s), i.e. a 1.42× speedup (≈ +42% throughput).

[Ideny42]

The current W&B timing metrics do not accurately describe the async execution path. install_phase_timing() instruments train_step(), but AsyncDiffusionTrainer runs through _train_async_loop(), so phase metrics such as generate/reward/train time are not recorded. The remaining perf/step_time_s mixes collect wait, refill, advantage computation, and training, while excluding quiesce/weight sync. It therefore cannot show the actual rollout/reward/training latency or their overlap. Could we add async-specific timing metrics for capacity release, reward completion, training, sync, and end-to-end round time?