Repository navigation
feat(rollout): bound driver task admission and support non-retaining collection - #3497
Conversation
Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
|
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. 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
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 I prepared a fix on The fix keeps the existing bounded-admission design. A failed completion now propagates immediately, matching 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]
Could you please cherry-pick this signed-off commit onto git cherry-pick 85fed73121763007930018e7aceef2b5213ecdf5Once it is on this PR, I can re-review the integrated revision so we can move it toward merge. |
ananthsub
left a comment
There was a problem hiding this comment.
requesting changes to address the potential deadlock
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…nto fix/bounded-rollout-admission
|
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. |
There was a problem hiding this comment.
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_parallelunset. Withnum_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=trueand 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
mainwith 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
…-admission # Conflicts: # nemo_gym/rollout_collection.py
There was a problem hiding this comment.
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_configis awaited inside a caller'sexceptblock. 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 twicemain. - 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_paralleldescription, andcli-commands.mdxnow 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_tasksstill 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.mdxstill calls the two settings independent.- The public
run_examplesdocstring does not describemax_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.
Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
|
@namtran1812 please address the code lint failure |
Signed-off-by: namtran1812 <158846154+namtran1812@users.noreply.github.com>
|
Thanks, @ananthsub . Fixed the Ruff formatting failure in |
|
/ok to test 4ea6ea8 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/36643880276 |
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>
…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>
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_parallelbounded 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:
max_resident_rollout_tasksas an independent driver-side admission boundnum_samples_in_parallelsemantics unchangedretain_results_in_memory=Falsebehavior; collection persists results incrementally and returns[]Correctness
Focused coverage includes:
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
The event-loop lag comparison used
num_samples_in_parallelunset. Withnum_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
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:
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_tasksandnum_samples_in_parallel.The default remains unset, preserving existing admission behavior unless a bound is configured.
retain_results_in_memory=Falsebounds 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.waitscans the resident window, so larger configured windows can add bookkeeping overhead.