From e727b9c51a89ea123bd896279bec12e70e9ba47f Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 21 Aug 2026 15:19:23 +0200 Subject: [PATCH 1/3] make sure workloads dependencies are configurable --- src/cloudai/workloads/ai_dynamo/ai_dynamo.py | 16 ++- .../workloads/aiconfig/aiconfigurator.py | 8 +- .../workloads/dynamo_mocker/dynamo_mocker.py | 6 +- .../megatron_bridge/megatron_bridge.py | 24 +++- .../slurm_command_gen_strategy.py | 6 +- .../workloads/nemo_launcher/nemo_launcher.py | 26 +++- .../slurm_command_gen_strategy.py | 1 + tests/ref_data/megatron-bridge.sbatch | 4 +- tests/test_acceptance.py | 120 +++++++++--------- .../test_command_gen_strategy_slurm.py | 28 ++-- .../test_command_gen_strategy_standalone.py | 18 +++ .../test_command_gen_strategy_standalone.py | 26 ++++ .../test_command_gen_strategy_slurm.py | 15 +++ .../test_command_gen_strategy_slurm.py | 11 ++ 14 files changed, 219 insertions(+), 90 deletions(-) diff --git a/src/cloudai/workloads/ai_dynamo/ai_dynamo.py b/src/cloudai/workloads/ai_dynamo/ai_dynamo.py index d3726fda1..f2195c528 100644 --- a/src/cloudai/workloads/ai_dynamo/ai_dynamo.py +++ b/src/cloudai/workloads/ai_dynamo/ai_dynamo.py @@ -394,6 +394,10 @@ class AIDynamoCmdArgs(CmdArgs): model_config = ConfigDict(extra="forbid") + dynamo_version: str = Field( + default="f7e468c7e8ff0d1426db987564e60572167e8464", + description="AI Dynamo Git commit, tag, or branch.", + ) docker_image_url: str startup_cmd: str | None = None startup_cmd_docker_image: str | None = None @@ -457,9 +461,7 @@ class AIDynamoTestDefinition(TestDefinition): _dcgm_exporter_image: Optional[DockerImage] = None _startup_cmd_docker_image: DockerImage | None = None script: File = File(Path(__file__).parent.parent / "ai_dynamo/ai_dynamo.sh") - repo: GitRepo = GitRepo( - url="https://github.com/ai-dynamo/dynamo.git", commit="f7e468c7e8ff0d1426db987564e60572167e8464" - ) + _repo: GitRepo | None = None _hf_model: HFModel | None = None constraints: Constraints = Constraints() @@ -483,6 +485,14 @@ def get_workload_map(self) -> dict[str, Workload]: self.cmd_args.aiperf.script.src.name: self.cmd_args.aiperf, } + @property + def repo(self) -> GitRepo: + """Return the AI Dynamo repository selected by ``cmd_args.dynamo_version``.""" + version = self.cmd_args.dynamo_version + if self._repo is None or self._repo.commit != version: + self._repo = GitRepo(url="https://github.com/ai-dynamo/dynamo.git", commit=version) + return self._repo + @property def docker_image(self) -> DockerImage: if not self._docker_image: diff --git a/src/cloudai/workloads/aiconfig/aiconfigurator.py b/src/cloudai/workloads/aiconfig/aiconfigurator.py index ef61047f9..76dad5805 100644 --- a/src/cloudai/workloads/aiconfig/aiconfigurator.py +++ b/src/cloudai/workloads/aiconfig/aiconfigurator.py @@ -18,7 +18,7 @@ from typing import List, Optional, Union -from pydantic import BaseModel, ConfigDict, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from cloudai.core import CmdArgs, Installable, PythonEnvironment, TestDefinition @@ -61,6 +61,10 @@ class Disagg(BaseModel): class AiconfiguratorCmdArgs(CmdArgs): """Command arguments for Aiconfigurator workload with nested agg/disagg configs.""" + requirements: str = Field( + default="aiconfigurator~=0.5.0", + description="Space-separated Python requirements installed into the Aiconfigurator environment.", + ) model_name: str system: str backend: str = "trtllm" @@ -91,7 +95,7 @@ def python_environment(self) -> PythonEnvironment: return PythonEnvironment( name="aiconfigurator", python_version="3.10", - requirements=["aiconfigurator~=0.5.0"], + requirements=self.cmd_args.requirements.split(), ) @property diff --git a/src/cloudai/workloads/dynamo_mocker/dynamo_mocker.py b/src/cloudai/workloads/dynamo_mocker/dynamo_mocker.py index 1ce6422e8..e48f1e510 100644 --- a/src/cloudai/workloads/dynamo_mocker/dynamo_mocker.py +++ b/src/cloudai/workloads/dynamo_mocker/dynamo_mocker.py @@ -218,6 +218,10 @@ def token_counts_positive(cls, v: Union[int, List[int]]) -> Union[int, List[int] class DynamoMockerCmdArgs(CmdArgs): """Top-level command arguments for the Dynamo Mocker workload.""" + requirements: str = Field( + default="ai-dynamo==1.3.0.post1 genai-perf==0.0.16 aiperf==0.11.0", + description="Space-separated Python requirements installed into the Dynamo Mocker environment.", + ) model_path: str = "Qwen/Qwen3-0.6B" nats_cmd: str = Field( default="nats-server -js", @@ -260,7 +264,7 @@ def python_environment(self) -> PythonEnvironment: self._python_environment = PythonEnvironment( name="dynamo-mocker", python_version="3.12", - requirements=["ai-dynamo", "genai-perf", "aiperf"], + requirements=self.cmd_args.requirements.split(), ) return self._python_environment diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 12f0c7178..627c010ad 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -66,6 +66,11 @@ class MegatronBridgeCmdArgs(CmdArgs): wandb_experiment_name: Optional[str] = Field(default=None) wandb_save_dir: Optional[str] = Field(default=None) wandb_version: str = Field(default="0.28.1", description="W&B version installed in the launcher environment.") + numpy_version: str = Field(default="1.26.4", description="NumPy version installed in the launcher environment.") + nemorun_version: str = Field( + default="v0.8.0", + description="NeMo Run Git commit, tag, or branch used by the launcher environment.", + ) # Retries max_retries: Optional[int] = Field(default=1) @@ -181,14 +186,18 @@ class MegatronBridgeTestDefinition(TestDefinition): cmd_args: MegatronBridgeCmdArgs - nemo_run_repo: GitRepo = GitRepo( - url="https://github.com/NVIDIA-NeMo/Run.git", - commit="main", - ) - _docker_image: Optional[DockerImage] = None _python_executable: Optional[PythonExecutable] = None _megatron_bridge_repo: Optional[GitRepo] = None + _nemo_run_repo: Optional[GitRepo] = None + + @property + def nemo_run_repo(self) -> GitRepo: + """Return the NeMo Run repository selected by ``cmd_args.nemorun_version``.""" + version = self.cmd_args.nemorun_version + if self._nemo_run_repo is None or self._nemo_run_repo.commit != version: + self._nemo_run_repo = GitRepo(url="https://github.com/NVIDIA-NeMo/Run.git", commit=version) + return self._nemo_run_repo @staticmethod def _select_megatron_bridge_repo(git_repos: list[GitRepo]) -> GitRepo | None: @@ -229,8 +238,9 @@ def docker_image(self) -> DockerImage: @property def python_executable(self) -> PythonExecutable: - if not self._python_executable: - self._python_executable = PythonExecutable(git_repo=self.nemo_run_repo) + nemo_run_repo = self.nemo_run_repo + if not self._python_executable or self._python_executable.git_repo != nemo_run_repo: + self._python_executable = PythonExecutable(git_repo=nemo_run_repo) return self._python_executable @property diff --git a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py b/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py index f34d2a65c..31cf500f4 100644 --- a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py @@ -90,6 +90,7 @@ def gen_exec_command(self) -> str: " ".join(parts), launcher_python, args.wandb_version, + args.numpy_version, pre_hook_sbatch_path=pre_hook_sbatch_path, base_slurm_params=base_slurm_params, capture_nodelist=capture_nodelist, @@ -302,6 +303,7 @@ def _wrap_launcher_for_job_id_and_quiet_output( launcher_cmd: str, launcher_python: str, wandb_version: str, + numpy_version: str, pre_hook_sbatch_path: Optional[Path] = None, base_slurm_params: str = "", capture_nodelist: bool = False, @@ -388,9 +390,9 @@ def _wrap_launcher_for_job_id_and_quiet_output( *pre_hook_lines, ': >"$LOG"', "WANDB_INSTALL_RC=0", - f'{shlex.quote(launcher_python)} -m pip install wandb=={wandb_version} numpy==1.26.4 >>"$LOG" 2>&1 || WANDB_INSTALL_RC=$?', # noqa: E501 + f'{shlex.quote(launcher_python)} -m pip install wandb=={wandb_version} numpy=={numpy_version} >>"$LOG" 2>&1 || WANDB_INSTALL_RC=$?', # noqa: E501 'if [ "${WANDB_INSTALL_RC}" -ne 0 ]; then', - f' echo "Failed to install runtime deps (wandb=={wandb_version}, numpy==1.26.4) in launcher venv (exit ${{WANDB_INSTALL_RC}})." >&2', # noqa: E501 + f' echo "Failed to install runtime deps (wandb=={wandb_version}, numpy=={numpy_version}) in launcher venv (exit ${{WANDB_INSTALL_RC}})." >&2', # noqa: E501 ' tail -n 40 "$LOG" >&2 || true', ' exit "${WANDB_INSTALL_RC}"', "fi", diff --git a/src/cloudai/workloads/nemo_launcher/nemo_launcher.py b/src/cloudai/workloads/nemo_launcher/nemo_launcher.py index 4c7ebbf67..25081c903 100644 --- a/src/cloudai/workloads/nemo_launcher/nemo_launcher.py +++ b/src/cloudai/workloads/nemo_launcher/nemo_launcher.py @@ -94,6 +94,10 @@ class Training(BaseModel): class NeMoLauncherCmdArgs(CmdArgs): """NeMoLauncher test command arguments.""" + launcher_version: str = Field( + default="599ecfcbbd64fd2de02f2cc093b1610d73854022", + description="NeMo Framework Launcher Git commit, tag, or branch.", + ) launcher_script: str = "launcher_scripts/main.py" docker_image_url: str = "nvcr.io/nvidia/nemo:24.12.01" stages: str = '["training"]' @@ -106,11 +110,20 @@ class NeMoLauncherTestDefinition(TestDefinition): """Test object for NeMoLauncher.""" cmd_args: NeMoLauncherCmdArgs - launcher_repo: GitRepo = GitRepo( - url="https://github.com/NVIDIA/NeMo-Framework-Launcher.git", commit="599ecfcbbd64fd2de02f2cc093b1610d73854022" - ) _docker_image: Optional[DockerImage] = None _python_executable: Optional[PythonExecutable] = None + _launcher_repo: Optional[GitRepo] = None + + @property + def launcher_repo(self) -> GitRepo: + """Return the launcher repository selected by ``cmd_args.launcher_version``.""" + version = self.cmd_args.launcher_version + if self._launcher_repo is None or self._launcher_repo.commit != version: + self._launcher_repo = GitRepo( + url="https://github.com/NVIDIA/NeMo-Framework-Launcher.git", + commit=version, + ) + return self._launcher_repo @property def docker_image(self) -> DockerImage: @@ -120,10 +133,9 @@ def docker_image(self) -> DockerImage: @property def python_executable(self) -> PythonExecutable: - if not self._python_executable: - self._python_executable = PythonExecutable( - GitRepo(url=self.launcher_repo.url, commit=self.launcher_repo.commit) - ) + launcher_repo = self.launcher_repo + if not self._python_executable or self._python_executable.git_repo != launcher_repo: + self._python_executable = PythonExecutable(launcher_repo) return self._python_executable @property diff --git a/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py b/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py index 9ed6484ce..0fbf74d6c 100644 --- a/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py @@ -47,6 +47,7 @@ def gen_exec_command(self) -> str: self.final_cmd_args["container"] = str(tdef.docker_image.installed_path) self.final_cmd_args.pop("docker_image_url", None) + self.final_cmd_args.pop("launcher_version", None) if self.job_prefix is None: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") diff --git a/tests/ref_data/megatron-bridge.sbatch b/tests/ref_data/megatron-bridge.sbatch index 3aabcca40..1141a5881 100644 --- a/tests/ref_data/megatron-bridge.sbatch +++ b/tests/ref_data/megatron-bridge.sbatch @@ -11,7 +11,7 @@ exec > >(tee -a "$WRAPPER_STDOUT") 2> >(tee -a "$WRAPPER_STDERR" >&2) : >"$LOG" WANDB_INSTALL_RC=0 -__INSTALL_DIR__/Run__main-venv/bin/python -m pip install wandb==0.28.1 numpy==1.26.4 >>"$LOG" 2>&1 || WANDB_INSTALL_RC=$? +__INSTALL_DIR__/Run__v0.8.0-venv/bin/python -m pip install wandb==0.28.1 numpy==1.26.4 >>"$LOG" 2>&1 || WANDB_INSTALL_RC=$? if [ "${WANDB_INSTALL_RC}" -ne 0 ]; then echo "Failed to install runtime deps (wandb==0.28.1, numpy==1.26.4) in launcher venv (exit ${WANDB_INSTALL_RC})." >&2 tail -n 40 "$LOG" >&2 || true @@ -19,7 +19,7 @@ if [ "${WANDB_INSTALL_RC}" -ne 0 ]; then fi LAUNCH_RC=0 -NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__main-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -t 00:20:00 -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 8 -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --save_config_filepath /nemo_run/configs/ConfigContainer.yaml --additional_slurm_params 'gpus-per-node=8;gres=gpu:8' logger.tensorboard_dir=/nemo_run/tb_logs logger.log_timers_to_tensorboard=true logger.log_throughput_to_tensorboard=true logger.log_memory_to_tensorboard=true >>"$LOG" 2>&1 || LAUNCH_RC=$? +NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__v0.8.0-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -t 00:20:00 -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 8 -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --save_config_filepath /nemo_run/configs/ConfigContainer.yaml --additional_slurm_params 'gpus-per-node=8;gres=gpu:8' logger.tensorboard_dir=/nemo_run/tb_logs logger.log_timers_to_tensorboard=true logger.log_throughput_to_tensorboard=true logger.log_memory_to_tensorboard=true >>"$LOG" 2>&1 || LAUNCH_RC=$? JOB_ID="" diff --git a/tests/test_acceptance.py b/tests/test_acceptance.py index 397c52f7e..080799d77 100644 --- a/tests/test_acceptance.py +++ b/tests/test_acceptance.py @@ -181,6 +181,12 @@ def create_test_run(partial_tr: partial[TestRun], name: str, test_definition: Te return tr +def with_installed_dynamo_repo(test_definition: AIDynamoTestDefinition, installed_path: Path) -> AIDynamoTestDefinition: + """Set the fake installed path used by acceptance command generation.""" + test_definition.repo.installed_path = installed_path + return test_definition + + def build_special_test_run( partial_tr: partial[TestRun], param: str, test_mapping: Dict[str, Callable[[], TestRun]] ) -> Tuple[TestRun, str, Optional[str]]: @@ -514,69 +520,67 @@ def test_req(request, slurm_system: SlurmSystem, partial_tr: partial[TestRun]) - "ai-dynamo": lambda: create_test_run( partial_tr, "ai-dynamo", - AIDynamoTestDefinition( - name="ai-dynamo", - description="AI Dynamo test", - test_template_name="ai-dynamo", - repo=GitRepo( - url="https://github.com/ai-dynamo/dynamo.git", - commit="f7e468c7e8ff0d1426db987564e60572167e8464", - installed_path=slurm_system.install_path, - ), - cmd_args=AIDynamoCmdArgs( - docker_image_url="nvcr.io/nvidia/ai-dynamo:24.09", - workloads="aiperf.sh", - dynamo=AIDynamoArgs( - model="model", - backend="vllm", - endpoint="v1/chat/completions", - workspace_path="/workspace", - dcgm_exporter=DCGMExporter(enabled=True, port=9501), - prefill_worker=WorkerConfig( - cmd="python3 -m dynamo.vllm --is-prefill-worker", - worker_initialized_regex="VllmWorker.*has.been.initialized", - **{ - "num-nodes": 1, - "args": WorkerBaseArgs(), - }, + with_installed_dynamo_repo( + AIDynamoTestDefinition( + name="ai-dynamo", + description="AI Dynamo test", + test_template_name="ai-dynamo", + cmd_args=AIDynamoCmdArgs( + docker_image_url="nvcr.io/nvidia/ai-dynamo:24.09", + workloads="aiperf.sh", + dynamo=AIDynamoArgs( + model="model", + backend="vllm", + endpoint="v1/chat/completions", + workspace_path="/workspace", + dcgm_exporter=DCGMExporter(enabled=True, port=9501), + prefill_worker=WorkerConfig( + cmd="python3 -m dynamo.vllm --is-prefill-worker", + worker_initialized_regex="VllmWorker.*has.been.initialized", + **{ + "num-nodes": 1, + "args": WorkerBaseArgs(), + }, + ), + decode_worker=WorkerConfig( + cmd="python3 -m dynamo.vllm", + worker_initialized_regex="VllmWorker.*has.been.initialized", + **{ + "num-nodes": 1, + "args": WorkerBaseArgs(), + }, + ), ), - decode_worker=WorkerConfig( - cmd="python3 -m dynamo.vllm", - worker_initialized_regex="VllmWorker.*has.been.initialized", + genai_perf=GenAIPerf( **{ - "num-nodes": 1, - "args": WorkerBaseArgs(), - }, + "streaming": True, + "extra-inputs": '{"temperature": 0.7, "max_tokens": 128}', + "output-tokens-mean": 128, + "random-seed": 42, + "request-count": 100, + "synthetic-input-tokens-mean": 550, + "warmup-request-count": 10, + } ), + aiperf=AIPerf.model_validate( + { + "extra-args": "--server-metrics-formats json csv", + "args": { + "concurrency": 2, + "request-count": 50, + "synthetic-input-tokens-mean": 300, + "output-tokens-mean": 500, + "server-metrics": "auto", + }, + } + ), + aiperf_phases=[ + AIPerfPhase.model_validate({"name": "round_1", "args": {"concurrency": 1}}), + AIPerfPhase.model_validate({"name": "round_2", "args": {"request-count": 10}}), + ], ), - genai_perf=GenAIPerf( - **{ - "streaming": True, - "extra-inputs": '{"temperature": 0.7, "max_tokens": 128}', - "output-tokens-mean": 128, - "random-seed": 42, - "request-count": 100, - "synthetic-input-tokens-mean": 550, - "warmup-request-count": 10, - } - ), - aiperf=AIPerf.model_validate( - { - "extra-args": "--server-metrics-formats json csv", - "args": { - "concurrency": 2, - "request-count": 50, - "synthetic-input-tokens-mean": 300, - "output-tokens-mean": 500, - "server-metrics": "auto", - }, - } - ), - aiperf_phases=[ - AIPerfPhase.model_validate({"name": "round_1", "args": {"concurrency": 1}}), - AIPerfPhase.model_validate({"name": "round_2", "args": {"request-count": 10}}), - ], ), + slurm_system.install_path, ), ), "moe-benchmark": lambda: create_test_run( diff --git a/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py b/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py index e0c3d8146..dce1df637 100644 --- a/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py +++ b/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py @@ -105,12 +105,8 @@ def test_run(tmp_path: Path, cmd_args: AIDynamoCmdArgs) -> TestRun: description="desc", test_template_name="template", cmd_args=cmd_args, - repo=GitRepo( - url="https://github.com/ai-dynamo/dynamo.git", - commit="f7e468c7e8ff0d1426db987564e60572167e8464", - installed_path=dynamo_repo_path, - ), ) + tdef.repo.installed_path = dynamo_repo_path return TestRun(name="run", test=tdef, nodes=["n0", "n1"], num_nodes=2, output_path=tmp_path) @@ -145,6 +141,18 @@ def test_installables_include_top_level_git_repos(cmd_args: AIDynamoCmdArgs) -> assert repo in tdef.installables +def test_repo_uses_configured_dynamo_version(cmd_args: AIDynamoCmdArgs) -> None: + cmd_args.dynamo_version = "release-branch" + tdef = AIDynamoTestDefinition( + name="test", + description="desc", + test_template_name="template", + cmd_args=cmd_args, + ) + + assert tdef.repo.commit == "release-branch" + + def test_startup_cmd_image_is_installable(cmd_args: AIDynamoCmdArgs) -> None: cmd_args.startup_cmd = "/mnt/discover.sh" cmd_args.startup_cmd_docker_image = "nvcr.io/test/discovery:latest" @@ -435,13 +443,14 @@ def test_dcgm_exporter_adds_configured_docker_image_installable(cmd_args: AIDyna def test_shared_node_disagg_preserves_explicit_smaller_node_count( slurm_system: SlurmSystem, tmp_path: Path, cmd_args: AIDynamoCmdArgs ) -> None: + cmd_args.dynamo_version = "main" tdef = AIDynamoTestDefinition( name="test", description="desc", test_template_name="template", cmd_args=cmd_args, - repo=GitRepo(url="https://github.com/ai-dynamo/dynamo.git", commit="main", installed_path=tmp_path), ) + tdef.repo.installed_path = tmp_path tr = TestRun( name="run", test=tdef, @@ -474,13 +483,14 @@ def test_explicit_overlapping_worker_nodes_are_allowed_for_shared_node( ) -> None: cmd_args.dynamo.prefill_worker.nodes = "n0" cmd_args.dynamo.decode_worker.nodes = "n0" + cmd_args.dynamo_version = "main" tdef = AIDynamoTestDefinition( name="test", description="desc", test_template_name="template", cmd_args=cmd_args, - repo=GitRepo(url="https://github.com/ai-dynamo/dynamo.git", commit="main", installed_path=tmp_path), ) + tdef.repo.installed_path = tmp_path tr = TestRun(name="run", test=tdef, nodes=[], num_nodes=1, output_path=tmp_path) strategy = AIDynamoSlurmCommandGenStrategy(slurm_system, tr) @@ -495,13 +505,14 @@ def test_explicit_overlapping_worker_nodes_reject_extra_allocated_nodes( ) -> None: cmd_args.dynamo.prefill_worker.nodes = "n0" cmd_args.dynamo.decode_worker.nodes = "n0" + cmd_args.dynamo_version = "main" tdef = AIDynamoTestDefinition( name="test", description="desc", test_template_name="template", cmd_args=cmd_args, - repo=GitRepo(url="https://github.com/ai-dynamo/dynamo.git", commit="main", installed_path=tmp_path), ) + tdef.repo.installed_path = tmp_path tr = TestRun( name="run", test=tdef, @@ -565,6 +576,7 @@ def test_aiperf_phase_roundtrip_does_not_emit_default_report_name(strategy: AIDy ] roundtripped = AIDynamoTestDefinition.model_validate(td.model_dump()) + roundtripped.repo.installed_path = td.repo.installed_path strategy.test_run.test = roundtripped assert roundtripped.cmd_args.aiperf_phases is not None diff --git a/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py b/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py index 6fc7198ae..6456b94ac 100644 --- a/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py +++ b/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py @@ -135,6 +135,24 @@ def test_installables_include_aiconfigurator_python_environment(): assert tdef.python_environment.requirements == ["aiconfigurator~=0.5.0"] +def test_python_environment_uses_configured_requirements() -> None: + tdef = AiconfiguratorTestDefinition( + name="aiconfig", + description="desc", + test_template_name="Aiconfigurator", + cmd_args=AiconfiguratorCmdArgs( + requirements="aiconfigurator==0.5.4", + model_name="LLAMA3.1_70B", + system="h200_sxm", + isl=4000, + osl=500, + agg=Agg(batch_size=8, ctx_tokens=16), + ), + ) + + assert tdef.python_environment.requirements == ["aiconfigurator==0.5.4"] + + def test_cmd_args_requires_exactly_one_mode() -> None: with pytest.raises(ValueError): AiconfiguratorCmdArgs( diff --git a/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py b/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py index 6d540f549..ec693934c 100644 --- a/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py +++ b/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py @@ -72,6 +72,32 @@ def _make_strategy( return DynamoMockerStandaloneCommandGenStrategy(system=system, test_run=tr) +def test_python_environment_uses_default_requirements() -> None: + tdef = DynamoMockerTestDefinition( + name="dynamo_mocker", + description="test", + test_template_name="DynamoMockerTest", + cmd_args=DynamoMockerCmdArgs(), + ) + + assert tdef.python_environment.requirements == [ + "ai-dynamo==1.3.0.post1", + "genai-perf==0.0.16", + "aiperf==0.11.0", + ] + + +def test_python_environment_uses_configured_requirements() -> None: + tdef = DynamoMockerTestDefinition( + name="dynamo_mocker", + description="test", + test_template_name="DynamoMockerTest", + cmd_args=DynamoMockerCmdArgs(requirements="ai-dynamo==1.2.1 aiperf==0.10.0"), + ) + + assert tdef.python_environment.requirements == ["ai-dynamo==1.2.1", "aiperf==0.10.0"] + + class TestBuildScriptArgsCombined: """_build_script_args in combined (disaggregation_mode=none) mode.""" diff --git a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py b/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py index bb16a13ee..7295d9d54 100644 --- a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py +++ b/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py @@ -321,6 +321,21 @@ def test_wrapper_uses_configured_wandb_version( assert "-m pip install wandb==0.27.2 numpy==1.26.4" in self._wrapper_content(cmd_gen) + def test_wrapper_uses_configured_numpy_version( + self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun] + ) -> None: + tr = make_test_run(cmd_args_overrides={"numpy_version": "2.2.1"}) + cmd_gen = MegatronBridgeSlurmCommandGenStrategy(configured_slurm_system, tr) + + assert "-m pip install wandb==0.28.1 numpy==2.2.1" in self._wrapper_content(cmd_gen) + + def test_nemo_run_repo_uses_configured_version(self, make_test_run: Callable[..., TestRun]) -> None: + tr = make_test_run(cmd_args_overrides={"nemorun_version": "feature/ref"}) + tdef = cast(MegatronBridgeTestDefinition, tr.test) + + assert tdef.nemo_run_repo.commit == "feature/ref" + assert tdef.python_executable.git_repo.commit == "feature/ref" + def test_wrapper_exits_when_wandb_install_fails( self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun] ) -> None: diff --git a/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py b/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py index d6c5138a1..ba289bc8a 100644 --- a/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py @@ -31,6 +31,17 @@ class TestNeMoLauncherSlurmCommandGenStrategy: + def test_launcher_repo_uses_configured_version(self) -> None: + tdef = NeMoLauncherTestDefinition( + name="t1", + description="desc1", + test_template_name="tt", + cmd_args=NeMoLauncherCmdArgs(launcher_version="release-branch"), + ) + + assert tdef.launcher_repo.commit == "release-branch" + assert tdef.python_executable.git_repo.commit == "release-branch" + @pytest.fixture def test_run(self, tmp_path: Path) -> TestRun: tdef = NeMoLauncherTestDefinition( From 16089d76b86c52083e392afa5e91273c6677951a Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Fri, 21 Aug 2026 23:29:20 +0200 Subject: [PATCH 2/3] remove redundant tests + change default nemorun version --- .../megatron_bridge/megatron_bridge.py | 2 +- .../workloads/nemo_launcher/nemo_launcher.py | 2 +- .../slurm_command_gen_strategy.py | 2 +- .../test_command_gen_strategy_standalone.py | 18 ------------- .../test_command_gen_strategy_standalone.py | 26 ------------------- .../test_command_gen_strategy_slurm.py | 15 ----------- .../test_command_gen_strategy_slurm.py | 11 -------- 7 files changed, 3 insertions(+), 73 deletions(-) diff --git a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py index 627c010ad..a3daa9f3e 100644 --- a/src/cloudai/workloads/megatron_bridge/megatron_bridge.py +++ b/src/cloudai/workloads/megatron_bridge/megatron_bridge.py @@ -68,7 +68,7 @@ class MegatronBridgeCmdArgs(CmdArgs): wandb_version: str = Field(default="0.28.1", description="W&B version installed in the launcher environment.") numpy_version: str = Field(default="1.26.4", description="NumPy version installed in the launcher environment.") nemorun_version: str = Field( - default="v0.8.0", + default="v0.10.0", description="NeMo Run Git commit, tag, or branch used by the launcher environment.", ) diff --git a/src/cloudai/workloads/nemo_launcher/nemo_launcher.py b/src/cloudai/workloads/nemo_launcher/nemo_launcher.py index 25081c903..4e4904dc6 100644 --- a/src/cloudai/workloads/nemo_launcher/nemo_launcher.py +++ b/src/cloudai/workloads/nemo_launcher/nemo_launcher.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py b/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py index 0fbf74d6c..d5413a690 100644 --- a/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py @@ -1,5 +1,5 @@ # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES -# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py b/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py index 6456b94ac..6fc7198ae 100644 --- a/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py +++ b/tests/workloads/aiconfig/test_command_gen_strategy_standalone.py @@ -135,24 +135,6 @@ def test_installables_include_aiconfigurator_python_environment(): assert tdef.python_environment.requirements == ["aiconfigurator~=0.5.0"] -def test_python_environment_uses_configured_requirements() -> None: - tdef = AiconfiguratorTestDefinition( - name="aiconfig", - description="desc", - test_template_name="Aiconfigurator", - cmd_args=AiconfiguratorCmdArgs( - requirements="aiconfigurator==0.5.4", - model_name="LLAMA3.1_70B", - system="h200_sxm", - isl=4000, - osl=500, - agg=Agg(batch_size=8, ctx_tokens=16), - ), - ) - - assert tdef.python_environment.requirements == ["aiconfigurator==0.5.4"] - - def test_cmd_args_requires_exactly_one_mode() -> None: with pytest.raises(ValueError): AiconfiguratorCmdArgs( diff --git a/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py b/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py index ec693934c..6d540f549 100644 --- a/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py +++ b/tests/workloads/dynamo_mocker/test_command_gen_strategy_standalone.py @@ -72,32 +72,6 @@ def _make_strategy( return DynamoMockerStandaloneCommandGenStrategy(system=system, test_run=tr) -def test_python_environment_uses_default_requirements() -> None: - tdef = DynamoMockerTestDefinition( - name="dynamo_mocker", - description="test", - test_template_name="DynamoMockerTest", - cmd_args=DynamoMockerCmdArgs(), - ) - - assert tdef.python_environment.requirements == [ - "ai-dynamo==1.3.0.post1", - "genai-perf==0.0.16", - "aiperf==0.11.0", - ] - - -def test_python_environment_uses_configured_requirements() -> None: - tdef = DynamoMockerTestDefinition( - name="dynamo_mocker", - description="test", - test_template_name="DynamoMockerTest", - cmd_args=DynamoMockerCmdArgs(requirements="ai-dynamo==1.2.1 aiperf==0.10.0"), - ) - - assert tdef.python_environment.requirements == ["ai-dynamo==1.2.1", "aiperf==0.10.0"] - - class TestBuildScriptArgsCombined: """_build_script_args in combined (disaggregation_mode=none) mode.""" diff --git a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py b/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py index 7295d9d54..bb16a13ee 100644 --- a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py +++ b/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py @@ -321,21 +321,6 @@ def test_wrapper_uses_configured_wandb_version( assert "-m pip install wandb==0.27.2 numpy==1.26.4" in self._wrapper_content(cmd_gen) - def test_wrapper_uses_configured_numpy_version( - self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun] - ) -> None: - tr = make_test_run(cmd_args_overrides={"numpy_version": "2.2.1"}) - cmd_gen = MegatronBridgeSlurmCommandGenStrategy(configured_slurm_system, tr) - - assert "-m pip install wandb==0.28.1 numpy==2.2.1" in self._wrapper_content(cmd_gen) - - def test_nemo_run_repo_uses_configured_version(self, make_test_run: Callable[..., TestRun]) -> None: - tr = make_test_run(cmd_args_overrides={"nemorun_version": "feature/ref"}) - tdef = cast(MegatronBridgeTestDefinition, tr.test) - - assert tdef.nemo_run_repo.commit == "feature/ref" - assert tdef.python_executable.git_repo.commit == "feature/ref" - def test_wrapper_exits_when_wandb_install_fails( self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun] ) -> None: diff --git a/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py b/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py index ba289bc8a..d6c5138a1 100644 --- a/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nemo_launcher/test_command_gen_strategy_slurm.py @@ -31,17 +31,6 @@ class TestNeMoLauncherSlurmCommandGenStrategy: - def test_launcher_repo_uses_configured_version(self) -> None: - tdef = NeMoLauncherTestDefinition( - name="t1", - description="desc1", - test_template_name="tt", - cmd_args=NeMoLauncherCmdArgs(launcher_version="release-branch"), - ) - - assert tdef.launcher_repo.commit == "release-branch" - assert tdef.python_executable.git_repo.commit == "release-branch" - @pytest.fixture def test_run(self, tmp_path: Path) -> TestRun: tdef = NeMoLauncherTestDefinition( From edb57115d5c682a35f164a7717857efd0a30c162 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 24 Aug 2026 11:46:47 +0200 Subject: [PATCH 3/3] update mbridge test with updated nemo-run version --- tests/ref_data/megatron-bridge.sbatch | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/ref_data/megatron-bridge.sbatch b/tests/ref_data/megatron-bridge.sbatch index 1141a5881..4f2e955f7 100644 --- a/tests/ref_data/megatron-bridge.sbatch +++ b/tests/ref_data/megatron-bridge.sbatch @@ -11,7 +11,7 @@ exec > >(tee -a "$WRAPPER_STDOUT") 2> >(tee -a "$WRAPPER_STDERR" >&2) : >"$LOG" WANDB_INSTALL_RC=0 -__INSTALL_DIR__/Run__v0.8.0-venv/bin/python -m pip install wandb==0.28.1 numpy==1.26.4 >>"$LOG" 2>&1 || WANDB_INSTALL_RC=$? +__INSTALL_DIR__/Run__v0.10.0-venv/bin/python -m pip install wandb==0.28.1 numpy==1.26.4 >>"$LOG" 2>&1 || WANDB_INSTALL_RC=$? if [ "${WANDB_INSTALL_RC}" -ne 0 ]; then echo "Failed to install runtime deps (wandb==0.28.1, numpy==1.26.4) in launcher venv (exit ${WANDB_INSTALL_RC})." >&2 tail -n 40 "$LOG" >&2 || true @@ -19,7 +19,7 @@ if [ "${WANDB_INSTALL_RC}" -ne 0 ]; then fi LAUNCH_RC=0 -NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__v0.8.0-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -t 00:20:00 -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 8 -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --save_config_filepath /nemo_run/configs/ConfigContainer.yaml --additional_slurm_params 'gpus-per-node=8;gres=gpu:8' logger.tensorboard_dir=/nemo_run/tb_logs logger.log_timers_to_tensorboard=true logger.log_throughput_to_tensorboard=true logger.log_memory_to_tensorboard=true >>"$LOG" 2>&1 || LAUNCH_RC=$? +NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__v0.10.0-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -t 00:20:00 -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 8 -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --save_config_filepath /nemo_run/configs/ConfigContainer.yaml --additional_slurm_params 'gpus-per-node=8;gres=gpu:8' logger.tensorboard_dir=/nemo_run/tb_logs logger.log_timers_to_tensorboard=true logger.log_throughput_to_tensorboard=true logger.log_memory_to_tensorboard=true >>"$LOG" 2>&1 || LAUNCH_RC=$? JOB_ID=""