fix(rollout): stop dynamic sampling from leaking generation past the rollout

[OnePunchMonk]

What

• When the dynamic sampling filter drops a group, fetch only as many replacement groups as were dropped, instead of a whole new over_sampling_batch_size batch.
• Before generate_rollout_async returns (or raises), cancel and await any tasks still running, then reset the shared state.

Why

With --dynamic-sampling-filter-path (including the stock check_reward_nonzero_std), a single dropped group made the loop fetch rollout_batch_size new groups. Once enough groups were kept, state.reset() discarded the running tasks without cancelling them. They run on the background event loop, which never stops, so they kept generating on the engines while the trainer offloaded them, trained on the same GPUs (colocated) or pushed new weights.

Measured with the real generate_rollout and a fake engine (8 groups per rollout, 1 in 4 groups with all-equal rewards, 3 rollouts), base vs this PR:

	base	this PR
groups still generating after each rollout returns	6, 3, 5	0, 0, 0
groups generated for 24 kept	48	31
prompts fetched but never trained on or filtered	12	0

The skipped prompts also advance the data-source position saved in the checkpoint, so a resume skipped them too. The kept set was whichever groups finished first, which biased training toward prompts that generate fast.

Unused groups never need data_source.add_samples now: after a successful rollout, fetched = kept + dropped.

Validation

• New tests/fast/rollout/test_dynamic_sampling_rollout.py (2 tests): no work in flight after the rollout returns, fetched = kept + dropped across rollouts, and a failing group cancels the rest and leaves clean state for the next rollout. Both fail on base and pass here.
• Modal CPU replica of the stage-a-cpu job (Ubuntu 24.04, sglang-miles): 496 passed.
• GPU, SD3.5-medium OCR GRPO with the filter on, 4 groups of 8 per rollout, 3 rollouts, 2x H100 (trainer and rollout engine on separate GPUs). Every 4th group was force-dropped on top of the stock filter, since OCR rewards rarely tie exactly. Exit 0.

rollout	started	kept	dropped	fetched	in flight at return
0	5	4	1	5	0
1	5	4	1	5	0
2	5	4	1	5	0

Notes on the GPU run:

• It is not the recipe's 2-GPU colocated setup. Modal blocks the CUDA IPC weight transfer that --colocate uses (pidfd_getfd: Operation not permitted), so I ran non-colocated.
• The image needed sglang-miles checked out over its sglang, as CI does.
• In a non-colocated run the CPU-only rollout manager crashes at import: fsdp_utils -> H3 weight bridge -> vsa calls torch.cuda.get_device_capability(0). I stubbed vsa in my probe harness only. This is separate from this PR.
• Base was not run on the GPU. The base behavior is shown by the CPU repro above.

Review: an independent Codex pass raised one point, that kept groups past the target could be dropped silently. That cannot happen here since kept <= submitted - dropped <= target; it was true on base and the test asserts it.

Tradeoff

Refills now arrive in rounds sized to the drops, so a rollout with many drops can wait for one extra group's generation per round. In exchange, no extra groups are generated.

Files

• miles/rollout/sglang_diffusion_rollout.py: top up only the dropped count, and cancel and await leftover tasks in a finally before resetting state.
• tests/fast/rollout/test_dynamic_sampling_rollout.py: regression tests.

Checklist

• ruff, black 24.3.0, isort, autoflake clean on the touched files
• Added tests for the new behaviour
• stage-a-cpu green on a Modal replica (496 passed)
• No launch flags changed
• No public flag added
• No example added

[OnePunchMonk]

Note:

1. The actual fix is a one-line fetch change: fetch only as many groups as the filter dropped, so nothing is left running or skipped.

2. The try/finally is just cleanup for the case where a group raises mid-rollout; it does nothing on a normal return.

[Rockdu]

Hi @OnePunchMonk,
Fixing the leak is right, but topping up exactly the dropped count serializes the tail. In-flight groups equal target - kept, so near the end only one group runs while the other engines sit idle.

Example: 4 engines, target 4 groups, one group takes 1 time unit, same keep/drop outcome for each group in both runs (green = kept, gray = dropped, red = extra).

[OnePunchMonk]

Added an abort endpoint to SGL diffusion here: sgl-project/sglang#44953