From 2249665c28ca503ebfea3549a7f8eb48c4398220 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Mon, 10 Aug 2026 22:19:07 +0200 Subject: [PATCH 1/9] nixl asio support --- conf/experimental/test/nixl-bench.toml | 36 +++++++ .../test_scenario/nixl-bench.toml | 57 +++++++++++ doc/workloads/nixl_bench.rst | 26 +++++ src/cloudai/workloads/common/nixl.py | 24 ++++- .../workloads/nixl_bench/nixl_bench.py | 34 ++++++- .../nixl_bench/slurm_command_gen_strategy.py | 37 +++++-- tests/ref_data/nixl_bench.sbatch | 1 + .../test_command_gen_strategy_slurm.py | 97 +++++++++++++++++++ 8 files changed, 298 insertions(+), 14 deletions(-) create mode 100644 conf/experimental/test/nixl-bench.toml create mode 100644 conf/experimental/test_scenario/nixl-bench.toml diff --git a/conf/experimental/test/nixl-bench.toml b/conf/experimental/test/nixl-bench.toml new file mode 100644 index 000000000..a78bcdb89 --- /dev/null +++ b/conf/experimental/test/nixl-bench.toml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name = "nixl-bench" +description = "Small NIXLBench smoke test using host memory" +test_template_name = "NIXLBench" + +[cmd_args] +docker_image_url = "nvcr.io#nvidia/cuda-dl-base:26.06-cuda13.3-inference-devel-ubuntu24.04" +path_to_benchmark = "nixlbench" +etcd_image_url = "sourcemation/etcd:3.6.13" +etcd_path = "/usr/local/bin/etcd" +backend = "UCX" +initiator_seg_type = "DRAM" +target_seg_type = "DRAM" +op_type = "WRITE" +total_buffer_size = "64mb" +start_block_size = 4096 +max_block_size = 1048576 +start_batch_size = 1 +max_batch_size = 1 +warmup_iter = 5 +num_iter = 20 diff --git a/conf/experimental/test_scenario/nixl-bench.toml b/conf/experimental/test_scenario/nixl-bench.toml new file mode 100644 index 000000000..2b1ce3010 --- /dev/null +++ b/conf/experimental/test_scenario/nixl-bench.toml @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name = "nixl-bench" + +[[Tests]] +id = "nixl-bench.etcd.1node" +num_nodes = 1 +time_limit = "00:05:00" +test_name = "nixl-bench" +description = "NIXLBench ETCD smoke test with two processes on one node" + + [Tests.cmd_args] + runtime_type = "ETCD" + +[[Tests]] +id = "nixl-bench.etcd.2nodes" +num_nodes = 2 +time_limit = "00:05:00" +test_name = "nixl-bench" +description = "NIXLBench ETCD smoke test with one process per node" + + [Tests.cmd_args] + runtime_type = "ETCD" + +[[Tests]] +id = "nixl-bench.asio.1node" +num_nodes = 1 +time_limit = "00:05:00" +test_name = "nixl-bench" +description = "NIXLBench ASIO smoke test with two processes on one node" + + [Tests.cmd_args] + runtime_type = "ASIO" + +[[Tests]] +id = "nixl-bench.asio.2nodes" +num_nodes = 2 +time_limit = "00:05:00" +test_name = "nixl-bench" +description = "NIXLBench ASIO smoke test with one process per node" + + [Tests.cmd_args] + runtime_type = "ASIO" diff --git a/doc/workloads/nixl_bench.rst b/doc/workloads/nixl_bench.rst index ec89b9713..196643da0 100644 --- a/doc/workloads/nixl_bench.rst +++ b/doc/workloads/nixl_bench.rst @@ -62,6 +62,32 @@ Test-in-Scenario example: target_seg_type = "DRAM" op_type = "WRITE" +Runtime Coordination +-------------------- + +NIXLBench uses ETCD by default. CloudAI starts ETCD from the benchmark image when +``etcd_image_url`` is omitted, or from the configured image otherwise. + +To use NIXLBench's direct two-process ASIO runtime instead, set: + +.. code-block:: toml + + runtime_type = "ASIO" + +CloudAI resolves ``asio_address`` to the first allocated node by default and does not +install or start ETCD. ``asio_address`` and ``asio_port`` can be overridden explicitly. +ASIO requires exactly two NIXLBench processes. For UCX, a one-node test runs both +processes locally, while a two-node test runs one process on each node. + +Storage backends can run without either runtime by using an empty ETCD endpoint: + +.. code-block:: toml + + backend = "POSIX" + etcd_endpoints = "" + +This null-runtime mode is limited to storage backends and launches one NIXLBench process. + API Documentation ----------------- diff --git a/src/cloudai/workloads/common/nixl.py b/src/cloudai/workloads/common/nixl.py index cbd64ed92..eb0f46c4a 100644 --- a/src/cloudai/workloads/common/nixl.py +++ b/src/cloudai/workloads/common/nixl.py @@ -145,10 +145,20 @@ def etcd_image(self) -> DockerImage | None: self._etcd_image = DockerImage(url=self.cmd_args.etcd_image_url) return self._etcd_image + @property + def uses_etcd(self) -> bool: + """Return whether this workload needs CloudAI to manage an ETCD server.""" + return True + + @property + def supports_asio(self) -> bool: + """Return whether this workload supports NIXLBench's ASIO runtime.""" + return False + @property def installables(self) -> list[Installable]: installables = [self.docker_image, *self.git_repos] - if self.etcd_image: + if self.uses_etcd and self.etcd_image: installables.append(self.etcd_image) return installables @@ -254,9 +264,15 @@ def _cleanup_targets(self) -> list[Path]: @property def final_env_vars(self) -> dict[str, str | list[str]]: env_vars = super().final_env_vars - env_vars["NIXL_ETCD_NAMESPACE"] = "/nixl/kvbench/$(uuidgen)" - env_vars["NIXL_ETCD_ENDPOINTS"] = '"$SLURM_JOB_MASTER_NODE:2379"' + tdef = cast(NIXLBaseTestDefinition[NIXLBaseCmdArgs], self.test_run.test) + if tdef.uses_etcd: + env_vars["NIXL_ETCD_NAMESPACE"] = "/nixl/kvbench/$(uuidgen)" + env_vars["NIXL_ETCD_ENDPOINTS"] = '"$SLURM_JOB_MASTER_NODE:2379"' env_vars["SLURM_JOB_MASTER_NODE"] = "$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1)" + if tdef.supports_asio: + env_vars["NIXL_ASIO_ADDRESS"] = ( + "$(getent ahostsv4 \"$SLURM_JOB_MASTER_NODE\" | awk 'NR == 1 {print $1; exit}')" + ) return env_vars @final_env_vars.setter @@ -309,7 +325,7 @@ def gen_nixlbench_srun_commands(self, test_cmd: list[str], backend: str) -> list def create_env_vars_file(self) -> None: with (self.test_run.output_path / "env_vars.sh").open("w") as f: for key, value in self.final_env_vars.items(): - if key in {"NIXL_ETCD_ENDPOINTS", "NIXL_ETCD_NAMESPACE"}: + if key in {"NIXL_ASIO_ADDRESS", "NIXL_ETCD_ENDPOINTS", "NIXL_ETCD_NAMESPACE"}: continue if key == "SLURM_JOB_MASTER_NODE": # this is an sbatch-level variable, not needed per-node continue diff --git a/src/cloudai/workloads/nixl_bench/nixl_bench.py b/src/cloudai/workloads/nixl_bench/nixl_bench.py index b13325434..928230dc8 100644 --- a/src/cloudai/workloads/nixl_bench/nixl_bench.py +++ b/src/cloudai/workloads/nixl_bench/nixl_bench.py @@ -16,6 +16,10 @@ from __future__ import annotations +from typing import Literal + +from pydantic import Field, field_validator + from cloudai.core import JobStatusResult, TestRun from cloudai.workloads.common.nixl import ( NIXLBaseCmdArgs, @@ -30,14 +34,33 @@ class NIXLBenchCmdArgs(NIXLBaseCmdArgs, NIXLExtendedCmdArgs): path_to_benchmark: str etcd_endpoints: str = "http://$NIXL_ETCD_ENDPOINTS" + runtime_type: Literal["ETCD", "ASIO"] = "ETCD" + asio_address: str = "$NIXL_ASIO_ADDRESS" + asio_port: int = Field(default=12345, ge=1, le=65535) + + @field_validator("runtime_type", mode="before") + @classmethod + def normalize_runtime_type(cls, value: str) -> str: + """Normalize the upstream NIXLBench runtime name.""" + return value.upper() class NIXLBenchTestDefinition(NIXLBaseTestDefinition[NIXLBenchCmdArgs]): """Test definition for a NIXL Bench test.""" + @property + def uses_etcd(self) -> bool: + """Return whether CloudAI should launch ETCD for this benchmark.""" + return self.cmd_args.runtime_type == "ETCD" and bool(self.cmd_args.etcd_endpoints) + + @property + def supports_asio(self) -> bool: + """Return whether this workload supports NIXLBench's ASIO runtime.""" + return True + @property def cmd_args_dict(self) -> dict[str, str | list[str]]: - return self.cmd_args.model_dump( + cmd_args = self.cmd_args.model_dump( exclude={ "docker_image_url", "path_to_benchmark", @@ -48,6 +71,15 @@ def cmd_args_dict(self) -> dict[str, str | list[str]]: }, exclude_none=True, ) + if self.cmd_args.runtime_type == "ETCD": + # ETCD is NIXLBench's default runtime, so keep existing commands concise. + cmd_args.pop("runtime_type") + cmd_args.pop("asio_address") + cmd_args.pop("asio_port") + else: + # ASIO performs direct peer-to-peer coordination and does not use ETCD endpoints. + cmd_args.pop("etcd_endpoints") + return cmd_args def was_run_successful(self, tr: TestRun) -> JobStatusResult: df = extract_nixlbench_data(tr.output_path / "stdout.txt") diff --git a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py index 7d0995e6f..3c38d6295 100644 --- a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py @@ -20,6 +20,8 @@ from .nixl_bench import NIXLBenchTestDefinition +STORAGE_BACKENDS = {"AZURE_BLOB", "GDS", "GDS_MT", "GUSLI", "HF3FS", "OBJ", "POSIX"} + class NIXLBenchSlurmCommandGenStrategy(NIXLCmdGenBase): """Command generation strategy for NIXL Bench tests.""" @@ -31,19 +33,36 @@ def tdef(self) -> NIXLBenchTestDefinition: def _gen_srun_command(self) -> str: self.create_env_vars_file() + backend = str(self.tdef.cmd_args_dict.get("backend", "unset")) + if ( + self.tdef.cmd_args.runtime_type == "ETCD" + and not self.tdef.cmd_args.etcd_endpoints + and backend.upper() not in STORAGE_BACKENDS + ): + raise ValueError("NIXLBench can run without ETCD or ASIO only with a storage backend.") + self._current_image_url = str(self.tdef.docker_image.installed_path) - etcd_command: list[str] = self.gen_etcd_srun_command(self.tdef.cmd_args.etcd_path) - nixl_commands = self.gen_nixlbench_srun_commands( - self.gen_nixlbench_command(), str(self.tdef.cmd_args_dict.get("backend", "unset")) - ) - self._current_image_url = None + try: + nixl_commands = self.gen_nixlbench_srun_commands(self.gen_nixlbench_command(), backend) + if self.tdef.cmd_args.runtime_type == "ASIO" and len(nixl_commands) != 2: + raise ValueError(f"ASIO runtime requires exactly two NIXLBench processes, got {len(nixl_commands)}.") + + commands = [ + *[" ".join(cmd) + " &\nsleep 15" for cmd in nixl_commands[:-1]], + " ".join(nixl_commands[-1]), + ] + if not self.tdef.uses_etcd: + return "\n".join(commands) + + etcd_command: list[str] = self.gen_etcd_srun_command(self.tdef.cmd_args.etcd_path) + finally: + self._current_image_url = None - commands: list[str] = [ + commands = [ " ".join(etcd_command), "etcd_pid=$!", - " ".join(self.gen_wait_for_etcd_command()), - *[" ".join(cmd) + " &\nsleep 15" for cmd in nixl_commands[:-1]], - " ".join(nixl_commands[-1]), + " ".join(self.gen_wait_for_etcd_command(self.tdef.cmd_args.wait_etcd_for)), + *commands, " ".join(self.gen_kill_and_wait_cmd("etcd_pid")), ] return "\n".join(commands) diff --git a/tests/ref_data/nixl_bench.sbatch b/tests/ref_data/nixl_bench.sbatch index cd59cb3f0..833220480 100644 --- a/tests/ref_data/nixl_bench.sbatch +++ b/tests/ref_data/nixl_bench.sbatch @@ -12,6 +12,7 @@ export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head export NIXL_ETCD_NAMESPACE=/nixl/kvbench/$(uuidgen) export NIXL_ETCD_ENDPOINTS="$SLURM_JOB_MASTER_NODE:2379" export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1) +export NIXL_ASIO_ADDRESS=$(getent ahostsv4 "$SLURM_JOB_MASTER_NODE" | awk 'NR == 1 {print $1; exit}') srun --export=ALL --mpi=pmix -N2 --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}." srun --export=ALL --mpi=pmix -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash __INSTALL_DIR__/slurm-metadata.sh diff --git a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py index 90c03d2eb..3b8cfda05 100644 --- a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py @@ -69,6 +69,30 @@ def test_can_set_any_cmd_arg(self, nixl_bench_tr: TestRun, slurm_system: SlurmSy for k, v in in_args.items(): assert f"--{k}={v}" in cmd + def test_asio_runtime_args(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + tdef.cmd_args.runtime_type = "ASIO" + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + assert strategy.gen_nixlbench_command() == [ + "./nixlbench", + "--runtime_type=ASIO", + "--asio_address=$NIXL_ASIO_ADDRESS", + "--asio_port=12345", + ] + assert not tdef.uses_etcd + + def test_runtime_type_is_case_insensitive(self): + cmd_args = NIXLBenchCmdArgs.model_validate( + { + "docker_image_url": "docker.io/library/ubuntu:22.04", + "path_to_benchmark": "nixlbench", + "runtime_type": "asio", + } + ) + + assert cmd_args.runtime_type == "ASIO" + def test_container_mounts(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): nixl_bench_tr.test.cmd_args = NIXLBenchCmdArgs.model_validate( { @@ -263,6 +287,79 @@ def test_gen_srun_command(nixl_bench_tr: TestRun, slurm_system: SlurmSystem): ] +@pytest.mark.parametrize("num_nodes", [1, 2]) +def test_asio_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem, num_nodes: int) -> None: + nixl_bench_tr.num_nodes = num_nodes + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + tdef.cmd_args.runtime_type = "ASIO" + nixl_bench_tr.test.cmd_args.backend = "UCX" + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + command = strategy.gen_srun_command() + + assert command.count("nixlbench --runtime_type=ASIO") == 2 + assert "--asio_address=$NIXL_ASIO_ADDRESS" in command + assert "etcd_pid" not in command + assert "until curl" not in command + if num_nodes == 1: + assert command.count("--nodelist=$SLURM_JOB_MASTER_NODE") == 2 + else: + assert "sed -n '1p'" in command + assert "sed -n '2p'" in command + + +def test_etcd_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + command = strategy.gen_srun_command() + + assert "etcd_pid=$!" in command + assert "until curl" in command + assert "kill -TERM $etcd_pid" in command + + +def test_storage_backend_without_runtime(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + nixl_bench_tr.test.cmd_args.backend = "POSIX" + tdef.cmd_args.etcd_endpoints = "" + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + command = strategy.gen_srun_command() + + assert command.count("nixlbench") == 1 + assert "--etcd-endpoints=" in command + assert "etcd_pid" not in command + assert "until curl" not in command + + +def test_network_backend_without_runtime_is_rejected(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + nixl_bench_tr.test.cmd_args.backend = "UCX" + tdef.cmd_args.etcd_endpoints = "" + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + with pytest.raises(ValueError, match="only with a storage backend"): + strategy.gen_srun_command() + + +def test_asio_does_not_install_custom_etcd_image(nixl_bench_tr: TestRun) -> None: + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + tdef.cmd_args.runtime_type = "ASIO" + tdef.cmd_args.etcd_image_url = "docker.io/library/etcd:latest" + + assert tdef.etcd_image not in tdef.installables + + +def test_asio_rejects_non_pairwise_process_shape(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + tdef.cmd_args.runtime_type = "ASIO" + nixl_bench_tr.test.cmd_args.backend = "POSIX" + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + with pytest.raises(ValueError, match="ASIO runtime requires exactly two NIXLBench processes"): + strategy.gen_srun_command() + + def test_gen_kill_and_wait_cmd(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) cmd = strategy.gen_kill_and_wait_cmd("PID", timeout=120) From 8e4dfa946e8bc7a2d80570ec04cc77efaca61ed5 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 12:58:52 +0200 Subject: [PATCH 2/9] fix asio startup --- .../workloads/nixl_bench/slurm_command_gen_strategy.py | 9 ++++++++- .../nixl_bench/test_command_gen_strategy_slurm.py | 5 +++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py index 3c38d6295..7f01931d3 100644 --- a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py @@ -21,6 +21,8 @@ from .nixl_bench import NIXLBenchTestDefinition STORAGE_BACKENDS = {"AZURE_BLOB", "GDS", "GDS_MT", "GUSLI", "HF3FS", "OBJ", "POSIX"} +ASIO_PROCESS_START_DELAY_SECONDS = 4 +ETCD_PROCESS_START_DELAY_SECONDS = 15 class NIXLBenchSlurmCommandGenStrategy(NIXLCmdGenBase): @@ -47,8 +49,13 @@ def _gen_srun_command(self) -> str: if self.tdef.cmd_args.runtime_type == "ASIO" and len(nixl_commands) != 2: raise ValueError(f"ASIO runtime requires exactly two NIXLBench processes, got {len(nixl_commands)}.") + process_start_delay = ( + ASIO_PROCESS_START_DELAY_SECONDS + if self.tdef.cmd_args.runtime_type == "ASIO" + else ETCD_PROCESS_START_DELAY_SECONDS + ) commands = [ - *[" ".join(cmd) + " &\nsleep 15" for cmd in nixl_commands[:-1]], + *[" ".join(cmd) + f" &\nsleep {process_start_delay}" for cmd in nixl_commands[:-1]], " ".join(nixl_commands[-1]), ] if not self.tdef.uses_etcd: diff --git a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py index 3b8cfda05..d2ca07d6b 100644 --- a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py @@ -301,6 +301,8 @@ def test_asio_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem, assert "--asio_address=$NIXL_ASIO_ADDRESS" in command assert "etcd_pid" not in command assert "until curl" not in command + assert "sleep 4" in command + assert "sleep 15" not in command if num_nodes == 1: assert command.count("--nodelist=$SLURM_JOB_MASTER_NODE") == 2 else: @@ -309,12 +311,15 @@ def test_asio_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem, def test_etcd_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + nixl_bench_tr.num_nodes = 2 + nixl_bench_tr.test.cmd_args.backend = "UCX" strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) command = strategy.gen_srun_command() assert "etcd_pid=$!" in command assert "until curl" in command + assert "sleep 15" in command assert "kill -TERM $etcd_pid" in command From 7dacb01db009ef3c2645cd3a60081d1a9e16f3e4 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 13:31:00 +0200 Subject: [PATCH 3/9] code cleanup --- src/cloudai/workloads/common/nixl.py | 6 +- .../workloads/nixl_bench/nixl_bench.py | 12 +-- .../nixl_bench/slurm_command_gen_strategy.py | 8 -- tests/ref_data/nixl_bench.sbatch | 1 - .../test_command_gen_strategy_slurm.py | 84 ++----------------- 5 files changed, 14 insertions(+), 97 deletions(-) diff --git a/src/cloudai/workloads/common/nixl.py b/src/cloudai/workloads/common/nixl.py index eb0f46c4a..39e96ab2a 100644 --- a/src/cloudai/workloads/common/nixl.py +++ b/src/cloudai/workloads/common/nixl.py @@ -151,8 +151,8 @@ def uses_etcd(self) -> bool: return True @property - def supports_asio(self) -> bool: - """Return whether this workload supports NIXLBench's ASIO runtime.""" + def uses_asio(self) -> bool: + """Return whether this workload uses NIXLBench's ASIO runtime.""" return False @property @@ -269,7 +269,7 @@ def final_env_vars(self) -> dict[str, str | list[str]]: env_vars["NIXL_ETCD_NAMESPACE"] = "/nixl/kvbench/$(uuidgen)" env_vars["NIXL_ETCD_ENDPOINTS"] = '"$SLURM_JOB_MASTER_NODE:2379"' env_vars["SLURM_JOB_MASTER_NODE"] = "$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1)" - if tdef.supports_asio: + if tdef.uses_asio: env_vars["NIXL_ASIO_ADDRESS"] = ( "$(getent ahostsv4 \"$SLURM_JOB_MASTER_NODE\" | awk 'NR == 1 {print $1; exit}')" ) diff --git a/src/cloudai/workloads/nixl_bench/nixl_bench.py b/src/cloudai/workloads/nixl_bench/nixl_bench.py index 928230dc8..462d1a500 100644 --- a/src/cloudai/workloads/nixl_bench/nixl_bench.py +++ b/src/cloudai/workloads/nixl_bench/nixl_bench.py @@ -18,7 +18,7 @@ from typing import Literal -from pydantic import Field, field_validator +import pydantic from cloudai.core import JobStatusResult, TestRun from cloudai.workloads.common.nixl import ( @@ -36,9 +36,9 @@ class NIXLBenchCmdArgs(NIXLBaseCmdArgs, NIXLExtendedCmdArgs): etcd_endpoints: str = "http://$NIXL_ETCD_ENDPOINTS" runtime_type: Literal["ETCD", "ASIO"] = "ETCD" asio_address: str = "$NIXL_ASIO_ADDRESS" - asio_port: int = Field(default=12345, ge=1, le=65535) + asio_port: int = pydantic.Field(default=12345, ge=1, le=65535) - @field_validator("runtime_type", mode="before") + @pydantic.field_validator("runtime_type", mode="before") @classmethod def normalize_runtime_type(cls, value: str) -> str: """Normalize the upstream NIXLBench runtime name.""" @@ -54,9 +54,9 @@ def uses_etcd(self) -> bool: return self.cmd_args.runtime_type == "ETCD" and bool(self.cmd_args.etcd_endpoints) @property - def supports_asio(self) -> bool: - """Return whether this workload supports NIXLBench's ASIO runtime.""" - return True + def uses_asio(self) -> bool: + """Return whether this benchmark uses NIXLBench's ASIO runtime.""" + return self.cmd_args.runtime_type == "ASIO" @property def cmd_args_dict(self) -> dict[str, str | list[str]]: diff --git a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py index 7f01931d3..32f2859f5 100644 --- a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py @@ -20,7 +20,6 @@ from .nixl_bench import NIXLBenchTestDefinition -STORAGE_BACKENDS = {"AZURE_BLOB", "GDS", "GDS_MT", "GUSLI", "HF3FS", "OBJ", "POSIX"} ASIO_PROCESS_START_DELAY_SECONDS = 4 ETCD_PROCESS_START_DELAY_SECONDS = 15 @@ -36,13 +35,6 @@ def _gen_srun_command(self) -> str: self.create_env_vars_file() backend = str(self.tdef.cmd_args_dict.get("backend", "unset")) - if ( - self.tdef.cmd_args.runtime_type == "ETCD" - and not self.tdef.cmd_args.etcd_endpoints - and backend.upper() not in STORAGE_BACKENDS - ): - raise ValueError("NIXLBench can run without ETCD or ASIO only with a storage backend.") - self._current_image_url = str(self.tdef.docker_image.installed_path) try: nixl_commands = self.gen_nixlbench_srun_commands(self.gen_nixlbench_command(), backend) diff --git a/tests/ref_data/nixl_bench.sbatch b/tests/ref_data/nixl_bench.sbatch index 833220480..cd59cb3f0 100644 --- a/tests/ref_data/nixl_bench.sbatch +++ b/tests/ref_data/nixl_bench.sbatch @@ -12,7 +12,6 @@ export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head export NIXL_ETCD_NAMESPACE=/nixl/kvbench/$(uuidgen) export NIXL_ETCD_ENDPOINTS="$SLURM_JOB_MASTER_NODE:2379" export SLURM_JOB_MASTER_NODE=$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1) -export NIXL_ASIO_ADDRESS=$(getent ahostsv4 "$SLURM_JOB_MASTER_NODE" | awk 'NR == 1 {print $1; exit}') srun --export=ALL --mpi=pmix -N2 --output=__OUTPUT_DIR__/output/mapping-stdout.txt --error=__OUTPUT_DIR__/output/mapping-stderr.txt bash -c "echo \$(date): \$(hostname):node \${SLURM_NODEID}:rank \${SLURM_PROCID}." srun --export=ALL --mpi=pmix -N2 --ntasks=2 --ntasks-per-node=1 --output=__OUTPUT_DIR__/output/metadata/node-%N.toml --error=__OUTPUT_DIR__/output/metadata/nodes.err bash __INSTALL_DIR__/slurm-metadata.sh diff --git a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py index d2ca07d6b..ef9da9c60 100644 --- a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py @@ -51,6 +51,8 @@ def test_default(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): cmd = strategy.gen_nixlbench_command() tdef: NIXLBenchTestDefinition = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) assert cmd == ["./nixlbench", f"--etcd-endpoints={tdef.cmd_args.etcd_endpoints}"] + assert not tdef.uses_asio + assert "NIXL_ASIO_ADDRESS" not in strategy.final_env_vars def test_can_set_any_cmd_arg(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): in_args = {"backend": "MPI", "dashed-opt": "DRAM", "under_score_opt": "VRAM"} @@ -81,6 +83,8 @@ def test_asio_runtime_args(self, nixl_bench_tr: TestRun, slurm_system: SlurmSyst "--asio_port=12345", ] assert not tdef.uses_etcd + assert tdef.uses_asio + assert "NIXL_ASIO_ADDRESS" in strategy.final_env_vars def test_runtime_type_is_case_insensitive(self): cmd_args = NIXLBenchCmdArgs.model_validate( @@ -203,31 +207,6 @@ def test_device_list_validation( assert cmd_args.total_buffer_size == expected_total_buffer_size -def test_gen_etcd_srun_command(nixl_bench_tr: TestRun, slurm_system: SlurmSystem): - strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) - tdef: NIXLBenchTestDefinition = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) - cmd = " ".join(strategy.gen_etcd_srun_command(tdef.cmd_args.etcd_path)) - assert ( - f"{tdef.cmd_args.etcd_path} --listen-client-urls=http://0.0.0.0:2379 --advertise-client-urls=http://$SLURM_JOB_MASTER_NODE:2379" - " --listen-peer-urls=http://0.0.0.0:2380 --initial-advertise-peer-urls=http://$SLURM_JOB_MASTER_NODE:2380" - ' --initial-cluster="default=http://$SLURM_JOB_MASTER_NODE:2380"' - ) in cmd - - tdef: NIXLBenchTestDefinition = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) - assert "--overlap" in cmd - assert "--ntasks-per-node=1" in cmd - assert "--ntasks=1" in cmd - assert "--nodelist=$SLURM_JOB_MASTER_NODE" in cmd - assert "-N1" in cmd - assert "--container-image=" not in cmd - assert "--container-mounts" not in cmd - - strategy._current_image_url = str(tdef.docker_image.installed_path) - cmd = " ".join(strategy.gen_etcd_srun_command(tdef.cmd_args.etcd_path)) - assert f"--container-image={tdef.docker_image.installed_path}" in cmd - assert "--container-mounts" in cmd - - def test_get_etcd_srun_command_with_etcd_image(nixl_bench_tr: TestRun, slurm_system: SlurmSystem): strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) tdef: NIXLBenchTestDefinition = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) @@ -272,21 +251,6 @@ def test_gen_nixl_srun_command( assert "--nodelist=$SLURM_JOB_MASTER_NODE" in cmd -def test_gen_srun_command(nixl_bench_tr: TestRun, slurm_system: SlurmSystem): - strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) - cmd = strategy.gen_wait_for_etcd_command() - assert cmd == [ - "timeout", - "60", - "bash", - "-c", - '"until curl -s $NIXL_ETCD_ENDPOINTS/health > /dev/null 2>&1; do sleep 1; done" || {\n', - ' echo "ETCD ($NIXL_ETCD_ENDPOINTS) was unreachable after 60 seconds";\n', - " exit 1\n", - "}", - ] - - @pytest.mark.parametrize("num_nodes", [1, 2]) def test_asio_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem, num_nodes: int) -> None: nixl_bench_tr.num_nodes = num_nodes @@ -297,6 +261,7 @@ def test_asio_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem, command = strategy.gen_srun_command() + assert "NIXL_ASIO_ADDRESS" in strategy.final_env_vars assert command.count("nixlbench --runtime_type=ASIO") == 2 assert "--asio_address=$NIXL_ASIO_ADDRESS" in command assert "etcd_pid" not in command @@ -310,19 +275,6 @@ def test_asio_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem, assert "sed -n '2p'" in command -def test_etcd_srun_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: - nixl_bench_tr.num_nodes = 2 - nixl_bench_tr.test.cmd_args.backend = "UCX" - strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) - - command = strategy.gen_srun_command() - - assert "etcd_pid=$!" in command - assert "until curl" in command - assert "sleep 15" in command - assert "kill -TERM $etcd_pid" in command - - def test_storage_backend_without_runtime(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) nixl_bench_tr.test.cmd_args.backend = "POSIX" @@ -337,16 +289,6 @@ def test_storage_backend_without_runtime(nixl_bench_tr: TestRun, slurm_system: S assert "until curl" not in command -def test_network_backend_without_runtime_is_rejected(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: - tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) - nixl_bench_tr.test.cmd_args.backend = "UCX" - tdef.cmd_args.etcd_endpoints = "" - strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) - - with pytest.raises(ValueError, match="only with a storage backend"): - strategy.gen_srun_command() - - def test_asio_does_not_install_custom_etcd_image(nixl_bench_tr: TestRun) -> None: tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) tdef.cmd_args.runtime_type = "ASIO" @@ -363,19 +305,3 @@ def test_asio_rejects_non_pairwise_process_shape(nixl_bench_tr: TestRun, slurm_s with pytest.raises(ValueError, match="ASIO runtime requires exactly two NIXLBench processes"): strategy.gen_srun_command() - - -def test_gen_kill_and_wait_cmd(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: - strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) - cmd = strategy.gen_kill_and_wait_cmd("PID", timeout=120) - assert cmd == [ - "kill -TERM $PID\n", - "timeout", - "120", - "bash", - "-c", - '"while kill -0 $PID 2>/dev/null; do sleep 1; done" || {\n', - ' echo "Failed to kill ETCD (pid=$PID) within 120 seconds";\n', - " exit 1\n", - "}", - ] From 9f3775b95d5ba0cb5af6060facf90512a1de85ee Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 15:43:38 +0200 Subject: [PATCH 4/9] restore backwards compatibility --- src/cloudai/workloads/common/nixl.py | 7 ++-- .../workloads/nixl_bench/nixl_bench.py | 9 +++-- .../nixl_bench/slurm_command_gen_strategy.py | 2 +- .../test_command_gen_strategy_slurm.py | 37 +++++++++++++++++-- 4 files changed, 44 insertions(+), 11 deletions(-) diff --git a/src/cloudai/workloads/common/nixl.py b/src/cloudai/workloads/common/nixl.py index 39e96ab2a..f750a26e1 100644 --- a/src/cloudai/workloads/common/nixl.py +++ b/src/cloudai/workloads/common/nixl.py @@ -39,6 +39,7 @@ # 8gb is the default value in the nixl itself # it's not set as a default in the model below to not propagate it into the srun if the user didn't explicitly set it DEFAULT_TOTAL_BUFFER_SIZE = 8 * 1024 * 1024 * 1024 +MANAGED_ETCD_ENDPOINTS = "http://$NIXL_ETCD_ENDPOINTS" class NIXLBaseCmdArgs(CmdArgs): @@ -146,7 +147,7 @@ def etcd_image(self) -> DockerImage | None: return self._etcd_image @property - def uses_etcd(self) -> bool: + def uses_managed_etcd(self) -> bool: """Return whether this workload needs CloudAI to manage an ETCD server.""" return True @@ -158,7 +159,7 @@ def uses_asio(self) -> bool: @property def installables(self) -> list[Installable]: installables = [self.docker_image, *self.git_repos] - if self.uses_etcd and self.etcd_image: + if self.uses_managed_etcd and self.etcd_image: installables.append(self.etcd_image) return installables @@ -265,7 +266,7 @@ def _cleanup_targets(self) -> list[Path]: def final_env_vars(self) -> dict[str, str | list[str]]: env_vars = super().final_env_vars tdef = cast(NIXLBaseTestDefinition[NIXLBaseCmdArgs], self.test_run.test) - if tdef.uses_etcd: + if tdef.uses_managed_etcd: env_vars["NIXL_ETCD_NAMESPACE"] = "/nixl/kvbench/$(uuidgen)" env_vars["NIXL_ETCD_ENDPOINTS"] = '"$SLURM_JOB_MASTER_NODE:2379"' env_vars["SLURM_JOB_MASTER_NODE"] = "$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1)" diff --git a/src/cloudai/workloads/nixl_bench/nixl_bench.py b/src/cloudai/workloads/nixl_bench/nixl_bench.py index 462d1a500..df91e911a 100644 --- a/src/cloudai/workloads/nixl_bench/nixl_bench.py +++ b/src/cloudai/workloads/nixl_bench/nixl_bench.py @@ -22,6 +22,7 @@ from cloudai.core import JobStatusResult, TestRun from cloudai.workloads.common.nixl import ( + MANAGED_ETCD_ENDPOINTS, NIXLBaseCmdArgs, NIXLBaseTestDefinition, NIXLExtendedCmdArgs, @@ -33,7 +34,7 @@ class NIXLBenchCmdArgs(NIXLBaseCmdArgs, NIXLExtendedCmdArgs): """Command line arguments for a NIXL Bench test.""" path_to_benchmark: str - etcd_endpoints: str = "http://$NIXL_ETCD_ENDPOINTS" + etcd_endpoints: str = MANAGED_ETCD_ENDPOINTS runtime_type: Literal["ETCD", "ASIO"] = "ETCD" asio_address: str = "$NIXL_ASIO_ADDRESS" asio_port: int = pydantic.Field(default=12345, ge=1, le=65535) @@ -49,9 +50,9 @@ class NIXLBenchTestDefinition(NIXLBaseTestDefinition[NIXLBenchCmdArgs]): """Test definition for a NIXL Bench test.""" @property - def uses_etcd(self) -> bool: + def uses_managed_etcd(self) -> bool: """Return whether CloudAI should launch ETCD for this benchmark.""" - return self.cmd_args.runtime_type == "ETCD" and bool(self.cmd_args.etcd_endpoints) + return self.cmd_args.runtime_type == "ETCD" and self.cmd_args.etcd_endpoints == MANAGED_ETCD_ENDPOINTS @property def uses_asio(self) -> bool: @@ -76,6 +77,8 @@ def cmd_args_dict(self) -> dict[str, str | list[str]]: cmd_args.pop("runtime_type") cmd_args.pop("asio_address") cmd_args.pop("asio_port") + if not self.cmd_args.etcd_endpoints: + cmd_args.pop("etcd_endpoints") else: # ASIO performs direct peer-to-peer coordination and does not use ETCD endpoints. cmd_args.pop("etcd_endpoints") diff --git a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py index 32f2859f5..578f9be92 100644 --- a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py @@ -50,7 +50,7 @@ def _gen_srun_command(self) -> str: *[" ".join(cmd) + f" &\nsleep {process_start_delay}" for cmd in nixl_commands[:-1]], " ".join(nixl_commands[-1]), ] - if not self.tdef.uses_etcd: + if not self.tdef.uses_managed_etcd: return "\n".join(commands) etcd_command: list[str] = self.gen_etcd_srun_command(self.tdef.cmd_args.etcd_path) diff --git a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py index ef9da9c60..8214953f3 100644 --- a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py @@ -50,7 +50,8 @@ def test_default(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) cmd = strategy.gen_nixlbench_command() tdef: NIXLBenchTestDefinition = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) - assert cmd == ["./nixlbench", f"--etcd-endpoints={tdef.cmd_args.etcd_endpoints}"] + assert cmd == ["./nixlbench", "--etcd-endpoints=http://$NIXL_ETCD_ENDPOINTS"] + assert tdef.uses_managed_etcd assert not tdef.uses_asio assert "NIXL_ASIO_ADDRESS" not in strategy.final_env_vars @@ -82,7 +83,7 @@ def test_asio_runtime_args(self, nixl_bench_tr: TestRun, slurm_system: SlurmSyst "--asio_address=$NIXL_ASIO_ADDRESS", "--asio_port=12345", ] - assert not tdef.uses_etcd + assert not tdef.uses_managed_etcd assert tdef.uses_asio assert "NIXL_ASIO_ADDRESS" in strategy.final_env_vars @@ -114,7 +115,7 @@ def test_container_mounts(self, nixl_bench_tr: TestRun, slurm_system: SlurmSyste "--filepath=/data", "--total_buffer_size=1024", "--device_list=11:K:/dev/nvme0n1,12:F:/p1/store0.bin,13:F:/p2/store0.bin", - f"--etcd-endpoints={nixl_bench_tr.test.cmd_args.etcd_endpoints}", + "--etcd-endpoints=http://$NIXL_ETCD_ENDPOINTS", "--backend=GUSLI", ] @@ -284,7 +285,35 @@ def test_storage_backend_without_runtime(nixl_bench_tr: TestRun, slurm_system: S command = strategy.gen_srun_command() assert command.count("nixlbench") == 1 - assert "--etcd-endpoints=" in command + assert "--etcd-endpoints" not in command + assert "etcd_pid" not in command + assert "until curl" not in command + + +def test_managed_etcd_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + nixl_bench_tr.test.cmd_args.backend = "UCX" + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + command = strategy.gen_srun_command() + + assert tdef.uses_managed_etcd + assert "--etcd-endpoints=http://$NIXL_ETCD_ENDPOINTS" in command + assert "etcd_pid=$!" in command + assert "until curl" in command + assert "kill -TERM $etcd_pid" in command + + +def test_external_etcd_is_passed_through(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + tdef = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) + nixl_bench_tr.test.cmd_args.backend = "UCX" + tdef.cmd_args.etcd_endpoints = "http://etcd.example:2379" + strategy = NIXLBenchSlurmCommandGenStrategy(slurm_system, nixl_bench_tr) + + command = strategy.gen_srun_command() + + assert not tdef.uses_managed_etcd + assert "--etcd-endpoints=http://etcd.example:2379" in command assert "etcd_pid" not in command assert "until curl" not in command From 2aee1178f290567aeb28069a7927b16097e0015f Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 15:46:56 +0200 Subject: [PATCH 5/9] remove redundant case handling --- src/cloudai/workloads/nixl_bench/nixl_bench.py | 6 ------ .../nixl_bench/test_command_gen_strategy_slurm.py | 11 ----------- 2 files changed, 17 deletions(-) diff --git a/src/cloudai/workloads/nixl_bench/nixl_bench.py b/src/cloudai/workloads/nixl_bench/nixl_bench.py index df91e911a..354b525db 100644 --- a/src/cloudai/workloads/nixl_bench/nixl_bench.py +++ b/src/cloudai/workloads/nixl_bench/nixl_bench.py @@ -39,12 +39,6 @@ class NIXLBenchCmdArgs(NIXLBaseCmdArgs, NIXLExtendedCmdArgs): asio_address: str = "$NIXL_ASIO_ADDRESS" asio_port: int = pydantic.Field(default=12345, ge=1, le=65535) - @pydantic.field_validator("runtime_type", mode="before") - @classmethod - def normalize_runtime_type(cls, value: str) -> str: - """Normalize the upstream NIXLBench runtime name.""" - return value.upper() - class NIXLBenchTestDefinition(NIXLBaseTestDefinition[NIXLBenchCmdArgs]): """Test definition for a NIXL Bench test.""" diff --git a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py index 8214953f3..66538fbb1 100644 --- a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py @@ -87,17 +87,6 @@ def test_asio_runtime_args(self, nixl_bench_tr: TestRun, slurm_system: SlurmSyst assert tdef.uses_asio assert "NIXL_ASIO_ADDRESS" in strategy.final_env_vars - def test_runtime_type_is_case_insensitive(self): - cmd_args = NIXLBenchCmdArgs.model_validate( - { - "docker_image_url": "docker.io/library/ubuntu:22.04", - "path_to_benchmark": "nixlbench", - "runtime_type": "asio", - } - ) - - assert cmd_args.runtime_type == "ASIO" - def test_container_mounts(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): nixl_bench_tr.test.cmd_args = NIXLBenchCmdArgs.model_validate( { From 1676f2a5195474cf8b5b5d16b908b27327b6733c Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 15:47:51 +0200 Subject: [PATCH 6/9] rename prop --- src/cloudai/workloads/common/nixl.py | 6 +++--- src/cloudai/workloads/nixl_bench/nixl_bench.py | 2 +- .../workloads/nixl_bench/slurm_command_gen_strategy.py | 2 +- .../nixl_bench/test_command_gen_strategy_slurm.py | 8 ++++---- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/cloudai/workloads/common/nixl.py b/src/cloudai/workloads/common/nixl.py index f750a26e1..6a8d06582 100644 --- a/src/cloudai/workloads/common/nixl.py +++ b/src/cloudai/workloads/common/nixl.py @@ -147,7 +147,7 @@ def etcd_image(self) -> DockerImage | None: return self._etcd_image @property - def uses_managed_etcd(self) -> bool: + def uses_etcd(self) -> bool: """Return whether this workload needs CloudAI to manage an ETCD server.""" return True @@ -159,7 +159,7 @@ def uses_asio(self) -> bool: @property def installables(self) -> list[Installable]: installables = [self.docker_image, *self.git_repos] - if self.uses_managed_etcd and self.etcd_image: + if self.uses_etcd and self.etcd_image: installables.append(self.etcd_image) return installables @@ -266,7 +266,7 @@ def _cleanup_targets(self) -> list[Path]: def final_env_vars(self) -> dict[str, str | list[str]]: env_vars = super().final_env_vars tdef = cast(NIXLBaseTestDefinition[NIXLBaseCmdArgs], self.test_run.test) - if tdef.uses_managed_etcd: + if tdef.uses_etcd: env_vars["NIXL_ETCD_NAMESPACE"] = "/nixl/kvbench/$(uuidgen)" env_vars["NIXL_ETCD_ENDPOINTS"] = '"$SLURM_JOB_MASTER_NODE:2379"' env_vars["SLURM_JOB_MASTER_NODE"] = "$(scontrol show hostname $SLURM_JOB_NODELIST | head -n 1)" diff --git a/src/cloudai/workloads/nixl_bench/nixl_bench.py b/src/cloudai/workloads/nixl_bench/nixl_bench.py index 354b525db..db7dd6b44 100644 --- a/src/cloudai/workloads/nixl_bench/nixl_bench.py +++ b/src/cloudai/workloads/nixl_bench/nixl_bench.py @@ -44,7 +44,7 @@ class NIXLBenchTestDefinition(NIXLBaseTestDefinition[NIXLBenchCmdArgs]): """Test definition for a NIXL Bench test.""" @property - def uses_managed_etcd(self) -> bool: + def uses_etcd(self) -> bool: """Return whether CloudAI should launch ETCD for this benchmark.""" return self.cmd_args.runtime_type == "ETCD" and self.cmd_args.etcd_endpoints == MANAGED_ETCD_ENDPOINTS diff --git a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py index 578f9be92..32f2859f5 100644 --- a/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/nixl_bench/slurm_command_gen_strategy.py @@ -50,7 +50,7 @@ def _gen_srun_command(self) -> str: *[" ".join(cmd) + f" &\nsleep {process_start_delay}" for cmd in nixl_commands[:-1]], " ".join(nixl_commands[-1]), ] - if not self.tdef.uses_managed_etcd: + if not self.tdef.uses_etcd: return "\n".join(commands) etcd_command: list[str] = self.gen_etcd_srun_command(self.tdef.cmd_args.etcd_path) diff --git a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py index 66538fbb1..e4f7858a8 100644 --- a/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py +++ b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py @@ -51,7 +51,7 @@ def test_default(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): cmd = strategy.gen_nixlbench_command() tdef: NIXLBenchTestDefinition = cast(NIXLBenchTestDefinition, nixl_bench_tr.test) assert cmd == ["./nixlbench", "--etcd-endpoints=http://$NIXL_ETCD_ENDPOINTS"] - assert tdef.uses_managed_etcd + assert tdef.uses_etcd assert not tdef.uses_asio assert "NIXL_ASIO_ADDRESS" not in strategy.final_env_vars @@ -83,7 +83,7 @@ def test_asio_runtime_args(self, nixl_bench_tr: TestRun, slurm_system: SlurmSyst "--asio_address=$NIXL_ASIO_ADDRESS", "--asio_port=12345", ] - assert not tdef.uses_managed_etcd + assert not tdef.uses_etcd assert tdef.uses_asio assert "NIXL_ASIO_ADDRESS" in strategy.final_env_vars @@ -286,7 +286,7 @@ def test_managed_etcd_lifecycle(nixl_bench_tr: TestRun, slurm_system: SlurmSyste command = strategy.gen_srun_command() - assert tdef.uses_managed_etcd + assert tdef.uses_etcd assert "--etcd-endpoints=http://$NIXL_ETCD_ENDPOINTS" in command assert "etcd_pid=$!" in command assert "until curl" in command @@ -301,7 +301,7 @@ def test_external_etcd_is_passed_through(nixl_bench_tr: TestRun, slurm_system: S command = strategy.gen_srun_command() - assert not tdef.uses_managed_etcd + assert not tdef.uses_etcd assert "--etcd-endpoints=http://etcd.example:2379" in command assert "etcd_pid" not in command assert "until curl" not in command From 9bf13b31a4960c8143f489cbd5f0b34dad6f990e Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 16:34:34 +0200 Subject: [PATCH 7/9] updated nixl configs --- conf/experimental/test/nixl-bench-posix.toml | 36 ++++++++++++++++++ .../{nixl-bench.toml => nixl-bench-ucx.toml} | 6 +-- ...ixl-bench.toml => nixl-bench-network.toml} | 20 +++++----- .../test_scenario/nixl-bench-storage.toml | 37 +++++++++++++++++++ 4 files changed, 86 insertions(+), 13 deletions(-) create mode 100644 conf/experimental/test/nixl-bench-posix.toml rename conf/experimental/test/{nixl-bench.toml => nixl-bench-ucx.toml} (89%) rename conf/experimental/test_scenario/{nixl-bench.toml => nixl-bench-network.toml} (79%) create mode 100644 conf/experimental/test_scenario/nixl-bench-storage.toml diff --git a/conf/experimental/test/nixl-bench-posix.toml b/conf/experimental/test/nixl-bench-posix.toml new file mode 100644 index 000000000..9dcad27a6 --- /dev/null +++ b/conf/experimental/test/nixl-bench-posix.toml @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name = "nixl-bench-posix" +description = "Small NIXLBench POSIX storage smoke test" +test_template_name = "NIXLBench" + +[cmd_args] +docker_image_url = "nvcr.io#nvidia/cuda-dl-base:26.06-cuda13.3-inference-devel-ubuntu24.04" +path_to_benchmark = "nixlbench" +etcd_image_url = "sourcemation/etcd:3.6.13" +etcd_path = "/usr/local/bin/etcd" +backend = "POSIX" +filepath = "/data" +posix_api_type = "AIO" +op_type = "WRITE" +total_buffer_size = "64mb" +start_block_size = 4096 +max_block_size = 1048576 +start_batch_size = 1 +max_batch_size = 1 +warmup_iter = 5 +num_iter = 20 diff --git a/conf/experimental/test/nixl-bench.toml b/conf/experimental/test/nixl-bench-ucx.toml similarity index 89% rename from conf/experimental/test/nixl-bench.toml rename to conf/experimental/test/nixl-bench-ucx.toml index a78bcdb89..5e23255d6 100644 --- a/conf/experimental/test/nixl-bench.toml +++ b/conf/experimental/test/nixl-bench-ucx.toml @@ -6,7 +6,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -name = "nixl-bench" -description = "Small NIXLBench smoke test using host memory" +name = "nixl-bench-ucx" +description = "Small NIXLBench UCX smoke test using host memory" test_template_name = "NIXLBench" [cmd_args] diff --git a/conf/experimental/test_scenario/nixl-bench.toml b/conf/experimental/test_scenario/nixl-bench-network.toml similarity index 79% rename from conf/experimental/test_scenario/nixl-bench.toml rename to conf/experimental/test_scenario/nixl-bench-network.toml index 2b1ce3010..d8e0d827d 100644 --- a/conf/experimental/test_scenario/nixl-bench.toml +++ b/conf/experimental/test_scenario/nixl-bench-network.toml @@ -6,7 +6,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, @@ -14,43 +14,43 @@ # See the License for the specific language governing permissions and # limitations under the License. -name = "nixl-bench" +name = "nixl-bench-network" [[Tests]] -id = "nixl-bench.etcd.1node" +id = "nixl-bench-network.etcd.1node" num_nodes = 1 time_limit = "00:05:00" -test_name = "nixl-bench" +test_name = "nixl-bench-ucx" description = "NIXLBench ETCD smoke test with two processes on one node" [Tests.cmd_args] runtime_type = "ETCD" [[Tests]] -id = "nixl-bench.etcd.2nodes" +id = "nixl-bench-network.etcd.2nodes" num_nodes = 2 time_limit = "00:05:00" -test_name = "nixl-bench" +test_name = "nixl-bench-ucx" description = "NIXLBench ETCD smoke test with one process per node" [Tests.cmd_args] runtime_type = "ETCD" [[Tests]] -id = "nixl-bench.asio.1node" +id = "nixl-bench-network.asio.1node" num_nodes = 1 time_limit = "00:05:00" -test_name = "nixl-bench" +test_name = "nixl-bench-ucx" description = "NIXLBench ASIO smoke test with two processes on one node" [Tests.cmd_args] runtime_type = "ASIO" [[Tests]] -id = "nixl-bench.asio.2nodes" +id = "nixl-bench-network.asio.2nodes" num_nodes = 2 time_limit = "00:05:00" -test_name = "nixl-bench" +test_name = "nixl-bench-ucx" description = "NIXLBench ASIO smoke test with one process per node" [Tests.cmd_args] diff --git a/conf/experimental/test_scenario/nixl-bench-storage.toml b/conf/experimental/test_scenario/nixl-bench-storage.toml new file mode 100644 index 000000000..e82f8c26e --- /dev/null +++ b/conf/experimental/test_scenario/nixl-bench-storage.toml @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name = "nixl-bench-storage" + +[[Tests]] +id = "nixl-bench-storage.posix.null-runtime" +num_nodes = 1 +time_limit = "00:05:00" +test_name = "nixl-bench-posix" +description = "NIXLBench POSIX smoke test without ETCD" + + [Tests.cmd_args] + etcd_endpoints = "" + +[[Tests]] +id = "nixl-bench-storage.posix.etcd" +num_nodes = 1 +time_limit = "00:05:00" +test_name = "nixl-bench-posix" +description = "NIXLBench POSIX smoke test with managed ETCD" + + [Tests.cmd_args] + runtime_type = "ETCD" From 1fe0be1f64177aee2f747ecc1df4ef7047c07659 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 16:36:10 +0200 Subject: [PATCH 8/9] simplify config --- conf/experimental/test_scenario/nixl-bench-storage.toml | 3 --- 1 file changed, 3 deletions(-) diff --git a/conf/experimental/test_scenario/nixl-bench-storage.toml b/conf/experimental/test_scenario/nixl-bench-storage.toml index e82f8c26e..ad1a52511 100644 --- a/conf/experimental/test_scenario/nixl-bench-storage.toml +++ b/conf/experimental/test_scenario/nixl-bench-storage.toml @@ -32,6 +32,3 @@ num_nodes = 1 time_limit = "00:05:00" test_name = "nixl-bench-posix" description = "NIXLBench POSIX smoke test with managed ETCD" - - [Tests.cmd_args] - runtime_type = "ETCD" From 9eafbf5857852d0a4f0b1e355befd2ca26e0db62 Mon Sep 17 00:00:00 2001 From: Ivan Podkidyshev Date: Tue, 11 Aug 2026 17:13:45 +0200 Subject: [PATCH 9/9] fix copyright headers --- conf/experimental/test/nixl-bench-posix.toml | 2 +- conf/experimental/test/nixl-bench-ucx.toml | 2 +- conf/experimental/test_scenario/nixl-bench-network.toml | 2 +- conf/experimental/test_scenario/nixl-bench-storage.toml | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/conf/experimental/test/nixl-bench-posix.toml b/conf/experimental/test/nixl-bench-posix.toml index 9dcad27a6..275341aa6 100644 --- a/conf/experimental/test/nixl-bench-posix.toml +++ b/conf/experimental/test/nixl-bench-posix.toml @@ -6,7 +6,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, diff --git a/conf/experimental/test/nixl-bench-ucx.toml b/conf/experimental/test/nixl-bench-ucx.toml index 5e23255d6..dab7492ff 100644 --- a/conf/experimental/test/nixl-bench-ucx.toml +++ b/conf/experimental/test/nixl-bench-ucx.toml @@ -6,7 +6,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, diff --git a/conf/experimental/test_scenario/nixl-bench-network.toml b/conf/experimental/test_scenario/nixl-bench-network.toml index d8e0d827d..02bfbb5fd 100644 --- a/conf/experimental/test_scenario/nixl-bench-network.toml +++ b/conf/experimental/test_scenario/nixl-bench-network.toml @@ -6,7 +6,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, diff --git a/conf/experimental/test_scenario/nixl-bench-storage.toml b/conf/experimental/test_scenario/nixl-bench-storage.toml index ad1a52511..bea8f25cd 100644 --- a/conf/experimental/test_scenario/nixl-bench-storage.toml +++ b/conf/experimental/test_scenario/nixl-bench-storage.toml @@ -6,7 +6,7 @@ # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # -# http://www.apache.org/licenses/LICENSE-2.0 +# http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS,