viettelcloud-aiplatform 0.3.0__py3-none-any.whl

viettelcloud/aiplatform/trainer/constants/constants.py

@@ -0,0 +1,179 @@
+# Copyright 2024 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.
+
+import os
+import textwrap
+
+# Common constants.
+GROUP = "trainer.kubeflow.org"
+VERSION = "v1alpha1"
+API_VERSION = f"{GROUP}/{VERSION}"
+
+# The Kind name for the ClusterTrainingRuntime.
+CLUSTER_TRAINING_RUNTIME_KIND = "ClusterTrainingRuntime"
+
+# The plural for the ClusterTrainingRuntime.
+CLUSTER_TRAINING_RUNTIME_PLURAL = "clustertrainingruntimes"
+
+# The Kind name for the TrainJob.
+TRAINJOB_KIND = "TrainJob"
+
+# The plural for the TrainJob.
+TRAINJOB_PLURAL = "trainjobs"
+
+# The default status for the TrainJob once users create it.
+TRAINJOB_CREATED = "Created"
+
+# The running status of the TrainJob, defined when all training node (e.g. Pods) are
+# running or succeeded.
+TRAINJOB_RUNNING = "Running"
+
+# The complete status of the TrainJob, defined when TrainJob CR has complete condition.
+TRAINJOB_COMPLETE = "Complete"
+
+# The failed status of the TrainJob, defined when TrainJob CR has failed condition.
+TRAINJOB_FAILED = "Failed"
+
+# The succeeded phase of the Pods.
+POD_SUCCEEDED = "Succeeded"
+
+# The label key to identify the relationship between TrainJob and Pod template in the runtime.
+# For example, what PodTemplate must be overridden by TrainJob's .spec.trainer APIs.
+TRAINJOB_ANCESTOR_LABEL = "trainer.kubeflow.org/trainjob-ancestor-step"
+
+# The label key to identify ML framework that runtime uses (e.g. torch, deepspeed, torchtune, etc.)
+RUNTIME_FRAMEWORK_LABEL = "trainer.kubeflow.org/framework"
+
+# The name of the ReplicatedJob and container of the dataset initializer.
+# Also, it represents the `trainjob-ancestor-step` label value for the dataset initializer step.
+DATASET_INITIALIZER = "dataset-initializer"
+
+# The name of the ReplicatedJob and container of the model initializer.
+# Also, it represents the `trainjob-ancestor-step` label value for the model initializer step.
+MODEL_INITIALIZER = "model-initializer"
+
+# The env name for the access token of dataset/model initializer.
+INITIALIZER_ENV_ACCESS_TOKEN = "ACCESS_TOKEN"
+
+# The default value for initializer to ignore files.
+INITIALIZER_DEFAULT_IGNORE_PATTERNS = ["*.msgpack", "*.h5", "*.bin", ".pt", ".pth"]
+
+# The default path to the users' workspace.
+# TODO (andreyvelich): Discuss how to keep this path is sync with pkg.initializers.constants
+WORKSPACE_PATH = "/workspace"
+
+# The path where initializer downloads dataset.
+DATASET_PATH = os.path.join(WORKSPACE_PATH, "dataset")
+
+# The path where initializer downloads model.
+MODEL_PATH = os.path.join(WORKSPACE_PATH, "model")
+
+# The name of the ReplicatedJob to launch mpirun.
+LAUNCHER = "launcher"
+
+# The name of the ReplicatedJob and container of the node. The node usually represents
+# single VM where distributed training code is executed.
+NODE = "node"
+
+# The label for cpu in the container resources.
+CPU_LABEL = "cpu"
+
+# The label for NVIDIA GPU in the container resources.
+GPU_LABEL = "nvidia.com/gpu"
+
+# The prefix for NVIDIA MIG device in the container resources.
+GPU_MIG_PREFIX = "nvidia.com/mig-"
+
+# The label for TPU in the container resources.
+TPU_LABEL = "google.com/tpu"
+
+# The label key to identify the JobSet name of the Pod.
+JOBSET_NAME_LABEL = "jobset.sigs.k8s.io/jobset-name"
+
+# The label key to identify the JobSet's ReplicatedJob of the Pod.
+JOBSET_RJOB_NAME_LABEL = "jobset.sigs.k8s.io/replicatedjob-name"
+
+# The label key to identify the Job completion index of the Pod.
+JOB_INDEX_LABEL = "batch.kubernetes.io/job-completion-index"
+
+# The Pod pending phase indicates that Pod has been accepted by the Kubernetes cluster,
+# but one or more of the containers has not been made ready to run.
+POD_PENDING = "Pending"
+
+# The label selector for Pods created by the TrainJob.
+# It checks the following rJob.name: dataset-initializer, model-initializer, launcher, node.
+POD_LABEL_SELECTOR = (
+    f"{JOBSET_NAME_LABEL}={{trainjob_name}},{JOBSET_RJOB_NAME_LABEL} "
+    f"in ({DATASET_INITIALIZER}, {MODEL_INITIALIZER}, {LAUNCHER}, {NODE})"
+)
+
+# Handle environment variable for multiple URLs (comma-separated).
+# The first URL will be the index-url, and remaining ones are extra-index-urls.
+DEFAULT_PIP_INDEX_URLS = os.getenv("DEFAULT_PIP_INDEX_URLS", "https://pypi.org/simple").split(",")
+
+# The exec script to embed training function into container command.
+# __ENTRYPOINT__ depends on the MLPolicy, func_code and func_file is substituted in the `train` API.
+EXEC_FUNC_SCRIPT = textwrap.dedent(
+    """
+        read -r -d '' SCRIPT << EOM\n
+        {func_code}
+        EOM
+        printf "%s" \"$SCRIPT\" > \"{func_file}\"
+        __ENTRYPOINT__ \"{func_file}\""""
+)
+
+# The default command for the PlainML CustomTrainer.
+DEFAULT_COMMAND = (
+    "bash",
+    "-c",
+    EXEC_FUNC_SCRIPT.replace("__ENTRYPOINT__", "python"),
+)
+
+# The default home directory for the MPI user.
+DEFAULT_MPI_USER_HOME = os.getenv("DEFAULT_MPI_USER_HOME", "/home/mpiuser")
+
+# The default command for the OpenMPI CustomTrainer.
+MPI_COMMAND = (
+    "mpirun",
+    "--hostfile",
+    "/etc/mpi/hostfile",
+    *DEFAULT_COMMAND,
+)
+
+# The default name for the ClusterTrainingRuntime.
+DEFAULT_TRAINING_RUNTIME = os.getenv("DEFAULT_TRAINING_RUNTIME", "torch-distributed")
+
+# The default container command for the Torch CustomTrainer
+TORCH_COMMAND = (
+    "bash",
+    "-c",
+    EXEC_FUNC_SCRIPT.replace("__ENTRYPOINT__", "torchrun"),
+)
+# The Torch env name for the number of procs per node (e.g. number of GPUs per Pod).
+TORCH_ENV_NUM_PROC_PER_NODE = "PET_NPROC_PER_NODE"
+
+# The default command for the TorchTune BuiltinTrainer.
+TORCH_TUNE_COMMAND = ("tune", "run")
+
+# The Instruct Datasets class in torchtune
+TORCH_TUNE_INSTRUCT_DATASET = "torchtune.datasets.instruct_dataset"
+
+# Default container images for each framework (used as fallback when runtime not provided)
+DEFAULT_FRAMEWORK_IMAGES = {
+    "torch": "pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime",
+}
+
+# The length of the UUID suffix for auto-generated job names.
+# Total name length = 1 (random letter) + 11 (UUID hex) = 12 characters
+JOB_NAME_UUID_LENGTH = 11

viettelcloud/aiplatform/trainer/options/__init__.py

@@ -0,0 +1,52 @@
+# 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.
+
+"""Training options for the Viettel Cloud AI Platform SDK.
+
+All options are available from this single import location:
+    from viettelcloud.aiplatform.trainer.options import Name, Labels, PodTemplateOverrides, ...
+
+Options self-validate their backend compatibility at runtime.
+Check each option's docstring for supported backends.
+"""
+
+from viettelcloud.aiplatform.trainer.options.common import Name
+from viettelcloud.aiplatform.trainer.options.kubernetes import (
+    Annotations,
+    ContainerOverride,
+    Labels,
+    PodSpecOverride,
+    PodTemplateOverride,
+    PodTemplateOverrides,
+    SpecAnnotations,
+    SpecLabels,
+    TrainerArgs,
+    TrainerCommand,
+)
+
+__all__ = [
+    # Common options (all backends)
+    "Name",
+    # Kubernetes options
+    "Annotations",
+    "ContainerOverride",
+    "Labels",
+    "PodSpecOverride",
+    "PodTemplateOverride",
+    "PodTemplateOverrides",
+    "SpecAnnotations",
+    "SpecLabels",
+    "TrainerArgs",
+    "TrainerCommand",
+]

viettelcloud/aiplatform/trainer/options/common.py

@@ -0,0 +1,55 @@
+# 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.
+
+"""Common options and helper classes used across multiple backends."""
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+from viettelcloud.aiplatform.trainer.backends.base import RuntimeBackend
+from viettelcloud.aiplatform.trainer.types.types import (
+    BuiltinTrainer,
+    CustomTrainer,
+    CustomTrainerContainer,
+)
+
+
+@dataclass
+class Name:
+    """Set a custom name for the TrainJob resource.
+
+    This option works with all backends.
+
+    Args:
+        name: Custom name for the job. Must be a valid identifier.
+    """
+
+    name: str
+
+    def __call__(
+        self,
+        job_spec: dict[str, Any],
+        trainer: Optional[Union[BuiltinTrainer, CustomTrainer, CustomTrainerContainer]],
+        backend: RuntimeBackend,
+    ) -> None:
+        """Apply custom name to the job specification.
+
+        Args:
+            job_spec: Job specification dictionary to modify.
+            trainer: Optional trainer instance for context.
+            backend: Backend instance for validation and context.
+        """
+        # Name option is generic - works with all backends
+        metadata = job_spec.setdefault("metadata", {})
+        metadata["name"] = self.name