nemo_gym.rollouts.rollout_collection

Call /aggregate_metrics on the environment server each rollout ran through.

Rows are grouped by the environment server stamped on them at preprocessing (_ng_environment_server); a row without the stamp is grouped by the environment server that fronts its agent, as before, so the identity decided at dispatch is the one aggregation uses, across shards and resumed runs alike. Writes a single _aggregate_metrics.json with one entry per environment server (same shape as the old _agent_metrics.json, plus the server name). Returns the file path. Failed requests retain an entry with empty metric collections and aggregation_error. By default, re-raise the first original exception after saving partial metrics, preserving its type and traceback for direct callers such as custom evaluation drivers. Collection and aggregation entrypoints pass raise_on_error=False to write batch status first, then call _raise_for_aggregation_errors for a summary error.

Internal dispatch shared by run_examples and Gym’s own collection paths.

record_unrouted_failures returns a failed /run that route_failures_to_sidecar would route as a failure row whose run_error holds the error, so the collection writer can keep evidence of it before it raises the error. Other callers keep receiving the exception.

When max_resident_tasks is set, at most that many rollout tasks are admitted at once. When unset, all examples are scheduled as before. The collection owner closes either iterator on completion, cancellation or error.

dispatch_budget_s stops starting rows that many seconds after this call, and drain_margin_s stops sooner for rows that would not have time to finish. A row drained this way resolves to a Gym-built _dispatch_drained_result.

Identical contract to run_examples, but each future resolves to a _CompletedRollout that carries rollout_latency_ms alongside the raw /run result instead of inside it, so internal-only timing never has to be smuggled through (and stripped back out of) a dict that a direct caller of run_examples could also observe.

interleave_by_agent is enabled by the full evaluation path. Lower-level callers retain the order they supplied.

Stamp compatibility-routed rows with the bound agent.

These rows never reach resolve_task_sources, so without this they carry no agent_ref and results, aggregate metrics and reward profiling lose the agent they ran on. A row that already names an agent is left alone; the name is validated against the environment server by _validate_environment_servers. Materialized tasks are skipped: they carry no agent by design, and their result projection is not the legacy shape this key belongs to.

Fail before any dispatch when a row names an agent absent from the running config.

Without this, the first bad row dies mid-collection with a raw omegaconf ConfigKeyError after valid rows have already been dispatched.

Fail before dispatch when a row points to agent incompatible with the resources server it runs on.

The verifier’s allowed_agents is read from the row’s task_source and from the resources server the agent binds; the dataset entry that declares the rows may also name the harnesses it admits (a Harbor dataset’s [gym] allowed_agents, carried into its generated datasets: entry, whose rows route by taskset to an environment server that binds no resources server). That list is checked against the agent’s own type, so a run config edited to another harness after gym eval run admitted one is refused here as well. A materialized row is matched to its entry by taskset, wherever the entry is declared; a flat row meets every entry the agent declares (see _dataset_admitted_agents).

Apply run-level routing and repetition to caller-held rows.

Public entry point for direct run_examples callers (e.g. trainer integrations that drive dispatch themselves): run_examples resolves task_sources and validates agent names, but agent_map, fan_out and num_repeats are applied only during preprocessing. Call this first, then pass the returned rows to run_examples.

Pass global_config_dict (the merged config) to also resolve task_source-only rows to their agents here; leave it None to defer that to run_examples, which does it against the head server’s config. Pass task_ids to keep only those tasks, selected by the id stored in each row (see nemo_gym.data.task_selection.select_tasks), each task once before repetition. Input rows are not mutated; the expanded, stamped copies are returned.

Stamp an agent_ref onto every row that carries only a task_source.

task_source names the config instance that declared the row’s dataset. Resolution is resolve_dataset_agent — the same rules benchmark discovery uses, so dispatch can never disagree with the listing. Conflicting agent: pins across one instance’s datasets are a hard error (rows carry only the instance name), as are unknown/non-routable instances; +agent_map is the disambiguator.

Rows that already have an agent_ref are left untouched, so this is a no-op on legacy datasets and on already-resolved (materialized) rows. Runs before any dispatch.

We provide this function as a lower level interface for running rollout collection.

Rows are dispatched as given: task_sources are resolved and agent names validated here, but run-level knobs (agent_map, fan_out, num_repeats) are NOT applied — call preprocess_examples first if you need them.

route_failures_to_sidecar makes a failed /run a failure row instead of an exception that ends every rollout still in flight. It defaults off because those rollouts then leave the score.

max_resident_tasks limits admitted tasks and therefore concurrent requests, even when semaphore allows more. Admission starts when the first returned awaitable is awaited. None schedules all examples up front, as with asyncio.as_completed.

To stop early, await the returned iterator’s aclose(): it cancels unfinished rollouts and waits for their client tasks in this process to end. Cancelling an in-flight /run closes its connection; a server with ClientDisconnectCancellationMiddleware subsequently cancels its handler. Server/agent cleanup is not awaited or acknowledged, depends on the agent, and may continue after return. Sandbox release is not guaranteed. Handed-out awaitables may return completed results or raise their exceptions, including CancelledError; bounded awaitables first awaited after closure raise CancelledError. The error propagates in the awaiting caller’s task even without cancelling that task, and asyncio.gather and TaskGroup treat it as cancellation. Without it, abandoned rollouts keep running to completion in the background.

To preserve completed but unconsumed work when stopping, use await completions.aclose_and_drain() instead. It stops admission before suspending, cancels and awaits unfinished client tasks, and returns settled (row, result) futures, including failures. Consume those futures and any outcome already claimed by an outstanding completion waiter; an unclaimed waiter raises CancelledError. Repeated drains return an empty list. This also works with bounded admission without starting more rows, and has the same server/agent cleanup limits as aclose().

Use finally to close the iterator if collection stops early:

completions = helper.run_examples(examples)
try:
    for future in completions:
        row, result = await future
        if should_stop(row, result):
            break
finally:
    await completions.aclose()

dispatch_budget_s stops starting rows that many seconds after this call, and drain_margin_s stops sooner for rows that would not have time to finish.

Rows start in the order given, once the returned iterator is first consumed.

A row that ran resolves to exactly the (row, result) pair Gym’s own /run endpoint returned — no Gym-private fields are added to result. A row that produced no /run result resolves to a Gym-built result carrying Gym-private fields instead: a row drained by the dispatch budget gets _ng_failure_class="cancelled" with the _ng_dispatch_drained and _ng_no_persist markers, and a failed /run under route_failures_to_sidecar gets a failure row with _ng_failure_* fields.

retry_requests=False disables automatic HTTP replay of a mutating /run. Connections that were never established can still be retried safely. It does not convert failures to successful result rows.

enrich_results runs enrich_run_result on each result before it is returned, so the caller receives what ng_collect_rollouts records: the row’s rollout identity, the captured model calls (ng_model_call_capture, with request and response bodies unless include_model_call_payloads is false), the standardized ng_trajectory and ng_perf. Model calls are captured only when the global config enables observability_enabled. delete_model_call_capture removes each rollout’s capture file after it is merged, and needs the payloads kept. With enrich_results the rows are dispatched as collection dispatches them, on copies that name this call’s run: results are validated, and under route_failures_to_sidecar a failed /run comes back as the typed RolloutFailureRecord, as run_outcomes returns it; rows without a rollout identity keep the plain dispatch. With enrich_results off, results are exactly what /run returned.

record_output_dir with record_shard also writes each rollout to that run folder through _RunRecordWriter, as ng_collect_rollouts would: rollouts-<shard>.jsonl, its failures sidecar with the typed failure rows collection writes, the rows as the shard’s materialized inputs, its run manifest and its saved progress for gym eval status (omitted-rollout rows with record_omitted_rollouts). It implies enrich_results, dispatches as collection does (typed failure records under route_failures_to_sidecar, validated results) on copies of the rows that name the shard’s run, and the Store, the run health check and gym eval aggregate read the shard like an eval run. Each copy’s rollout id is the row’s own id (its _ng_rollout_id, or <task>-<rollout>) prefixed with the shard’s run id, so the shard’s model calls are captured under keys of its own; a caller that looks a capture up reads its key from the returned row (maybe_rollout_id_from_run_body). gym eval run with +input_jsonl_fpath=<dir>/rollouts-<shard>_materialized_inputs.jsonl +output_jsonl_fpath=<dir>/rollouts-<shard>.jsonl +resume_from_cache=true resumes a shard and runs only what it still lacks.

Collect rollouts for a whole config. Wrapped in the run-scoped job span.

This is the driver side of an evaluation run and the outermost span Gym produces. Rollouts are not its descendants: each dispatch opens its own trace, linked to this span (see _dispatch_span), so a run’s traces stay bounded by one rollout each while job keeps the run’s own lifecycle. job is in the default preset and not in per_rollout.

Yield (row, completed_dict | RolloutFailureRecord) in completion order.

Expected HTTP, JSON and result-shape failures become bounded records. Programming errors and cancellation still propagate. Call preprocess_examples first, or provide stable explicit rollout ids. This does not persist or retry work; a direct caller owns those decisions. Supply run_id to associate direct calls with a run, overriding any row stamp. Without it, existing row stamps are preserved; unstamped rows share a new invocation ID. Returned rows are invocation-owned copies; input rows are not modified. run_examples retains its behavior.

On early stopping, await the iterator’s aclose() to cancel and join unfinished client tasks, or aclose_and_drain() to also retrieve settled, unconsumed outcomes. Both have the same server/agent cleanup limitations documented by run_examples.

Return a client for the head server, reusing the cached one when the helper opted in.

Without opt-in, every call fetches the configuration from the head server again. An opted-in helper keeps one client and fetches again only when the head config differs from the cached one. Build a new helper after the configuration behind the same head changes.

One sample of requests waiting and running, summed over every server.

Raise the strict run’s rejection, if any; see completion_rejection for the arguments.

The error that rejects unfinished work, or None; omissions alongside completed tasks are allowed.

Collection decides this before it computes the aggregate, so a rejected run publishes no score. Drivers must supply whether any tasks can still be retried and whether this invocation saved outcomes that made progress (not just interruptions). Assuming either can cause endless automatic restarts when a driver has no work left or makes no progress.

final_zero_rows are failure rows that settle their rollout as a scored zero (see _final_zero_failure_rows): resume does not run them again, so they are completed samples. A caller that passes completed has already counted them.

Run the close hook once, after the outcomes a close still delivers have been enriched.

Cancel unfinished rollouts and await the end of their client tasks in this process.

Cancelling an in-flight /run request closes its connection. A server using ClientDisconnectCancellationMiddleware subsequently cancels its handler. This method does not wait for or acknowledge server/agent cleanup, which depends on the agent and may continue after return; it does not guarantee sandbox release or retire an attempt.

A handed-out awaitable may return an already completed result or raise its exception, including CancelledError for cancelled work. In bounded mode, an awaitable first awaited after closure raises CancelledError even if a result was ready. This error propagates in the awaiting caller’s task without a call to that task’s cancel(); asyncio.gather and TaskGroup treat it as cancellation. Iteration stops after closure. Safe to call again after closure completes, or before first use (no rollouts start).

Stop admission, cancel unfinished client tasks, and preserve completed outcomes.

Before its first suspension, this operation stops admitting rollouts and snapshots completed outcomes not yet claimed by a completion awaitable. It then cancels and awaits unfinished client tasks, and returns settled futures yielding (row, result) or raising the original exception. No new rollouts start while draining, including in bounded mode. Outcomes are transferred once; repeated calls return an empty list.

A handed-out awaitable that already claimed an outcome retains it. Other issued awaitables raise CancelledError. Callers with an outstanding waiter must also await that waiter and consume its outcome if it was not cancelled. The returned futures have their exceptions retrieved even if abandoned; awaiting them still raises. An outcome already cancelled at the snapshot is returned as a cancelled future and raises CancelledError when awaited. Tasks cancelled by this operation are excluded.

Calling this before first use starts nothing. Calling it after aclose() returns nothing. Like aclose(), it waits only for local client tasks, not server/agent cleanup, and does not guarantee sandbox release.

Hold the shard’s lock across processes, so opening and each write never interleave with another’s.

The lock file is never unlinked or replaced, so every writer of a label locks the same object; it also holds the token of the writer that opened the shard last.

The manifest gym eval run computes when it resumes this shard, with every row’s attempt reserved.

That resume names the shard’s materialized inputs as its input and the shard as its output, and leaves the other collection settings at their defaults: the same source file, the same rows and, with the same servers, the same configuration. Collection numbers each task’s rollouts 0, 1, … in input order, so rows numbered otherwise are recorded too, and that resume refuses them as other materialized inputs. A configuration this call cannot take an identity from still records the shard, whose manifest then verifies its own inputs only, so a resume refuses it and says why.

Write the shard’s status as it stands, for a caller that stops consuming before every row is in.

Append one finished rollout’s record; returns the file written, or None for an omitted row not kept.

Raises:

• RuntimeError: Another writer opened the same shard since this one did.