[NIXL] ASIO support

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

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

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"

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"

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
 -----------------
 

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

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,15 @@ def cmd_args_dict(self) -> dict[str, str | list[str]]:
             },
             exclude_none=True,
         )
+        if self.cmd_args.runtime_type == "ETCD":
+            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:
+            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")

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)