nemo_gym.rollout_collection

View as Markdown

Module Contents

Classes

NameDescription
E2ERolloutCollectionConfigSpin up all necessary servers and perform a batch of rollout collection using each dataset inside the provided configs.
RolloutAggregationConfigAggregate metrics across rollout shards produced by gym eval run --no-serve +disable_aggregation=true.
RolloutAggregationHelper-
RolloutCollectionConfigPerform a batch of rollout collection.
RolloutCollectionHelper-
SharedRolloutCollectionConfig-

Functions

NameDescription
_agent_request_failure_rowOne sidecar row for a /run call that came back without a result.
_attach_ng_perf-
_attach_trajectory_record-
_build_ng_perfAssemble the per-rollout ng_perf summary from ng_trajectory.
_build_trajectory_record-
_expand_input_globExpand a glob-or-comma-separated-globs string into a sorted, deduplicated list of paths.
_failure_rows_counted_as_zeroSidecar rows the caller opted to count in the metrics denominator.
_get_max_rollout_attemptsRead NEMO_GYM_MAX_ROLLOUT_ATTEMPTS (positive int) or default to 3.
_has_observation_gap-
_nonnegative_int-
_normalize_health_check_ignored_checks-
_rollout_for_exportReturn an exporter view without the complete trajectory or raw capture payloads.
_rollout_request_debug_summary-
_strip_capture_payloads-
_trajectory_identity-
_truncated_bodyDecode at most the kept prefix, so a huge error page is never decoded in full.
loads_jsonl_lineParse one JSONL line, raising a clean ConfigError (naming file + line) on malformed JSON.

Data

AGENT_REQUEST_FAILED_FAILURE_CLASS

AGENT_RUN_ERROR_FAILURE_CLASS

NG_FAILURE_CLASS_KEY

NG_NO_PERSIST_KEY

NG_PERF_KEY

NG_TERMINAL_KEY

NG_TRAJECTORY_KEY

_AGENT_DID_NOT_RUN_STATUSES

_DEFAULT_MAX_ROLLOUT_ATTEMPTS

_MAX_FAILURE_BODY_CHARS

_MODEL_CALL_PAYLOAD_KEYS

_NG_ROLLOUT_LATENCY_MS_KEY

_NO_RESULT_FAILURE_CLASSES

_RUN_FAILURE_ERRORS

__getattr__

_failures_path_for

logger

API

class nemo_gym.rollout_collection.E2ERolloutCollectionConfig()

Bases: SharedRolloutCollectionConfig

Spin up all necessary servers and perform a batch of rollout collection using each dataset inside the provided configs.

Examples:

gym eval run +output_jsonl_fpath=weather_rollouts.jsonl +num_samples_in_parallel=10
reuse_existing_data_preparation
bool = False
split
Union[Literal['train'], Literal['validation'], Literal['benchmark']]
nemo_gym.rollout_collection.E2ERolloutCollectionConfig._reject_input_jsonl_fpath(
data
)
classmethod
class nemo_gym.rollout_collection.RolloutAggregationConfig()

Bases: BaseNeMoGymCLIConfig

Aggregate metrics across rollout shards produced by gym eval run --no-serve +disable_aggregation=true.

Reads every JSONL file matching input_glob, computes aggregate metrics by POSTing to each agent server’s /aggregate_metrics endpoint over the global union of records, and writes a single <output_jsonl_fpath stem>_aggregate_metrics.json next to the rollouts. By default also concatenates all shards into output_jsonl_fpath.

Examples:

gym eval aggregate "+config_paths=[benchmarks/aime24/config.yaml,responses_api_models/vllm_model/configs/vllm_model.yaml]" +input_glob='results/rollouts-rs*-chunk*.jsonl' +output_jsonl_fpath=results/rollouts.jsonl
count_failure_classes_as_zero
List[str]
disable_health_check
bool
health_check_ignored_checks
List[str]
health_check_workers
Optional[int]
input_glob
str
merge_shards
bool
output_jsonl_fpath
str
nemo_gym.rollout_collection.RolloutAggregationConfig._validate_health_check_ignored_checks(
value
)
classmethod
class nemo_gym.rollout_collection.RolloutAggregationHelper()

Bases: BaseModel

nemo_gym.rollout_collection.RolloutAggregationHelper.run_from_config(
config: nemo_gym.rollout_collection.RolloutAggregationConfig
) -> typing.Optional[pathlib.Path]
async
class nemo_gym.rollout_collection.RolloutCollectionConfig()

Bases: SharedRolloutCollectionConfig

Perform a batch of rollout collection.

Examples:

gym eval run --no-serve +agent_name=example_single_tool_call_simple_agent +input_jsonl_fpath=weather_query.jsonl +output_jsonl_fpath=weather_rollouts.jsonl +limit=100 +num_repeats=4 +num_samples_in_parallel=10
agent_map
Optional[Dict[str, str]]
agent_name
Optional[str]
fan_out
Optional[Dict[str, List[str]]]
input_jsonl_fpath
str
limit
Optional[int]
materialized_jsonl_fpath
Path
num_repeats
Union[int, Dict[str, int]]
num_repeats_add_seed
bool
prompt_config
Optional[str]
resume_from_cache
bool
skills
Optional[SkillsConfig]
nemo_gym.rollout_collection.RolloutCollectionConfig._coerce_null_num_repeats(
v
)
classmethod
class nemo_gym.rollout_collection.RolloutCollectionHelper()

Bases: BaseModel

nemo_gym.rollout_collection.RolloutCollectionHelper._call_aggregate_metrics(
results: typing.List[typing.Dict],
rows: typing.List[typing.Dict],
output_fpath: pathlib.Path
) -> typing.Optional[pathlib.Path]
async

Call /aggregate_metrics on each agent server after rollouts complete.

Writes a single _aggregate_metrics.json with one entry per agent (same shape as the old _agent_metrics.json). Returns the file path.

nemo_gym.rollout_collection.RolloutCollectionHelper._load_from_cache(
config: nemo_gym.rollout_collection.RolloutCollectionConfig
) -> typing.Tuple[typing.List[typing.Dict], typing.List[typing.Dict], typing.List[typing.Dict], typing.List[typing.List[str]]]
nemo_gym.rollout_collection.RolloutCollectionHelper._preprocess_raw_rows(
raw_rows: typing.List[typing.Tuple[int, str, typing.Dict]],
config: nemo_gym.rollout_collection.RolloutCollectionConfig
) -> typing.List[typing.Dict]
staticmethod
nemo_gym.rollout_collection.RolloutCollectionHelper._preprocess_rows_from_config(
config: nemo_gym.rollout_collection.RolloutCollectionConfig
) -> typing.List[typing.Dict]
nemo_gym.rollout_collection.RolloutCollectionHelper._run_from_config(
config: nemo_gym.rollout_collection.RolloutCollectionConfig
) -> typing.Tuple[typing.List[typing.Dict]]
async
nemo_gym.rollout_collection.RolloutCollectionHelper._validate_agent_names(
examples: typing.List[typing.Dict],
global_config_dict: omegaconf.DictConfig
) -> None
staticmethod

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.

nemo_gym.rollout_collection.RolloutCollectionHelper._validate_agent_pairings(
examples: typing.List[typing.Dict],
global_config_dict: omegaconf.DictConfig
) -> None
staticmethod

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

nemo_gym.rollout_collection.RolloutCollectionHelper.preprocess_examples(
examples: typing.List[typing.Dict],
agent_map: typing.Optional[typing.Dict[str, str]] = None,
fan_out: typing.Optional[typing.Dict[str, typing.List[str]]] = None,
num_repeats: typing.Union[int, typing.Dict[str, int]] = 1,
num_repeats_add_seed: bool = False,
global_config_dict: typing.Optional[omegaconf.DictConfig] = None
) -> typing.List[typing.Dict]

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. Input rows are not mutated; the expanded, stamped copies are returned.

nemo_gym.rollout_collection.RolloutCollectionHelper.resolve_task_sources(
examples: typing.List[typing.Dict],
global_config_dict: omegaconf.DictConfig
) -> None
staticmethod

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 against the merged config, first match wins:

  • the instance is an agent (self-contained environment) -> route to it directly;
  • the instance is a resources server -> route to the unique agent whose resources_server.name edge points at it (inversion of the edge every agent config already declares);
  • zero or 2+ candidate agents, or an unknown/non-routable instance -> hard error (2+ names +agent_map as 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.

nemo_gym.rollout_collection.RolloutCollectionHelper.run_examples(
examples: typing.List[typing.Dict],
head_server_config: typing.Optional[nemo_gym.config_types.BaseServerConfig] = None,
semaphore: typing.Optional[asyncio.Semaphore] = None,
route_failures_to_sidecar: bool = False
) -> typing.Iterator[asyncio.Future]

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 is set by managed collection (run_from_config), where a failed /run becomes a failure row instead of ending every rollout still in flight. Direct callers such as NeMo-RL leave it off and keep receiving the exception.

nemo_gym.rollout_collection.RolloutCollectionHelper.run_from_config(
config: nemo_gym.rollout_collection.RolloutCollectionConfig
) -> typing.Tuple[typing.List[typing.Dict]]
async

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, so every rollout it dispatches is a descendant of it. job is in the default preset but deliberately not in per_rollout, where each rollout is meant to be its own bounded root trace.

nemo_gym.rollout_collection.RolloutCollectionHelper.setup_server_client(
head_server_config: typing.Optional[nemo_gym.config_types.BaseServerConfig] = None
) -> nemo_gym.server_utils.ServerClient
class nemo_gym.rollout_collection.SharedRolloutCollectionConfig()

Bases: UploadRolloutsConfigMixin, BaseNeMoGymCLIConfig

count_failure_classes_as_zero
List[str]
disable_aggregation
bool
disable_health_check
bool
health_check_ignored_checks
List[str]
health_check_workers
Optional[int]
num_samples_in_parallel
Optional[int]
output_jsonl_fpath
str
responses_create_params
Dict[str, Any]
rollout_collection_driver
Optional[str]
nemo_gym.rollout_collection.SharedRolloutCollectionConfig._validate_health_check_ignored_checks(
value
)
classmethod
nemo_gym.rollout_collection._agent_request_failure_row(
exc: BaseException,
status: typing.Optional[int]
) -> typing.Dict[str, typing.Any]

One sidecar row for a /run call that came back without a result.

No reward and no response: an infrastructure failure is not a verifier score of zero, and a placeholder would read as real generation data to token capture, aggregation and trainers. The class says whether the rollout ran. A NeMo Gym agent answers 500 when its own handler raises, so any status it answered with means the agent ran and broke, which is also how a model server rejecting the model’s own output arrives here. A gateway status, or no reply to take a status from, says nothing about the rollout. Neither class carries a reward; an evaluation that wants the first counted names it in count_failure_classes_as_zero.

nemo_gym.rollout_collection._attach_ng_perf(
result: dict[str, typing.Any],
observability_enabled: bool
) -> None
nemo_gym.rollout_collection._attach_trajectory_record(
row: dict[str, typing.Any],
result: dict[str, typing.Any]
) -> None
nemo_gym.rollout_collection._build_ng_perf(
result: dict[str, typing.Any],
rollout_latency_ms: typing.Optional[float]
) -> typing.Optional[dict[str, typing.Any]]

Assemble the per-rollout ng_perf summary from ng_trajectory.

Returns None (ng_perf stays absent) unless at least one reasoning turn was observed: per-turn evidence is needed rather than just raw model-call capture, so a rollout collected with observability disabled produces no ng_perf at all.

Token fields are summed only over model calls owned by a reasoning-turn AgentInvocation, excluding compaction calls — mixing in compaction overhead would skew the token efficiency signal.

num_turns counts reasoning turns summed across all invocations (an AgentInvocation is one root-agent or subagent conversation that may span many turns). Each invocation contributes its explicit TrajectoryTurn count when the harness emits turn records, falling back to its owned model-call count (one assistant response per turn), then to 1 (an invocation that ran had at least one turn) — so hybrid trajectories where only some invocations report turns still count every conversation.

token_observability_coverage reports what fraction of those turns actually resolved to a captured call: a turn whose ModelCallRef was unmatched or ambiguous silently loses its tokens from the sums below, and this is the only signal that it happened.

nemo_gym.rollout_collection._build_trajectory_record(
row: dict[str, typing.Any],
result: dict[str, typing.Any]
) -> nemo_gym.rollout_observability.TrajectoryRecord
nemo_gym.rollout_collection._expand_input_glob(
input_glob: str
) -> typing.List[str]

Expand a glob-or-comma-separated-globs string into a sorted, deduplicated list of paths.

Examples:

'results/rollouts.jsonl' -> ['results/rollouts.jsonl'] (if it exists)
'a/*.jsonl, b/*.jsonl' -> matches of both patterns, deduplicated
nemo_gym.rollout_collection._failure_rows_counted_as_zero(
failures_fpaths: typing.List[pathlib.Path],
failure_classes: typing.List[str],
scored_keys: set
) -> typing.List[typing.Dict[str, typing.Any]]

Sidecar rows the caller opted to count in the metrics denominator.

The last attempt of a rollout is the one that stands, so it is selected across every failure class before the wanted classes are picked out. Selecting the other way round would let a stale attempt be counted after a later one landed in a class the caller did not ask for.

A row that already carries a reward is counted as it stands. A row that carries none records that no rollout happened, so it is counted as a zero here and only here: the score enters the metric input, never the sidecar or the rollouts jsonl, which keeps the artifacts free of a verdict no verifier gave. A rollout that also succeeded is never counted.

nemo_gym.rollout_collection._get_max_rollout_attempts() -> int

Read NEMO_GYM_MAX_ROLLOUT_ATTEMPTS (positive int) or default to 3.

nemo_gym.rollout_collection._has_observation_gap(
result: dict[str, typing.Any],
code: str
) -> bool
nemo_gym.rollout_collection._nonnegative_int(
value: typing.Any
) -> typing.Optional[int]
nemo_gym.rollout_collection._normalize_health_check_ignored_checks(
value
) -> typing.List[str]
nemo_gym.rollout_collection._rollout_for_export(
result: dict[str, typing.Any]
) -> dict[str, typing.Any]

Return an exporter view without the complete trajectory or raw capture payloads.

nemo_gym.rollout_collection._rollout_request_debug_summary(
row: typing.Dict[str, typing.Any]
) -> typing.Dict[str, typing.Any]
nemo_gym.rollout_collection._strip_capture_payloads(
result: dict[str, typing.Any]
) -> None
nemo_gym.rollout_collection._trajectory_identity(
row: dict[str, typing.Any]
) -> tuple[str, str]
nemo_gym.rollout_collection._truncated_body(
body: typing.Optional[bytes]
) -> typing.Optional[str]

Decode at most the kept prefix, so a huge error page is never decoded in full.

nemo_gym.rollout_collection.loads_jsonl_line(
raw,
fpath,
line_no: int
)

Parse one JSONL line, raising a clean ConfigError (naming file + line) on malformed JSON.

nemo_gym.rollout_collection.AGENT_REQUEST_FAILED_FAILURE_CLASS = 'agent_request_failed'
nemo_gym.rollout_collection.AGENT_RUN_ERROR_FAILURE_CLASS = 'agent_run_error'
nemo_gym.rollout_collection.NG_FAILURE_CLASS_KEY = '_ng_failure_class'
nemo_gym.rollout_collection.NG_NO_PERSIST_KEY = '_ng_no_persist'
nemo_gym.rollout_collection.NG_PERF_KEY = 'ng_perf'
nemo_gym.rollout_collection.NG_TERMINAL_KEY = '_ng_failure_terminal'
nemo_gym.rollout_collection.NG_TRAJECTORY_KEY = 'ng_trajectory'
nemo_gym.rollout_collection._AGENT_DID_NOT_RUN_STATUSES = frozenset({429, 502, 503, 504})
nemo_gym.rollout_collection._DEFAULT_MAX_ROLLOUT_ATTEMPTS = 3
nemo_gym.rollout_collection._MAX_FAILURE_BODY_CHARS = 2000
nemo_gym.rollout_collection._MODEL_CALL_PAYLOAD_KEYS = ('request', 'response', 'request_raw', 'response_raw')
nemo_gym.rollout_collection._NG_ROLLOUT_LATENCY_MS_KEY = '_ng_rollout_latency_ms'
nemo_gym.rollout_collection._NO_RESULT_FAILURE_CLASSES = frozenset({AGENT_REQUEST_FAILED_FAILURE_CLASS, AGENT_RUN_ERROR_FAILURE_CLASS})
nemo_gym.rollout_collection._RUN_FAILURE_ERRORS = (ClientError, orjson.JSONDecodeError, TimeoutError)
nemo_gym.rollout_collection.__getattr__ = moved_attr_getter(__name__, {'collect_rollouts': 'nemo_gym.cli.eval', 'aggregate...
nemo_gym.rollout_collection._failures_path_for = failures_path_for
nemo_gym.rollout_collection.logger = logging.getLogger(__name__)