Merge pull request #752 from NVIDIA/am/k8s-dse

Enhancements for Dynamo with k8s

Author: amaslenn

src/cloudai/_core/json_gen_strategy.py

@@ -18,6 +18,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict
 
+import toml
+
 from .system import System
 from .test_scenario import TestRun
 
@@ -29,6 +31,8 @@ class JsonGenStrategy(ABC):
     It specifies how to generate JSON job specifications based on system and test parameters.
     """
 
+    TEST_RUN_DUMP_FILE_NAME: str = "test-run.toml"
+
     def __init__(self, system: System, test_run: TestRun) -> None:
         self.system = system
         self.test_run = test_run
@@ -54,6 +58,14 @@ def sanitize_k8s_job_name(self, job_name: str) -> str:
         sanitized_name = re.sub(r"[^a-z0-9]+$", "", sanitized_name)
         return sanitized_name[:253]
 
+    def store_test_run(self) -> None:
+        from cloudai.models.scenario import TestRunDetails
+
+        test_cmd, srun_cmd = ("", "n/a")
+        with (self.test_run.output_path / self.TEST_RUN_DUMP_FILE_NAME).open("w") as f:
+            trd = TestRunDetails.from_test_run(self.test_run, test_cmd=test_cmd, full_cmd=srun_cmd)
+            toml.dump(trd.model_dump(), f)
+
     @abstractmethod
     def gen_json(self) -> Dict[Any, Any]:
         """

src/cloudai/cli/cli.py

@@ -74,6 +74,10 @@ def setup_logging(log_file: str, log_level: str) -> None:
                 "handlers": ["debug_file"],
                 "propagate": False,
             },
+            "kubernetes": {
+                "handlers": [],
+                "propagate": False,
+            },
         },
     }
     logging.config.dictConfig(LOGGING_CONFIG)

src/cloudai/systems/kubernetes/kubernetes_runner.py

@@ -42,6 +42,10 @@ def _submit_test(self, tr: TestRun) -> KubernetesJob:
 
         return job
 
+    def on_job_submit(self, tr: TestRun) -> None:
+        json_gen = self.get_json_gen_strategy(self.system, tr)
+        json_gen.store_test_run()
+
     def on_job_completion(self, job: BaseJob) -> None:
         k8s_system: KubernetesSystem = cast(KubernetesSystem, self.system)
         k_job = cast(KubernetesJob, job)

src/cloudai/systems/kubernetes/kubernetes_system.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import subprocess
 import time
@@ -43,7 +42,6 @@ class KubernetesSystem(System):
     _core_v1: Optional[k8s.client.CoreV1Api] = None
     _batch_v1: Optional[k8s.client.BatchV1Api] = None
     _custom_objects_api: Optional[k8s.client.CustomObjectsApi] = None
-    _port_forward_process: subprocess.Popen | None = None
     _genai_perf_completed: bool = False
 
     def __getstate__(self) -> dict[str, Any]:
@@ -279,60 +277,15 @@ def are_vllm_pods_ready(self, job: KubernetesJob) -> bool:
 
         return all_ready
 
-    def _setup_port_forward(self, job: KubernetesJob) -> None:
-        if self._port_forward_process and self._port_forward_process.poll() is None:
-            logging.debug("Port forwarding is already running")
-            return
-
-        if not self.are_vllm_pods_ready(job):
-            logging.debug("Pods are not ready yet, skipping port forward")
-            return
-
-        cmd = f"kubectl port-forward svc/{job.name}-frontend 8000:8000 -n {self.default_namespace}"
-        logging.debug("Starting port forwarding")
-        self._port_forward_process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-        logging.debug(f"Port forwarding started (pid={self._port_forward_process.pid})")
-
-    def _check_model_server(self) -> bool:
-        if not self._port_forward_process:
-            logging.debug("Port forward process is not running")
-            return False
-
-        server = "localhost:8000"
-        cmd = f"curl -s http://{server}/v1/models"
-        logging.debug(f"Checking if model server is up at {server}: {cmd}")
-        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
-
-        if result.returncode != 0:
-            logging.debug(
-                f"Failed to connect to model server={server}, "
-                f"output={result.stdout.strip()}, "
-                f"error={result.stderr.strip()}"
-            )
-            return False
-
-        try:
-            response = json.loads(result.stdout)
-            if response.get("data") and len(response["data"]) > 0:
-                logging.debug(f"Model server is running. Response: {result.stdout}")
-                return True
-            else:
-                logging.debug("Model server is up but no models are loaded yet")
-                return False
-        except json.JSONDecodeError:
-            logging.warning("Invalid JSON response from model server")
-            return False
-
-    def _get_frontend_pod_name(self) -> str:
+    def _get_dynamo_pod_by_role(self, role: str) -> str:
         for pod in self.core_v1.list_namespaced_pod(namespace=self.default_namespace).items:
             labels = pod.metadata.labels
             logging.debug(f"Found pod: {pod.metadata.name} with labels: {labels}")
-            if labels and str(labels.get("nvidia.com/dynamo-component", "")).lower() == "frontend":  # v0.6.x
+            if labels and str(labels.get("nvidia.com/dynamo-component", "")).lower() == role.lower():  # v0.6.x
                 return pod.metadata.name
-            if labels and str(labels.get("nvidia.com/dynamo-component-type", "")).lower() == "frontend":  # v0.7.x
+            if labels and str(labels.get("nvidia.com/dynamo-component-type", "")).lower() == role.lower():  # v0.7.x
                 return pod.metadata.name
-        raise RuntimeError("No frontend pod found for the job")
+        raise RuntimeError(f"No pod found for the role '{role}'")
 
     def _run_genai_perf(self, job: KubernetesJob) -> None:
         from cloudai.workloads.ai_dynamo.ai_dynamo import AIDynamoTestDefinition
@@ -352,7 +305,7 @@ def _run_genai_perf(self, job: KubernetesJob) -> None:
             genai_perf_cmd.extend(extra_args.split())
         logging.debug(f"GenAI perf arguments: {genai_perf_cmd=}")
 
-        frontend_pod = self._get_frontend_pod_name()
+        frontend_pod = self._get_dynamo_pod_by_role(role="frontend")
 
         logging.debug(f"Executing genai-perf in pod={frontend_pod} cmd={genai_perf_cmd}")
         try:
@@ -402,12 +355,20 @@ def _is_dynamo_graph_deployment_running(self, job: KubernetesJob) -> bool:
             return False
 
         if self.are_vllm_pods_ready(job):
-            self._setup_port_forward(job)
-            if self._port_forward_process and self._check_model_server():
-                logging.debug("vLLM server is up and models are loaded")
-                self._run_genai_perf(job)
-                self._genai_perf_completed = True
-                return False
+            self._run_genai_perf(job)
+            self._genai_perf_completed = True
+
+            for pod_role in {"decode", "prefill", "frontend"}:
+                try:
+                    pod_name = self._get_dynamo_pod_by_role(pod_role)
+                    logging.debug(f"Fetching logs for {pod_role=} {pod_name=}")
+                    logs = self.core_v1.read_namespaced_pod_log(name=pod_name, namespace=self.default_namespace)
+                    with (job.test_run.output_path / f"{pod_role}_pod.log").open("w") as f:
+                        f.write(logs)
+                except Exception as e:
+                    logging.debug(f"Error fetching logs for role '{pod_role}': {e}")
+
+            return False
 
         deployment = cast(
             dict,
@@ -485,9 +446,7 @@ def _delete_dynamo_graph_deployment(self, job_name: str) -> None:
         if result.returncode != 0:
             logging.debug(f"Failed to delete DynamoGraphDeployment: {result.stderr}")
 
-        if self._port_forward_process and self._port_forward_process.poll() is None:
-            self._port_forward_process.kill()
-        self._port_forward_process = None
+        self._genai_perf_completed = False
 
     def create_job(self, job_spec: Dict[Any, Any], timeout: int = 60, interval: int = 1) -> str:
         """
@@ -562,6 +521,10 @@ def _create_mpi_job(self, job_spec: Dict[Any, Any]) -> str:
         return job_name
 
     def _create_dynamo_graph_deployment(self, job_spec: Dict[Any, Any]) -> str:
+        logging.debug(f"Attempting to delete existing job='{job_spec['metadata']['name']}' before creation.")
+        self._delete_dynamo_graph_deployment(job_spec["metadata"]["name"])
+
+        logging.debug("Creating DynamoGraphDeployment with spec")
         try:
             api_response = self.custom_objects_api.create_namespaced_custom_object(
                 group="nvidia.com",

src/cloudai/workloads/ai_dynamo/ai_dynamo.py

@@ -154,7 +154,7 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult:
         output_path = tr.output_path
         csv_files = list(output_path.rglob(CSV_FILES_PATTERN))
         json_files = list(output_path.rglob(JSON_FILES_PATTERN))
-        logging.debug(f"Found CSV files: {csv_files}, JSON files: {json_files}")
+        logging.debug(f"Found CSV files in {output_path.absolute()}: {csv_files}, JSON files: {json_files}")
         has_results = len(csv_files) > 0 and len(json_files) > 0
         if not has_results:
             return JobStatusResult(False, "No result files found in the output directory.")