viettelcloud-aiplatform 0.3.0__py3-none-any.whl

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

@@ -0,0 +1,668 @@
+# 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.
+
+"""
+ContainerBackend
+----------------
+
+Unified local execution backend for `CustomTrainer` jobs using containers.
+
+This backend automatically detects and uses either Docker or Podman.
+It provides a single interface regardless of the underlying container runtime.
+
+Key behaviors:
+- Auto-detection: Tries Docker first, then Podman. Can be overridden via config.
+- Multi-node jobs: one container per node connected via a per-job network.
+- Entry script generation: we serialize the user's training function and embed it
+  inline in the container command using a heredoc (no file I/O on the host). The
+  script is created inside the container at /tmp/train.py and invoked using
+  `torchrun` (preferred) or `python` as a fallback.
+- Runtimes: we use `config/training_runtimes` to define runtime images and
+  characteristics (e.g., torch). Defaults to `torch-distributed` if no runtime
+  is provided.
+- Image pulling: controlled via `pull_policy` and performed automatically if
+  needed.
+- Logs and lifecycle: streaming logs and deletion semantics similar to the
+  Docker/Podman backends, but with automatic runtime detection.
+"""
+
+from collections.abc import Callable, Iterator
+from datetime import datetime
+import logging
+import os
+import random
+import shutil
+import string
+from typing import Optional, Union
+import uuid
+
+from viettelcloud.aiplatform.trainer.backends.base import RuntimeBackend
+from viettelcloud.aiplatform.trainer.backends.container import utils as container_utils
+from viettelcloud.aiplatform.trainer.backends.container.adapters.base import (
+    BaseContainerClientAdapter,
+)
+from viettelcloud.aiplatform.trainer.backends.container.adapters.docker import DockerClientAdapter
+from viettelcloud.aiplatform.trainer.backends.container.adapters.podman import PodmanClientAdapter
+from viettelcloud.aiplatform.trainer.backends.container.runtime_loader import (
+    get_training_runtime_from_sources,
+    list_training_runtimes_from_sources,
+)
+from viettelcloud.aiplatform.trainer.backends.container.types import ContainerBackendConfig
+from viettelcloud.aiplatform.trainer.constants import constants
+from viettelcloud.aiplatform.trainer.types import types
+
+logger = logging.getLogger(__name__)
+
+
+class ContainerBackend(RuntimeBackend):
+    """
+    Unified container backend that auto-detects Docker or Podman.
+
+    This backend uses the adapter pattern to abstract away differences between
+    Docker and Podman, providing a single consistent interface.
+    """
+
+    def __init__(self, cfg: ContainerBackendConfig):
+        self.cfg = cfg
+        self.label_prefix = "trainer.kubeflow.org"
+
+        # Initialize the container client adapter
+        self._adapter = self._create_adapter()
+
+    def _get_common_socket_locations(self, runtime_name: str) -> list[Optional[str]]:
+        """
+        Get common socket locations to try for the given runtime.
+
+        Args:
+            runtime_name: "docker" or "podman"
+
+        Returns:
+            List of socket URLs to try, including None (for default)
+        """
+        import os
+        from pathlib import Path
+
+        locations = [self.cfg.container_host] if self.cfg.container_host else []
+
+        if runtime_name == "docker":
+            # Common Docker socket locations
+            colima_sock = Path.home() / ".colima/default/docker.sock"
+            if colima_sock.exists():
+                locations.append(f"unix://{colima_sock}")
+            # Standard Docker socket
+            locations.append(None)  # Use docker.from_env() default
+
+        elif runtime_name == "podman":
+            # Common Podman socket locations on macOS
+            uid = os.getuid() if hasattr(os, "getuid") else None
+            if uid:
+                user_sock = f"/run/user/{uid}/podman/podman.sock"
+                if Path(user_sock).exists():
+                    locations.append(f"unix://{user_sock}")
+            # Standard Podman socket
+            locations.append(None)  # Use PodmanClient() default
+
+        # Remove duplicates while preserving order
+        seen = set()
+        unique_locations = []
+        for loc in locations:
+            if loc not in seen:
+                unique_locations.append(loc)
+                seen.add(loc)
+
+        return unique_locations
+
+    def _create_adapter(self) -> BaseContainerClientAdapter:
+        """
+        Create the appropriate container client adapter.
+
+        Tries Docker first, then Podman if Docker fails, unless a specific
+        runtime is requested in the config. Automatically tries common socket
+        locations (e.g., Colima for Docker on macOS, user socket for Podman).
+
+        Raises RuntimeError if neither Docker nor Podman are available.
+        """
+        runtime_map = {
+            "docker": DockerClientAdapter,
+            "podman": PodmanClientAdapter,
+        }
+
+        # Determine which runtimes to try
+        runtimes_to_try = (
+            [self.cfg.container_runtime] if self.cfg.container_runtime else ["docker", "podman"]
+        )
+
+        attempted_connections = []
+        last_error = None
+
+        for runtime_name in runtimes_to_try:
+            if runtime_name not in runtime_map:
+                continue
+
+            # Try common socket locations for this runtime
+            socket_locations = self._get_common_socket_locations(runtime_name)
+
+            for host in socket_locations:
+                try:
+                    adapter = runtime_map[runtime_name](host)
+                    adapter.ping()
+                    host_display = host or "default"
+                    logger.debug(
+                        f"Using {runtime_name} as container runtime (host: {host_display})"
+                    )
+                    return adapter
+                except Exception as e:
+                    host_str = host or "default"
+                    logger.debug(f"{runtime_name} initialization failed at {host_str}: {e}")
+                    attempted_connections.append(f"{runtime_name} at {host_str}")
+                    last_error = e
+
+        # Build helpful error message
+        import platform
+
+        system = platform.system()
+
+        attempted = ", ".join(attempted_connections)
+        error_msg = f"Could not connect to Docker or Podman (tried: {attempted}).\n"
+
+        if system == "Darwin":  # macOS
+            error_msg += (
+                "Ensure Docker/Podman is running "
+                "(e.g., 'colima start' or 'podman machine start').\n"
+            )
+        else:
+            error_msg += "Ensure Docker/Podman is installed and running.\n"
+
+        error_msg += (
+            "To specify a custom socket: ContainerBackendConfig(container_host='unix:///path/to/socket')\n"
+            "Or use LocalProcessBackendConfig for non-containerized execution."
+        )
+
+        raise RuntimeError(error_msg) from last_error
+
+    @property
+    def _runtime_type(self) -> str:
+        """Get the runtime type for debugging/logging."""
+        return self._adapter._runtime_type
+
+    # ---- Runtime APIs ----
+    def list_runtimes(self) -> list[types.Runtime]:
+        return list_training_runtimes_from_sources(self.cfg.runtime_source.sources)
+
+    def get_runtime(self, name: str) -> types.Runtime:
+        return get_training_runtime_from_sources(name, self.cfg.runtime_source.sources)
+
+    def get_runtime_packages(self, runtime: types.Runtime):
+        """
+        Spawn a short-lived container to report Python version, pip list, and nvidia-smi.
+        """
+        container_utils.maybe_pull_image(self._adapter, runtime.trainer.image, self.cfg.pull_policy)
+
+        command = [
+            "bash",
+            "-lc",
+            "python -c \"import sys; print(f'Python: {sys.version}')\" && "
+            "(pip list || echo 'pip not found') && "
+            "(nvidia-smi || echo 'nvidia-smi not found')",
+        ]
+
+        logs = self._adapter.run_oneoff_container(image=runtime.trainer.image, command=command)
+        print(logs)
+
+    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:
+        if runtime is None:
+            runtime = self.get_runtime(constants.DEFAULT_TRAINING_RUNTIME)
+        elif isinstance(runtime, str):
+            runtime = self.get_runtime(runtime)
+
+        # Process options to extract configuration
+        name = None
+        if options:
+            job_spec = {}
+            for option in options:
+                option(job_spec, trainer, self)
+
+            metadata_section = job_spec.get("metadata", {})
+            name = metadata_section.get("name")
+
+        if not isinstance(trainer, types.CustomTrainer):
+            raise ValueError(f"{self.__class__.__name__} supports only CustomTrainer in v1")
+
+        # Generate train job name if not provided via options
+        trainjob_name = name or (
+            random.choice(string.ascii_lowercase)
+            + uuid.uuid4().hex[: constants.JOB_NAME_UUID_LENGTH]
+        )
+
+        logger.debug(f"Starting training job: {trainjob_name}")
+        try:
+            # Create per-job working directory on host (for outputs, checkpoints, etc.)
+            workdir = container_utils.create_workdir(trainjob_name)
+            logger.debug(f"Created working directory: {workdir}")
+
+            # Generate training script code (inline, not written to disk)
+            training_script_code = container_utils.get_training_script_code(trainer)
+            logger.debug("Generated training script code")
+
+            # Get the image from the trainer or runtime.
+            image = trainer.image if trainer.image else runtime.trainer.image
+            logger.debug(f"Using image: {image}")
+
+            container_utils.maybe_pull_image(self._adapter, image, self.cfg.pull_policy)
+            logger.debug(f"Image ready: {image}")
+
+            # Build base environment
+            env = container_utils.build_environment(trainer)
+
+            # Construct pre-run command to install packages
+            pre_install_cmd = container_utils.build_pip_install_cmd(trainer)
+
+            # Create network for multi-node communication
+            num_nodes = trainer.num_nodes or runtime.trainer.num_nodes or 1
+            logger.debug(f"Creating network for {num_nodes} nodes")
+
+            # Determine number of processes per node from GPU count
+            # For GPU training: spawn one process per GPU for optimal utilization
+            # For CPU training: use single process (PyTorch parallelizes internally via threads)
+            nproc_per_node = 1  # Default for CPU training
+            if trainer.resources_per_node and "gpu" in trainer.resources_per_node:
+                try:
+                    nproc_per_node = int(trainer.resources_per_node["gpu"])
+                    logger.debug(f"Using {nproc_per_node} processes per node (1 per GPU)")
+                except (ValueError, TypeError):
+                    logger.warning(
+                        f"Invalid GPU count in resources_per_node: "
+                        f"{trainer.resources_per_node['gpu']}, defaulting to 1 process per node"
+                    )
+            else:
+                logger.debug("No GPU specified, using 1 process per node")
+
+            network_id = self._adapter.create_network(
+                name=f"{trainjob_name}-net",
+                labels={
+                    f"{self.label_prefix}/trainjob-name": trainjob_name,
+                    f"{self.label_prefix}/runtime-name": runtime.name,
+                    f"{self.label_prefix}/workdir": workdir,
+                },
+            )
+            logger.debug(f"Created network: {network_id}")
+
+            # Create N containers (one per node)
+            container_ids: list[str] = []
+            master_container_id = None
+            master_ip = None
+
+            for rank in range(num_nodes):
+                container_name = f"{trainjob_name}-node-{rank}"
+
+                # Get master address and port for torchrun
+                master_port = 29500
+
+                # For Podman: use IP address to avoid DNS timing issues
+                # For Docker: use hostname (DNS is reliable)
+                if rank == 0:
+                    # Master node - will be created first
+                    master_addr = f"{trainjob_name}-node-0"
+                else:
+                    # Worker nodes - determine master address based on runtime
+                    if self._runtime_type == "podman" and master_ip:
+                        master_addr = master_ip
+                        logger.debug(f"Using master IP address for Podman: {master_ip}")
+                    else:
+                        master_addr = f"{trainjob_name}-node-0"
+                        logger.debug(f"Using master hostname: {master_addr}")
+
+                # Prefer torchrun; fall back to python if torchrun is unavailable
+                # For worker nodes, wait for master to be reachable before starting torchrun
+                wait_for_master = ""
+                if rank > 0:
+                    wait_for_master = (
+                        f"echo 'Waiting for master node {master_addr}:{master_port}...'; "
+                        f"for i in {{1..60}}; do "
+                        f"  if timeout 1 bash -c 'cat < /dev/null > "
+                        f"/dev/tcp/{master_addr}/{master_port}' 2>/dev/null; then "
+                        f"    echo 'Master node is reachable'; break; "
+                        f"  fi; "
+                        f"  if [ $i -eq 60 ]; then "
+                        f"echo 'Timeout waiting for master node'; exit 1; fi; "
+                        f"  sleep 2; "
+                        f"done; "
+                    )
+
+                # Embed training script inline using heredoc (no file I/O on host)
+                entry_cmd = (
+                    f"{pre_install_cmd}"
+                    f"{wait_for_master}"
+                    f"cat > /tmp/train.py << 'TRAINING_SCRIPT_EOF'\n"
+                    f"{training_script_code}\n"
+                    f"TRAINING_SCRIPT_EOF\n"
+                    "if command -v torchrun >/dev/null 2>&1; then "
+                    f"  torchrun --nproc_per_node={nproc_per_node} --nnodes={num_nodes} "
+                    f"  --node-rank={rank} --rdzv-backend=static "
+                    f"  --rdzv-endpoint={master_addr}:{master_port} "
+                    f"  /tmp/train.py; "
+                    "else "
+                    f"  python /tmp/train.py; "
+                    "fi"
+                )
+
+                full_cmd = ["bash", "-lc", entry_cmd]
+
+                labels = {
+                    f"{self.label_prefix}/trainjob-name": trainjob_name,
+                    f"{self.label_prefix}/step": f"node-{rank}",
+                    f"{self.label_prefix}/network-id": network_id,
+                }
+
+                volumes = {
+                    workdir: {
+                        "bind": constants.WORKSPACE_PATH,
+                        "mode": "rw",
+                    }
+                }
+
+                logger.debug(f"Creating container {rank}/{num_nodes}: {container_name}")
+
+                container_id = self._adapter.create_and_start_container(
+                    image=image,
+                    command=full_cmd,
+                    name=container_name,
+                    network_id=network_id,
+                    environment=env,
+                    labels=labels,
+                    volumes=volumes,
+                    working_dir=constants.WORKSPACE_PATH,
+                )
+
+                logger.debug(f"Started container {container_name} (ID: {container_id[:12]})")
+                container_ids.append(container_id)
+
+                # If this is the master node and we're using Podman, get its IP address
+                if rank == 0:
+                    master_container_id = container_id
+                    if self._runtime_type == "podman":
+                        # Get master IP for worker nodes to use
+                        master_ip = self._adapter.get_container_ip(master_container_id, network_id)
+                        if master_ip:
+                            logger.debug(f"Master node IP address: {master_ip}")
+                        else:
+                            logger.warning(
+                                "Could not retrieve master IP address. "
+                                "Worker nodes will fall back to DNS resolution."
+                            )
+
+            logger.debug(
+                f"Training job {trainjob_name} created successfully with "
+                f"{len(container_ids)} container(s)"
+            )
+            return trainjob_name
+
+        except Exception as e:
+            # Clean up on failure
+            logger.error(f"Failed to create training job {trainjob_name}: {e}")
+            logger.exception("Full traceback:")
+
+            # Try to clean up any resources that were created
+            from contextlib import suppress
+
+            try:
+                # Stop and remove any containers that were created
+                if "container_ids" in locals():
+                    for container_id in container_ids:
+                        with suppress(Exception):
+                            self._adapter.stop_container(container_id, timeout=5)
+                            self._adapter.remove_container(container_id, force=True)
+
+                # Remove network if it was created
+                if "network_id" in locals():
+                    with suppress(Exception):
+                        self._adapter.delete_network(network_id)
+
+                # Remove working directory if it was created
+                if "workdir" in locals() and os.path.isdir(workdir):
+                    shutil.rmtree(workdir, ignore_errors=True)
+
+            except Exception as cleanup_error:
+                logger.error(f"Error during cleanup: {cleanup_error}")
+
+            # Re-raise the original exception
+            raise
+
+    def _get_job_containers(self, name: str) -> list[dict]:
+        """
+        Get containers for a specific training job.
+
+        Args:
+            name: Name of the training job
+
+        Returns:
+            List of container dictionaries for this job
+
+        Raises:
+            ValueError: If no containers found for the job
+        """
+        filters = {"label": [f"{self.label_prefix}/trainjob-name={name}"]}
+        containers = self._adapter.list_containers(filters=filters)
+
+        if not containers:
+            raise ValueError(f"No TrainJob with name {name}")
+
+        return containers
+
+    def __get_trainjob_from_containers(
+        self, job_name: str, containers: list[dict]
+    ) -> types.TrainJob:
+        """
+        Build a TrainJob object from a list of containers.
+
+        Args:
+            job_name: Name of the training job
+            containers: List of container dictionaries for this job
+
+        Returns:
+            TrainJob object
+
+        Raises:
+            ValueError: If network metadata is missing or runtime not found
+        """
+        if not containers:
+            raise ValueError(f"No containers found for TrainJob {job_name}")
+
+        # Get metadata from network
+        network_id = containers[0]["labels"].get(f"{self.label_prefix}/network-id")
+        if not network_id:
+            raise ValueError(f"TrainJob {job_name} is missing network metadata")
+
+        network_info = self._adapter.get_network(network_id)
+        if not network_info:
+            raise ValueError(f"TrainJob {job_name} network not found")
+
+        network_labels = network_info.get("labels", {})
+        runtime_name = network_labels.get(f"{self.label_prefix}/runtime-name")
+
+        # Get runtime object
+        try:
+            job_runtime = self.get_runtime(runtime_name) if runtime_name else None
+        except Exception as e:
+            raise ValueError(f"Runtime {runtime_name} not found for job {job_name}") from e
+
+        if not job_runtime:
+            raise ValueError(f"Runtime {runtime_name} not found for job {job_name}")
+
+        # Parse creation timestamp from first container
+        created_str = containers[0].get("created", "")
+        try:
+            from dateutil import parser
+
+            creation_timestamp = parser.isoparse(created_str)
+        except Exception:
+            creation_timestamp = datetime.now()
+
+        # Build steps from containers
+        steps = []
+        for container in sorted(containers, key=lambda c: c["name"]):
+            step_name = container["labels"].get(f"{self.label_prefix}/step", "")
+            steps.append(
+                types.Step(
+                    name=step_name,
+                    pod_name=container["name"],
+                    status=container_utils.get_container_status(self._adapter, container["id"]),
+                )
+            )
+
+        # Get num_nodes from container count
+        num_nodes = len(containers)
+
+        return types.TrainJob(
+            name=job_name,
+            creation_timestamp=creation_timestamp,
+            runtime=job_runtime,
+            steps=steps,
+            num_nodes=num_nodes,
+            status=container_utils.aggregate_container_statuses(self._adapter, containers),
+        )
+
+    def list_jobs(self, runtime: Optional[types.Runtime] = None) -> list[types.TrainJob]:
+        """List all training jobs by querying container runtime."""
+        # Get all containers with our label prefix
+        filters = {"label": [f"{self.label_prefix}/trainjob-name"]}
+        containers = self._adapter.list_containers(filters=filters)
+
+        # Group containers by job name
+        jobs_map: dict[str, list[dict]] = {}
+        for container in containers:
+            job_name = container["labels"].get(f"{self.label_prefix}/trainjob-name")
+            if job_name:
+                if job_name not in jobs_map:
+                    jobs_map[job_name] = []
+                jobs_map[job_name].append(container)
+
+        result: list[types.TrainJob] = []
+        for job_name, job_containers in jobs_map.items():
+            # Skip jobs with no containers
+            if not job_containers:
+                continue
+
+            # Filter by runtime if specified
+            if runtime:
+                network_id = job_containers[0]["labels"].get(f"{self.label_prefix}/network-id")
+                if network_id:
+                    network_info = self._adapter.get_network(network_id)
+                    if network_info:
+                        network_labels = network_info.get("labels", {})
+                        runtime_name = network_labels.get(f"{self.label_prefix}/runtime-name")
+                        if runtime_name != runtime.name:
+                            continue
+
+            # Build TrainJob from containers
+            try:
+                result.append(self.__get_trainjob_from_containers(job_name, job_containers))
+            except Exception as e:
+                logger.warning(f"Failed to get TrainJob {job_name}: {e}")
+                continue
+
+        return result
+
+    def get_job(self, name: str) -> types.TrainJob:
+        """Get a specific training job by querying container runtime."""
+        containers = self._get_job_containers(name)
+        return self.__get_trainjob_from_containers(name, containers)
+
+    def get_job_logs(
+        self,
+        name: str,
+        follow: bool = False,
+        step: str = constants.NODE + "-0",
+    ) -> Iterator[str]:
+        """Get logs for a training job by querying container runtime."""
+        containers = self._get_job_containers(name)
+
+        want_all = step == constants.NODE + "-0"
+        for container in sorted(containers, key=lambda c: c["name"]):
+            container_step = container["labels"].get(f"{self.label_prefix}/step", "")
+            if not want_all and container_step != step:
+                continue
+            try:
+                yield from self._adapter.container_logs(container["id"], follow)
+            except Exception as e:
+                logger.warning(f"Failed to get logs for {container['name']}: {e}")
+                yield f"Error getting logs: {e}\n"
+
+    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:
+        import time
+
+        end = time.time() + timeout
+        while time.time() < end:
+            tj = self.get_job(name)
+            logger.debug(f"TrainJob {name}, status {tj.status}")
+
+            # Invoke callbacks if provided
+            if callbacks:
+                for callback in callbacks:
+                    callback(tj)
+
+            if tj.status in status:
+                return tj
+            if constants.TRAINJOB_FAILED not in status and tj.status == constants.TRAINJOB_FAILED:
+                raise RuntimeError(f"TrainJob {name} is Failed")
+            time.sleep(polling_interval)
+        raise TimeoutError(f"Timeout waiting for TrainJob {name} to reach status: {status}")
+
+    def delete_job(self, name: str):
+        """Delete a training job by querying container runtime."""
+        containers = self._get_job_containers(name)
+
+        # Get network_id and workdir from labels
+        network_id = containers[0]["labels"].get(f"{self.label_prefix}/network-id")
+
+        # Get workdir from network labels
+        workdir_host = None
+        if network_id:
+            network_info = self._adapter.get_network(network_id)
+            if network_info:
+                network_labels = network_info.get("labels", {})
+                workdir_host = network_labels.get(f"{self.label_prefix}/workdir")
+
+        # Stop containers and remove
+        from contextlib import suppress
+
+        for container in containers:
+            with suppress(Exception):
+                self._adapter.stop_container(container["id"], timeout=10)
+            with suppress(Exception):
+                self._adapter.remove_container(container["id"], force=True)
+
+        # Remove network (best-effort)
+        if network_id:
+            with suppress(Exception):
+                self._adapter.delete_network(network_id)
+
+        # Remove working directory if configured
+        if self.cfg.auto_remove and workdir_host and os.path.isdir(workdir_host):
+            shutil.rmtree(workdir_host, ignore_errors=True)