nemo_gym.environment.benchmark_preparation

Prepare a resolved config’s benchmark datasets by running their prepare scripts, and plan a benchmark.

The gym eval prepare and gym eval plan commands and callers that hold a resolved config in Python share this code. It sits in environment because it builds on BenchmarkConfig. It is not nemo_gym.data.dataset_preparation, which prepares training data.

Module Contents

Classes

Name	Description
`PreparationFailedError`	A benchmark’s data preparation failed in its own code or environment.
`PrepareBenchmarkConfig`	Prepare benchmark data by running the benchmark’s prepare.py script.

Functions

Name	Description
`_install_prepare_dependencies`	Install what a benchmark’s prepare script imports, before importing it.
`_multiprocess_benchmark_prepare_fn`	-
`_prepare_environment`	Use explicit dependency arguments as before, or sync an isolated locked project.
`_prepare_json_default`	Serialize path arguments without silently coercing unsupported values.
`_run_isolated_prepare`	Keep preparer stdout visible and pass structured results over a separate file.
`datasets_to_prepare`	Every dataset the resolved config prepares: benchmark datasets and other datasets with a prepare_script.
`gym_relative_path`	A config path relative to the Gym checkout when it lies inside it, so plans match across machines.
`plan_for_config`	Prepare the one benchmark a resolved config declares and build its plan (shared by `gym eval plan`).
`prepare_datasets`	Run each dataset’s prepare script, after validating all of them.

Data

logger

API

class nemo_gym.environment.benchmark_preparation.PreparationFailedError(
    module: str,
    action: str,
    reason: str
)

Bases: RuntimeError

A benchmark’s data preparation failed in its own code or environment.

Not resumable: the CLI exits with stop_exit_code (76) and records the kind for the guard of a job chain’s next segment, which would fail the same way.

_init_args

= (module, action, reason)

stop_exit_code

= 76

nemo_gym.environment.benchmark_preparation.PreparationFailedError.__reduce__()

class nemo_gym.environment.benchmark_preparation.PrepareBenchmarkConfig()

Bases: BaseNeMoGymCLIConfig

Prepare benchmark data by running the benchmark’s prepare.py script.

The benchmark is identified from a config_paths entry pointing to a benchmarks/*/config.yaml file.

With use_cached_prepared_benchmarks=true, an existing prepared file is reused. A prepare.py whose output depends on its own code (for example, settings written into each row) can define is_prepared_data_current(fpath: Path) -> bool; when it returns False, the cached file is prepared again.

Examples:

gym eval prepare --benchmark aime24

num_prepare_benchmark_processes

int

prepare_script_args

Dict[str, Any]

use_cached_prepared_benchmarks

bool

uv_venv_dir

Path

nemo_gym.environment.benchmark_preparation._install_prepare_dependencies(
    benchmark_config: nemo_gym.environment.benchmarks.BenchmarkConfig
) -> None

Install what a benchmark’s prepare script imports, before importing it.

Gym cannot depend on every benchmark’s data-prep requirements, so a benchmark needing something extra had to shell out to pip from inside the prepare script itself. Declaring it on the dataset puts it in the config instead.

nemo_gym.environment.benchmark_preparation._multiprocess_benchmark_prepare_fn(
    args
)

nemo_gym.environment.benchmark_preparation._prepare_environment(
    benchmark_config: nemo_gym.environment.benchmarks.BenchmarkConfig,
    venv_root: pathlib.Path | None = None
) -> pathlib.Path | None

Use explicit dependency arguments as before, or sync an isolated locked project.

nemo_gym.environment.benchmark_preparation._prepare_json_default(
    value: typing.Any
) -> typing.Any

Serialize path arguments without silently coercing unsupported values.

nemo_gym.environment.benchmark_preparation._run_isolated_prepare(
    python: pathlib.Path,
    module: str,
    action: str,
    expected: pathlib.Path,
    args: dict,
    config: omegaconf.DictConfig | dict | None = None
) -> dict

Keep preparer stdout visible and pass structured results over a separate file.

nemo_gym.environment.benchmark_preparation.datasets_to_prepare(
    global_config_dict: omegaconf.DictConfig
) -> typing.List[nemo_gym.environment.benchmarks.BenchmarkConfig]

Every dataset the resolved config prepares: benchmark datasets and other datasets with a prepare_script.

nemo_gym.environment.benchmark_preparation.gym_relative_path(
    path: str
) -> str

A config path relative to the Gym checkout when it lies inside it, so plans match across machines.

nemo_gym.environment.benchmark_preparation.plan_for_config(
    global_config_dict: omegaconf.DictConfig,
    prepare_config: nemo_gym.environment.benchmark_preparation.PrepareBenchmarkConfig,
    task_ids: typing.Sequence[str] | None = None
) -> nemo_gym.eval.eval_plan.EvalPlan

Prepare the one benchmark a resolved config declares and build its plan (shared by gym eval plan).

nemo_gym.environment.benchmark_preparation.prepare_datasets(
    datasets: typing.Sequence[nemo_gym.environment.benchmarks.BenchmarkConfig],
    prepare_benchmark_config: nemo_gym.environment.benchmark_preparation.PrepareBenchmarkConfig
) -> None

Run each dataset’s prepare script, after validating all of them.

nemo_gym.environment.benchmark_preparation.logger = logging.getLogger(__name__)