Nemo Run#

This workload (test_template_name is NemoRun) executes NeMo training and fine-tuning tasks using the NeMo Run framework.

Usage Examples#

Test TOML example:

name = "my_nemo_test"
description = "Example NeMo Run test"
test_template_name = "NemoRun"

[cmd_args]
recipe = "llama3_8b"
task = "pretrain"

Test Scenario example:

name = "nemo-run-test"

[[Tests]]
id = "nemo.1"
num_nodes = 4
time_limit = "02:00:00"

test_name = "my_nemo_test"

Test-in-Scenario example:

name = "nemo-run-test"

[[Tests]]
id = "nemo.1"
num_nodes = 4
time_limit = "02:00:00"

name = "my_nemo_test"
description = "Example NeMo Run test"
test_template_name = "NemoRun"

  [Tests.cmd_args]
  recipe = "llama3_8b"
  task = "pretrain"

API Documentation#

Command Arguments#

class cloudai.workloads.nemo_run.nemo_run.NeMoRunCmdArgs(

*,

docker_image_url: str,

task: str,

recipe_name: str,

num_layers: int | None = None,

trainer: ~cloudai.workloads.nemo_run.nemo_run.Trainer = <factory>,

log: ~cloudai.workloads.nemo_run.nemo_run.Log = <factory>,

data: ~cloudai.workloads.nemo_run.nemo_run.Data = <factory>,

optim: ~cloudai.workloads.nemo_run.nemo_run.Optim = <factory>,

**extra_data: ~typing.Any,

)[source]#

Bases: CmdArgs

NeMoRun test command arguments.

Test Definition#

class cloudai.workloads.nemo_run.nemo_run.NeMoRunTestDefinition(*, name: str, description: str, test_template_name: str, cmd_args: ~cloudai.workloads.nemo_run.nemo_run.NeMoRunCmdArgs, dse_excluded_args: list[str] = <factory>, extra_env_vars: dict[str, str | ~typing.List[str]] = {}, extra_cmd_args: dict[str, str] = {}, extra_container_mounts: list[str] = [], git_repos: list[~cloudai._core.installables.git_repo.GitRepo] = [], nsys: ~cloudai.models.workload.NsysConfiguration | None = None, predictor: ~cloudai.models.workload.PredictorConfig | None = None, agent: str = 'grid_search', agent_steps: int = 1, agent_metrics: list[str] = ['default'], agent_reward_function: str = 'inverse', agent_config: dict[str, ~typing.Any] | None = None, env_params: dict[str, ~cloudai.configurator.env_params.EnvParamSpec] = <factory>, script: ~cloudai._core.installables.file.File = File(src=PosixPath('/home/runner/work/cloudai/cloudai/src/cloudai/workloads/nemo_run/cloudai_nemorun.py')))[source]#

Bases: TestDefinition

NeMoRun test definition.

property installables: list[Installable]#: Get list of installable objects.

constraint_check( tr: TestRun, system: System | None, ) → bool[source]#: Check constraints for NeMoRun.

property update_num_train_samples: int | None#: Calculate num_train_samples based on global_batch_size and max_steps.

property is_domain_randomization_enabled: bool#

at least one env_params annotation.

Type:: Whether the config declares domain randomization

is_dse_excluded_arg(path: str) → bool#: Return whether a dot-separated cmd_args path should be ignored by DSE.

is_env_sampled(cmd_args_path: str) → bool#: Whether a cmd_args field is env-sampled (env draws it per trial, not the agent).

validator validate_env_params » all fields#

Validate env_params annotations against cmd_args.

env_params is an annotation: each key names a cmd_args field whose value is the candidate set (the single source of truth), and the entry carries only how to sample. So each key must name a real cmd_args field whose value is a candidate list; a scalar is already fixed, so annotating it is a meaningless label and is rejected here. When weights are declared, the list needs >= 2 values and the weights must align 1:1 with it. Sampling, persistence, the per-trial cmd_args overlay, and the cache key all live in CloudAIGymEnv; keeping this shape check in core lets the overlay stay agent- and workload-agnostic rather than re-implemented per workload.