fix(async): gate rollout batches on validation pause

[jepio]

What does this PR do ?

Prevents async training rollouts from starting during validation.

On validation steps, set_weight_version() and resume_after_refit() could wake a collector below its loop-top pause check. The collector could then start a rollout batch during validation.

This change pauses the collector before refit can wake it and checks the manual pause again before starting a batch. In-flight rollouts still finish.

Issues

None.

Usage

No user-facing changes.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs. No documentation change is needed because this fixes internal scheduling behavior.

Additional Information

Validation:

• Focused regression test passed in Lima: uv run --frozen --group test pytest tests/unit/algorithms/test_async_utils.py -k collection_loop_defers_new_batch_while_manually_paused -q
• All applicable pre-commit hooks passed.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[jepio]

Generated by Claude Code

[jepio]

grpo.py:5417

1 action item. Pre-existing bug, not introduced by this PR — but the PR description describes the fix generally ("prevents async training rollouts from starting during validation"), and the same race is still open in PPO's async path.

grpo.py and ppo.py share the same AsyncTrajectoryCollector (trajectory_collector.py:126), the same _manual_pause_cleared/_refit_pause_cleared Events, and the same _run_collection_loop. This PR reorders grpo.py so pause.remote() runs before resume_after_refit()/set_weight_version(), closing the race here. ppo.py's async loop still has the old ordering:

ray.get(trajectory_collector.resume_after_refit.remote())   # ppo.py:2790 — wakes the collector
...
if (val_period > 0 and (step + 1) % val_period == 0) or (val_at_end and is_last_step):
    with timer.time("idle/validation"):
        ray.get(trajectory_collector.pause.remote())          # ppo.py:2804 — pause lands after the wake

The new re-check added to trajectory_collector.py only blocks the race if the caller's pause() already landed before the collector wakes — true for grpo.py now, still false for ppo.py. Not dead code either: examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async.yaml runs async PPO with val_period: 1.

Action: move the should_run_validation predicate and ray.get(trajectory_collector.pause.remote()) in ppo.py to before the prepare_for_refit/resume_after_refit block, mirroring this PR's grpo.py structure — or scope this PR's description to GRPO and file a fast-follow for PPO.

[jepio]

Confirmed — I checked ppo.py at HEAD (commit 25c4b60) and the ordering bug is real, not stale:

• ppo.py:2790 calls resume_after_refit.remote() unconditionally after refit.
• ppo.py:2799-2804 computes the validation gate and calls pause.remote() after that resume — same ordering this PR fixes in grpo.py.
• It's not dead code: examples/configs/recipes/llm/ppo-qwen2.5-1.5b-gsm8k-2n8g-megatron-valuetp2sp-dynbatch-noncolocated-async.yaml sets val_period: 1 and runs async PPO through this exact path (also the -automodel- and -single-controller variants).

Since this PR's AsyncTrajectoryCollector/re-check changes are shared infra, the fix mirrors cleanly: move the should_run_validation predicate + pause.remote() in ppo.py to before the resume_after_refit/set_weight_version block, same as the grpo.py reorder here.

Would you like to extend this fix to ppo.py in this PR, or should we scope the description to GRPO and track PPO as a fast-follow? Happy to help draft the ppo.py change if useful.