Skip to content

[Installables] configurable workloads dependencies #1009

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
podkidyshev wants to merge 3 commits into
base: main
Choose a base branch
from
Open

[Installables] configurable workloads dependencies #1009

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e727b9c
make sure workloads dependencies are configurable
podkidyshev Aug 21, 2026
16089d7
remove redundant tests + change default nemorun version
podkidyshev Aug 21, 2026
edb5711
update mbridge test with updated nemo-run version
podkidyshev Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions src/cloudai/workloads/ai_dynamo/ai_dynamo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand All @@ -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:
Expand Down
8 changes: 6 additions & 2 deletions src/cloudai/workloads/aiconfig/aiconfigurator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.",
)
Comment on lines +64 to +67

@jj10306 jj10306 Aug 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

any reason to use a str here instead of List[str]? There are two potential issues with str:

  1. afaik, "aiconfigurator >= 0.5.0" is valid (despite spaces around the operator) - if this is the case, this would break this logic
  2. an empty string would fail at runtime because "" gets split into [] which would cause an empty venv to be created and lead to imports failing at runtime.

I believe using a List here would naturally be spported by TOML arrays and you can enforce the non-empty requirement via min_length

Suggested change
requirements: str = Field(
default="aiconfigurator~=0.5.0",
description="Space-separated Python requirements installed into the Aiconfigurator environment.",
)
requirements: List[str] = Field(
default=["aiconfigurator~=0.5.0"],
min_length=1,
description="Space-separated Python requirements installed into the Aiconfigurator environment.",
)

@podkidyshev podkidyshev Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

any list object in our TOMLs is by default a DSE sweep. to avoid that, one must also set dse_excluded_args param to include this one which I thought won't be the best UX

long story short, this str typing is a trade-off

@jj10306 jj10306 Aug 25, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

got it. in this case, should the str splitting logic be updated in order to address the 2 potential issues listed above?

@podkidyshev podkidyshev Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. to address the first - the split logic would be rather complicated (iterating chunks manually, checking for characters, etc.) yet I don't see much profit
  2. well, user's run will fail early anyway if that's the case. or maybe user's environment already has all the reqs installed so they actually don't need any reqs installed at all

in general I try to trust users and don't limit them with what they want to do. we have so many config options for so many workloads, ain't no way we verify everything....

model_name: str
system: str
backend: str = "trtllm"
Expand Down Expand Up @@ -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
Expand Down
6 changes: 5 additions & 1 deletion src/cloudai/workloads/dynamo_mocker/dynamo_mocker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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

Expand Down
24 changes: 17 additions & 7 deletions src/cloudai/workloads/megatron_bridge/megatron_bridge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.10.0",
description="NeMo Run Git commit, tag, or branch used by the launcher environment.",
)

# Retries
max_retries: Optional[int] = Field(default=1)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Comment thread
coderabbitai[bot] marked this conversation as resolved.
' tail -n 40 "$LOG" >&2 || true',
' exit "${WANDB_INSTALL_RC}"',
"fi",
Expand Down
28 changes: 20 additions & 8 deletions src/cloudai/workloads/nemo_launcher/nemo_launcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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");
Expand Down Expand Up @@ -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"]'
Expand All @@ -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:
Expand All @@ -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
Expand Down
3 changes: 2 additions & 1 deletion src/cloudai/workloads/nemo_launcher/slurm_command_gen_strategy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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");
Expand Down Expand Up @@ -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")
Expand Down
4 changes: 2 additions & 2 deletions tests/ref_data/megatron-bridge.sbatch
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,15 @@ 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.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
exit "${WANDB_INSTALL_RC}"
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.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=""
Expand Down
120 changes: 62 additions & 58 deletions tests/test_acceptance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
Expand Down Expand Up @@ -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(
Expand Down
Loading
Loading