viettelcloud-aiplatform 0.3.0__py3-none-any.whl

viettelcloud/aiplatform/trainer/api/trainer_client.py

@@ -0,0 +1,277 @@
+# 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.
+
+from collections.abc import Callable, Iterator
+import logging
+from typing import Optional, Union
+
+from viettelcloud.aiplatform.common.types import KubernetesBackendConfig
+from viettelcloud.aiplatform.trainer.backends.container.backend import ContainerBackend
+from viettelcloud.aiplatform.trainer.backends.container.types import ContainerBackendConfig
+from viettelcloud.aiplatform.trainer.backends.kubernetes.backend import KubernetesBackend
+from viettelcloud.aiplatform.trainer.backends.localprocess.backend import (
+    LocalProcessBackend,
+    LocalProcessBackendConfig,
+)
+from viettelcloud.aiplatform.trainer.constants import constants
+from viettelcloud.aiplatform.trainer.types import types
+
+logger = logging.getLogger(__name__)
+
+
+class TrainerClient:
+    def __init__(
+        self,
+        backend_config: Optional[
+            Union[
+                KubernetesBackendConfig,
+                LocalProcessBackendConfig,
+                ContainerBackendConfig,
+            ]
+        ] = None,
+    ):
+        """Initialize a Kubeflow Trainer client.
+
+        Args:
+            backend_config: Backend configuration. Either KubernetesBackendConfig,
+                            LocalProcessBackendConfig, ContainerBackendConfig,
+                            or None to use the backend's default config class.
+                            Defaults to KubernetesBackendConfig.
+
+        Raises:
+            ValueError: Invalid backend configuration.
+
+        """
+        # Set the default backend config.
+        if not backend_config:
+            backend_config = KubernetesBackendConfig()
+
+        if isinstance(backend_config, KubernetesBackendConfig):
+            self.backend = KubernetesBackend(backend_config)
+        elif isinstance(backend_config, LocalProcessBackendConfig):
+            self.backend = LocalProcessBackend(backend_config)
+        elif isinstance(backend_config, ContainerBackendConfig):
+            self.backend = ContainerBackend(backend_config)
+        else:
+            raise ValueError(f"Invalid backend config '{backend_config}'")
+
+    def list_runtimes(self) -> list[types.Runtime]:
+        """List of the available runtimes.
+
+        Returns:
+            A list of available training runtimes. If no runtimes exist, an empty list is returned.
+
+        Raises:
+            TimeoutError: Timeout to list runtimes.
+            RuntimeError: Failed to list runtimes.
+        """
+        return self.backend.list_runtimes()
+
+    def get_runtime(self, name: str) -> types.Runtime:
+        """Get the runtime object
+        Args:
+            name: Name of the runtime.
+
+        Returns:
+            A runtime object.
+        """
+        return self.backend.get_runtime(name=name)
+
+    def get_runtime_packages(self, runtime: types.Runtime):
+        """Print the installed Python packages for the given runtime. If a runtime has GPUs it also
+        prints available GPUs on the single training node.
+
+        Args:
+            runtime: Reference to one of existing runtimes.
+
+        Raises:
+            ValueError: Input arguments are invalid.
+            RuntimeError: Failed to get Runtime.
+
+        """
+        return self.backend.get_runtime_packages(runtime=runtime)
+
+    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:
+        """Create a TrainJob. You can configure the TrainJob using one of these trainers:
+
+        - CustomTrainer: Runs training with a user-defined function that fully encapsulates the
+            training process.
+        - CustomTrainerContainer: Runs training with a user-defined image that fully encapsulates
+            the training process.
+        - BuiltinTrainer: Uses a predefined trainer with built-in post-training logic, requiring
+            only parameter configuration.
+
+        Args:
+            runtime: Optional reference to one of the existing runtimes. It can accept the runtime
+                name or Runtime object from the `get_runtime()` API.
+                Defaults to the torch-distributed runtime if not provided.
+            initializer: Optional configuration for the dataset and model initializers.
+            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 viettelcloud.aiplatform.trainer.options.
+
+        Returns:
+            The unique name of the TrainJob that has been generated.
+
+        Raises:
+            ValueError: Input arguments are invalid.
+            TimeoutError: Timeout to create TrainJobs.
+            RuntimeError: Failed to create TrainJobs.
+        """
+        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
+        that runtime are returned.
+
+        Args:
+            runtime: Reference to one of the existing runtimes.
+
+        Returns:
+            List of created TrainJobs. If no TrainJob exist, an empty list is returned.
+
+        Raises:
+            TimeoutError: Timeout to list TrainJobs.
+            RuntimeError: Failed to list TrainJobs.
+        """
+        return self.backend.list_jobs(runtime=runtime)
+
+    def get_job(self, name: str) -> types.TrainJob:
+        """Get the TrainJob object
+
+        Args:
+            name: Name of the TrainJob.
+
+        Returns:
+            A TrainJob object.
+
+        Raises:
+            TimeoutError: Timeout to get a TrainJob.
+            RuntimeError: Failed to get a TrainJob.
+        """
+
+        return self.backend.get_job(name=name)
+
+    def get_job_logs(
+        self,
+        name: str,
+        step: str = constants.NODE + "-0",
+        follow: Optional[bool] = False,
+    ) -> Iterator[str]:
+        """Get logs from a specific step of a TrainJob.
+
+        You can watch for the logs in realtime as follows:
+        ```python
+        from viettelcloud.aiplatform.trainer import TrainerClient
+
+        for logline in TrainerClient().get_job_logs(name="s8d44aa4fb6d", follow=True):
+            print(logline)
+        ```
+
+        Args:
+            name: Name of the TrainJob.
+            step: Step of the TrainJob to collect logs from, like dataset-initializer or node-0.
+            follow: Whether to stream logs in realtime as they are produced.
+
+        Returns:
+            Iterator of log lines.
+
+
+        Raises:
+            TimeoutError: Timeout to get a TrainJob.
+            RuntimeError: Failed to get a TrainJob.
+        """
+        return self.backend.get_job_logs(name=name, follow=follow, step=step)
+
+    def get_job_events(self, name: str) -> list[types.Event]:
+        """Get events for a TrainJob.
+
+        This provides additional clarity about the state of the TrainJob
+        when logs alone are not sufficient. Events include information about
+        pod state changes, errors, and other significant occurrences.
+
+        Args:
+            name: Name of the TrainJob.
+
+        Returns:
+            A list of Event objects associated with the TrainJob.
+
+        Raises:
+            TimeoutError: Timeout to get a TrainJob events.
+            RuntimeError: Failed to get a TrainJob events.
+        """
+        return self.backend.get_job_events(name=name)
+
+    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:
+        """Wait for a TrainJob to reach a desired status.
+
+        Args:
+            name: Name of the TrainJob.
+            status: Expected statuses. Must be a subset of Created, Running, Complete, and
+                Failed statuses.
+            timeout: Maximum number of seconds to wait for the TrainJob to reach one of the
+                expected statuses.
+            polling_interval: The polling interval in seconds to check TrainJob status.
+            callbacks: Optional list of callback functions to be invoked after each polling
+                interval. Each callback should accept a single argument: the TrainJob object.
+
+        Returns:
+            A TrainJob object that reaches the desired status.
+
+        Raises:
+            ValueError: The input values are incorrect.
+            RuntimeError: Failed to get TrainJob or TrainJob reaches unexpected Failed status.
+            TimeoutError: Timeout to wait for TrainJob status.
+        """
+        return self.backend.wait_for_job_status(
+            name=name,
+            status=status,
+            timeout=timeout,
+            polling_interval=polling_interval,
+            callbacks=callbacks,
+        )
+
+    def delete_job(self, name: str):
+        """Delete the TrainJob.
+
+        Args:
+            name: Name of the TrainJob.
+
+        Raises:
+            TimeoutError: Timeout to delete TrainJob.
+            RuntimeError: Failed to delete TrainJob.
+        """
+        return self.backend.delete_job(name=name)

viettelcloud/aiplatform/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 viettelcloud.aiplatform.common.types import KubernetesBackendConfig
+from viettelcloud.aiplatform.trainer.api.trainer_client import TrainerClient
+from viettelcloud.aiplatform.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"]

viettelcloud/aiplatform/trainer/backends/base.py

@@ -0,0 +1,94 @@
+# 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.
+
+import abc
+from collections.abc import Callable, Iterator
+from typing import Optional, Union
+
+from viettelcloud.aiplatform.trainer.constants import constants
+from viettelcloud.aiplatform.trainer.types import types
+
+
+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()
+
+    @abc.abstractmethod
+    def get_runtime(self, name: str) -> types.Runtime:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_runtime_packages(self, runtime: types.Runtime):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    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:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def list_jobs(self, runtime: Optional[types.Runtime] = None) -> list[types.TrainJob]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_job(self, name: str) -> types.TrainJob:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_job_logs(
+        self,
+        name: str,
+        follow: bool = False,
+        step: str = constants.NODE + "-0",
+    ) -> Iterator[str]:
+        raise NotImplementedError()
+
+    def get_job_events(self, name: str) -> list[types.Event]:
+        """Get events for a TrainJob.
+
+        Args:
+            name: Name of the TrainJob.
+
+        Returns:
+            A list of Event objects associated with the TrainJob.
+        """
+        return []
+
+    @abc.abstractmethod
+    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:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def delete_job(self, name: str):
+        raise NotImplementedError()

viettelcloud/aiplatform/trainer/backends/container/adapters/base.py

@@ -0,0 +1,195 @@
+# 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.
+
+"""
+Container client adapters for Docker and Podman.
+
+This module implements the adapter pattern to abstract away differences between
+Docker and Podman APIs, allowing the backend to work with either runtime through
+a common interface.
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Iterator
+from typing import Optional
+
+
+class BaseContainerClientAdapter(abc.ABC):
+    """
+    Abstract adapter interface for container clients.
+
+    This adapter abstracts the container runtime API, allowing the backend
+    to work with Docker and Podman through a unified interface.
+    """
+
+    @abc.abstractmethod
+    def ping(self):
+        """Test the connection to the container runtime."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def create_network(
+        self,
+        name: str,
+        labels: dict[str, str],
+    ) -> str:
+        """
+        Create a container network.
+
+        Args:
+            name: Network name
+            labels: Labels to attach to the network
+
+        Returns:
+            Network ID or name
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def delete_network(self, network_id: str):
+        """Delete a network."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def create_and_start_container(
+        self,
+        image: str,
+        command: list[str],
+        name: str,
+        network_id: str,
+        environment: dict[str, str],
+        labels: dict[str, str],
+        volumes: dict[str, dict[str, str]],
+        working_dir: str,
+    ) -> str:
+        """
+        Create and start a container.
+
+        Args:
+            image: Container image
+            command: Command to run
+            name: Container name
+            network_id: Network to attach to
+            environment: Environment variables
+            labels: Container labels
+            volumes: Volume mounts
+            working_dir: Working directory
+
+        Returns:
+            Container ID
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_container(self, container_id: str):
+        """Get container object by ID."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def container_logs(self, container_id: str, follow: bool) -> Iterator[str]:
+        """Stream logs from a container."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def stop_container(self, container_id: str, timeout: int = 10):
+        """Stop a container."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def remove_container(self, container_id: str, force: bool = True):
+        """Remove a container."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def pull_image(self, image: str):
+        """Pull an image."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def image_exists(self, image: str) -> bool:
+        """Check if an image exists locally."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def run_oneoff_container(self, image: str, command: list[str]) -> str:
+        """
+        Run a short-lived container and return its output.
+
+        Args:
+            image: Container image
+            command: Command to run
+
+        Returns:
+            Container output as string
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def container_status(self, container_id: str) -> tuple[str, Optional[int]]:
+        """
+        Get container status.
+
+        Returns:
+            Tuple of (status_string, exit_code)
+            Status strings: "running", "created", "exited", etc.
+            Exit code is None if container hasn't exited
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_container_ip(self, container_id: str, network_id: str) -> Optional[str]:
+        """
+        Get container's IP address on a specific network.
+
+        Args:
+            container_id: Container ID
+            network_id: Network name or ID
+
+        Returns:
+            IP address string or None if not found
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def list_containers(self, filters: Optional[dict[str, list[str]]] = None) -> list[dict]:
+        """
+        List containers, optionally filtered by labels.
+
+        Args:
+            filters: Dictionary of filters (e.g., {"label": ["key=value"]})
+
+        Returns:
+            List of container info dictionaries with keys:
+            - id: Container ID
+            - name: Container name
+            - labels: Dictionary of labels
+            - status: Container status
+            - created: Creation timestamp
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_network(self, network_id: str) -> Optional[dict]:
+        """
+        Get network information by ID or name.
+
+        Args:
+            network_id: Network ID or name
+
+        Returns:
+            Dictionary with network info including labels, or None if not found
+        """
+        raise NotImplementedError()