Merge pull request #739 from NVIDIA/am/dynamo

Expand K8s Dynamo support to disagg and multinode

Author: amaslenn

conf/experimental/ai_dynamo/test/vllm.toml

@@ -29,20 +29,10 @@ docker_image_url = "nvcr.io/nvidia/ai-dynamo/vllm-runtime:0.6.1.post1"
   decode-cmd = 'python3 -m dynamo.vllm'
 
     [cmd_args.dynamo.prefill_worker]
-    num-nodes = 2
-    gpu-memory-utilization = 0.95
-    tensor-parallel-size = 8
     pipeline-parallel-size = 1
-    data-parallel-size = 1
-    extra-args = "--no-enable-expert-parallel"
 
     [cmd_args.dynamo.decode_worker]
-    num-nodes = 2
-    gpu-memory-utilization = 0.95
-    tensor-parallel-size = 8
     pipeline-parallel-size = 1
-    data-parallel-size = 1
-    extra-args = "--no-enable-expert-parallel"
 
   [cmd_args.genai_perf]
   model = "Qwen/Qwen3-0.6B"
@@ -52,10 +42,10 @@ docker_image_url = "nvcr.io/nvidia/ai-dynamo/vllm-runtime:0.6.1.post1"
   output-tokens-mean = 500
   output-tokens-stddev = 0
   random-seed = 123
-  request-count = 2
+  request-count = 128
   synthetic-input-tokens-mean = 300
   synthetic-input-tokens-stddev = 0
-  warmup-request-count = 1
+  warmup-request-count = 10
   concurrency = 1
   extra-args = "--streaming -- -v --async"
 

conf/experimental/ai_dynamo/test_scenario/vllm_k8s.toml

@@ -24,5 +24,7 @@ test_name = "vLLM-Qwen3-0.6B"
     [Tests.cmd_args.dynamo]
       [Tests.cmd_args.dynamo.prefill_worker]
       num-nodes = 1
+      tensor-parallel-size = 8
       [Tests.cmd_args.dynamo.decode_worker]
       num-nodes = 1
+      tensor-parallel-size = 8

conf/experimental/ai_dynamo/test_scenario/vllm_slurm.toml

@@ -18,20 +18,13 @@ name = "dynamo-vllm-slurm"
 
 [[Tests]]
 id = "qwen3-0.6B"
-num_nodes = 3
+test_name = "vLLM-Qwen3-0.6B"
+num_nodes = 4
 time_limit = "00:20:00"
 
-name = "vllm"
-description = "vllm"
-test_template_name = "AIDynamo"
-
   [Tests.cmd_args]
-  docker_image_url = "nvcr.io/nvidia/ai-dynamo/vllm-runtime:0.6.1.post1"
 
     [Tests.cmd_args.dynamo]
-    backend = "vllm"
-    model = "Qwen/Qwen3-0.6B"
-    decode-cmd = 'python3 -m dynamo.vllm'
     decode-initialized-regex = 'VllmWorker.*has.been.initialized'
     etcd-cmd = "etcd --log-level debug"
     etcd-port = 2379
@@ -41,41 +34,19 @@ test_template_name = "AIDynamo"
     nats-port = 4222
     node-setup-cmd = "apt-get update -o APT::Sandbox::User=root && apt-get install -y curl libibverbs1 rdma-core ibverbs-utils libibumad3 libnuma1 librdmacm1 ibverbs-providers; /usr/local/ucx/bin/ucx_info -d |grep Transport | sort -u;"
     port = 8787
-    prefill-cmd = 'python3 -m dynamo.vllm --is-prefill-worker'
     prefill-initialized-regex = 'VllmWorker.*has.been.initialized'
-    workspace-path = "/workspace/"
 
       [Tests.cmd_args.dynamo.prefill_worker]
-      data-parallel-size = 1
       gpu-memory-utilization = 0.90
       max_model_len = 19280
       num-nodes = 2
-      pipeline-parallel-size = 1
-      tensor-parallel-size = 2
-      extra-args = "--no-enable-expert-parallel"
+      tensor-parallel-size = 4
 
       [Tests.cmd_args.dynamo.decode_worker]
-      data-parallel-size = 1
       gpu-memory-utilization = 0.90
       max_model_len = 19280
-      num-nodes = 1
-      pipeline-parallel-size = 1
-      tensor-parallel-size = 2
-      extra-args = "--no-enable-expert-parallel"
-
-    [Tests.cmd_args.genai_perf]
-    concurrency = 8
-    endpoint = "v1/chat/completions"
-    endpoint-type = "chat"
-    extra-inputs = 'min_tokens:10'
-    output-tokens-mean = 150
-    output-tokens-stddev = 0
-    random-seed = 123
-    request-count = 128
-    synthetic-input-tokens-mean = 3000
-    synthetic-input-tokens-stddev = 0
-    warmup-request-count = 8
-    extra-args = "--streaming -- -v --async"
+      num-nodes = 2
+      tensor-parallel-size = 4
 
   [Tests.extra_env_vars]
   UCX_LOG_LEVEL = "warn"

src/cloudai/workloads/ai_dynamo/ai_dynamo.py

@@ -16,7 +16,7 @@
 
 import logging
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional
 
 from pydantic import AliasChoices, BaseModel, ConfigDict, Field
 
@@ -31,10 +31,36 @@ class WorkerBaseArgs(BaseModel):
 
     model_config = ConfigDict(extra="allow", populate_by_name=True)
 
-    num_nodes: Union[int, list[int]] = Field(
-        serialization_alias="num-nodes", validation_alias=AliasChoices("num-nodes", "num_nodes")
+    num_nodes: int | list[int] = Field(
+        default=1, serialization_alias="num-nodes", validation_alias=AliasChoices("num-nodes", "num_nodes")
+    )
+    nodes: str | None = Field(default=None)
+
+    data_parallel_size: int | list[int] | None = Field(
+        default=None,
+        serialization_alias="data-parallel-size",
+        validation_alias=AliasChoices("data-parallel-size", "data_parallel_size"),
+    )
+    gpu_memory_utilization: float | list[float] | None = Field(
+        default=None,
+        serialization_alias="gpu-memory-utilization",
+        validation_alias=AliasChoices("gpu-memory-utilization", "gpu_memory_utilization"),
+    )
+    pipeline_parallel_size: int | list[int] | None = Field(
+        default=None,
+        serialization_alias="pipeline-parallel-size",
+        validation_alias=AliasChoices("pipeline-parallel-size", "pipeline_parallel_size"),
+    )
+    tensor_parallel_size: int | list[int] | None = Field(
+        default=None,
+        serialization_alias="tensor-parallel-size",
+        validation_alias=AliasChoices("tensor-parallel-size", "tensor_parallel_size"),
+    )
+    extra_args: str | list[str] | None = Field(
+        default=None,
+        serialization_alias="extra-args",
+        validation_alias=AliasChoices("extra-args", "extra_args"),
     )
-    nodes: Optional[str] = Field(default=None)
 
 
 class PrefillWorkerArgs(WorkerBaseArgs):
@@ -57,17 +83,22 @@ class AIDynamoArgs(BaseModel):
     model: str = "Qwen/Qwen3-0.6B"
     backend: str = "vllm"
     workspace_path: str = Field(
+        default="/workspace",
         serialization_alias="workspace-path",
         validation_alias=AliasChoices("workspace-path", "workspace_path"),
-        default="/workspace",
     )
-    decode_worker: DecodeWorkerArgs
+    decode_worker: DecodeWorkerArgs = Field(default_factory=DecodeWorkerArgs)
     decode_cmd: str = Field(
+        default="python3 -m dynamo.vllm",
         serialization_alias="decode-cmd",
         validation_alias=AliasChoices("decode-cmd", "decode_cmd"),
+    )
+    prefill_worker: PrefillWorkerArgs | None = None
+    prefill_cmd: str = Field(
         default="python3 -m dynamo.vllm",
+        serialization_alias="prefill-cmd",
+        validation_alias=AliasChoices("prefill-cmd", "prefill_cmd"),
     )
-    prefill_worker: PrefillWorkerArgs
 
 
 class GenAIPerfArgs(BaseModel):

src/cloudai/workloads/ai_dynamo/kubernetes_json_gen_strategy.py

@@ -19,15 +19,19 @@
 from pathlib import Path
 from typing import Any, Dict, cast
 
+import yaml
+
 from cloudai.core import JsonGenStrategy
 from cloudai.systems.kubernetes import KubernetesSystem
 
-from .ai_dynamo import AIDynamoTestDefinition
+from .ai_dynamo import AIDynamoTestDefinition, WorkerBaseArgs
 
 
 class AIDynamoKubernetesJsonGenStrategy(JsonGenStrategy):
     """JSON generation strategy for AI Dynamo on Kubernetes systems."""
 
+    DEPLOYMENT_FILE_NAME = "deployment.yaml"
+
     def _install_python_packages(self, repo_root: Path, venv_pip: Path) -> None:
         installs = [
             ("perf_analyzer", repo_root),
@@ -68,30 +72,50 @@ def gen_frontend_dict(self) -> dict[str, Any]:
         }
 
     def gen_decode_dict(self) -> dict[str, Any]:
-        system = cast(KubernetesSystem, self.system)
         tdef = cast(AIDynamoTestDefinition, self.test_run.test)
-        return {
-            "dynamoNamespace": system.default_namespace,
-            "componentType": "worker",
-            "replicas": 1,
-            "resources": {"limits": {"gpu": f"{system.gpus_per_node}"}},
-            "extraPodSpec": {
-                "mainContainer": {
-                    "image": tdef.cmd_args.docker_image_url,
-                    "workingDir": tdef.cmd_args.dynamo.workspace_path,
-                    "command": tdef.cmd_args.dynamo.decode_cmd.split(),
-                    "args": ["--model", tdef.cmd_args.dynamo.model],
-                }
-            },
-        }
+
+        decode_cfg = self._get_base_service_dict()
+        decode_cfg["extraPodSpec"]["mainContainer"]["command"] = tdef.cmd_args.dynamo.decode_cmd.split()
+
+        args = ["--model", tdef.cmd_args.dynamo.model]
+        if tdef.cmd_args.dynamo.prefill_worker:
+            decode_cfg["subComponentType"] = "decode-worker"
+            args.append("--is-decode-worker")
+        args.extend(self._args_from_worker_config(tdef.cmd_args.dynamo.decode_worker))
+
+        decode_cfg["extraPodSpec"]["mainContainer"]["args"] = args
+
+        self._set_multinode_if_needed(decode_cfg, tdef.cmd_args.dynamo.decode_worker)
+
+        return decode_cfg
+
+    def gen_prefill_dict(self) -> dict[str, Any]:
+        tdef = cast(AIDynamoTestDefinition, self.test_run.test)
+        if not tdef.cmd_args.dynamo.prefill_worker:
+            raise ValueError("Prefill worker configuration is not defined in the test definition.")
+
+        prefill_cfg = self._get_base_service_dict()
+        prefill_cfg["subComponentType"] = "prefill"
+        prefill_cfg["extraPodSpec"]["mainContainer"]["command"] = tdef.cmd_args.dynamo.prefill_cmd.split()
+
+        prefill_cfg["extraPodSpec"]["mainContainer"]["args"] = [
+            "--model",
+            tdef.cmd_args.dynamo.model,
+            "--is-prefill-worker",
+            *self._args_from_worker_config(tdef.cmd_args.dynamo.prefill_worker),
+        ]
+
+        self._set_multinode_if_needed(prefill_cfg, tdef.cmd_args.dynamo.prefill_worker)
+
+        return prefill_cfg
 
     def gen_json(self) -> Dict[Any, Any]:
         td = cast(AIDynamoTestDefinition, self.test_run.test)
         k8s_system = cast(KubernetesSystem, self.system)
 
         self._setup_genai(td)
 
-        return {
+        deployment = {
             "apiVersion": "nvidia.com/v1alpha1",
             "kind": "DynamoGraphDeployment",
             "metadata": {"name": k8s_system.default_namespace},
@@ -102,3 +126,44 @@ def gen_json(self) -> Dict[Any, Any]:
                 },
             },
         }
+        if td.cmd_args.dynamo.prefill_worker:
+            deployment["spec"]["services"]["VllmPrefillWorker"] = self.gen_prefill_dict()
+
+        with (self.test_run.output_path / self.DEPLOYMENT_FILE_NAME).open("w") as f:
+            yaml.safe_dump(deployment, f)
+
+        return deployment
+
+    def _get_base_service_dict(self) -> dict[str, Any]:
+        system = cast(KubernetesSystem, self.system)
+        tdef = cast(AIDynamoTestDefinition, self.test_run.test)
+        return {
+            "dynamoNamespace": system.default_namespace,
+            "componentType": "worker",
+            "replicas": 1,
+            "resources": {"limits": {"gpu": f"{system.gpus_per_node}"}},
+            "extraPodSpec": {
+                "mainContainer": {
+                    "image": tdef.cmd_args.docker_image_url,
+                    "workingDir": tdef.cmd_args.dynamo.workspace_path,
+                }
+            },
+        }
+
+    def _to_dynamo_arg(self, arg_name: str) -> str:
+        return "--" + arg_name.replace("_", "-")
+
+    def _dynamo_args_dict(self, model: WorkerBaseArgs) -> dict:
+        return model.model_dump(exclude={"num_nodes", "extra_args", "nodes"}, exclude_none=True)
+
+    def _args_from_worker_config(self, worker: WorkerBaseArgs) -> list[str]:
+        args = []
+        for arg, value in self._dynamo_args_dict(worker).items():
+            args.extend([self._to_dynamo_arg(arg), str(value)])
+        if worker.extra_args:
+            args.append(f"{worker.extra_args}")
+        return args
+
+    def _set_multinode_if_needed(self, cfg: dict[str, Any], worker: WorkerBaseArgs) -> None:
+        if cast(int, worker.num_nodes) > 1:
+            cfg["multinode"] = {"nodeCount": worker.num_nodes}

src/cloudai/workloads/ai_dynamo/report_generation_strategy.py

@@ -20,11 +20,15 @@
 import logging
 import shutil
 from pathlib import Path
-from typing import ClassVar
+from typing import TYPE_CHECKING, ClassVar, cast
 
 from cloudai.core import METRIC_ERROR, ReportGenerationStrategy
+from cloudai.systems.kubernetes.kubernetes_system import KubernetesSystem
 from cloudai.systems.slurm.slurm_system import SlurmSystem
 
+if TYPE_CHECKING:
+    from .ai_dynamo import AIDynamoTestDefinition
+
 CSV_FILES_PATTERN = "profile*_genai_perf.csv"
 JSON_FILES_PATTERN = "profile*_genai_perf.json"
 
@@ -118,16 +122,19 @@ def get_metric(self, metric: str) -> float:
 
     def _calculate_total_gpus(self) -> int | None:
         gpus_per_node = None
-        if isinstance(self.system, SlurmSystem):
+        if isinstance(self.system, (SlurmSystem, KubernetesSystem)):
             gpus_per_node = self.system.gpus_per_node
 
         if gpus_per_node is None:
             return None
 
-        num_frontend_nodes = 1
-        num_prefill_nodes = self.test_run.test.cmd_args.dynamo.prefill_worker.num_nodes
-        num_decode_nodes = self.test_run.test.cmd_args.dynamo.decode_worker.num_nodes
+        tdef = cast("AIDynamoTestDefinition", self.test_run.test)
 
+        num_frontend_nodes = 1
+        num_prefill_nodes = (
+            cast(int, tdef.cmd_args.dynamo.prefill_worker.num_nodes) if tdef.cmd_args.dynamo.prefill_worker else 0
+        )
+        num_decode_nodes = cast(int, tdef.cmd_args.dynamo.decode_worker.num_nodes)
         return (num_frontend_nodes + num_prefill_nodes + num_decode_nodes) * gpus_per_node
 
     def _read_csv_sections(self, source_csv: Path) -> list[list[list[str]]]:

src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py

@@ -77,7 +77,14 @@ def _gen_script_args(self, td: AIDynamoTestDefinition) -> List[str]:
             self._get_toml_args(
                 td.cmd_args.dynamo,
                 "--dynamo-",
-                exclude=["prefill_worker", "decode_worker", "genai_perf", "workspace_path", "decode_cmd"],
+                exclude=[
+                    "prefill_worker",
+                    "decode_worker",
+                    "genai_perf",
+                    "workspace_path",
+                    "decode_cmd",
+                    "prefill_cmd",
+                ],
             )
         )
 
@@ -106,7 +113,8 @@ def _gen_script_args(self, td: AIDynamoTestDefinition) -> List[str]:
                 ]
             )
 
-        args.extend(self._get_toml_args(td.cmd_args.dynamo.prefill_worker, "--prefill-"))
+        if td.cmd_args.dynamo.prefill_worker:
+            args.extend(self._get_toml_args(td.cmd_args.dynamo.prefill_worker, "--prefill-"))
         args.extend(self._get_toml_args(td.cmd_args.dynamo.decode_worker, "--decode-"))
         args.extend(self._get_toml_args(td.cmd_args.genai_perf, "--genai-perf-"))
 
@@ -194,9 +202,11 @@ def get_cached_nodes_spec(self) -> tuple[int, list[str]]:
             return self._node_spec_cache[cache_key]
 
         td = cast(AIDynamoTestDefinition, self.test_run.test)
-        prefill_n = td.cmd_args.dynamo.prefill_worker.num_nodes
+        prefill_n, prefill_nodes = 0, ""
+        if td.cmd_args.dynamo.prefill_worker:
+            prefill_n = cast(int, td.cmd_args.dynamo.prefill_worker.num_nodes)
+            prefill_nodes = td.cmd_args.dynamo.prefill_worker.nodes
         decode_n = td.cmd_args.dynamo.decode_worker.num_nodes
-        prefill_nodes = td.cmd_args.dynamo.prefill_worker.nodes
         decode_nodes = td.cmd_args.dynamo.decode_worker.nodes
 
         assert isinstance(prefill_n, int), "prefill_worker.num_nodes must be an integer"