feat(backend): Support namespaced TrainingRuntime in the SDK

Author: shaikmoeed

kubeflow/trainer/backends/kubernetes/backend.py

@@ -60,21 +60,33 @@ def __init__(
     def list_runtimes(self) -> list[types.Runtime]:
         result = []
         try:
-            thread = self.custom_api.list_cluster_custom_object(
+            cluster_thread = self.custom_api.list_cluster_custom_object(
                 constants.GROUP,
                 constants.VERSION,
                 constants.CLUSTER_TRAINING_RUNTIME_PLURAL,
                 async_req=True,
             )
 
-            runtime_list = models.TrainerV1alpha1ClusterTrainingRuntimeList.from_dict(
-                thread.get(constants.DEFAULT_TIMEOUT)
+            namespace_thread = self.custom_api.list_namespaced_custom_object(
+                constants.GROUP,
+                constants.VERSION,
+                self.namespace,
+                constants.TRAINING_RUNTIME_PLURAL,
+                async_req=True,
+            )
+
+            cluster_runtime_list = models.TrainerV1alpha1ClusterTrainingRuntimeList.from_dict(
+                cluster_thread.get(constants.DEFAULT_TIMEOUT)
+            )
+
+            namespace_runtime_list = models.TrainerV1alpha1TrainingRuntimeList.from_dict(
+                namespace_thread.get(constants.DEFAULT_TIMEOUT)
             )
 
-            if not runtime_list:
+            if not (cluster_runtime_list and namespace_runtime_list):
                 return result
 
-            for runtime in runtime_list.items:
+            for runtime in namespace_runtime_list.items + cluster_runtime_list.items:
                 if not (
                     runtime.metadata
                     and runtime.metadata.labels
@@ -89,33 +101,55 @@ def list_runtimes(self) -> list[types.Runtime]:
 
         except multiprocessing.TimeoutError as e:
             raise TimeoutError(
-                f"Timeout to list {constants.CLUSTER_TRAINING_RUNTIME_KIND}s "
+                "Timeout to list "
+                f"{constants.CLUSTER_TRAINING_RUNTIME_KIND}s/{constants.TRAINING_RUNTIME_KIND}s "
                 f"in namespace: {self.namespace}"
             ) from e
         except Exception as e:
             raise RuntimeError(
-                f"Failed to list {constants.CLUSTER_TRAINING_RUNTIME_KIND}s "
+                "Failed to list "
+                f"{constants.CLUSTER_TRAINING_RUNTIME_KIND}s/{constants.TRAINING_RUNTIME_KIND}s "
                 f"in namespace: {self.namespace}"
             ) from e
 
         return result
 
     def get_runtime(self, name: str) -> types.Runtime:
-        """Get the the Runtime object"""
+        """Get the the Runtime object prefer namespaced, fall-back to cluster-scoped"""
 
         try:
-            thread = self.custom_api.get_cluster_custom_object(
+            cluster_thread = self.custom_api.get_cluster_custom_object(
                 constants.GROUP,
                 constants.VERSION,
                 constants.CLUSTER_TRAINING_RUNTIME_PLURAL,
                 name,
                 async_req=True,
             )
 
-            runtime = models.TrainerV1alpha1ClusterTrainingRuntime.from_dict(
-                thread.get(constants.DEFAULT_TIMEOUT)  # type: ignore
+            namespace_thread = self.custom_api.get_namespaced_custom_object(
+                constants.GROUP,
+                constants.VERSION,
+                self.namespace,
+                constants.TRAINING_RUNTIME_PLURAL,
+                name,
+                async_req=True,
             )
 
+            # Try namespaced runtime first, fall back to cluster-scoped one
+            try:
+                runtime = models.TrainerV1alpha1TrainingRuntime.from_dict(
+                    namespace_thread.get(constants.DEFAULT_TIMEOUT)  # type: ignore
+                )
+            except Exception as e:
+                logger.warning(
+                    f"Namespaced TrainingRuntime '{self.namespace}/{name}' not found "
+                    f"({type(e).__name__}: {e}); falling back to cluster-scoped runtime."
+                )
+
+                runtime = models.TrainerV1alpha1ClusterTrainingRuntime.from_dict(
+                    cluster_thread.get(constants.DEFAULT_TIMEOUT)  # type: ignore
+                )
+
         except multiprocessing.TimeoutError as e:
             raise TimeoutError(
                 f"Timeout to get {constants.CLUSTER_TRAINING_RUNTIME_PLURAL}: "
@@ -433,8 +467,13 @@ def delete_job(self, name: str):
 
     def __get_runtime_from_crd(
         self,
-        runtime_crd: models.TrainerV1alpha1ClusterTrainingRuntime,
+        runtime_crd: Union[
+            models.TrainerV1alpha1ClusterTrainingRuntime, models.TrainerV1alpha1TrainingRuntime
+        ],
     ) -> types.Runtime:
+        crd_kind = getattr(runtime_crd, "kind", "UnknownKind")
+        crd_name = getattr(runtime_crd.metadata, "name", "UnknownName")
+
         if not (
             runtime_crd.metadata
             and runtime_crd.metadata.name
@@ -443,7 +482,11 @@ def __get_runtime_from_crd(
             and runtime_crd.spec.template.spec
             and runtime_crd.spec.template.spec.replicated_jobs
         ):
-            raise Exception(f"ClusterTrainingRuntime CRD is invalid: {runtime_crd}")
+            raise Exception(
+                f"{crd_kind} '{crd_name}' is invalid — missing one or more required fields: "
+                f"metadata.name, spec.mlPolicy, spec.template.spec.replicatedJobs.\n"
+                f"Full object: {runtime_crd}"
+            )
 
         if not (
             runtime_crd.metadata.labels

kubeflow/trainer/constants/constants.py

@@ -32,6 +32,12 @@
 # The plural for the ClusterTrainingRuntime.
 CLUSTER_TRAINING_RUNTIME_PLURAL = "clustertrainingruntimes"
 
+# The Kind name for the TrainingRuntime.
+TRAINING_RUNTIME_KIND = "ClusterTrainingRuntime"
+
+# The plural for the ClusterTrainingRuntime.
+TRAINING_RUNTIME_PLURAL = "trainingruntimes"
+
 # The Kind name for the TrainJob.
 TRAINJOB_KIND = "TrainJob"