feat: route NeMo-Gym rollout rows across sharded actors

[ananthsub]

What does this PR do ?

Routes rollout rows to the right actor and merges the results back into one stream. With K actors now buildable (PR 5 of this stack), a training step has to decide which rows go where and then reassemble what comes back.

Dispatch is by whole prompt group, never by row. _bucket_nemo_gym_rows_by_shard walks the batch in groups of num_generations, resolves the group's agent_ref to a shard, and sends the entire group to one actor. Splitting a group across actors would deadlock any verifier that scores rollouts against their peers — genrm_compare buffers a whole group in one process and blocks until it is complete, so each half would wait forever for rollouts sitting in the other. A group whose rows name agents in different shards is rejected rather than split; group members inherit the prompt's agent_ref so this should not happen, but the check is cheap and catches upstream regressions.

Rows keep the _rowidx stamped over the whole batch before bucketing, so the accumulator downstream never learns that shards exist.

Fan-in preserves streaming. run_rollouts is an async generator, so there is nothing to gather: _merge_nemo_gym_shard_streams replaces one stream with an interleave of K, yielding rows as they complete and feeding them to the existing accumulator unchanged. With one bucket it degenerates to that single stream, which is what keeps the unsharded path identical. Leaving the merge early — a downstream error, or the consumer closing the generator — cancels the in-flight anext tasks rather than stranding them.

Timing metrics are reported per shard and rolled up. _merge_nemo_gym_timing_metrics emits each shard's numbers under {timer_prefix}/shard/{name}/... and takes the maximum for the rolled-up value rather than the sum, because shards run concurrently and the step is gated by the slowest one.

Two smaller changes fall out of this:

• task_to_env["nemo_gym"] now accepts either a NemoGymShardSet or a bare actor handle, the latter read as the one-shard, one-replica case (as_nemo_gym_shard_set). This is what lets the router land before any entrypoint builds a real shard set.
• shutdown_environments had assumed every environment is a Ray actor handle. A shard set is not — it owns actors — so teardown now asks an owner object to shut itself down and does not escalate to ray.kill, since there is no single handle to kill and the owner has already reported which shard would not stop.

Issues

None.

Usage

No user-facing change. Unsharded jobs take the same path as before: one bucket over one stream.

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

• Which shard a group goes to is fixed by the data (its agent_ref) and cannot change mid-job — shard contents freeze when RunHelper serializes each child's config at Popen time. Which replica within that shard is chosen at dispatch, round-robin.
• Taking the maximum on roll-up makes the rolled-up postprocess_results_pct the worst shard's share rather than a batch-wide average. A true batch-wide percentage needs each shard's rollout total, which the actor pops before returning; exposing it is a small upstream Gym change and is not worth coupling to this stack. Per-shard percentages are reported and are exact.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.