Backfill async rollouts on sample completions

[ys-2020]

Motivation

In group-based RL (GRPO-style, n_samples_per_prompt samples per prompt) with long-horizon agentic rollouts, the async schedulers currently backfill at group granularity: a replacement group is only submitted when an entire group task returns. When individual trials run 20-120 minutes (e.g. SWE-agent tasks at 64K context), every group is gated by its slowest sibling, and engine capacity idles while stragglers finish. In our production runs (n=8, 32 concurrent groups on SWE tasks) this straggler gating was a primary rollout-throughput limiter.

What this PR does

• Adds an opt-in flag --rollout-sample-completion-backfill (default off — the legacy group-level scheduler remains the default path and behavior is unchanged unless the flag is set).
• When enabled, each finished sample contributes a completion credit; once n_samples_per_prompt credits accumulate, one replacement prompt group is submitted immediately — keeping in-flight sample concurrency saturated instead of waiting for whole groups to return.
• Wires an optional sample_done_callback through generate_and_rm_group and both inference-rollout implementations, plus the fully-async example worker loop (refactored around a submit_one_group() helper).

Testing

• New CPU CI tests (tests/fast/rollout/inference_rollout/test_sample_completion_backfill.py):
  • default-disabled routing goes to the legacy group-level scheduler;
  • replacement-group scheduling triggers after n_samples_per_prompt sample completions.
• Full tests/fast/rollout/inference_rollout suite: 45 passed.

🤖 Generated with Claude Code

https://claude.ai/code/session_013vyY58s2TFpCzqEP4KaSxC

[gemini-code-assist]

Code Review

This pull request introduces a sample-completion backfill feature for training rollouts, allowing new prompt groups to be scheduled as individual samples complete rather than waiting for entire groups. Feedback on the changes highlights several concurrency and robustness issues, including potential event loop blocking when using synchronous queues, deadlocks upon data source exhaustion, incorrect completion credits for failed or cancelled tasks, and lost concurrency slots when group tasks raise exceptions.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Using self.output_queue.put directly inside the asyncio task done callback can block the asyncio event loop thread if the queue is full. Since self.output_queue is a synchronous queue.Queue, blocking here will freeze the entire event loop, preventing other rollout tasks from running. To avoid this, try a non-blocking put_nowait first, and fallback to running the blocking put in an executor.

Suggested change

	def task_done_callback(done_task):
	result = done_task.result()
	self.output_queue.put((gid, result))
	def task_done_callback(done_task):
	result = done_task.result()
	try:
	self.output_queue.put_nowait((gid, result))
	except queue.Full:
	asyncio.get_running_loop().run_in_executor(
	None, self.output_queue.put, (gid, result)
	)

[ys-2020]

This pattern is pre-existing on main (the same output_queue.put in the task callback, this PR only moved it into submit_one_group()), and the full-queue condition is unreachable in practice: in-flight groups are capped at rollout_batch_size (typically 32-64) while the queue's maxsize is 1000 and the consumer drains continuously. Happy to switch to put_nowait with an executor fallback in a follow-up if maintainers prefer, but keeping this PR's surface minimal.

[gemini-code-assist]

If the data source is exhausted, submit_groups will return 0 and pendings will remain empty. Without breaking out of the loop, the code will proceed to wait on sample_done_queue.get(), which will block forever since no more tasks are running to populate the queue. Breaking when no new groups can be submitted prevents this deadlock.

Suggested change

	while len(data) < target_data_size:
	if not pendings:
	# Defensive fallback for group-level task exceptions. Normal flow keeps
	# pending sample slots replenished from sample completion credits.
	submit_groups(max(1, target_data_size - len(data)))
	while len(data) < target_data_size:
	if not pendings:
	# Defensive fallback for group-level task exceptions. Normal flow keeps
	# pending sample slots replenished from sample completion credits.
	if submit_groups(max(1, target_data_size - len(data))) == 0:
	break

[ys-2020]

Good catch — fixed in f814b44: submit_groups returning 0 with nothing in flight now breaks out of the loop instead of blocking on the sample-completion queue.

[gemini-code-assist]

Only trigger the sample_done_callback if the sample task completed successfully (i.e., was not cancelled and did not raise an exception). This ensures that failed or cancelled samples do not incorrectly contribute to completion credits, which could otherwise lead to premature or double backfilling when combined with group-level failure recovery.

Suggested change

	if sample_done_callback is not None:
	task.add_done_callback(lambda _task: sample_done_callback())
	if sample_done_callback is not None:
	task.add_done_callback(
	lambda t: sample_done_callback() if not t.cancelled() and t.exception() is None else None
	)

[ys-2020]

This is deliberate, and conditioning on success would introduce a subtler bug: the callback fires exactly once per submitted sample task (success, exception, or cancellation), so completion credits conserve in-flight sample concurrency exactly — every n-credit batch corresponds to exactly one freed group slot. If failed samples were skipped, credits would leak and concurrency would decay monotonically with every failure (unless paired with group-level replenish, which then over-submits on partial group failures — see the reply on the other thread). A failed sample does free engine capacity, so counting it is the semantically correct choice for a throughput scheduler. Documented the invariant at this call site in f814b44.

[gemini-code-assist]

Only trigger the sample_done_callback if the sample task completed successfully (i.e., was not cancelled and did not raise an exception). This ensures that failed or cancelled samples do not incorrectly contribute to completion credits.

Suggested change

	if sample_done_callback is not None:
	task.add_done_callback(lambda _task: sample_done_callback())
	if sample_done_callback is not None:
	task.add_done_callback(
	lambda t: sample_done_callback() if not t.cancelled() and t.exception() is None else None
	)

[ys-2020]

Same rationale as in inference_rollout_common.py: credits deliberately fire for every sample outcome so that in-flight concurrency is conserved exactly; conditioning on success would leak credits. Comment added in f814b44.

[gemini-code-assist]

When a group task fails with an exception, its samples do not complete successfully and won't trigger on_sample_done. This causes the pending slot to be permanently lost, reducing the rollout concurrency until pendings is completely empty. Replenish the failed group immediately by calling submit_groups(1) to maintain saturated concurrency.

Suggested change

	except Exception as e:
	logger.error(f"[rollout] Task raised exception: {e!r}", exc_info=True)
	continue
	except Exception as e:
	logger.error(f"[rollout] Task raised exception: {e!r}", exc_info=True)
	submit_groups(1)
	continue

[ys-2020]

The premise doesn't hold for the current implementation: when a group task raises, its per-sample done callbacks still fire (they're attached to the individual sample tasks, which run to completion independently of the group-level gather), so the failed group's n credits still arrive and trigger exactly one replacement group — the slot is not lost. Adding submit_groups(1) here on top of that would double-replenish: e.g. a group where 1 of 8 samples raises would emit 7-8 credits and an immediate extra group, over-submitting by up to ~2x per failed group. Keeping replenishment purely credit-driven preserves the exact conservation invariant.

[autoweave-bot]

This looks like a really neat way to keep things running smoothly! ✨ 😄
Explore here →

_{Maintainer? Turn off weaves from non-maintainers →
Carefully crafted by Subweave · 🧶 used ~356k LLM tokens}