Skip to content

feat(rollout): bound driver task admission and support non-retaining collection - #3497

Merged
ananthsub merged 10 commits into
NVIDIA-NeMo:mainfrom
namtran1812:fix/bounded-rollout-admission
Sep 29, 2026
Merged

ananthsub merged 10 commits into
NVIDIA-NeMo:mainfrom
namtran1812:fix/bounded-rollout-admission

Conversation

@namtran1812

@namtran1812 namtran1812 commented Sep 17, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Closes #3006.

This bounds driver-side rollout task admission independently of request concurrency and adds an explicit non-retaining collection mode for artifact-oriented workloads.

Previously, asyncio.as_completed() scheduled the full iterable of rollout coroutines up front. num_samples_in_parallel bounded execution inside the request semaphore, but did not bound the number of asyncio tasks resident in the driver. Completed rows/results were also retained for the duration of collection even when the persisted artifacts were the desired output.

This PR:

  • adds max_resident_rollout_tasks as an independent driver-side admission bound
  • adds a bounded completion iterator that admits new rollout tasks as completed tasks are consumed
  • keeps num_samples_in_parallel semantics unchanged
  • adds explicit retain_results_in_memory=False behavior; collection persists results incrementally and returns []
  • preserves per-completion append + flush for main and failure-sidecar artifacts
  • preserves resume, upload, aggregation, health/coverage, failure routing, deterministic final ordering, and token-capture finalization
  • cancels only the bounded resident task set during exceptional shutdown
  • preserves upload semantics in non-retaining mode using a temporary spool, including resumed successful rollouts

Correctness

Focused coverage includes:

  • admission bound independent of request concurrency
  • simultaneous task completions without loss or duplication
  • cancellation of the bounded resident set
  • processing failures after rollout completion
  • success / failure-sidecar / no-persist routing
  • retaining vs non-retaining artifact equivalence
  • partial resume without duplicate execution, including counted failure classes in non-retaining mode
  • actual resident task peak and release of completed result objects
  • upload equivalence, including resume
  • aggregation and health/coverage behavior
  • owned vs externally installed token-capture source lifecycle

The rollout collection unit suite passes:

236 passed (7 expected warnings)

All pre-commit hooks pass for the changed files after merging upstream main.

Driver benchmark

I used a deterministic local synthetic driver benchmark against the real run_from_config() persistence path. The benchmark removes model/network latency so task admission, Python object residency, event-loop behavior, and persistence overhead are visible directly.

Environment: macOS / Apple Silicon, CPython 3.13.14. Bounded runs use a resident window of 64.

131,072 small rollouts

Mode Peak rollout tasks Peak Python heap p99 event-loop lag Rollouts/s
Existing unbounded 131,073 316.4 MB 2229.9 ms 5,052
Bounded, retaining 65 141.2 MB 10.4 ms 5,329
Bounded, non-retaining 65 89.5 MB 10.6 ms 5,416

The event-loop lag comparison used num_samples_in_parallel unset. With num_samples_in_parallel=64, the base p99 lag was already about 12 ms; the bound still removes the multi-second task-creation stall when it is unset.

At a configured window of 64, peak rollout task residency falls from 131,073 to 65 (the 64 admitted rollout tasks plus the surrounding benchmark task).

No throughput regression was observed in this synthetic benchmark; the throughput numbers above are not intended as production performance claims.

8,192 rollouts with 64-KiB response-text payloads

Mode Peak rollout tasks Peak Python heap
Existing unbounded 8,193 562.6 MB
Bounded, retaining 65 553.2 MB
Bounded, non-retaining 65 12.0 MB

With aggregation disabled and upload off, the non-retaining path reduced peak traced Python heap by ~97.9% versus the existing path for this large-result workload. Enabling aggregation or upload can reload the complete persisted results and increase final peak memory; this figure does not describe those configurations.

All three modes produced equivalent canonical 8,192-record main artifacts (~538.6 MB) after parsing and sorting by rollout identity. Raw JSONL byte order can vary with completion scheduling and is not guaranteed to match across modes.

This workload uses 64-KiB response text per rollout; it is not intended to represent a literal 64K-token model response.

Persistence and resume checks

A mixed 100-rollout workload containing successes, sidecar failures, and no-persist results produced:

  • 70 main-artifact records
  • 20 failure-sidecar records
  • 10 no-persist results

Canonical main and failure artifacts were identical across existing, bounded-retaining, and bounded-non-retaining collection.

For partial resume, I persisted 37 of 100 deterministic rollouts and resumed against the full materialized input. The resumed run dispatched exactly the 63 missing rollout identities, produced no duplicates, and the final canonical artifact matched an uninterrupted 100-rollout run.

Notes

When both limits are set, active requests cannot exceed the smaller of max_resident_rollout_tasks and num_samples_in_parallel.

The default remains unset, preserving existing admission behavior unless a bound is configured.

retain_results_in_memory=False bounds completed row/result residency during collection. Operations that explicitly require the complete result set, such as aggregation, may reload the persisted artifact after collection.

Input materialization itself remains proportional to the number of offered rows; this change specifically bounds rollout task residency and optionally removes completed-result residency during collection. asyncio.wait scans the resident window, so larger configured windows can add bookkeeping overhead.

Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 17, 2026

Copy link
Copy Markdown

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.

@nemo-automation-bot nemo-automation-bot Bot added the community-request Issue reported or requested by someone from the community label Sep 17, 2026
@yaoyu-33 yaoyu-33 added area:evaluation Evaluation, rollout collection, reward profiling, comparison, and exporters complexity:medium Single-domain change with interacting parts or a moderate review surface feature New capabilities, enhancements, or enablement work needs-review PR is ready for code review and waiting on a reviewer labels Sep 17, 2026
@ananthsub

ananthsub commented Sep 18, 2026 •

Copy link
Copy Markdown
Contributor

Thank you for taking this on and for the careful correctness and benchmark work. Bounding driver-side admission independently from request concurrency is the right direction for #3006, and the retaining and non-retaining split is useful.

I found two lifecycle problems that need to be addressed before this can merge.

First, the bounded iterator can deadlock when callers await multiple yielded completions concurrently. _next_completed() holds _lock while waiting for a rollout. When one rollout raises, its consumer calls aclose(), but another consumer can acquire that same lock and hold it while waiting indefinitely for another rollout. The original exception then cannot finish propagating because cleanup is blocked on the lock.

sequenceDiagram
    participant ConsumerA
    participant ConsumerB
    participant Iterator
    participant FailedRollout
    participant BlockedRollout

    ConsumerA->>Iterator: Await next completion
    ConsumerB->>Iterator: Await next completion
    Iterator->>FailedRollout: Wait while holding lock
    FailedRollout-->>ConsumerA: Raise exception
    ConsumerB->>Iterator: Acquire lock and wait
    Iterator->>BlockedRollout: Wait while holding lock
    ConsumerA->>Iterator: aclose needs lock
    Note over ConsumerA,BlockedRollout: Exception propagation and cancellation are stuck
Loading

Second, resource cleanup covers only exceptions raised inside the collection loop. Synchronous dispatch validation happens after output files and the upload spool are opened but before the protected block. The zero-success error also occurs after that block. Those paths can leave file handles or the .upload.tmp spool behind, including rollout data that may be large or sensitive.

I prepared a fix on ananthsub/fix-bounded-rollout-correctness, in commit 85fed7312.

The fix keeps the existing bounded-admission design. A failed completion now propagates immediately, matching asyncio.as_completed() semantics. _run_from_config() remains the owner that cancels and gathers resident tasks in its finally block. Output files, the upload spool, and the framework-owned token source are registered under one cleanup boundary, and the spool is removed on both successful upload and failure. It also documents that non-retaining mode bounds memory during collection, while aggregation and upload can still materialize their inputs.

flowchart TD
    Completion[Rollout completion] --> Outcome{Result or exception}
    Outcome -->|Result| Persist[Append and flush selected artifact]
    Outcome -->|Exception| Propagate[Propagate immediately to collection owner]
    Persist --> Continue[Admit next rollout within resident bound]
    Propagate --> OwnerFinally[run_from_config finally]
    OwnerFinally --> Cancel[Cancel and gather resident tasks]
    OwnerFinally --> Close[Close output files and token source]
    OwnerFinally --> Remove[Remove upload spool on failure]
    Continue --> Postprocess[Upload or aggregate when configured]
    Postprocess --> RemoveSuccess[Close and remove upload spool]
Loading

Could you please cherry-pick this signed-off commit onto fix/bounded-rollout-admission?

git cherry-pick 85fed73121763007930018e7aceef2b5213ecdf5

Once it is on this PR, I can re-review the integrated revision so we can move it toward merge.

@ananthsub ananthsub left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

requesting changes to address the potential deadlock

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Sep 18, 2026
ananthsub and others added 2 commits September 18, 2026 18:01
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@namtran1812

Copy link
Copy Markdown
Contributor Author

Thanks Ananth — I cherry-picked 85fed73 and integrated it with the latest PR branch. I re-ran the rollout collection suite (205 passed) and all pre-commit checks successfully. The updated revision is pushed and ready for re-review. Thanks again for catching these lifecycle issues and preparing the fix.

@yaoyu-33 yaoyu-33 removed the needs-review PR is ready for code review and waiting on a reviewer label Sep 18, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label Sep 19, 2026
@yaoyu-33 yaoyu-33 added the needs-review PR is ready for code review and waiting on a reviewer label Sep 19, 2026
@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Sep 21, 2026
@yaoyu-33 yaoyu-33 removed the needs-review PR is ready for code review and waiting on a reviewer label Sep 21, 2026
@github-actions github-actions Bot added the sla:review-overdue Review response is over the one-business-day SLA label Sep 21, 2026

@ananthsub ananthsub left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The inline comment on the zero-success check explains why the rebase needs a small semantic change, not just re-indentation.

Benchmark results

The benchmark script is not part of this PR, so we rebuilt both benchmarks. Each runs the real run_from_config() path with only the agent network call stubbed. They ran on Linux x86 with CPython 3.13.14. The hardware differs from yours, so the ratios matter more than the absolute sizes.

131,072 small rollouts, window 64, num_samples_in_parallel unset, aggregation off

Mode Peak rollout tasks Peak Python heap p99 event-loop lag Rollouts/s
Base (before this PR) 131,072 472 MB 7,134 ms 14,194
This PR, window unset 131,072 472 MB 7,217 ms 14,106
Window 64, retaining 64 276 MB 319 ms 22,796
Window 64, non-retaining 64 187 MB 2.3 ms 28,988

These claims reproduce:

  • Resident rollout tasks drop to the window size.
  • The default path shows no throughput regression.
  • Resume and the mixed success, sidecar, and no-persist routing produce the same sorted artifacts in every mode.

Some claims need qualifying:

  • Event-loop lag. The drop to about 10 ms holds for non-retaining mode against a run with num_samples_in_parallel unset. With num_samples_in_parallel=64, base p99 lag is already 11.7 ms. The bounded modes still remove the multi-second stall when all tasks are created up front. Retaining mode sits around 319 ms p99 here, because retained results cause long garbage-collection pauses.
  • Throughput. The bounded modes were faster in this synthetic run mainly because they spend less time in garbage collection. As the PR says, this is not a production throughput claim.
  • Large-payload memory. The 97.9% heap reduction reproduces only with disable_aggregation=true and upload off. With default aggregation, non-retaining mode peaks at the same heap as base. With upload also on, it peaks at about twice base. The inline comments on the aggregation reload and the spool upload give the numbers.
  • Raw artifacts. Retaining and non-retaining runs did not produce byte-identical raw files in general. The line order depends on which rollouts finish in the same event-loop turn, and that varies from run to run. The sorted, parsed artifacts matched in every mode, so "canonically identical" is the accurate description.

Could you update the PR description with these caveats? The rollout collection suite now has 205 tests, not 201.

Before merge

  • Rebase onto main with the persisted-key change described inline.
  • Replace the sys.exc_info() failure check. That bug came from my earlier fix.
  • Release the upload list before the aggregation reload.
  • Restore the progress output formatting.
  • Add the suggested tests. Two mutations that reintroduce result retention currently pass the whole suite.

resources_servers/gdpval/multistage_orchestrator.py also calls _run_examples_with_metadata without passing the new bound, so this setting is silently ignored there. That can be a follow-up rather than part of this PR.

Generated by Claude Code

Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
@svcnvidia-nemo-ci svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Sep 29, 2026
…-admission

# Conflicts:
#	nemo_gym/rollout_collection.py

@ananthsub ananthsub left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the quick turnaround.

These items are now fixed:

  • Rebase. The failure-counted-as-zero exclusion is correct after the rebase. A failure that later succeeds on resume is not scored twice, in either retaining or non-retaining mode. The result matches main.
  • Upload spool. It is no longer deleted when run_from_config is awaited inside a caller's except block. The upload now completes.
  • Upload memory. With upload and aggregation both enabled, non-retaining mode peaked at 539 MB against 542 MB on main. Before this revision it was about twice main.
  • Tests. The new tests catch both mutations that brought back result retention. There is also a check that the resident-task bound holds during a real collection run.
  • Documentation and description. The PR description, the num_samples_in_parallel description, and cli-commands.mdx now state the caveats accurately.
  • Headline claims. Resident tasks still peak at the window size plus one. The non-retaining heap stays small when aggregation and upload are off.

A few items from the previous review are still open, and I would like all of them fixed before this merges. Each has an inline comment with a verified change:

  • aclose() can still deadlock when two consumers await completions concurrently.
  • The existing concurrent-consumer test hides that deadlock, and two aclose() calls have no timeout.
  • test_run_from_config_processing_failure_cancels_bounded_resident_tasks still hangs, rather than fails, if the bound stops reaching the iterator.
  • The "Found ... rollouts" progress block still prints extra indentation in the default mode.
  • system-design.mdx still calls the two settings independent.
  • The public run_examples docstring does not describe max_resident_tasks.

I applied the code and test changes to the current head, and all 232 tests passed.

One more item for the PR description: the event-loop lag comparison should say it was measured with num_samples_in_parallel unset. With num_samples_in_parallel=64, main already has a p99 lag of about 12 ms. The bound still removes the multi-second stall that happens when every task is created up front.

Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Comment thread nemo_gym/rollout_collection.py Resolved
Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
ananthsub
ananthsub previously approved these changes Sep 29, 2026

@ananthsub ananthsub left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for working through all of the review feedback!

ananthsub
ananthsub previously approved these changes Sep 29, 2026
@ananthsub

Copy link
Copy Markdown
Contributor

@namtran1812 please address the code lint failure

@ananthsub
ananthsub disabled auto-merge September 29, 2026 20:24
Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
@namtran1812

Copy link
Copy Markdown
Contributor Author

Thanks, @ananthsub . Fixed the Ruff formatting failure in 4ea6ea8b6. The new workflows are awaiting approval; could you approve them when you have a chance?

@ananthsub

Copy link
Copy Markdown
Contributor

/ok to test 4ea6ea8

@ananthsub
ananthsub added this pull request to the merge queue Sep 29, 2026
@github-actions

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/36643880276

Merged via the queue into NVIDIA-NeMo:main with commit 6f308e1 Sep 29, 2026
41 checks passed
ananthsub added a commit that referenced this pull request Sep 29, 2026
Resolve conflicts with bounded rollout admission (#3497). The collection loop keeps
main's bounded, optionally non-retaining structure and applies this branch's changes
inside it: episode replies become records, records carry the Environment Server name
and type, rows are counted by the server they dispatch to, metrics are averaged over
the results that report them, and token capture resolves the agent through the row's
server. run_examples accepts both environment_server_name and max_resident_tasks.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@svcnvidia-nemo-ci svcnvidia-nemo-ci removed the waiting-on-customer Waiting on the original author to respond label Sep 30, 2026
OliviaViessmann pushed a commit to OliviaViessmann/Gym that referenced this pull request Oct 15, 2026
…eMo#4646)

## What changed

After the rollout loop finishes,
`RolloutCollectionHelper._run_from_config` (the `gym eval run` path)
reads several whole artifacts before it reports metrics. With
`retain_results_in_memory=false`, collection keeps no results in memory,
so it parses the entire rollouts JSONL again with `_read_jsonl` for
aggregation and for `check_completion`, as NVIDIA-NeMo#3497 designed. It also reads
the failure sidecar in `_failure_rows_counted_as_zero`, the materialized
inputs in `_missing_rollout_rows_counted_as_zero` and
`_fill_task_fields`, the upload spool, and the materialized inputs again
to count the expected rollouts.

All of these reads ran on the asyncio event loop without yielding. For a
large run the loop served nothing else until they finished, so telemetry
sampling, any remaining tasks, and shutdown all waited. Rollouts were
not slowed, because the reads start after collection ends.

This change takes two approaches, depending on what each read keeps:

- The two reloads that keep every row, the rollouts JSONL and the upload
spool, go through an async `_read_jsonl`. It parses on the event loop
thread and yields to other tasks every few milliseconds. It stays on the
loop thread on purpose: a worker thread allocates from its own malloc
arena, so it cannot reuse the memory the loop thread freed during
collection, and the driver would hold both copies. A worker-thread
version of this reload raised the simulator driver's memory from about 6
GB to about 10 GB.
- The scans that keep only a small result run in a worker thread through
`asyncio.to_thread`: the failure sidecar scan, the missing-rollout scan,
the task-field fill, and the expected-count read. The expected-count
read moves into a small `_count_lines` helper, which also closes the
file it opens.

The functions, their inputs, and their order are unchanged, so the
aggregate metrics, the counted failure and missing rows, the uploaded
rows, the returned value, and the raised errors stay the same. Sorting
and the `/aggregate_metrics` requests stay as they were. A full cyclic
garbage collection while the reload builds many objects can still pause
the loop briefly, as it would in any thread holding the GIL. Peak memory
during the reload is unchanged, because the whole rollouts file is still
loaded for aggregation, as NVIDIA-NeMo#3497 documented.

No issue: this is a focused fix for an event-loop stall found with the
CPU capacity simulator.

## Validation

- `PYTHONPATH=. .venv/bin/python -m pytest -n 16 -q
tests/unit_tests/test_rollout_collection.py`
- `PYTHONPATH=. .venv/bin/python -m pytest -n 16 -q tests/unit_tests
--ignore=tests/unit_tests/test_exporters.py` in a venv with the `dev`
and `telemetry` extras. The only failures are in
`test_opensandbox_compose.py`, `test_reward_profile.py`,
`test_dataset_orchestrator.py`, and `test_sandbox_stop_retry.py`; they
depend on the local environment and fail the same way on `main`.
- `pre-commit run --files nemo_gym/rollout_collection.py
tests/unit_tests/test_rollout_collection.py`

`test_read_jsonl_yields_to_other_tasks_between_rows` checks that a
coroutine on the same loop runs between parsed rows and that the parsed
rows are unchanged.
`test_run_from_config_non_retaining_reads_artifacts_off_the_event_loop`
runs a non-retaining collection with one success, one sidecar failure
counted as zero, one missing rollout counted as zero, and upload
enabled. Each scan waits for a coroutine on the loop to advance before
it starts. On the loop thread that wait blocks the coroutine, so the
scan is recorded as stalling the loop. Against `main`, the end-to-end
test fails on every scan, and the reader test fails because `main`'s
`_read_jsonl` cannot be awaited. The end-to-end test also checks the
rows sent to aggregation, including the task field filled into the
counted zeros, and the rows sent to the exporter.

On one simulator run's 4.8 GB rollouts file, the new `_read_jsonl`
returned the same rows as the previous synchronous reader.

## Rollout evidence

N/A: this change only changes how the driver reads its own artifacts
after collection. It does not change agent, model, or verifier behavior,
so no model rollouts were run.

The CPU capacity simulator compared `main` (`974d1a921b`) with this
branch on the production topology: a long-context eval with 4,096
rollouts, 1,024 in flight, 80 history messages of 2,000 characters each,
rollout observability on, console metrics export, and
`++retain_results_in_memory=false`, with two interleaved repeats per
variant. The driver's `gym.event_loop.lag_ms` histogram is cumulative,
so its max in the final export includes any end-of-run stall, and the
export before it covers only collection. `process.memory.usage` is the
driver's peak and final sample.

| Run | Lag max before the final export | Lag max in the final export |
Driver memory, peak / end |
| --- | ---: | ---: | ---: |
| `main`, repeat 0 | 1,747 ms | 8,400 ms | 5.99 / 5.55 GB |
| `main`, repeat 1 | 1,701 ms | 6,073 ms | 5.99 / 5.52 GB |
| this branch, repeat 0 | 2,244 ms | 2,244 ms | 5.99 / 5.56 GB |
| this branch, repeat 1 | 1,143 ms | 1,143 ms | 5.99 / 5.74 GB |

On `main` the end-of-run read raises the lag max by 4 to 7 seconds. On
this branch the max does not move after collection; the values in the
table were reached during collection and are unaffected by this change.
Driver memory is the same. Throughput was within run-to-run noise
(median 51.0 versus 48.8 rollouts per second, with overlapping ranges).

The aggregate metrics files agree between variants on every reward
metric. They differ only in `input_tokens` and `total_tokens`, which
also differ between two `main` runs, because the simulator's token
counts vary from run to run.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

This branch was previously deployed

1 inactive deployment
public — 4ea6ea8b Deployed Sep 29, 2026 by copy-pr-bot[bot] via release / finalize / notify #3979
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:evaluation Evaluation, rollout collection, reward profiling, comparison, and exporters community-request Issue reported or requested by someone from the community complexity:medium Single-domain change with interacting parts or a moderate review surface feature New capabilities, enhancements, or enablement work sla:review-overdue Review response is over the one-business-day SLA

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Bound rollout task admission and persist completed results incrementally

4 participants