diff --git a/conf/experimental/test/nixl-bench-posix.toml b/conf/experimental/test/nixl-bench-posix.toml new file mode 100644 index 000000000..275341aa6 --- /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-ucx.toml b/conf/experimental/test/nixl-bench-ucx.toml new file mode 100644 index 000000000..dab7492ff --- /dev/null +++ b/conf/experimental/test/nixl-bench-ucx.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-ucx" +description = "Small NIXLBench UCX 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-network.toml b/conf/experimental/test_scenario/nixl-bench-network.toml new file mode 100644 index 000000000..02bfbb5fd --- /dev/null +++ b/conf/experimental/test_scenario/nixl-bench-network.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-network" + +[[Tests]] +id = "nixl-bench-network.etcd.1node" +num_nodes = 1 +time_limit = "00:05:00" +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-network.etcd.2nodes" +num_nodes = 2 +time_limit = "00:05:00" +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-network.asio.1node" +num_nodes = 1 +time_limit = "00:05:00" +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-network.asio.2nodes" +num_nodes = 2 +time_limit = "00:05:00" +test_name = "nixl-bench-ucx" +description = "NIXLBench ASIO smoke test with one process per node" + + [Tests.cmd_args] + runtime_type = "ASIO" 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..bea8f25cd --- /dev/null +++ b/conf/experimental/test_scenario/nixl-bench-storage.toml @@ -0,0 +1,34 @@ +# 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" 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..6a8d06582 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): @@ -145,10 +146,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 uses_asio(self) -> bool: + """Return whether this workload uses 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 +265,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.uses_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 +326,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 5c524d60f..68e3f1208 100644 --- a/src/cloudai/workloads/nixl_bench/nixl_bench.py +++ b/src/cloudai/workloads/nixl_bench/nixl_bench.py @@ -16,12 +16,15 @@ from __future__ import annotations -from typing import Any, cast +from typing import Any, Literal, cast + +import pydantic import cloudai.metrics from cloudai.core import JobStatusResult, System, TestRun from cloudai.util.lazy_imports import lazy from cloudai.workloads.common.nixl import ( + MANAGED_ETCD_ENDPOINTS, NIXLBaseCmdArgs, NIXLBaseTestDefinition, NIXLExtendedCmdArgs, @@ -33,15 +36,28 @@ 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) 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 self.cmd_args.etcd_endpoints == MANAGED_ETCD_ENDPOINTS + + @property + 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]]: - return self.cmd_args.model_dump( + cmd_args = self.cmd_args.model_dump( exclude={ "docker_image_url", "path_to_benchmark", @@ -52,6 +68,17 @@ 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") + 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") + 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..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,6 +20,9 @@ from .nixl_bench import NIXLBenchTestDefinition +ASIO_PROCESS_START_DELAY_SECONDS = 4 +ETCD_PROCESS_START_DELAY_SECONDS = 15 + class NIXLBenchSlurmCommandGenStrategy(NIXLCmdGenBase): """Command generation strategy for NIXL Bench tests.""" @@ -31,19 +34,34 @@ 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")) 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)}.") + + 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) + f" &\nsleep {process_start_delay}" 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/workloads/nixl_bench/test_command_gen_strategy_slurm.py b/tests/workloads/nixl_bench/test_command_gen_strategy_slurm.py index 90c03d2eb..e4f7858a8 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,10 @@ 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_etcd + 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"} @@ -69,6 +72,21 @@ 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 + assert tdef.uses_asio + assert "NIXL_ASIO_ADDRESS" in strategy.final_env_vars + def test_container_mounts(self, nixl_bench_tr: TestRun, slurm_system: SlurmSystem): nixl_bench_tr.test.cmd_args = NIXLBenchCmdArgs.model_validate( { @@ -86,7 +104,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", ] @@ -179,31 +197,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) @@ -248,32 +241,85 @@ 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): +@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 "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 + 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: + assert "sed -n '1p'" in command + assert "sed -n '2p'" 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" 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) - 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", - "}", - ] - - -def test_gen_kill_and_wait_cmd(nixl_bench_tr: TestRun, slurm_system: SlurmSystem) -> None: + + command = strategy.gen_srun_command() + + assert tdef.uses_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) - 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", - "}", - ] + + command = strategy.gen_srun_command() + + 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 + + +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()