viettelcloud-aiplatform 0.3.0__py3-none-any.whl

viettelcloud/aiplatform/trainer/backends/kubernetes/backend.py

@@ -0,0 +1,710 @@
+# 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.
+
+from collections.abc import Callable, Iterator
+import copy
+import logging
+import multiprocessing
+import random
+import re
+import string
+import time
+from typing import Any, Optional, Union
+import uuid
+
+from kubeflow_trainer_api import models
+from kubernetes import client, config, watch
+
+import viettelcloud.aiplatform.common.constants as common_constants
+from viettelcloud.aiplatform.common.types import KubernetesBackendConfig
+import viettelcloud.aiplatform.common.utils as common_utils
+from viettelcloud.aiplatform.trainer.backends.base import RuntimeBackend
+import viettelcloud.aiplatform.trainer.backends.kubernetes.utils as utils
+from viettelcloud.aiplatform.trainer.constants import constants
+from viettelcloud.aiplatform.trainer.types import types
+
+logger = logging.getLogger(__name__)
+
+
+class KubernetesBackend(RuntimeBackend):
+    def __init__(self, cfg: KubernetesBackendConfig):
+        if cfg.namespace is None:
+            cfg.namespace = common_utils.get_default_target_namespace(cfg.context)
+
+        # If client configuration is not set, use kube-config to access Kubernetes APIs.
+        if cfg.client_configuration is None:
+            # Load kube-config or in-cluster config.
+            if cfg.config_file or not common_utils.is_running_in_k8s():
+                config.load_kube_config(config_file=cfg.config_file, context=cfg.context)
+            else:
+                config.load_incluster_config()
+
+        k8s_client = client.ApiClient(cfg.client_configuration)
+        self.custom_api = client.CustomObjectsApi(k8s_client)
+        self.core_api = client.CoreV1Api(k8s_client)
+
+        self.namespace = cfg.namespace
+
+    def list_runtimes(self) -> list[types.Runtime]:
+        result = []
+        try:
+            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(common_constants.DEFAULT_TIMEOUT)
+            )
+
+            if not runtime_list:
+                return result
+
+            for runtime in runtime_list.items:
+                if not (
+                    runtime.metadata
+                    and runtime.metadata.labels
+                    and constants.RUNTIME_FRAMEWORK_LABEL in runtime.metadata.labels
+                ):
+                    logger.warning(
+                        f"Runtime {runtime.metadata.name} must have "  # type: ignore
+                        f"{constants.RUNTIME_FRAMEWORK_LABEL} label."
+                    )
+                    continue
+                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") from e
+        except Exception as e:
+            raise RuntimeError(f"Failed to list {constants.CLUSTER_TRAINING_RUNTIME_KIND}s") from e
+
+        return result
+
+    def get_runtime(self, name: str) -> types.Runtime:
+        """Get the Runtime object"""
+
+        try:
+            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(common_constants.DEFAULT_TIMEOUT)  # type: ignore
+            )
+
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout to get {constants.CLUSTER_TRAINING_RUNTIME_PLURAL}: {name}"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get {constants.CLUSTER_TRAINING_RUNTIME_PLURAL}: {name}"
+            ) from e
+
+        return self.__get_runtime_from_cr(runtime)  # type: ignore
+
+    def get_runtime_packages(self, runtime: types.Runtime):
+        if runtime.trainer.trainer_type == types.TrainerType.BUILTIN_TRAINER:
+            raise ValueError("Cannot get Runtime packages for BuiltinTrainer")
+
+        # Create a deepcopy of the runtime to avoid modifying the original command.
+        runtime_copy = copy.deepcopy(runtime)
+
+        # Run mpirun only within the single process.
+        if runtime_copy.trainer.command[0] == "mpirun":
+            mpi_command = list(constants.MPI_COMMAND)
+            mpi_command[1:3] = ["-np", "1"]
+            runtime_copy.trainer.set_command(tuple(mpi_command))
+
+        def print_packages():
+            import shutil
+            import subprocess
+            import sys
+
+            # Print Python version.
+            print(f"Python: {sys.version}")
+
+            # Print Python packages.
+            if shutil.which("pip"):
+                pip_list = subprocess.run(["pip", "list"], capture_output=True, text=True)
+                print(pip_list.stdout)
+            else:
+                print("Unable to get installed packages: pip command not found")
+
+            # Print nvidia-smi if GPUs are available.
+            if shutil.which("nvidia-smi"):
+                print("Available GPUs on the single training node")
+                nvidia_smi = subprocess.run(["nvidia-smi"], capture_output=True, text=True)
+                print(nvidia_smi.stdout)
+
+        # Create the TrainJob and wait until it completes.
+        # If Runtime trainer has GPU resources use them, otherwise run TrainJob with 1 CPU.
+        job_name = self.train(
+            runtime=runtime_copy,
+            trainer=types.CustomTrainer(
+                func=print_packages,
+                num_nodes=1,
+                resources_per_node=({"cpu": 1} if runtime_copy.trainer.device != "gpu" else None),
+            ),
+        )
+
+        self.wait_for_job_status(job_name)
+        print("\n".join(self.get_job_logs(name=job_name)))
+        self.delete_job(job_name)
+
+    def train(
+        self,
+        runtime: Optional[Union[str, types.Runtime]] = None,
+        initializer: Optional[types.Initializer] = None,
+        trainer: Optional[
+            Union[types.CustomTrainer, types.CustomTrainerContainer, types.BuiltinTrainer]
+        ] = None,
+        options: Optional[list] = None,
+    ) -> str:
+        # 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, labels=labels, annotations=annotations
+            ),
+            spec=trainjob_spec,
+        )
+
+        # Create the TrainJob.
+        try:
+            self.custom_api.create_namespaced_custom_object(
+                constants.GROUP,
+                constants.VERSION,
+                self.namespace,
+                constants.TRAINJOB_PLURAL,
+                train_job.to_dict(),
+            )
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout to create {constants.TRAINJOB_KIND}: {self.namespace}/{train_job_name}"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to create {constants.TRAINJOB_KIND}: {self.namespace}/{train_job_name}"
+            ) from e
+
+        logger.debug(
+            f"{constants.TRAINJOB_KIND} {self.namespace}/{train_job_name} has been created"
+        )
+
+        return train_job_name
+
+    def list_jobs(self, runtime: Optional[types.Runtime] = None) -> list[types.TrainJob]:
+        result = []
+        try:
+            thread = self.custom_api.list_namespaced_custom_object(
+                constants.GROUP,
+                constants.VERSION,
+                self.namespace,
+                constants.TRAINJOB_PLURAL,
+                async_req=True,
+            )
+
+            trainjob_list = models.TrainerV1alpha1TrainJobList.from_dict(
+                thread.get(common_constants.DEFAULT_TIMEOUT)
+            )
+
+            if not trainjob_list:
+                return result
+
+            for trainjob in trainjob_list.items:
+                # If runtime object is set, we check the TrainJob's runtime reference.
+                if (
+                    runtime is not None
+                    and trainjob.spec
+                    and trainjob.spec.runtime_ref
+                    and trainjob.spec.runtime_ref.name != runtime.name
+                ):
+                    continue
+
+                result.append(self.__get_trainjob_from_cr(trainjob))
+
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout to list {constants.TRAINJOB_KIND}s in namespace: {self.namespace}"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to list {constants.TRAINJOB_KIND}s in namespace: {self.namespace}"
+            ) from e
+
+        return result
+
+    def get_job(self, name: str) -> types.TrainJob:
+        """Get the TrainJob object"""
+
+        try:
+            thread = self.custom_api.get_namespaced_custom_object(
+                constants.GROUP,
+                constants.VERSION,
+                self.namespace,
+                constants.TRAINJOB_PLURAL,
+                name,
+                async_req=True,
+            )
+
+            trainjob = models.TrainerV1alpha1TrainJob.from_dict(
+                thread.get(common_constants.DEFAULT_TIMEOUT)  # type: ignore
+            )
+
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout to get {constants.TRAINJOB_KIND}: {self.namespace}/{name}"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to get {constants.TRAINJOB_KIND}: {self.namespace}/{name}"
+            ) from e
+
+        return self.__get_trainjob_from_cr(trainjob)  # type: ignore
+
+    def get_job_logs(
+        self,
+        name: str,
+        follow: bool = False,
+        step: str = constants.NODE + "-0",
+    ) -> Iterator[str]:
+        """Get the TrainJob logs"""
+        # Get the TrainJob Pod name.
+        pod_name = None
+        for c in self.get_job(name).steps:
+            if c.status != constants.POD_PENDING and c.name == step:
+                pod_name = c.pod_name
+                break
+        if pod_name is None:
+            return
+
+        # Remove the number for the node step.
+        container_name = re.sub(r"-\d+$", "", step)
+        yield from self._read_pod_logs(
+            pod_name=pod_name, container_name=container_name, follow=follow
+        )
+
+    def wait_for_job_status(
+        self,
+        name: str,
+        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        timeout: int = 600,
+        polling_interval: int = 2,
+        callbacks: Optional[list[Callable[[types.TrainJob], None]]] = None,
+    ) -> types.TrainJob:
+        job_statuses = {
+            constants.TRAINJOB_CREATED,
+            constants.TRAINJOB_RUNNING,
+            constants.TRAINJOB_COMPLETE,
+            constants.TRAINJOB_FAILED,
+        }
+        if not status.issubset(job_statuses):
+            raise ValueError(f"Expected status {status} must be a subset of {job_statuses}")
+
+        if polling_interval > timeout:
+            raise ValueError(
+                f"Polling interval {polling_interval} must be less than timeout: {timeout}"
+            )
+
+        for _ in range(round(timeout / polling_interval)):
+            # Check the status after event is generated for the TrainJob's Pods.
+            trainjob = self.get_job(name)
+            logger.debug(f"TrainJob {name}, status {trainjob.status}")
+
+            # Invoke callbacks if provided
+            if callbacks:
+                for callback in callbacks:
+                    callback(trainjob)
+
+            # Raise an error if TrainJob is Failed and it is not the expected status.
+            if (
+                constants.TRAINJOB_FAILED not in status
+                and trainjob.status == constants.TRAINJOB_FAILED
+            ):
+                raise RuntimeError(f"TrainJob {name} is Failed")
+
+            # Return the TrainJob if it reaches the expected status.
+            if trainjob.status in status:
+                return trainjob
+
+            time.sleep(polling_interval)
+
+        raise TimeoutError(f"Timeout waiting for TrainJob {name} to reach status: {status} status")
+
+    def delete_job(self, name: str):
+        try:
+            self.custom_api.delete_namespaced_custom_object(
+                constants.GROUP,
+                constants.VERSION,
+                self.namespace,
+                constants.TRAINJOB_PLURAL,
+                name=name,
+            )
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout to delete {constants.TRAINJOB_KIND}: {self.namespace}/{name}"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to delete {constants.TRAINJOB_KIND}: {self.namespace}/{name}"
+            ) from e
+
+        logger.debug(f"{constants.TRAINJOB_KIND} {self.namespace}/{name} has been deleted")
+
+    def get_job_events(self, name: str) -> list[types.Event]:
+        # Get all pod names related to this TrainJob
+        trainjob = self.get_job(name)
+
+        # Create set of all TrainJob-related resource names
+        trainjob_resources = {name}
+        for step in trainjob.steps:
+            trainjob_resources.add(step.name)
+            if step.pod_name:
+                trainjob_resources.add(step.pod_name)
+
+        events = []
+        try:
+            # Retrieve events from the namespace
+            event_response: models.IoK8sApiCoreV1EventList = self.core_api.list_namespaced_event(
+                namespace=self.namespace,
+                async_req=True,
+            ).get(common_constants.DEFAULT_TIMEOUT)
+
+            # Filter events related to this TrainJob or its pods
+            for event in event_response.items:
+                if not (event.metadata and event.involved_object and event.first_timestamp):
+                    continue
+
+                involved_object = event.involved_object
+
+                # Check if event is related to TrainJob resources
+                if (
+                    involved_object.kind in {constants.TRAINJOB_KIND, "JobSet", "Job", "Pod"}
+                    and involved_object.name in trainjob_resources
+                ):
+                    events.append(
+                        types.Event(
+                            involved_object_kind=involved_object.kind,
+                            involved_object_name=involved_object.name,
+                            message=event.message or "",
+                            reason=event.reason or "",
+                            event_time=event.first_timestamp,
+                        )
+                    )
+
+            # Sort events by first occurrence time
+            events.sort(key=lambda e: e.event_time)
+
+            return events
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout getting {constants.TRAINJOB_KIND} events: {self.namespace}/{name}"
+            ) from e
+
+    def __get_runtime_from_cr(
+        self,
+        runtime_cr: models.TrainerV1alpha1ClusterTrainingRuntime,
+    ) -> types.Runtime:
+        if not (
+            runtime_cr.metadata
+            and runtime_cr.metadata.name
+            and runtime_cr.spec
+            and runtime_cr.spec.ml_policy
+            and runtime_cr.spec.template.spec
+            and runtime_cr.spec.template.spec.replicated_jobs
+        ):
+            raise Exception(f"ClusterTrainingRuntime CR is invalid: {runtime_cr}")
+
+        if not (
+            runtime_cr.metadata.labels
+            and constants.RUNTIME_FRAMEWORK_LABEL in runtime_cr.metadata.labels
+        ):
+            raise Exception(
+                f"Runtime {runtime_cr.metadata.name} must have "
+                f"{constants.RUNTIME_FRAMEWORK_LABEL} label"
+            )
+
+        return types.Runtime(
+            name=runtime_cr.metadata.name,
+            trainer=utils.get_runtime_trainer(
+                runtime_cr.metadata.labels[constants.RUNTIME_FRAMEWORK_LABEL],
+                runtime_cr.spec.template.spec.replicated_jobs,
+                runtime_cr.spec.ml_policy,
+            ),
+        )
+
+    def _read_pod_logs(self, pod_name: str, container_name: str, follow: bool) -> Iterator[str]:
+        """Read logs from a pod container."""
+        try:
+            if follow:
+                log_stream = watch.Watch().stream(
+                    self.core_api.read_namespaced_pod_log,
+                    name=pod_name,
+                    namespace=self.namespace,
+                    container=container_name,
+                    follow=True,
+                )
+
+                # Stream logs incrementally.
+                yield from log_stream  # type: ignore
+            else:
+                logs = self.core_api.read_namespaced_pod_log(
+                    name=pod_name,
+                    namespace=self.namespace,
+                    container=container_name,
+                )
+
+                yield from logs.splitlines()
+
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to read logs for the pod {self.namespace}/{pod_name}"
+            ) from e
+
+    def __get_trainjob_from_cr(
+        self,
+        trainjob_cr: models.TrainerV1alpha1TrainJob,
+    ) -> types.TrainJob:
+        if not (
+            trainjob_cr.metadata
+            and trainjob_cr.metadata.name
+            and trainjob_cr.metadata.namespace
+            and trainjob_cr.spec
+            and trainjob_cr.metadata.creation_timestamp
+        ):
+            raise Exception(f"TrainJob CR is invalid: {trainjob_cr}")
+
+        name = trainjob_cr.metadata.name
+        namespace = trainjob_cr.metadata.namespace
+
+        runtime = self.get_runtime(trainjob_cr.spec.runtime_ref.name)
+
+        # Construct the TrainJob from the CR.
+        trainjob = types.TrainJob(
+            name=name,
+            creation_timestamp=trainjob_cr.metadata.creation_timestamp,
+            runtime=runtime,
+            steps=[],
+            # Number of nodes is taken from TrainJob or TrainingRuntime
+            num_nodes=(
+                trainjob_cr.spec.trainer.num_nodes
+                if trainjob_cr.spec.trainer and trainjob_cr.spec.trainer.num_nodes
+                else runtime.trainer.num_nodes
+            ),
+            status=constants.TRAINJOB_CREATED,  # The default TrainJob status.
+        )
+
+        # Add the TrainJob components, e.g. trainer nodes and initializer.
+        try:
+            response = self.core_api.list_namespaced_pod(
+                namespace,
+                label_selector=constants.POD_LABEL_SELECTOR.format(trainjob_name=name),
+                async_req=True,
+            ).get(common_constants.DEFAULT_TIMEOUT)
+
+            # Convert Pod to the correct format.
+            # This is required to convert Pod's container resources into API object from str
+            pod_list = models.IoK8sApiCoreV1PodList.from_dict(response.to_dict())
+            if not pod_list:
+                return trainjob
+
+            for pod in pod_list.items:
+                # Pod must have labels to detect the TrainJob step.
+                # Every Pod always has a single TrainJob step.
+                if not (pod.metadata and pod.metadata.name and pod.metadata.labels and pod.spec):
+                    raise Exception(f"TrainJob Pod is invalid: {pod}")
+
+                # Get the Initializer step.
+                if pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL] in {
+                    constants.DATASET_INITIALIZER,
+                    constants.MODEL_INITIALIZER,
+                }:
+                    trainjob.steps.append(
+                        utils.get_trainjob_initializer_step(
+                            pod.metadata.name,
+                            pod.spec,
+                            pod.status,
+                        )
+                    )
+                # Get the Node step.
+                elif pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL] in {
+                    constants.LAUNCHER,
+                    constants.NODE,
+                }:
+                    trainjob.steps.append(
+                        utils.get_trainjob_node_step(
+                            pod.metadata.name,
+                            pod.spec,
+                            pod.status,
+                            trainjob.runtime,
+                            pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL],
+                            int(pod.metadata.labels[constants.JOB_INDEX_LABEL]),
+                        )
+                    )
+        except multiprocessing.TimeoutError as e:
+            raise TimeoutError(
+                f"Timeout to list {constants.TRAINJOB_KIND}'s steps: {namespace}/{name}"
+            ) from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to list {constants.TRAINJOB_KIND}'s steps: {namespace}/{name}"
+            ) from e
+
+        # Update the TrainJob status from its conditions.
+        if trainjob_cr.status and trainjob_cr.status.conditions:
+            for c in trainjob_cr.status.conditions:
+                if (
+                    c.type == constants.TRAINJOB_COMPLETE
+                    and c.status == "True"
+                    or c.type == constants.TRAINJOB_FAILED
+                    and c.status == "True"
+                ):
+                    trainjob.status = c.type
+        else:
+            # The TrainJob running status is defined when all training node (e.g. Pods) are
+            # running or succeeded.
+            num_running_nodes = sum(
+                1
+                for step in trainjob.steps
+                if step.name.startswith(constants.NODE)
+                and (
+                    step.status == constants.TRAINJOB_RUNNING
+                    or step.status == constants.POD_SUCCEEDED
+                )
+            )
+
+            if trainjob.num_nodes == num_running_nodes:
+                trainjob.status = constants.TRAINJOB_RUNNING
+
+        return trainjob
+
+    def _get_trainjob_spec(
+        self,
+        runtime: Optional[Union[str, types.Runtime]] = None,
+        initializer: Optional[types.Initializer] = None,
+        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:
+            runtime = self.get_runtime(constants.DEFAULT_TRAINING_RUNTIME)
+        elif isinstance(runtime, str):
+            runtime = self.get_runtime(runtime)
+
+        # Build the Trainer.
+        trainer_cr = models.TrainerV1alpha1Trainer()
+
+        if trainer:
+            # If users choose to use a custom training script.
+            if isinstance(trainer, (types.CustomTrainer, types.CustomTrainerContainer)):
+                if runtime.trainer.trainer_type != types.TrainerType.CUSTOM_TRAINER:
+                    raise ValueError(f"CustomTrainer can't be used with {runtime} runtime")
+                trainer_cr = utils.get_trainer_cr_from_custom_trainer(runtime, trainer)
+
+            # If users choose to use a builtin trainer for post-training.
+            elif isinstance(trainer, types.BuiltinTrainer):
+                if runtime.trainer.trainer_type != types.TrainerType.BUILTIN_TRAINER:
+                    raise ValueError(f"BuiltinTrainer can't be used with {runtime} runtime")
+                trainer_cr = utils.get_trainer_cr_from_builtin_trainer(
+                    runtime, trainer, initializer
+                )
+
+            else:
+                raise ValueError(
+                    f"The trainer type {type(trainer)} is not supported. "
+                    "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"]
+
+        trainjob_spec = models.TrainerV1alpha1TrainJobSpec(
+            runtimeRef=models.TrainerV1alpha1RuntimeRef(name=runtime.name),
+            trainer=trainer_cr if trainer_cr != models.TrainerV1alpha1Trainer() else None,
+            labels=spec_labels,
+            annotations=spec_annotations,
+            pod_template_overrides=pod_template_overrides,
+        )
+
+        # Add initializer if users define it.
+        if initializer and (initializer.dataset or initializer.model):
+            trainjob_spec.initializer = models.TrainerV1alpha1Initializer(
+                dataset=utils.get_dataset_initializer(initializer.dataset)
+                if initializer.dataset
+                else None,
+                model=utils.get_model_initializer(initializer.model) if initializer.model else None,
+            )
+
+        return trainjob_spec