nemo_gym.rollout_collection

Module Contents

Classes

Name	Description
`E2ERolloutCollectionConfig`	Spin up all necessary servers and perform a batch of rollout collection using each dataset inside the provided configs.
`RolloutAggregationConfig`	Aggregate metrics across rollout shards produced by `ng_collect_rollouts +disable_aggregation=true`.
`RolloutAggregationHelper`	-
`RolloutCollectionConfig`	Perform a batch of rollout collection.
`RolloutCollectionHelper`	-
`SharedRolloutCollectionConfig`	-

Functions

Name	Description
`_expand_input_glob`	Expand a glob-or-comma-separated-globs string into a sorted, deduplicated list of paths.
`_rollout_request_debug_summary`	-
`aggregate_rollouts`	-
`collect_rollouts`	-

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:

ng_collect_rollouts         +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']]

class nemo_gym.rollout_collection.RolloutAggregationConfig()

Bases: BaseNeMoGymCLIConfig

Aggregate metrics across rollout shards produced by ng_collect_rollouts +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:

ng_aggregate_rollouts         "+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

input_glob

str

merge_shards

bool

output_jsonl_fpath

str

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:

ng_collect_rollouts         +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_name

Optional[str]

input_jsonl_fpath

str

limit

Optional[int]

materialized_jsonl_fpath

Path

num_repeats

Optional[int]

num_repeats_add_seed

bool

prompt_config

Optional[str]

resume_from_cache

bool

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_rows_from_config(
    config: nemo_gym.rollout_collection.RolloutCollectionConfig
) -> typing.List[typing.Dict]

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
) -> typing.Iterator[asyncio.Future]

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

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.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: BaseNeMoGymCLIConfig

disable_aggregation

bool

num_samples_in_parallel

Optional[int]

output_jsonl_fpath

str

responses_create_params

Dict[str, Any]

upload_rollouts_to_wandb

bool

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._rollout_request_debug_summary(
    row: typing.Dict[str, typing.Any]
) -> typing.Dict[str, typing.Any]

nemo_gym.rollout_collection.aggregate_rollouts()

nemo_gym.rollout_collection.collect_rollouts()