feat: Add simplified Training Options for Kubernetes backend type with podTemplateOverrides

Signed-off-by: Abhijeet Dhumal <[email protected]>

Author: abhijeet-dhumal

Makefile

@@ -75,12 +75,8 @@ release: install-dev
 .PHONY: test-python
 test-python: uv-venv  ## Run Python unit tests
 	@uv sync
-	@uv run coverage run --source=kubeflow.trainer.backends.kubernetes.backend,kubeflow.trainer.utils.utils -m pytest \
-		./kubeflow/trainer/backends/kubernetes/backend_test.py \
-		./kubeflow/trainer/backends/kubernetes/utils_test.py
-	@uv run coverage report -m \
-		kubeflow/trainer/backends/kubernetes/backend.py \
-		kubeflow/trainer/backends/kubernetes/utils.py
+	@uv run coverage run --source=kubeflow -m pytest ./kubeflow/
+	@uv run coverage report --omit='*_test.py' --skip-covered --skip-empty
 ifeq ($(report),xml)
 	@uv run coverage xml
 else

kubeflow/trainer/api/trainer_client.py

@@ -107,6 +107,7 @@ def train(
         trainer: Optional[
             Union[types.CustomTrainer, types.CustomTrainerContainer, types.BuiltinTrainer]
         ] = None,
+        options: Optional[list] = None,
     ) -> str:
         """Create a TrainJob. You can configure the TrainJob using one of these trainers:
 
@@ -124,6 +125,8 @@ def train(
             trainer: Optional configuration for a CustomTrainer, CustomTrainerContainer, or
                 BuiltinTrainer. If not specified, the TrainJob will use the
                 runtime's default values.
+            options: Optional list of configuration options to apply to the TrainJob.
+                Options can be imported from kubeflow.trainer.options.
 
         Returns:
             The unique name of the TrainJob that has been generated.
@@ -133,7 +136,12 @@ def train(
             TimeoutError: Timeout to create TrainJobs.
             RuntimeError: Failed to create TrainJobs.
         """
-        return self.backend.train(runtime=runtime, initializer=initializer, trainer=trainer)
+        return self.backend.train(
+            runtime=runtime,
+            initializer=initializer,
+            trainer=trainer,
+            options=options,
+        )
 
     def list_jobs(self, runtime: Optional[types.Runtime] = None) -> list[types.TrainJob]:
         """List of the created TrainJobs. If a runtime is specified, only TrainJobs associated with

kubeflow/trainer/api/trainer_client_test.py

@@ -0,0 +1,72 @@
+# Copyright 2025 The Kubeflow Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Unit tests for TrainerClient backend selection.
+"""
+
+from unittest.mock import Mock, patch
+
+import pytest
+
+from kubeflow.common.types import KubernetesBackendConfig
+from kubeflow.trainer.api.trainer_client import TrainerClient
+from kubeflow.trainer.backends.localprocess.types import LocalProcessBackendConfig
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        {
+            "name": "default_backend_is_kubernetes",
+            "backend_config": None,
+            "expected_backend": "KubernetesBackend",
+            "use_k8s_mocks": True,
+        },
+        {
+            "name": "local_process_backend_selection",
+            "backend_config": LocalProcessBackendConfig(),
+            "expected_backend": "LocalProcessBackend",
+            "use_k8s_mocks": False,
+        },
+        {
+            "name": "kubernetes_backend_selection",
+            "backend_config": KubernetesBackendConfig(),
+            "expected_backend": "KubernetesBackend",
+            "use_k8s_mocks": True,
+        },
+    ],
+)
+def test_backend_selection(test_case):
+    """Test TrainerClient backend selection logic."""
+    if test_case["use_k8s_mocks"]:
+        with (
+            patch("kubernetes.config.load_kube_config"),
+            patch("kubernetes.client.CustomObjectsApi") as mock_custom_api,
+            patch("kubernetes.client.CoreV1Api") as mock_core_api,
+        ):
+            mock_custom_api.return_value = Mock()
+            mock_core_api.return_value = Mock()
+
+            if test_case["backend_config"]:
+                client = TrainerClient(backend_config=test_case["backend_config"])
+            else:
+                client = TrainerClient()
+
+            backend_name = client.backend.__class__.__name__
+            assert backend_name == test_case["expected_backend"]
+    else:
+        client = TrainerClient(backend_config=test_case["backend_config"])
+        backend_name = client.backend.__class__.__name__
+        assert backend_name == test_case["expected_backend"]

kubeflow/trainer/backends/base.py

@@ -21,6 +21,11 @@
 
 
 class RuntimeBackend(abc.ABC):
+    """Base class for runtime backends.
+
+    Options self-validate by checking the backend instance type in their __call__ method.
+    """
+
     @abc.abstractmethod
     def list_runtimes(self) -> list[types.Runtime]:
         raise NotImplementedError()
@@ -41,6 +46,7 @@ def train(
         trainer: Optional[
             Union[types.CustomTrainer, types.CustomTrainerContainer, types.BuiltinTrainer]
         ] = None,
+        options: Optional[list] = None,
     ) -> str:
         raise NotImplementedError()
 

kubeflow/trainer/backends/kubernetes/backend.py

@@ -20,7 +20,7 @@
 import re
 import string
 import time
-from typing import Optional, Union
+from typing import Any, Optional, Union
 import uuid
 
 from kubeflow_trainer_api import models
@@ -87,15 +87,9 @@ def list_runtimes(self) -> list[types.Runtime]:
                 result.append(self.__get_runtime_from_cr(runtime))
 
         except multiprocessing.TimeoutError as e:
-            raise TimeoutError(
-                f"Timeout to list {constants.CLUSTER_TRAINING_RUNTIME_KIND}s "
-                f"in namespace: {self.namespace}"
-            ) from e
+            raise TimeoutError(f"Timeout to list {constants.CLUSTER_TRAINING_RUNTIME_KIND}s") from e
         except Exception as e:
-            raise RuntimeError(
-                f"Failed to list {constants.CLUSTER_TRAINING_RUNTIME_KIND}s "
-                f"in namespace: {self.namespace}"
-            ) from e
+            raise RuntimeError(f"Failed to list {constants.CLUSTER_TRAINING_RUNTIME_KIND}s") from e
 
         return result
 
@@ -184,16 +178,62 @@ def train(
         trainer: Optional[
             Union[types.CustomTrainer, types.CustomTrainerContainer, types.BuiltinTrainer]
         ] = None,
+        options: Optional[list] = None,
     ) -> str:
-        # Generate unique name for the TrainJob.
-        train_job_name = random.choice(string.ascii_lowercase) + uuid.uuid4().hex[:11]
+        if runtime is None:
+            runtime = self.get_runtime(constants.TORCH_RUNTIME)
+
+        # Process options to extract configuration
+        job_spec = {}
+        labels = None
+        annotations = None
+        name = None
+        spec_labels = None
+        spec_annotations = None
+        trainer_overrides = {}
+        pod_template_overrides = None
+
+        if options:
+            for option in options:
+                option(job_spec, trainer, self)
+
+            metadata_section = job_spec.get("metadata", {})
+            labels = metadata_section.get("labels")
+            annotations = metadata_section.get("annotations")
+            name = metadata_section.get("name")
+
+            # Extract spec-level labels/annotations and other spec configurations
+            spec_section = job_spec.get("spec", {})
+            spec_labels = spec_section.get("labels")
+            spec_annotations = spec_section.get("annotations")
+            trainer_overrides = spec_section.get("trainer", {})
+            pod_template_overrides = spec_section.get("podTemplateOverrides")
+
+        # Generate unique name for the TrainJob if not provided
+        train_job_name = name or (
+            random.choice(string.ascii_lowercase)
+            + uuid.uuid4().hex[: constants.JOB_NAME_UUID_LENGTH]
+        )
+
+        # Build the TrainJob spec using the common _get_trainjob_spec method
+        trainjob_spec = self._get_trainjob_spec(
+            runtime=runtime,
+            initializer=initializer,
+            trainer=trainer,
+            trainer_overrides=trainer_overrides,
+            spec_labels=spec_labels,
+            spec_annotations=spec_annotations,
+            pod_template_overrides=pod_template_overrides,
+        )
 
         # Build the TrainJob.
         train_job = models.TrainerV1alpha1TrainJob(
             apiVersion=constants.API_VERSION,
             kind=constants.TRAINJOB_KIND,
-            metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(name=train_job_name),
-            spec=self._get_trainjob_spec(runtime, initializer, trainer),
+            metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(
+                name=train_job_name, labels=labels, annotations=annotations
+            ),
+            spec=trainjob_spec,
         )
 
         # Create the TrainJob.
@@ -549,6 +589,10 @@ def _get_trainjob_spec(
         trainer: Optional[
             Union[types.CustomTrainer, types.CustomTrainerContainer, types.BuiltinTrainer]
         ] = None,
+        trainer_overrides: Optional[dict[str, Any]] = None,
+        spec_labels: Optional[dict[str, str]] = None,
+        spec_annotations: Optional[dict[str, str]] = None,
+        pod_template_overrides: Optional[models.IoK8sApiCoreV1PodTemplateSpec] = None,
     ) -> models.TrainerV1alpha1TrainJobSpec:
         """Get TrainJob spec from the given parameters"""
         if runtime is None:
@@ -575,9 +619,16 @@ def _get_trainjob_spec(
             else:
                 raise ValueError(
                     f"The trainer type {type(trainer)} is not supported. "
-                    "Please use CustomTrainer or BuiltinTrainer."
+                    "Please use CustomTrainer, CustomTrainerContainer, or BuiltinTrainer."
                 )
 
+        # Apply trainer overrides if trainer was not provided but overrides exist
+        if trainer_overrides:
+            if "command" in trainer_overrides:
+                trainer_cr.command = trainer_overrides["command"]
+            if "args" in trainer_overrides:
+                trainer_cr.args = trainer_overrides["args"]
+
         return models.TrainerV1alpha1TrainJobSpec(
             runtimeRef=models.TrainerV1alpha1RuntimeRef(name=runtime.name),
             trainer=(trainer_cr if trainer_cr != models.TrainerV1alpha1Trainer() else None),
@@ -589,4 +640,7 @@ def _get_trainjob_spec(
                 if isinstance(initializer, types.Initializer)
                 else None
             ),
+            labels=spec_labels,
+            annotations=spec_annotations,
+            pod_template_overrides=pod_template_overrides,
         )

kubeflow/trainer/backends/kubernetes/backend_test.py

@@ -35,6 +35,12 @@
 from kubeflow.trainer.backends.kubernetes.backend import KubernetesBackend
 import kubeflow.trainer.backends.kubernetes.utils as utils
 from kubeflow.trainer.constants import constants
+from kubeflow.trainer.options import (
+    Annotations,
+    Labels,
+    SpecAnnotations,
+    SpecLabels,
+)
 from kubeflow.trainer.test.common import (
     DEFAULT_NAMESPACE,
     FAILED,
@@ -274,17 +280,27 @@ def get_train_job(
     runtime_name: str,
     train_job_name: str = BASIC_TRAIN_JOB_NAME,
     train_job_trainer: Optional[models.TrainerV1alpha1Trainer] = None,
+    labels: Optional[dict[str, str]] = None,
+    annotations: Optional[dict[str, str]] = None,
+    spec_labels: Optional[dict[str, str]] = None,
+    spec_annotations: Optional[dict[str, str]] = None,
 ) -> models.TrainerV1alpha1TrainJob:
     """
     Create a mock TrainJob object with optional trainer configurations.
     """
     train_job = models.TrainerV1alpha1TrainJob(
         apiVersion=constants.API_VERSION,
         kind=constants.TRAINJOB_KIND,
-        metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(name=train_job_name),
+        metadata=models.IoK8sApimachineryPkgApisMetaV1ObjectMeta(
+            name=train_job_name,
+            labels=labels,
+            annotations=annotations,
+        ),
         spec=models.TrainerV1alpha1TrainJobSpec(
             runtimeRef=models.TrainerV1alpha1RuntimeRef(name=runtime_name),
             trainer=train_job_trainer,
+            labels=spec_labels,
+            annotations=spec_annotations,
         ),
     )
 
@@ -879,6 +895,58 @@ def test_get_runtime_packages(kubernetes_backend, test_case):
             },
             expected_error=ValueError,
         ),
+        TestCase(
+            name="train with metadata labels and annotations",
+            expected_status=SUCCESS,
+            config={
+                "options": [
+                    Labels({"team": "ml-platform"}),
+                    Annotations({"created-by": "sdk"}),
+                ],
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                labels={"team": "ml-platform"},
+                annotations={"created-by": "sdk"},
+            ),
+        ),
+        TestCase(
+            name="train with spec labels and annotations",
+            expected_status=SUCCESS,
+            config={
+                "options": [
+                    SpecLabels({"app": "training", "version": "v1.0"}),
+                    SpecAnnotations({"prometheus.io/scrape": "true"}),
+                ],
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                spec_labels={"app": "training", "version": "v1.0"},
+                spec_annotations={"prometheus.io/scrape": "true"},
+            ),
+        ),
+        TestCase(
+            name="train with both metadata and spec labels/annotations",
+            expected_status=SUCCESS,
+            config={
+                "options": [
+                    Labels({"owner": "ml-team"}),
+                    Annotations({"description": "Fine-tuning job"}),
+                    SpecLabels({"app": "training", "version": "v1.0"}),
+                    SpecAnnotations({"prometheus.io/scrape": "true"}),
+                ],
+            },
+            expected_output=get_train_job(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=BASIC_TRAIN_JOB_NAME,
+                labels={"owner": "ml-team"},
+                annotations={"description": "Fine-tuning job"},
+                spec_labels={"app": "training", "version": "v1.0"},
+                spec_annotations={"prometheus.io/scrape": "true"},
+            ),
+        ),
     ],
 )
 def test_train(kubernetes_backend, test_case):
@@ -888,8 +956,12 @@ def test_train(kubernetes_backend, test_case):
         kubernetes_backend.namespace = test_case.config.get("namespace", DEFAULT_NAMESPACE)
         runtime = kubernetes_backend.get_runtime(test_case.config.get("runtime", TORCH_RUNTIME))
 
+        options = test_case.config.get("options", [])
+
         train_job_name = kubernetes_backend.train(
-            runtime=runtime, trainer=test_case.config.get("trainer", None)
+            runtime=runtime,
+            trainer=test_case.config.get("trainer", None),
+            options=options,
         )
 
         assert test_case.expected_status == SUCCESS