Enable AI Dynamo w/ K8S SPCx (Alpha) - Rel 2B (#667)

Author: TaekyungHeo

conf/common/system/kubernetes_cluster.toml

@@ -1,5 +1,5 @@
 # SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
-# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

src/cloudai/registration.py

@@ -55,6 +55,7 @@ def register_all():
     from cloudai.systems.slurm import SlurmInstaller, SlurmRunner, SlurmSystem
     from cloudai.systems.standalone import StandaloneInstaller, StandaloneRunner, StandaloneSystem
     from cloudai.workloads.ai_dynamo import (
+        AIDynamoKubernetesJsonGenStrategy,
         AIDynamoReportGenerationStrategy,
         AIDynamoSlurmCommandGenStrategy,
         AIDynamoTestDefinition,
@@ -183,6 +184,9 @@ def register_all():
     Registry().add_command_gen_strategy(SlurmSystem, NixlPerftestTestDefinition, NixlPerftestSlurmCommandGenStrategy)
 
     Registry().add_command_gen_strategy(SlurmSystem, AIDynamoTestDefinition, AIDynamoSlurmCommandGenStrategy)
+    Registry().add_strategy(
+        JsonGenStrategy, [KubernetesSystem], [AIDynamoTestDefinition], AIDynamoKubernetesJsonGenStrategy
+    )
     Registry().add_command_gen_strategy(SlurmSystem, BashCmdTestDefinition, BashCmdCommandGenStrategy)
     Registry().add_command_gen_strategy(SlurmSystem, NIXLKVBenchTestDefinition, NIXLKVBenchSlurmCommandGenStrategy)
 

src/cloudai/systems/kubernetes/kubernetes_installer.py

@@ -20,8 +20,9 @@
 import shutil
 import subprocess
 from pathlib import Path
+from shutil import rmtree
 
-from cloudai.core import BaseInstaller, DockerImage, GitRepo, Installable, InstallStatusResult
+from cloudai.core import BaseInstaller, DockerImage, File, GitRepo, Installable, InstallStatusResult, PythonExecutable
 from cloudai.util.lazy_imports import lazy
 
 
@@ -55,83 +56,36 @@ def _check_prerequisites(self) -> InstallStatusResult:
             logging.error(message)
             return InstallStatusResult(False, message)
 
-        # Check MPIJob-related prerequisites
-        mpi_job_result = self._check_mpi_job_prerequisites()
-        if not mpi_job_result.success:
-            return mpi_job_result
-
         logging.info("All prerequisites are met. Proceeding with installation.")
         return InstallStatusResult(True)
 
-    def _check_mpi_job_prerequisites(self) -> InstallStatusResult:
-        """
-        Check if the MPIJob CRD is installed and if MPIJob kind is supported in the Kubernetes cluster.
-
-        This ensures that the system is ready for MPI-based operations.
-
-        Returns
-            InstallStatusResult: Result containing the status of the MPIJob prerequisite check and any error message.
-        """
-        # Check if MPIJob CRD is installed
-        try:
-            custom_api = lazy.k8s.client.CustomObjectsApi()
-            custom_api.get_cluster_custom_object(group="kubeflow.org", version="v1", plural="mpijobs", name="mpijobs")
-        except lazy.k8s.client.ApiException as e:
-            if e.status == 404:
-                message = (
-                    "Installation failed during prerequisite checking stage because MPIJob CRD is not installed on "
-                    "this Kubernetes cluster. Please ensure that the MPI Operator is installed and MPIJob kind is "
-                    "supported. You can follow the instructions in the MPI Operator repository to install it: "
-                    "https://github.com/kubeflow/mpi-operator"
-                )
-                logging.error(message)
-                return InstallStatusResult(False, message)
-            else:
-                message = (
-                    f"Installation failed during prerequisite checking stage due to an error while checking for MPIJob "
-                    f"CRD. Original error: {e!r}. Please ensure that the Kubernetes cluster is accessible and the "
-                    f"MPI Operator is correctly installed."
-                )
-                logging.error(message)
-                return InstallStatusResult(False, message)
-
-        # Check if MPIJob kind is supported
-        try:
-            api_resources = lazy.k8s.client.ApiextensionsV1Api().list_custom_resource_definition()
-            mpi_job_supported = any(item.metadata.name == "mpijobs.kubeflow.org" for item in api_resources.items)
-        except lazy.k8s.client.ApiException as e:
-            message = (
-                f"Installation failed during prerequisite checking stage due to an error while checking for MPIJob "
-                f"kind support. Original error: {e!r}. Please ensure that the Kubernetes cluster is accessible and "
-                f"the MPI Operator is correctly installed."
-            )
-            logging.error(message)
-            return InstallStatusResult(False, message)
-
-        if not mpi_job_supported:
-            message = (
-                "Installation failed during prerequisite checking stage because MPIJob kind is not supported on this "
-                "Kubernetes cluster. Please ensure that the MPI Operator is installed and MPIJob kind is supported. "
-                "You can follow the instructions in the MPI Operator repository to install it: "
-                "https://github.com/kubeflow/mpi-operator"
-            )
-            logging.error(message)
-            return InstallStatusResult(False, message)
-
-        return InstallStatusResult(True)
-
     def install_one(self, item: Installable) -> InstallStatusResult:
         if isinstance(item, DockerImage):
             return InstallStatusResult(True, f"Docker image {item} installed")
         elif isinstance(item, GitRepo):
             return self._install_one_git_repo(item)
+        elif isinstance(item, PythonExecutable):
+            return self._install_python_executable(item)
+        elif isinstance(item, File):
+            item.installed_path = self.system.install_path / item.src.name
+            shutil.copyfile(item.src, item.installed_path, follow_symlinks=False)
+            return InstallStatusResult(True)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def uninstall_one(self, item: Installable) -> InstallStatusResult:
         if isinstance(item, DockerImage):
             return InstallStatusResult(True, f"Docker image {item} uninstalled")
         elif isinstance(item, GitRepo):
             return self._uninstall_git_repo(item)
+        elif isinstance(item, PythonExecutable):
+            return self._uninstall_python_executable(item)
+        elif isinstance(item, File):
+            if item.installed_path != item.src:
+                item.installed_path.unlink()
+                item._installed_path = None
+                return InstallStatusResult(True)
+            logging.debug(f"File {item.installed_path} does not exist.")
+            return InstallStatusResult(True)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def is_installed_one(self, item: Installable) -> InstallStatusResult:
@@ -143,6 +97,8 @@ def is_installed_one(self, item: Installable) -> InstallStatusResult:
                 item.installed_path = repo_path
                 return InstallStatusResult(True)
             return InstallStatusResult(False, f"Git repository {item.url} not cloned")
+        elif isinstance(item, PythonExecutable):
+            return self._is_python_executable_installed(item)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
@@ -151,6 +107,10 @@ def mark_as_installed_one(self, item: Installable) -> InstallStatusResult:
         elif isinstance(item, GitRepo):
             item.installed_path = self.system.install_path / item.repo_name
             return InstallStatusResult(True)
+        elif isinstance(item, PythonExecutable):
+            item.git_repo.installed_path = self.system.install_path / item.git_repo.repo_name
+            item.venv_path = self.system.install_path / item.venv_name
+            return InstallStatusResult(True)
         return InstallStatusResult(False, f"Unsupported item type: {type(item)}")
 
     def _install_one_git_repo(self, item: GitRepo) -> InstallStatusResult:
@@ -172,6 +132,45 @@ def _install_one_git_repo(self, item: GitRepo) -> InstallStatusResult:
         item.installed_path = repo_path
         return InstallStatusResult(True)
 
+    def _install_python_executable(self, item: PythonExecutable) -> InstallStatusResult:
+        res = self._install_one_git_repo(item.git_repo)
+        if not res.success:
+            return res
+
+        res = self._create_venv(item)
+        if not res.success:
+            return res
+
+        return InstallStatusResult(True)
+
+    def _install_dependencies(self, item: PythonExecutable) -> InstallStatusResult:
+        venv_path = self.system.install_path / item.venv_name
+
+        if not item.git_repo.installed_path:
+            return InstallStatusResult(False, "Git repository must be installed before creating virtual environment.")
+
+        project_dir = item.git_repo.installed_path
+
+        if item.project_subpath:
+            project_dir = project_dir / item.project_subpath
+
+        pyproject_toml = project_dir / "pyproject.toml"
+        requirements_txt = project_dir / "requirements.txt"
+
+        if pyproject_toml.exists() and requirements_txt.exists():
+            if item.dependencies_from_pyproject:
+                res = self._install_pyproject(venv_path, project_dir)
+            else:
+                res = self._install_requirements(venv_path, requirements_txt)
+        elif pyproject_toml.exists():
+            res = self._install_pyproject(venv_path, project_dir)
+        elif requirements_txt.exists():
+            res = self._install_requirements(venv_path, requirements_txt)
+        else:
+            return InstallStatusResult(False, "No pyproject.toml or requirements.txt found for installation.")
+
+        return res
+
     def _clone_repository(self, git_url: str, path: Path) -> InstallStatusResult:
         logging.debug(f"Cloning repository {git_url} into {path}")
         clone_cmd = ["git", "clone"]
@@ -195,6 +194,58 @@ def _checkout_commit(self, commit_hash: str, path: Path) -> InstallStatusResult:
             return InstallStatusResult(False, f"Failed to checkout commit: {result.stderr}")
         return InstallStatusResult(True)
 
+    def _create_venv(self, item: PythonExecutable) -> InstallStatusResult:
+        venv_path = self.system.install_path / item.venv_name
+        logging.debug(f"Creating virtual environment in {venv_path}")
+        if venv_path.exists():
+            msg = f"Virtual environment already exists at {venv_path}."
+            logging.debug(msg)
+            return InstallStatusResult(True, msg)
+
+        cmd = ["python", "-m", "venv", str(venv_path)]
+        logging.debug(f"Creating venv using cmd: {' '.join(cmd)}")
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        logging.debug(f"venv creation STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}")
+        if result.returncode != 0:
+            if venv_path.exists():
+                rmtree(venv_path)
+            return InstallStatusResult(
+                False, f"Failed to create venv:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
+            )
+
+        res = self._install_dependencies(item)
+        if not res.success:
+            if venv_path.exists():
+                rmtree(venv_path)
+            return res
+
+        item.venv_path = self.system.install_path / item.venv_name
+
+        return InstallStatusResult(True)
+
+    def _install_pyproject(self, venv_dir: Path, project_dir: Path) -> InstallStatusResult:
+        install_cmd = [str(venv_dir / "bin" / "python"), "-m", "pip", "install", str(project_dir)]
+        logging.debug(f"Installing dependencies using: {' '.join(install_cmd)}")
+        result = subprocess.run(install_cmd, capture_output=True, text=True)
+
+        if result.returncode != 0:
+            return InstallStatusResult(False, f"Failed to install {project_dir} using pip: {result.stderr}")
+
+        return InstallStatusResult(True)
+
+    def _install_requirements(self, venv_dir: Path, requirements_txt: Path) -> InstallStatusResult:
+        if not requirements_txt.is_file():
+            return InstallStatusResult(False, f"Requirements file is invalid or does not exist: {requirements_txt}")
+
+        install_cmd = [str(venv_dir / "bin" / "python"), "-m", "pip", "install", "-r", str(requirements_txt)]
+        logging.debug(f"Installing dependencies using: {' '.join(install_cmd)}")
+        result = subprocess.run(install_cmd, capture_output=True, text=True)
+
+        if result.returncode != 0:
+            return InstallStatusResult(False, f"Failed to install dependencies from requirements.txt: {result.stderr}")
+
+        return InstallStatusResult(True)
+
     def _uninstall_git_repo(self, item: GitRepo) -> InstallStatusResult:
         logging.debug(f"Uninstalling git repository at {item.installed_path=}")
         repo_path = item.installed_path if item.installed_path else self.system.install_path / item.repo_name
@@ -207,3 +258,37 @@ def _uninstall_git_repo(self, item: GitRepo) -> InstallStatusResult:
         item.installed_path = None
 
         return InstallStatusResult(True)
+
+    def _uninstall_python_executable(self, item: PythonExecutable) -> InstallStatusResult:
+        res = self._uninstall_git_repo(item.git_repo)
+        if not res.success:
+            return res
+
+        logging.debug(f"Uninstalling virtual environment at {item.venv_path=}")
+        venv_path = item.venv_path if item.venv_path else self.system.install_path / item.venv_name
+        if not venv_path.exists():
+            msg = f"Virtual environment {item.venv_name} is not created."
+            return InstallStatusResult(True, msg)
+
+        logging.debug(f"Removing folder {venv_path}")
+        rmtree(venv_path)
+        item.venv_path = None
+
+        return InstallStatusResult(True)
+
+    def _is_python_executable_installed(self, item: PythonExecutable) -> InstallStatusResult:
+        repo_path = (
+            item.git_repo.installed_path
+            if item.git_repo.installed_path
+            else self.system.install_path / item.git_repo.repo_name
+        )
+        if not repo_path.exists():
+            return InstallStatusResult(False, f"Git repository {item.git_repo.url} not cloned")
+        item.git_repo.installed_path = repo_path
+
+        venv_path = item.venv_path if item.venv_path else self.system.install_path / item.venv_name
+        if not venv_path.exists():
+            return InstallStatusResult(False, f"Virtual environment not created for {item.git_repo.url}")
+        item.venv_path = venv_path
+
+        return InstallStatusResult(True, "Python executable installed")

src/cloudai/systems/kubernetes/kubernetes_runner.py

@@ -38,13 +38,18 @@ def _submit_test(self, tr: TestRun) -> KubernetesJob:
             k8s_system: KubernetesSystem = cast(KubernetesSystem, self.system)
             job_name = k8s_system.create_job(job_spec)
 
-        return KubernetesJob(tr, id=job_name, name=job_name, kind=job_kind)
+        job = KubernetesJob(tr, id=job_name, name=job_name, kind=job_kind)
+
+        return job
 
     def on_job_completion(self, job: BaseJob) -> None:
         k8s_system: KubernetesSystem = cast(KubernetesSystem, self.system)
         k_job = cast(KubernetesJob, job)
-        k8s_system.store_logs_for_job(k_job.name, k_job.test_run.output_path)
-        k8s_system.delete_job(k_job.name, k_job.kind)
+        if k_job.kind == "dynamographdeployment":
+            k8s_system._delete_dynamo_graph_deployment(k_job.name)
+        else:
+            k8s_system.store_logs_for_job(k_job.name, k_job.test_run.output_path)
+            k8s_system.delete_job(k_job.name, k_job.kind)
 
     def kill_job(self, job: BaseJob) -> None:
         k8s_system: KubernetesSystem = cast(KubernetesSystem, self.system)