feat(grpo): stream rollout batches by prompt group

[yfw]

What does this PR do ?

Stream NeMo-Gym rollout results back to NeMo RL as they complete, assemble them into prompt groups, and enqueue each complete group without waiting for the slowest prompt in the batch.

Thanks to @bxyu-nvidia for writing the initial version of this change.

What changed

• Stream (row index, result, timing) tuples from the NeMo-Gym Ray actor.
• Validate streamed rows and restore input order within each prompt group.
• Emit groups based on num_generations; no Gym-side num_groups_nemo_rl support is required.
• Use one batched collector worker per reserved target for both Gym and native rollouts.
• Enqueue completed groups concurrently with bounded replay-buffer backoff.
• Retry failed Gym streams without re-enqueuing groups already accepted by the replay buffer.
• Persist monotonic _ng_task_index values so retries and checkpoint restores do not reuse task identities.
• Preserve existing behavior for synchronous GRPO, PPO, distillation, and native environments.
• Keep native rollout scheduling unchanged while reporting group-local metrics.
• Compute teacher logprobs before replay-buffer insertion without blocking the async event loop.
• Exclude consumer-side pauses between generator pulls from rollout timing metrics.
• Update the NeMo-Gym integration documentation.
• Add logger.wandb.log_nemo_gym_full_result_tables flag to make logging full gym results to wandb optional (and default to false).
  • Migration Note: we were previously logging these results to wandb by default. In this PR, we set the default of this flag to false so after this change, we will no longer log these results to wandb by default. To enable logging of these results to match the previous behavior, set this flag to true.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

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.

Additional Information

• ...

[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.

[yfw]

/ok to test e02c941

[yfw]

/ok to test fa88050

[yfw]

/ok to test 9924919

[yfw]

/ok to test 2d61051

[yfw]

/ok to test c8d51f6

[yfw]

/ok to test 4fa8108

[terrykong]

TL;DR — what / why / how

What: Async GRPO previously generated rollouts prompt-group by prompt-group but consumed them batch by batch — nothing reached the replay buffer until every prompt in a batch finished, so one slow prompt held all completed groups hostage. This PR streams results back as each prompt group completes, so finished groups become trainable immediately and the slowest prompt only delays itself.

Why: Throughput. The trainer eats from the replay buffer while the collector fills it; head-of-line blocking on the slowest prompt starves the buffer and inflates step time — worst for NeMo-Gym workloads with heavy tail latency (agentic/SWE tasks).

How (bottom-up):

1. Gym actor — run_rollouts becomes an async generator yielding (row_index, result, timing) per completed task, consumed via Ray num_returns="streaming".
2. Rollout assembly — _NemoGymStreamAccumulator buckets streamed rows by row_index // num_generations and emits each complete prompt group in restored input order, fail-louding on duplicate/out-of-range/mixed rows. run_nemo_gym_rollout_sync drains the stream so sync GRPO/PPO/distillation keep whole-batch semantics (parity tests unchanged).
3. Collector — one batched worker per reserved target (replacing thread-per-group + semaphore) runs an asyncio loop in _collect_rollout_batch: streams groups, enqueues each concurrently with bounded backoff, computes teacher logprobs off-loop, and retries failed Gym streams up to 3× without re-enqueuing accepted groups.
4. Identity — a monotonic _ng_task_index is stamped per prompt, checkpointed to rollouts.pt, and restored as max(saved counter, max buffered index + 1), so retries and restores never reuse a task identity.

Also riding along: NaN generation-logprobs now raise instead of silently passing through (rollout_max_attempts_to_avoid_lp_nan removed), and W&B full_result tables get their own opt-in flag (logger.wandb.log_nemo_gym_full_result_tables).

---

Team review (5 agents: RL-guidelines, async-loop audit, independent bug scan, test review, adversarial verification) with extra scrutiny on the async/streaming semantics.

Overall: the async machinery held up very well under a dedicated 9-point concurrency audit. Verified with evidence: group ordering/identity (duplicate/out-of-range/mixed-agent rows all rejected), retry semantics (no group enqueued twice, no silent loss, failures raise), monotonic _ng_task_index checkpointing (restore takes max(rollouts.pt, max buffered index + 1)), event-loop hygiene (teacher logprobs via asyncio.to_thread + per-teacher locks; all push tasks gathered), correct Ray streaming-generator usage, and the efficiency-timer thread-seconds vs wall-clock split. Both parity tests binding the new implementations to the legacy ones are byte-identical to main, and replay_buffer.py is untouched. 37 unit tests from this PR pass locally; the vLLM/nemo_gym-marked E2E parity tests were left to CI.

The inline comments cover: the W&B full_result logging behavior change and a ruff-format failure (also failing the PR's Lint check in CI).

Two informational notes (no action needed):

• NemoGym is @ray.remote(max_restarts=-1, max_task_retries=-1); on actor restart Ray may replay already-yielded stream rows. This is safely converted into a clean batch retry by the accumulator's duplicate-row rejection — a one-line code comment noting the interaction would help future readers.
• timing/rollout/* metrics now ride on one group per batched worker instead of every group; the driver merge averages only over present keys, so values shift from per-prompt-group to per-worker-batch semantics without dilution.

Generated by Claude Code

[terrykong]

@yfw generally lgtm (also thanks @bxyu-nvidia for the change).

two extremely minor comments

[yfw]

/ok to test 07601e9

[yfw]

/ok to test 5c602b3