fix: Minor improvemnts in imports and trainer types

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

Author: abhijeet-dhumal

kubeflow/trainer/api/trainer_client.py

@@ -112,7 +112,7 @@ def train(
             trainer: Optional configuration for a CustomTrainer 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. Use
-                WithLabels and WithAnnotations for basic metadata configuration.
+                Labels and Annotations for basic metadata configuration.
 
         Returns:
             The unique name of the TrainJob that has been generated.

kubeflow/trainer/options/__init__.py

@@ -20,13 +20,13 @@
 
 from kubeflow.trainer.options.common import (
     ContainerOverride,
+    Name,
     PodTemplateOverride,
     PodTemplateSpecOverride,
 )
 from kubeflow.trainer.options.kubernetes import (
     Annotations,
     Labels,
-    Name,
     PodTemplateOverrides,
     SpecAnnotations,
     SpecLabels,
@@ -39,12 +39,12 @@
 __all__ = [
     # Common options
     "ContainerOverride",
+    "Name",
     "PodTemplateOverride",
     "PodTemplateSpecOverride",
     # Kubernetes options
     "Annotations",
     "Labels",
-    "Name",
     "PodTemplateOverrides",
     "SpecAnnotations",
     "SpecLabels",

kubeflow/trainer/options/common.py

@@ -14,16 +14,40 @@
 
 """Common options and helper classes used across multiple backends."""
 
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import Optional
+from typing import TYPE_CHECKING, Any
+
+from kubeflow.trainer.backends.base import UniversalCompatible
+
+if TYPE_CHECKING:
+    from kubeflow.trainer.types.types import BuiltinTrainer, CustomTrainer
 
 __all__ = [
     "ContainerOverride",
+    "Name",
     "PodTemplateOverride",
     "PodTemplateSpecOverride",
 ]
 
 
+@dataclass
+class Name(UniversalCompatible):
+    """Set a custom name for the TrainJob resource."""
+
+    name: str
+
+    def __call__(
+        self,
+        job_spec: dict[str, Any],
+        trainer: BuiltinTrainer | CustomTrainer | None = None,
+    ) -> None:
+        """Apply custom name to the job specification."""
+        metadata = job_spec.setdefault("metadata", {})
+        metadata["name"] = self.name
+
+
 @dataclass
 class ContainerOverride:
     """Configuration for overriding a specific container in a pod.
@@ -37,8 +61,8 @@ class ContainerOverride:
     """
 
     name: str
-    env: Optional[list[dict]] = None
-    volume_mounts: Optional[list[dict]] = None
+    env: list[dict] | None = None
+    volume_mounts: list[dict] | None = None
 
     def __post_init__(self):
         """Validate the container override configuration."""
@@ -108,15 +132,15 @@ class PodTemplateSpecOverride:
         image_pull_secrets: Image pull secrets for the pods.
     """
 
-    service_account_name: Optional[str] = None
-    node_selector: Optional[dict[str, str]] = None
-    affinity: Optional[dict] = None
-    tolerations: Optional[list[dict]] = None
-    volumes: Optional[list[dict]] = None
-    init_containers: Optional[list[ContainerOverride]] = None
-    containers: Optional[list[ContainerOverride]] = None
-    scheduling_gates: Optional[list[dict]] = None
-    image_pull_secrets: Optional[list[dict]] = None
+    service_account_name: str | None = None
+    node_selector: dict[str, str] | None = None
+    affinity: dict | None = None
+    tolerations: list[dict] | None = None
+    volumes: list[dict] | None = None
+    init_containers: list[ContainerOverride] | None = None
+    containers: list[ContainerOverride] | None = None
+    scheduling_gates: list[dict] | None = None
+    image_pull_secrets: list[dict] | None = None
 
 
 @dataclass
@@ -130,5 +154,5 @@ class PodTemplateOverride:
     """
 
     target_jobs: list[str]
-    metadata: Optional[dict] = None
-    spec: Optional[PodTemplateSpecOverride] = None
+    metadata: dict | None = None
+    spec: PodTemplateSpecOverride | None = None

kubeflow/trainer/options/kubernetes.py

@@ -14,8 +14,10 @@
 
 """Kubernetes-specific training options for the Kubeflow Trainer SDK."""
 
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any
 
 from kubeflow.trainer.backends.base import KubernetesCompatible
 from kubeflow.trainer.options.common import PodTemplateOverride
@@ -33,7 +35,7 @@ class Labels(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply labels to the job specification."""
         metadata = job_spec.setdefault("metadata", {})
@@ -49,7 +51,7 @@ class Annotations(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply annotations to the job specification."""
         metadata = job_spec.setdefault("metadata", {})
@@ -69,7 +71,7 @@ class SpecLabels(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply spec-level labels to the job specification."""
         spec = job_spec.setdefault("spec", {})
@@ -89,29 +91,13 @@ class SpecAnnotations(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply spec-level annotations to the job specification."""
         spec = job_spec.setdefault("spec", {})
         spec["annotations"] = self.annotations
 
 
-@dataclass
-class Name(KubernetesCompatible):
-    """Set a custom name for the TrainJob resource (.metadata.name)."""
-
-    name: str
-
-    def __call__(
-        self,
-        job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
-    ) -> None:
-        """Apply custom name to the job specification."""
-        metadata = job_spec.setdefault("metadata", {})
-        metadata["name"] = self.name
-
-
 @dataclass
 class PodTemplateOverrides(KubernetesCompatible):
     """Add pod template overrides to the TrainJob (.spec.podTemplateOverrides)."""
@@ -121,7 +107,7 @@ class PodTemplateOverrides(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply pod template overrides to the job specification."""
         spec = job_spec.setdefault("spec", {})
@@ -187,7 +173,7 @@ class TrainerImage(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply trainer image override to the job specification.
 
@@ -209,7 +195,7 @@ class TrainerCommand(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply trainer command override to the job specification.
 
@@ -243,7 +229,7 @@ class TrainerArgs(KubernetesCompatible):
     def __call__(
         self,
         job_spec: dict[str, Any],
-        trainer: Optional[Union["CustomTrainer", "BuiltinTrainer"]] = None,
+        trainer: CustomTrainer | BuiltinTrainer | None = None,
     ) -> None:
         """Apply trainer args override to the job specification.
 

kubeflow/trainer/options/localprocess.py

@@ -14,8 +14,10 @@
 
 """LocalProcess-specific training options for the Kubeflow Trainer SDK."""
 
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 from kubeflow.trainer.backends.base import LocalProcessCompatible
 
@@ -32,7 +34,7 @@ class ProcessTimeout(LocalProcessCompatible):
     def __call__(
         self,
         job_spec: dict,
-        trainer: Optional[Union["BuiltinTrainer", "CustomTrainer"]] = None,
+        trainer: BuiltinTrainer | CustomTrainer | None = None,
     ) -> None:
         """Apply timeout to local process configuration.
 
@@ -52,7 +54,7 @@ class WorkingDirectory(LocalProcessCompatible):
     def __call__(
         self,
         job_spec: dict,
-        trainer: Optional[Union["BuiltinTrainer", "CustomTrainer"]] = None,
+        trainer: BuiltinTrainer | CustomTrainer | None = None,
     ) -> None:
         """Apply working directory to local process configuration.