nemo_gym.rollout_collection#

Module Contents#

Classes#

SharedRolloutCollectionConfig

E2ERolloutCollectionConfig

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

RolloutCollectionConfig

Perform a batch of rollout collection.

RolloutCollectionHelper

Functions#

collect_rollouts

API#

class nemo_gym.rollout_collection.SharedRolloutCollectionConfig(/, **data: typing.Any)[source]#

Bases: nemo_gym.config_types.BaseNeMoGymCLIConfig

output_jsonl_fpath: str#

‘Field(…)’

num_samples_in_parallel: Optional[int]#

‘Field(…)’

responses_create_params: Dict[str, Any]#

‘Field(…)’

upload_rollouts_to_wandb: bool#

‘Field(…)’

class nemo_gym.rollout_collection.E2ERolloutCollectionConfig(/, **data: typing.Any)[source]#

Bases: nemo_gym.rollout_collection.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

Initialization

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

split: Union[Literal[train], Literal[validation], Literal[benchmark]]#

None

reuse_existing_data_preparation: bool#

False

class nemo_gym.rollout_collection.RolloutCollectionConfig(/, **data: typing.Any)[source]#

Bases: nemo_gym.rollout_collection.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

Initialization

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

agent_name: Optional[str]#

‘Field(…)’

input_jsonl_fpath: str#

‘Field(…)’

limit: Optional[int]#

‘Field(…)’

num_repeats: Optional[int]#

‘Field(…)’

num_repeats_add_seed: bool#

‘Field(…)’

resume_from_cache: bool#

‘Field(…)’

prompt_config: Optional[str]#

‘Field(…)’

property materialized_jsonl_fpath: pathlib.Path#
class nemo_gym.rollout_collection.RolloutCollectionHelper(/, **data: typing.Any)[source]#

Bases: pydantic.BaseModel

_preprocess_rows_from_config(
config: nemo_gym.rollout_collection.RolloutCollectionConfig,
) → List[Dict][source]#
_load_from_cache(
config: nemo_gym.rollout_collection.RolloutCollectionConfig,
) → Tuple[List[Dict], List[Dict], List[Dict], List[List[str]]][source]#
async run_from_config(
config: nemo_gym.rollout_collection.RolloutCollectionConfig,
) → Tuple[List[Dict]][source]#
async _call_aggregate_metrics(
results: List[Dict],
rows: List[Dict],
output_fpath: pathlib.Path,
) → Optional[pathlib.Path][source]#

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.

run_examples(
examples: List[Dict],
head_server_config: Optional[nemo_gym.config_types.BaseServerConfig] = None,
semaphore: Optional[asyncio.Semaphore] = None,
) → Iterator[asyncio.Future][source]#

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

setup_server_client(
head_server_config: Optional[nemo_gym.config_types.BaseServerConfig] = None,
) → nemo_gym.server_utils.ServerClient[source]#
nemo_gym.rollout_collection.collect_rollouts()[source]#