nemo_gym.environment.in_process

Run a Gym benchmark from Python: plan it in one call, then run single tasks against running servers.

planned = plan_benchmark("ifeval")                        # resolve, prepare, plan and collate
runner = TaskRunner(planned.rows, registry_server=BaseServerConfig(host="127.0.0.1", port=11000))
result = await runner.run_task(planned.plan.tasks[0].task_id.task_id)
result.rewards                                            # one reward per repeat

To start the servers from Python too, serve the benchmark and give the yielded registry server to the runner:

overrides = {"policy_base_url": url, "policy_api_key": key, "policy_model_name": model}
with serve_benchmark("ifeval", overrides=overrides) as registry_server:
    runner = TaskRunner(planned.rows, registry_server=registry_server)
    ...

plan_benchmark builds the plan gym eval plan --benchmark NAME writes and collates the benchmark’s rows. It reads no command-line arguments and leaves Gym’s process-wide configuration untouched, so a long-lived process can call it. serve_benchmark starts the servers gym env start --benchmark NAME starts, from the same kind of config dict; each server runs in its own process and environment either way. The registry server is the server that serves the run’s config and the other servers’ addresses; a TaskRunner talks to it.

This is the API for callers outside Gym. They import the names nemo_gym.environment exports (plan_benchmark, serve_benchmark, TaskRunner, PlannedBenchmark, TaskResult and EvalPlan) and ConfigError from nemo_gym.config. Signatures and result fields change only additively, and tests pin them. TaskResult.results holds Gym’s raw /run responses, whose keys can change: depend on rewards, and use results for debugging.

Module Contents

Classes

Name	Description
`PlannedBenchmark`	A prepared benchmark: its plan (task ids and input digests) and the collated rows to run.
`TaskResult`	One task’s rollouts: a reward per repeat and Gym’s raw `/run` results, in the same order.
`TaskRunner`	Runs single tasks of collated rows against Gym servers that are already running.

Functions

Name	Description
`_benchmark_config`	The resolved config for one benchmark, built without reading sys.argv or caching it process-wide.
`_checked_overrides`	The overrides as a dict, refusing a dotted key at any depth.
`_collated_rows`	-
`_first_row_of_each_task`	-
`_refuse_dotted_keys`	-
`_with_rollout_id`	-
`plan_benchmark`	Resolve the benchmark `NAME[/FLAVOR]`, prepare its data, build its plan and collate its rows.
`serve_benchmark`	Start the benchmark’s servers, and stop them on exit. Yields the registry server a `TaskRunner` talks to.

API

class nemo_gym.environment.in_process.PlannedBenchmark(
    plan: nemo_gym.eval.eval_plan.EvalPlan,
    rows: list[dict[str, typing.Any]]
)

Dataclass

A prepared benchmark: its plan (task ids and input digests) and the collated rows to run.

gym_commit

Optional[str]

The Gym commit the plan was made with, or None when Gym was not installed from a git checkout.

plan

EvalPlan

rows

list[dict[str, Any]]

nemo_gym.environment.in_process.PlannedBenchmark.from_json(
    text: str
) -> nemo_gym.environment.in_process.PlannedBenchmark

classmethod

nemo_gym.environment.in_process.PlannedBenchmark.task_digests() -> dict[str, str]

Each planned task’s id mapped to the digest of its input, which changes when the input changes.

nemo_gym.environment.in_process.PlannedBenchmark.task_ids() -> list[str]

nemo_gym.environment.in_process.PlannedBenchmark.to_json() -> str

The planned benchmark as JSON text, for handing it to another process; from_json reads it back.

class nemo_gym.environment.in_process.TaskResult(
    task_id: str,
    rewards: list[float],
    results: list[dict[str, typing.Any]]
)

Dataclass

One task’s rollouts: a reward per repeat and Gym’s raw /run results, in the same order.

results

list[dict[str, Any]]

rewards

list[float]

task_id

str

class nemo_gym.environment.in_process.TaskRunner(
    rows: collections.abc.Sequence[collections.abc.Mapping[str, typing.Any]],
    registry_server: nemo_gym.config.config_types.BaseServerConfig
)

Runs single tasks of collated rows against Gym servers that are already running.

The rows are checked and indexed by task id once. Each run_task call then selects one task’s row, repeats it, sends it through the agent’s /run and returns the rewards. Concurrent calls in one event loop are supported.

_helper

= RolloutCollectionHelper()

_row_of_task

= _first_row_of_each_task(rows)

nemo_gym.environment.in_process.TaskRunner.run_task(
    task_id: str,
    num_repeats: int = 1,
    rollout_id: typing.Optional[str] = None
) -> nemo_gym.environment.in_process.TaskResult

async

Run num_repeats rollouts of one task.

rollout_id names the rollouts for Gym’s rollout correlation (e.g. verification events); repeats get <rollout_id>-r<n>. It must differ between concurrent calls; a fresh one is made when it is None.

nemo_gym.environment.in_process.TaskRunner.task_ids() -> list[str]

nemo_gym.environment.in_process._benchmark_config(
    config_path: str,
    overrides: typing.Optional[collections.abc.Mapping[str, typing.Any]]
) -> omegaconf.DictConfig

The resolved config for one benchmark, built without reading sys.argv or caching it process-wide.

nemo_gym.environment.in_process._checked_overrides(
    overrides: typing.Optional[collections.abc.Mapping[str, typing.Any]]
) -> dict[str, typing.Any]

The overrides as a dict, refusing a dotted key at any depth.

Merged as a dict, "judge.model" would be one key with a dot in its name, while ++judge.model=x on the command line sets model inside judge. Refusing it keeps the dict and the command line from disagreeing.

nemo_gym.environment.in_process._collated_rows(
    global_config_dict: omegaconf.DictConfig,
    output_dir: pathlib.Path
) -> list[dict[str, typing.Any]]

nemo_gym.environment.in_process._first_row_of_each_task(
    rows: collections.abc.Sequence[collections.abc.Mapping[str, typing.Any]]
) -> dict[str, collections.abc.Mapping[str, typing.Any]]

nemo_gym.environment.in_process._refuse_dotted_keys(
    node: typing.Any
) -> None

nemo_gym.environment.in_process._with_rollout_id(
    row: dict[str, typing.Any],
    rollout_id: str,
    index: int,
    num_repeats: int
) -> dict[str, typing.Any]

nemo_gym.environment.in_process.plan_benchmark(
    benchmark: str,
    task_ids: typing.Optional[collections.abc.Sequence[str]] = None,
    overrides: typing.Optional[collections.abc.Mapping[str, typing.Any]] = None,
    work_dir: typing.Optional[pathlib.Path] = None
) -> nemo_gym.environment.in_process.PlannedBenchmark

Resolve the benchmark NAME[/FLAVOR], prepare its data, build its plan and collate its rows.

overrides are config keys merged on top, as a nested dict like the config file’s; a dotted key is refused (ConfigError), because it would be one literal key here but a path on the command line. work_dir receives the collated files; a temporary directory is used when it is None.

The caller’s command line and active exporters are left alone. The benchmark’s prepare script sees overrides as its Gym configuration, which is this process’s configuration only while it runs.

nemo_gym.environment.in_process.serve_benchmark(
    benchmark: str,
    model_type: str = 'openai_model',
    overrides: typing.Optional[collections.abc.Mapping[str, typing.Any]] = None,
    registry_port: typing.Optional[int] = None
) -> collections.abc.Iterator[nemo_gym.config.config_types.BaseServerConfig]

Start the benchmark’s servers, and stop them on exit. Yields the registry server a TaskRunner talks to.

The benchmark NAME[/FLAVOR] and the model type are named as in gym env start --benchmark NAME --model-type TYPE. overrides are config keys merged on top, as in plan_benchmark, so planning and serving read the same dict: nested dicts, and a dotted key is refused (ConfigError). The model’s policy_base_url, policy_api_key and policy_model_name are among them. registry_port fixes the registry server’s port, so a caller that serves from another process can know the address before the servers start; Gym chooses the port when it is None.

Returns once every server is ready; raises ConfigError if a server or a model endpoint is not, or if registry_port is in use. Like gym env start, it sets this process’s Gym configuration, so serve one benchmark per process.