Support separate ETCD container for NIXL workloads

doc/conf.py

@@ -75,11 +75,15 @@ def setup(app):
     "members": True,
     "member-order": "bysource",
     "special-members": "__init__",
+    "inherited-members": "BaseModel",
 }
 
 autodoc_pydantic_model_show_json = False
 autodoc_pydantic_model_show_field_summary = False
 autodoc_pydantic_model_show_config_summary = False
+autodoc_pydantic_field_list_validators = False
+autodoc_pydantic_model_show_validator_summary = False
+autodoc_pydantic_model_show_validator_members = False
 
 # Generate autosummary even if no references
 autosummary_generate = True

doc/workloads/nixl_bench.rst

@@ -64,9 +64,8 @@ API Documentation
 Command Arguments
 ~~~~~~~~~~~~~~~~~
 
-.. autoclass:: cloudai.workloads.nixl_bench.nixl_bench.NIXLBenchCmdArgs
+.. autopydantic_model:: cloudai.workloads.nixl_bench.nixl_bench.NIXLBenchCmdArgs
    :members:
-   :show-inheritance:
 
 Test Definition
 ~~~~~~~~~~~~~~~

doc/workloads/nixl_kvbench.rst

@@ -65,9 +65,8 @@ API Documentation
 Command Arguments
 ~~~~~~~~~~~~~~~~~
 
-.. autoclass:: cloudai.workloads.nixl_kvbench.nixl_kvbench.NIXLKVBenchCmdArgs
+.. autopydantic_model:: cloudai.workloads.nixl_kvbench.nixl_kvbench.NIXLKVBenchCmdArgs
    :members:
-   :show-inheritance:
 
 Test Definition
 ~~~~~~~~~~~~~~~

doc/workloads/nixl_perftest.rst

@@ -75,9 +75,8 @@ API Documentation
 Command Arguments
 ~~~~~~~~~~~~~~~~~
 
-.. autoclass:: cloudai.workloads.nixl_perftest.nixl_perftest.NixlPerftestCmdArgs
+.. autopydantic_model:: cloudai.workloads.nixl_perftest.nixl_perftest.NixlPerftestCmdArgs
    :members:
-   :show-inheritance:
 
 Test Definition
 ~~~~~~~~~~~~~~~

src/cloudai/workloads/common/nixl.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,18 +18,77 @@
 import logging
 from functools import cache
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
+from pydantic import Field
+
+from cloudai.core import DockerImage, Installable, TestRun
+from cloudai.models.workload import CmdArgs, TestDefinition
 from cloudai.systems.slurm import SlurmCommandGenStrategy
+from cloudai.systems.slurm.slurm_system import SlurmSystem
 from cloudai.util.lazy_imports import lazy
 
 if TYPE_CHECKING:
     import pandas as pd
 
 
+class NIXLBaseCmdArgs(CmdArgs):
+    """Command line arguments for a NIXL workloads."""
+
+    docker_image_url: str = Field(description="URL of the Docker image to use for the benchmark.")
+    etcd_path: str = Field(default="etcd", description="Path to the etcd executable.")
+    wait_etcd_for: int = Field(default=60, description="Number of seconds to wait for etcd to become healthy.")
+    etcd_image_url: str | None = Field(
+        default=None,
+        description=(
+            "Optional URL of the Docker image to use for etcd, by default etcd will be run from the same image "
+            "as the benchmark."
+        ),
+    )
+
+
+NIXLCmdArgsT = TypeVar("NIXLCmdArgsT", bound=NIXLBaseCmdArgs)
+
+
+class NIXLBaseTestDefinition(TestDefinition, Generic[NIXLCmdArgsT]):
+    """Test definition for a NIXL workloads."""
+
+    cmd_args: NIXLCmdArgsT
+    _nixl_image: DockerImage | None = None
+    _etcd_image: DockerImage | None = None
+
+    @property
+    def docker_image(self) -> DockerImage:
+        if not self._nixl_image:
+            self._nixl_image = DockerImage(url=self.cmd_args.docker_image_url)
+        return self._nixl_image
+
+    @property
+    def etcd_image(self) -> DockerImage | None:
+        if not self.cmd_args.etcd_image_url:
+            return None
+        if not self._etcd_image:
+            self._etcd_image = DockerImage(url=self.cmd_args.etcd_image_url)
+        return self._etcd_image
+
+    @property
+    def installables(self) -> list[Installable]:
+        installables = [self.docker_image, *self.git_repos]
+        if self.etcd_image:
+            installables.append(self.etcd_image)
+        return installables
+
+
 class NIXLCmdGenBase(SlurmCommandGenStrategy):
     """Base command generation strategy for NIXL-based workloads."""
 
+    def __init__(self, system: SlurmSystem, test_run: TestRun) -> None:
+        super().__init__(system, test_run)
+        self._current_image_url: str | None = None
+
+    def image_path(self) -> str | None:
+        return self._current_image_url
+
     @property
     def final_env_vars(self) -> dict[str, str | list[str]]:
         env_vars = super().final_env_vars
@@ -52,6 +111,10 @@ def gen_etcd_srun_command(self, etcd_path: str) -> list[str]:
             '--initial-cluster="default=http://$SLURM_JOB_MASTER_NODE:2380"',
             "--initial-cluster-state=new",
         ]
+        tdef = cast(NIXLBaseTestDefinition[NIXLBaseCmdArgs], self.test_run.test)
+        curr_image = self._current_image_url
+        if tdef.etcd_image:
+            self._current_image_url = str(tdef.etcd_image.installed_path)
         cmd = [
             *self.gen_srun_prefix(with_num_nodes=False),
             f"--output={self.test_run.output_path.absolute() / 'etcd.log'}",
@@ -63,6 +126,7 @@ def gen_etcd_srun_command(self, etcd_path: str) -> list[str]:
             *etcd_cmd,
             " &",
         ]
+        self._current_image_url = curr_image
         return cmd
 
     def gen_wait_for_etcd_command(self, timeout: int = 60) -> list[str]:

src/cloudai/workloads/nixl_bench/nixl_bench.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,39 +16,32 @@
 
 from __future__ import annotations
 
-from cloudai.core import DockerImage, Installable, JobStatusResult, TestRun
-from cloudai.models.workload import CmdArgs, TestDefinition
-from cloudai.workloads.common.nixl import extract_nixlbench_data
+from cloudai.core import JobStatusResult, TestRun
+from cloudai.workloads.common.nixl import NIXLBaseCmdArgs, NIXLBaseTestDefinition, extract_nixlbench_data
 
 
-class NIXLBenchCmdArgs(CmdArgs):
+class NIXLBenchCmdArgs(NIXLBaseCmdArgs):
     """Command line arguments for a NIXL Bench test."""
 
-    docker_image_url: str
     path_to_benchmark: str
-    etcd_path: str = "etcd"
     etcd_endpoints: str = "http://$NIXL_ETCD_ENDPOINTS"
 
 
-class NIXLBenchTestDefinition(TestDefinition):
+class NIXLBenchTestDefinition(NIXLBaseTestDefinition[NIXLBenchCmdArgs]):
     """Test definition for a NIXL Bench test."""
 
-    cmd_args: NIXLBenchCmdArgs
-    _nixl_image: DockerImage | None = None
-
-    @property
-    def docker_image(self) -> DockerImage:
-        if not self._nixl_image:
-            self._nixl_image = DockerImage(url=self.cmd_args.docker_image_url)
-        return self._nixl_image
-
-    @property
-    def installables(self) -> list[Installable]:
-        return [self.docker_image, *self.git_repos]
-
     @property
     def cmd_args_dict(self) -> dict[str, str | list[str]]:
-        return self.cmd_args.model_dump(exclude={"docker_image_url", "path_to_benchmark", "cmd_args", "etcd_path"})
+        return self.cmd_args.model_dump(
+            exclude={
+                "docker_image_url",
+                "path_to_benchmark",
+                "cmd_args",
+                "etcd_path",
+                "wait_etcd_for",
+                "etcd_image_url",
+            }
+        )
 
     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

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,8 +16,6 @@
 
 from typing import cast
 
-from cloudai.core import TestRun
-from cloudai.systems.slurm import SlurmSystem
 from cloudai.workloads.common.nixl import NIXLCmdGenBase
 
 from .nixl_bench import NIXLBenchTestDefinition
@@ -26,14 +24,6 @@
 class NIXLBenchSlurmCommandGenStrategy(NIXLCmdGenBase):
     """Command generation strategy for NIXL Bench tests."""
 
-    def __init__(self, system: SlurmSystem, test_run: TestRun) -> None:
-        super().__init__(system, test_run)
-
-        self._current_image_url: str | None = None
-
-    def image_path(self) -> str | None:
-        return self._current_image_url
-
     def _container_mounts(self) -> list[str]:
         return []
 

src/cloudai/workloads/nixl_kvbench/nixl_kvbench.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -18,18 +18,15 @@
 
 from typing import Literal
 
-from cloudai.core import CmdArgs, DockerImage, Installable, JobStatusResult, TestDefinition, TestRun
-from cloudai.workloads.common.nixl import extract_nixlbench_data
+from cloudai.core import JobStatusResult, TestRun
+from cloudai.workloads.common.nixl import NIXLBaseCmdArgs, NIXLBaseTestDefinition, extract_nixlbench_data
 
 
-class NIXLKVBenchCmdArgs(CmdArgs):
-    """Command line arguments for NIXLKVBench."""
+class NIXLKVBenchCmdArgs(NIXLBaseCmdArgs):
+    """Command line arguments for NIXL KVBench."""
 
     command: Literal["profile"] = "profile"
-    etcd_path: str = "etcd"
-    wait_etcd_for: int = 60
 
-    docker_image_url: str
     kvbench_script: str = "/workspace/nixl/benchmark/kvbench/main.py"
     python_executable: str = "python"
 
@@ -39,21 +36,8 @@ class NIXLKVBenchCmdArgs(CmdArgs):
     backend: str | list[str] | None = None
 
 
-class NIXLKVBenchTestDefinition(TestDefinition):
-    """Test definition for NIXLKVBench."""
-
-    _docker_image: DockerImage | None = None
-    cmd_args: NIXLKVBenchCmdArgs
-
-    @property
-    def docker_image(self) -> DockerImage:
-        if not self._docker_image:
-            self._docker_image = DockerImage(url=self.cmd_args.docker_image_url)
-        return self._docker_image
-
-    @property
-    def installables(self) -> list[Installable]:
-        return [*self.git_repos, self.docker_image]
+class NIXLKVBenchTestDefinition(NIXLBaseTestDefinition[NIXLKVBenchCmdArgs]):
+    """Test definition for NIXL KVBench."""
 
     @property
     def cmd_args_dict(self) -> dict[str, str | list[str]]:
@@ -65,6 +49,7 @@ def cmd_args_dict(self) -> dict[str, str | list[str]]:
                 "wait_etcd_for",
                 "docker_image_url",
                 "command",
+                "etcd_image_url",
             },
         )
 

src/cloudai/workloads/nixl_kvbench/slurm_command_gen_strategy.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -38,6 +38,8 @@ def tdef(self) -> NIXLKVBenchTestDefinition:
         return cast(NIXLKVBenchTestDefinition, self.test_run.test)
 
     def image_path(self) -> str | None:
+        if self._current_image_url is not None:
+            return self._current_image_url
         return str(self.tdef.docker_image.installed_path)
 
     def _gen_srun_command(self) -> str:

src/cloudai/workloads/nixl_perftest/nixl_perftest.py

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,11 +14,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Literal, Optional
+from typing import Literal
 
 from pydantic import Field, model_validator
 
-from cloudai.core import CmdArgs, DockerImage, Installable, TestDefinition, TestRun
+from cloudai.core import CmdArgs, TestRun
+from cloudai.workloads.common.nixl import NIXLBaseCmdArgs, NIXLBaseTestDefinition
 
 
 class MatgenCmdArgs(CmdArgs):
@@ -27,17 +28,13 @@ class MatgenCmdArgs(CmdArgs):
     ppn: int | None = None
 
 
-class NixlPerftestCmdArgs(CmdArgs):
-    """CmdArgs for NixlPerftestTestDefinition."""
-
-    docker_image_url: str
+class NixlPerftestCmdArgs(NIXLBaseCmdArgs):
+    """CmdArgs for NIXL Perftest."""
 
     subtest: Literal["sequential-ct-perftest"]
     perftest_script: str = "/workspace/nixl/benchmark/kvbench/main.py"
     matgen_script: str = "/workspace/nixl/benchmark/kvbench/test/inference_workload_matgen.py"
     python_executable: str = "python"
-    etcd_path: str = "etcd"
-    wait_etcd_for: int = 60
 
     num_user_requests: int | list[int]
     batch_size: int | list[int]
@@ -91,22 +88,9 @@ def model_vs_custom(self):
         return self
 
 
-class NixlPerftestTestDefinition(TestDefinition):
+class NixlPerftestTestDefinition(NIXLBaseTestDefinition[NixlPerftestCmdArgs]):
     """TestDefinition for NixlPerftest."""
 
-    _docker_image: Optional[DockerImage] = None
-    cmd_args: NixlPerftestCmdArgs
-
-    @property
-    def docker_image(self) -> DockerImage:
-        if not self._docker_image:
-            self._docker_image = DockerImage(url=self.cmd_args.docker_image_url)
-        return self._docker_image
-
-    @property
-    def installables(self) -> list[Installable]:
-        return [*self.git_repos, self.docker_image]
-
     def constraint_check(self, tr: TestRun) -> bool:
         decode_tp = int(tr.test.cmd_args.decode_tp)
         decode_nodes = int(tr.test.cmd_args.num_decode_nodes)