AI-Dynamo updates.

Signed-off-by: Kapil Arya <kapila@nvidia.com>

Author: karya0

src/cloudai/systems/kubernetes/kubernetes_system.py

@@ -296,25 +296,73 @@ def _run_genai_perf(self, job: KubernetesJob) -> None:
             raise TypeError("Test definition must be an instance of AIDynamoTestDefinition")
 
         genai_perf_results_path = "/tmp/cloudai/genai-perf"
+        frontend_pod = self._get_dynamo_pod_by_role(role="frontend")
+
+        # Copy wrapper script and calc_percentile_csv script to the pod
+        wrapper_script_path = tdef.genai_perf_wrapper_script.installed_path
+        calc_csv_script_path = tdef.calc_percentile_csv.installed_path
+
+        pod_wrapper_path = "/tmp/genai_perf_wrapper.sh"
+        pod_calc_csv_path = "/tmp/calc_percentile_csv.py"
+
+        logging.debug(f"Copying wrapper script {wrapper_script_path} to pod {frontend_pod}")
+        cp_wrapper_cmd = f"kubectl cp {wrapper_script_path} {self.default_namespace}/{frontend_pod}:{pod_wrapper_path}"
+        subprocess.run(cp_wrapper_cmd, shell=True, capture_output=True, text=True, check=True)
+
+        logging.debug(f"Copying calc_percentile_csv script {calc_csv_script_path} to pod {frontend_pod}")
+        cp_calc_cmd = f"kubectl cp {calc_csv_script_path} {self.default_namespace}/{frontend_pod}:{pod_calc_csv_path}"
+        subprocess.run(cp_calc_cmd, shell=True, capture_output=True, text=True, check=True)
 
-        genai_perf_cmd = ["genai-perf", "profile", f"--artifact-dir={genai_perf_results_path}"]
+        # Make wrapper script executable
+        chmod_cmd = ["chmod", "+x", pod_wrapper_path]
+        logging.debug(f"Making wrapper script executable in pod {frontend_pod}")
+        try:
+            lazy.k8s.stream.stream(
+                self.core_v1.connect_get_namespaced_pod_exec,
+                name=frontend_pod,
+                namespace=self.default_namespace,
+                command=chmod_cmd,
+                stderr=True,
+                stdin=False,
+                stdout=True,
+                tty=False,
+            )
+        except lazy.k8s.client.ApiException as e:
+            logging.error(f"Error making wrapper script executable in pod '{frontend_pod}': {e}")
+
+        # Build genai-perf command arguments
+        genai_perf_cmd_parts = ["genai-perf", "profile", f"--artifact-dir={genai_perf_results_path}"]
         for k, v in tdef.cmd_args.genai_perf.model_dump(
             exclude={"extra_args", "extra-args"}, exclude_none=True
         ).items():
-            genai_perf_cmd.append(f"--{k}={v}")
+            genai_perf_cmd_parts.append(f"--{k}={v}")
         if extra_args := tdef.cmd_args.genai_perf.extra_args:
-            genai_perf_cmd.extend(extra_args.split())
-        logging.debug(f"GenAI perf arguments: {genai_perf_cmd=}")
-
-        frontend_pod = self._get_dynamo_pod_by_role(role="frontend")
-
-        logging.debug(f"Executing genai-perf in pod={frontend_pod} cmd={genai_perf_cmd}")
+            genai_perf_cmd_parts.extend(extra_args.split())
+
+        # Build wrapper command with proper parameters
+        report_file = "genai_perf_report.csv"
+        wrapper_cmd = [
+            "/bin/bash",
+            pod_wrapper_path,
+            "--result_dir",
+            genai_perf_results_path,
+            "--report_file",
+            report_file,
+            "--calc_percentile_csv_script",
+            pod_calc_csv_path,
+            "--gpus_per_node",
+            str(self.gpus_per_node),
+            "--",
+            *genai_perf_cmd_parts,
+        ]
+
+        logging.debug(f"Executing genai-perf wrapper in pod={frontend_pod} cmd={wrapper_cmd}")
         try:
             genai_results = lazy.k8s.stream.stream(
                 self.core_v1.connect_get_namespaced_pod_exec,
                 name=frontend_pod,
                 namespace=self.default_namespace,
-                command=genai_perf_cmd,
+                command=wrapper_cmd,
                 stderr=True,
                 stdin=False,
                 stdout=True,
@@ -324,7 +372,7 @@ def _run_genai_perf(self, job: KubernetesJob) -> None:
             with (job.test_run.output_path / "genai_perf.log").open("w") as f:
                 f.write(genai_results)
         except lazy.k8s.client.ApiException as e:
-            logging.error(f"Error executing genai-perf command in pod '{frontend_pod}': {e}")
+            logging.error(f"Error executing genai-perf wrapper command in pod '{frontend_pod}': {e}")
 
         cp_logs_cmd = " ".join(
             [

src/cloudai/workloads/ai_dynamo/__init__.py

@@ -19,7 +19,7 @@
     AIDynamoCmdArgs,
     AIDynamoTestDefinition,
     DecodeWorkerArgs,
-    GenAIPerfArgs,
+    GenAIPerf,
     PrefillWorkerArgs,
 )
 from .kubernetes_json_gen_strategy import AIDynamoKubernetesJsonGenStrategy
@@ -34,6 +34,6 @@
     "AIDynamoSlurmCommandGenStrategy",
     "AIDynamoTestDefinition",
     "DecodeWorkerArgs",
-    "GenAIPerfArgs",
+    "GenAIPerf",
     "PrefillWorkerArgs",
 ]

src/cloudai/workloads/ai_dynamo/ai_dynamo.py

@@ -26,6 +26,39 @@
 from .report_generation_strategy import CSV_FILES_PATTERN, JSON_FILES_PATTERN
 
 
+class BenchmarkArgs(BaseModel):
+    """Arguments for custom benchmarks."""
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    extra_args: str | list[str] | None = Field(
+        default=None,
+        serialization_alias="extra-args",
+        validation_alias=AliasChoices("extra-args", "extra_args"),
+    )
+
+
+class Benchmark(BaseModel):
+    """Arguments for custom benchmarks."""
+
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
+
+    name: str
+    cmd: str
+    repo: Optional[GitRepo] = None
+    enabled: bool = False
+    args: Optional[BenchmarkArgs] = None
+    extra_args: str | list[str] | None = Field(
+        default=None,
+        serialization_alias="extra-args",
+        validation_alias=AliasChoices("extra-args", "extra_args"),
+    )
+
+    @property
+    def report_name(self) -> str:
+        return f"{self.name}_report.csv"
+
+
 class WorkerBaseArgs(BaseModel):
     """Base arguments for VLLM workers."""
 
@@ -78,10 +111,11 @@ class DecodeWorkerArgs(WorkerBaseArgs):
 class AIDynamoArgs(BaseModel):
     """Arguments for AI Dynamo setup."""
 
-    model_config = ConfigDict(extra="allow")
+    model_config = ConfigDict(extra="allow", populate_by_name=True)
 
     model: str = "Qwen/Qwen3-0.6B"
     backend: str = "vllm"
+    connector: Optional[str] = None  # none, lmcache, kvbm
     workspace_path: str = Field(
         default="/workspace",
         serialization_alias="workspace-path",
@@ -95,32 +129,92 @@ class AIDynamoArgs(BaseModel):
     )
     prefill_worker: PrefillWorkerArgs | None = None
     prefill_cmd: str = Field(
-        default="python3 -m dynamo.vllm",
+        default="python3 -m dynamo.vllm --is-prefill-worker",
         serialization_alias="prefill-cmd",
         validation_alias=AliasChoices("prefill-cmd", "prefill_cmd"),
     )
 
 
-class GenAIPerfArgs(BaseModel):
-    """Arguments for GenAI performance profiling."""
+class LMCacheArgs(BaseModel):
+    """Arguments for LMCache."""
+
+    model_config = ConfigDict(extra="allow")
+
+    chunk_size: int = 256
+    local_cpu: bool = False
+    nixl_buffer_size: int = 10737418240
+    nixl_buffer_device: str = "cuda"
+    extra_config_enable_nixl_storage: bool = True
+    extra_config_nixl_backend: str = "GDS_MT"
+    extra_config_nixl_file_pool_size: int = 64
+    extra_config_nixl_path: str = "%CACHEDIR%"
+
+    # LMCache controller configuration
+    enable_controller: bool = True
+    lmcache_instance_id: str = "lmcache_default_instance"
+    controller_url: str = "localhost:9001"
+    lmcache_worker_port: int = 8788
+    distributed_url: str = "localhost:8789"
+
+
+class LMCache(BaseModel):
+    """LMCache configuration."""
 
     model_config = ConfigDict(extra="allow")
 
-    extra_args: str | None = Field(
+    controller_cmd: str = "lmcache_controller --host localhost --port 9000 --monitor-port 9001"
+    repo: Optional[GitRepo] = GitRepo(
+        url="git@github.com:LMCache/LMCache.git", commit="ab8530993992db873869ba882320953582d94309"
+    )
+
+    args: Optional[LMCacheArgs] = None
+    extra_args: str | list[str] | None = Field(
         default=None,
         serialization_alias="extra-args",
         validation_alias=AliasChoices("extra-args", "extra_args"),
     )
 
 
+class GenAIPerf(Benchmark):
+    """Benchmark configuration for GenAI performance profiling."""
+
+    model_config = ConfigDict(extra="allow")
+
+    cmd: str = "genai-perf profile"
+
+
+class LMBench(Benchmark):
+    """Benchmark configuration for LMBench."""
+
+    model_config = ConfigDict(extra="allow")
+
+    repo: Optional[GitRepo] = GitRepo(
+        url="git@github.com:LMCache/LMBenchmark.git", commit="e1406623c5e88878cf2b7fbd64fe6c47f7dcb66f"
+    )
+
+    cmd: str = "python3 ./synthetic-multi-round-qa/multi-round-qa.py"
+
+
+class Constraints(BaseModel):
+    """Constraints for validation of AI Dynamo configurations when using DSE."""
+
+    model_config = ConfigDict(extra="allow")
+
+    prefill_tp_le_decode_tp: bool = True
+    tp_times_pp_le_gpus_per_node: bool = True
+    prefill_decode_nodes_le_total_nodes: bool = True
+
+
 class AIDynamoCmdArgs(CmdArgs):
     """Arguments for AI Dynamo."""
 
     docker_image_url: str
-    huggingface_home_container_path: Path = Path("/root/.cache/huggingface")
+    storage_cache_dir: Optional[str] = None
     dynamo: AIDynamoArgs
-    genai_perf: GenAIPerfArgs
-    run_script: str = ""
+    lmcache: LMCacheArgs
+    genai_perf: GenAIPerf
+    lmbench: LMBench
+    custom_bench: Optional[Benchmark] = None
 
 
 class AIDynamoTestDefinition(TestDefinition):
@@ -129,10 +223,14 @@ class AIDynamoTestDefinition(TestDefinition):
     cmd_args: AIDynamoCmdArgs
     _docker_image: Optional[DockerImage] = None
     script: File = File(Path(__file__).parent.parent / "ai_dynamo/ai_dynamo.sh")
+    genai_perf_wrapper_script: File = File(Path(__file__).parent.parent / "ai_dynamo/genai_perf_wrapper.sh")
+    calc_percentile_csv: File = File(Path(__file__).parent.parent / "ai_dynamo/calc_percentile_csv.py")
     dynamo_repo: GitRepo = GitRepo(
         url="https://github.com/ai-dynamo/dynamo.git", commit="f7e468c7e8ff0d1426db987564e60572167e8464"
     )
     _hf_model: HFModel | None = None
+    benchmarks: str = "genai_perf"
+    constraints: Constraints = Constraints()
 
     @property
     def docker_image(self) -> DockerImage:
@@ -148,7 +246,14 @@ def hf_model(self) -> HFModel:
 
     @property
     def installables(self) -> list[Installable]:
-        return [self.docker_image, self.script, self.dynamo_repo, self.hf_model]
+        return [
+            self.docker_image,
+            self.script,
+            self.genai_perf_wrapper_script,
+            self.dynamo_repo,
+            self.hf_model,
+            self.calc_percentile_csv,
+        ]
 
     def was_run_successful(self, tr: TestRun) -> JobStatusResult:
         output_path = tr.output_path
@@ -159,3 +264,30 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult:
         if not has_results:
             return JobStatusResult(False, "No result files found in the output directory.")
         return JobStatusResult(True)
+
+    def constraint_check(self, tr: TestRun) -> bool:
+        if self.constraints.prefill_tp_le_decode_tp:
+            prefill_tp = self.cmd_args.dynamo.prefill_worker.tensor_parallel_size
+            decode_tp = self.cmd_args.dynamo.decode_worker.tensor_parallel_size
+            if prefill_tp > decode_tp:
+                logging.info("constraint_check failed for: prefill_tp_le_decode_tp")
+                return False
+            logging.info("constraint_check passed for: prefill_tp_le_decode_tp")
+
+        if self.constraints.tp_times_pp_le_gpus_per_node:
+            tp = self.cmd_args.dynamo.prefill_worker.tensor_parallel_size
+            pp = self.cmd_args.dynamo.prefill_worker.pipeline_parallel_size
+            if tp * pp > self.cmd_args.gpus_per_node:
+                logging.info("constraint_check failed for: tp_times_pp_le_gpus_per_node")
+                return False
+            logging.info("constraint_check passed for: tp_times_pp_le_gpus_per_node")
+
+        if self.constraints.prefill_decode_nodes_le_total_nodes:
+            prefill_nodes = self.cmd_args.dynamo.prefill_worker.num_nodes
+            decode_nodes = self.cmd_args.dynamo.decode_worker.num_nodes
+            if prefill_nodes + decode_nodes > self.cmd_args.num_nodes:
+                logging.info("constraint_check failed for: prefill_decode_nodes_le_total_nodes")
+                return False
+            logging.info("constraint_check passed for: prefill_decode_nodes_le_total_nodes")
+
+        return True