What does this PR do?

Moves the sampler weight transfer in WorkerDispatch.save_weights_for_sampler off the trainer's event loop.

TLDR: _broadcast_and_finish() is a blocking ray.get called inline from an async def. Under the fully-async trainer, that freezes every generator coroutine on the loop for the whole transfer. It now runs through asyncio.to_thread, and a cancellation waits for it, so generation cannot resume mid-transfer. The offload-KV path also wakes the KV cache before resuming after a failed or cancelled sync.

Why

1. All six call sites in save_weights_for_sampler call _broadcast_and_finish() synchronously. _broadcast_to_inference_engines is deadline.ray_get(...) over the policy workers, so the event loop is blocked for the whole broadcast.
2. In the fully-async trainer, the _run_generate_for_a_group_loop workers, and whatever HTTP clients and environments they drive, run on that same loop. The engines are paused during the sync, but non-engine work (environment steps, tool calls, reward computation) stalls with them, so the overlap the fully-async trainer is meant to provide is lost for the length of the transfer. For a large MoE policy we measured timing/sync_weights_only_transfer at 38–43 s on every step.
3. A timeout armed by a coroutine just before the freeze fires on the first tick after it, even though the awaited operation was healthy. In-tree, that includes the 10 s bound on finish_session. In our case it was a custom generator's 10 s HTTP read. The resulting TimeoutError became the generator's BackgroundFailure, was re-raised at the next _drain_next_group, and killed the run. That failure is timing-dependent, so it can strike on any step once the transfer is longer than the shortest in-flight timeout.

Solution

_broadcast_and_finish is now async. It runs the broadcast plus _finish_weight_sync via asyncio.to_thread and is awaited at every call site. The pause, sleep, wake and resume brackets around it, their ordering, and the trailing increment_weight_version() are unchanged.

• Engine safety is unchanged. Requests are kept off half-updated weights by the engine-side bracket (pause_generation in KEEP mode, or sleep/wake), not by the frozen driver loop. Generators that run outside the driver process already send requests during a sync. Only the trainer coroutine issues engine control calls, and it still awaits the sync before moving on.
• Cancellation. A cancellation that arrives mid-transfer, for example on shutdown, waits for the transfer to finish (asyncio.wait on the thread's future) and then propagates, chained to the transfer's error if it failed. Without that wait, cleanup_preserving_primary(resume_generation) would resume generation on the NCCL and offload-KV paths while the policy workers were still broadcasting. The blocking version could not be cancelled mid-transfer at all, and this keeps that guarantee.
• Deadlines. step_deadline state is ContextVar-based and to_thread copies the context, so the threaded deadline.ray_get still gets remaining() as its timeout. On expiry, asyncio.timeout cancels the awaiting task and the thread's ray.get times out at the same instant, so the wait adds nothing. The await is wrapped in deadline.operation("broadcast_to_inference_engines"), so a StepTimeoutError names the same operation whichever side fires first.
• Thread safety. The thread makes only Ray calls (actor-method dispatch plus ray.get, both thread-safe). last_weight_sync_seconds is set after the await. The colocated _finish_weight_sync → _offload path is likewise ray.get only.
• KV cache on an interrupted sync (pre-existing on main). On the fully-async offload-KV path, a failed or cancelled broadcast skipped wake_for_weight_sync(kv_cache), and the resume cleanup then scheduled the frozen requests onto a sleeping KV cache. The KV wake is now its own cleanup_preserving_primary bracket inside the resume bracket, as the synchronous-trainer path already does for wake_up(kv_cache).
• Not done: per-caller fixes, such as a longer timeout in each HTTP client. Any coroutine on the trainer loop can be hit, and the freeze grows with model size.

Test plan

New CPU tests in TestSaveWeights, with a fake time.sleep broadcast:

• test_broadcast_does_not_block_event_loop uses a 0.5 s broadcast. It arms a 0.2 s asyncio.timeout around a 0.05 s sleep while the sync runs, then asserts that the timeout did not fire, that the sync was still in progress, and that last_weight_sync_seconds >= 0.5. On main the timeout fires, which is the production failure mode.
• test_cancellation_waits_for_the_transfer_before_resuming cancels the sync during a 0.3 s broadcast and asserts the order broadcast_done → resumed. With to_thread but no wait it fails with ['resumed']. On main it fails with DID NOT RAISE, because the blocking sync completes before the cancellation can land.
• test_cancellation_survives_a_failed_transfer cancels, then makes the broadcast fail. It asserts that CancelledError propagates with the RuntimeError as its __cause__.
• test_offload_kv_wakes_kv_cache_before_resuming[failure|cancellation] runs the fully-async offload-KV path and asserts the order wake_weights → broadcast_done → wake_kv_cache → resumed. On main the failure case resumes without the KV wake.

All three fail on this PR's first commit.

The existing TestSaveWeights cases run unchanged. They cover bracket ordering, resume after a broadcast failure, and the delta, LoRA and colocated paths.

No GPU CI was run. The broadcast itself is unchanged; only the thread it runs on moved.

🤖 Generated with Claude Code