manta-node 0.5b0.dev9__py3-none-any.whl

manta_node/infrastructure/container/manager.py

@@ -0,0 +1,536 @@
+import os
+import subprocess
+import sys
+from datetime import datetime
+from math import nan
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union, cast
+
+import aiodocker
+import docker
+from aiodocker.docker import DockerContainer
+
+from manta_common.build.common.tasks import TaskStatus
+from manta_common.conversions import ID
+from manta_common.errors import MantaDockerError
+from manta_common.traces import Tracer
+from ...domain.task_lifecycle import TaskLifecycleService
+from .docker_adapter import DockerRaiser
+
+
+class ContainerManager(DockerRaiser):
+    """
+    ContainerManager class to run tasks in docker containers
+
+    Parameters
+    ----------
+    node_id: ID
+        Node ID
+    docker_client: DockerClient
+        Docker client
+    async_client: AiodockerClient
+        Aiodocker client
+    containers: Dict[ID, DockerContainer]
+        Containers
+    gpu_config: dict
+        GPU configuration
+    """
+
+    tracer: Tracer
+
+    def __init__(self):
+        self.docker_client = docker.from_env()
+
+        # Windows-specific: Use TCP connection instead of Named Pipes to avoid event loop conflict
+        # SECURITY NOTE: Port 2375 is typically unencrypted. For production deployments on Windows:
+        #   1. Enable TLS on Docker daemon (port 2376)
+        #   2. Configure client certificates
+        #   3. Update connection URL to tcp://localhost:2376 with tls=True
+        # See: https://docs.docker.com/engine/security/protect-access/
+        if sys.platform.lower() == "win32" or os.name.lower() == "nt":
+            try:
+                # Docker Desktop on Windows exposes the API on tcp://localhost:2375
+                # WARNING: This connection is unencrypted. Use only in development.
+                self.async_client = aiodocker.Docker(url="tcp://localhost:2375")
+                self.tracer.warning(
+                    "Using unencrypted Docker TCP connection on Windows (port 2375). "
+                    "For production, configure TLS on port 2376."
+                )
+            except Exception as e:
+                self.tracer.error(
+                    f"Failed to connect to Docker on Windows via TCP: {e}. "
+                    "Ensure Docker Desktop is running and exposing TCP port 2375."
+                )
+                raise
+        else:
+            # Unix systems use the default socket
+            self.async_client = aiodocker.Docker()
+
+        self.containers: Dict[Tuple[ID, ID], DockerContainer] = {}
+        self.gpu_config = self.get_gpu_config()
+        if self.gpu_config is not None:
+            self.tracer.debug(
+                f"GPU support detected: \n{self.gpu_config}"
+            )  # pragma: no cover
+        else:
+            self.tracer.warning("No GPU support detected.")
+
+    @property
+    def gpu_supported(self):
+        return self.gpu_config is not None
+
+    def get_gpu_config(self) -> Optional[Dict[str, Any]]:
+        """
+        Check if GPU is supported and return the appropriate config field for aiodocker.
+
+        Returns
+        -------
+        Optional[Dict[str, Any]]
+            GPU configuration
+        """
+        # Get Docker version
+        version_info = self.docker_client.version()
+        docker_version = tuple(
+            map(int, version_info["Version"].split(".")[:2])
+        )  # Extract major.minor
+
+        # Check runtimes if Docker < 19.03
+        if docker_version < (19, 3):
+            if self._is_nvidia_runtime_installed():
+                return {"Runtime": "nvidia"}  # Old way
+            return None  # No GPU support
+
+        # Check for NVIDIA Container Toolkit (Docker >= 19.03)
+        if self._is_nvidia_toolkit_installed():
+            return {
+                "HostConfig": {
+                    "DeviceRequests": [
+                        {"Driver": "nvidia", "Count": -1, "Capabilities": [["gpu"]]}
+                    ]
+                }
+            }
+        elif self._is_nvidia_runtime_installed():
+            return {"Runtime": "nvidia"}
+
+        return None  # No GPU support detected
+
+    def _is_nvidia_runtime_installed(self):
+        """Method 1: Check if nvidia runtime is installed."""
+        info = self.docker_client.info()
+        return "nvidia" in info.get("Runtimes", {})
+
+    def _is_nvidia_toolkit_installed(self):
+        """Method 2: Check if nvidia-container-toolkit is installed."""
+        try:
+            subprocess.run(
+                ["nvidia-container-cli", "-V"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                check=True,
+            )
+            return True
+        except (FileNotFoundError, subprocess.CalledProcessError) as e:
+            self.tracer.debug(f"nvidia-container-cli not found: {e}")
+
+        return False
+
+    async def ensure_tailscale_service(
+        self,
+        swarm_id: ID,
+        authkey: str,
+        login_server: str,
+        hostname: str,
+        temp_folder: Path,
+    ):
+        """
+        Ensure the tailscale service is running
+
+        Parameters
+        ----------
+        swarm_id : ID
+            Swarm ID
+        authkey : str
+            Tailscale auth key
+        login_server : str
+            Tailscale login server
+        hostname : str
+            Tailscale hostname
+        temp_folder : Path
+            Temporary folder
+
+        Returns
+        -------
+        datetime
+            Start time of the container
+        """
+        # network = f"tailscale-{swarm_id}"
+        # Crée le réseau Docker s’il n’existe pas
+        # if not self.docker_client.networks.list(names=[network]):
+        #     self.docker_client.networks.create(network, driver="bridge", attachable=True)
+        # Lancement du conteneur Tailscale
+        config = {
+            "Image": "tailscale/tailscale:stable",
+            "HostConfig": {
+                "Hostname": hostname,
+                "Binds": [
+                    f"tailscale-{swarm_id}-data:/var/lib/tailscale",
+                    "/dev/net/tun:/dev/net/tun",
+                ],
+                "CapAdd": ["NET_ADMIN", "SYS_MODULE"],
+                "RestartPolicy": {"Name": "unless-stopped"},
+                "NetworkMode": "host",
+            },
+            "Env": [
+                f"TS_AUTHKEY=tskey-client-{authkey}?ephemeral=false",
+                "TS_SERVE_CONFIG=/config/config.json",
+                "TS_STATE_DIR=/var/lib/tailscale",
+                "TS_USERSPACE=false",
+                f"TS_EXTRA_ARGS=--login-server={login_server} --advertise-tags=tag:container --reset",
+            ],
+            "Volumes": {f"tailscale-{swarm_id}-data": {}},
+            "Detach": True,
+            "Labels": {
+                "com.docker.compose.project": "manta_tasks",
+                "com.docker.compose.service": "stirling-ts",  # optionnel
+            },
+        }
+        self.tracer.debug("Running tailscale with following configuration:")
+        for key, value in config.items():
+            self.tracer.debug(f"      {key}: {value}")
+        self.containers[
+            (swarm_id, ID(f"tailscale-{swarm_id}"))
+        ] = await self.run_or_raise(config)
+        return f"service:tailscale-{swarm_id}"
+
+    async def run_container(
+        self,
+        swarm_id: ID,
+        task_id: ID,
+        image: str,
+        command: str,
+        volumes: dict,
+        environment: Optional[dict] = None,
+        memory: Optional[Union[int, str]] = None,
+        cpus: Optional[str] = None,
+        gpu: Optional[bool] = False,
+        network: Optional[str] = None,
+        alias: Optional[str] = None,
+        force_pull: bool = False,
+    ) -> datetime:
+        """
+        Run a container with the given image, command, volumes and environment
+
+        Parameters
+        ----------
+        swarm_id : ID
+            ID of the swarm
+        task_id : ID
+            ID of the task
+        image : str
+            Image to run the container
+        command : str
+            Command to run in the container
+        volumes : list
+            Volumes to mount in the container
+        environment : Optional[dict]
+            Environment variables
+        memory : Optional[Union[int, str]]
+            Memory limit.
+            Accepts float values (which represent the memory limit of
+            the created container in bytes) or a string with a units
+            identification char (100000b, 1000k, 128m, 1g).
+            If a string is specified without a units character, bytes
+            are assumed as an intended unit.
+        cpus : Optional[str]
+            CPU limit. CPUs in which to allow execution (0-3, 0, 1)
+            or a percentage of the available CPUs
+        gpu : Optional[bool]
+            GPU support
+        network : Optional[str]
+            Network to attach to the container
+        alias : Optional[str]
+            Alias of the container
+        force_pull : bool
+            Force pull the image before running the container
+
+        Returns
+        -------
+        datetime
+            Start time of the container
+        """
+        if gpu and not self.gpu_supported:
+            raise MantaDockerError(
+                "GPU support is not enabled for this container manager.",
+                metadata=self.tracer.metadata,
+            )
+
+        host_config = (
+            self.gpu_config.get("HostConfig", {}) if (gpu and self.gpu_config) else {}
+        )
+        host_config["Binds"] = [
+            f"{str(host_path) if isinstance(host_path, Path) else host_path}:{bind['bind']}:{bind['mode']}"
+            for host_path, bind in volumes.items()
+        ]
+        host_config["NetworkMode"] = network if network else "host"
+        host_config["Hostname"] = alias if alias else None
+        host_config["ExtraHosts"] = ["host.docker.internal:host-gateway"]
+        env = (
+            [f"{key}={value}" for key, value in environment.items()]
+            if environment
+            else None
+        )
+        log_config = {
+            "type": "json-file",
+            "config": {"max-size": "10m", "max-file": "3"},
+        }
+
+        config = {
+            "Cmd": command.split(),
+            "Image": image,
+            "HostConfig": host_config,
+            "Env": env,
+            "Detach": True,
+            "Memory": memory,
+            "CpuShares": cpus,
+            "LogConfig": log_config,
+            "Runtime": (
+                self.gpu_config.get("Runtime") if (gpu and self.gpu_config) else None
+            ),
+            "Labels": {
+                "com.docker.compose.project": "manta_tasks",
+            },
+        }
+
+        start_time = datetime.now()
+        self.tracer.debug("Running container with following configuration:")
+        for key, value in config.items():
+            self.tracer.debug(f"      {key}: {value}")
+
+        if force_pull:
+            self.tracer.info(f"Pulling image {image}...")
+            await self.pull_image_or_raise(image)
+
+        self.containers[(swarm_id, task_id)] = await self.run_or_raise(config)
+
+        self.tracer.info(f"Task {task_id} container started.")
+        return start_time
+
+    async def restart_container(self, swarm_id: ID, task_id: ID) -> datetime:
+        """
+        Restart a task's container
+
+        Parameters
+        ----------
+        swarm_id : ID
+            Swarm ID
+        task_id : ID
+            Task ID
+
+        Returns
+        -------
+        datetime
+            Start time of the container
+
+        Raises
+        ------
+        MantaDockerError
+            Task not found
+        """
+        if container := self.containers.get((swarm_id, task_id)):
+            start_time = datetime.now()
+            await self.restart_or_raise(container)
+            self.tracer.debug(f"Task {task_id} container restarted.")
+            return start_time
+        message = f"Task {task_id} not found."
+        self.tracer.error(message)
+        raise MantaDockerError(message, metadata=self.tracer.metadata)
+
+    async def stop_container(self, swarm_id: ID, task_id: ID):
+        """
+        Stop a running container
+
+        Parameters
+        ----------
+        swarm_id : ID
+            Swarm ID
+        task_id : ID
+            Task ID
+
+        Raises
+        ------
+        MantaDockerError
+            Task not found
+        """
+        if container := self.containers.get((swarm_id, task_id)):
+            await self.stop_or_raise(container)
+            self.tracer.debug(f"Task {task_id} container stopped.")
+        else:
+            message = f"Task {task_id} not found."
+            self.tracer.error(message)
+            raise MantaDockerError(message, metadata=self.tracer.metadata)
+
+    async def remove_container(self, swarm_id: ID, task_id: ID):
+        """
+        Remove a container
+
+        Parameters
+        ----------
+        swarm_id : ID
+            Swarm ID
+        task_id : ID
+            Task ID
+
+        Raises
+        ------
+        MantaDockerError
+            Task not found
+        """
+        if container := self.containers.pop((swarm_id, task_id), None):
+            await self.remove_or_raise(container)
+            self.tracer.debug(f"Task {task_id} container removed.")
+        else:
+            message = f"Task {task_id} not found."
+            self.tracer.error(message)
+            raise MantaDockerError(message, metadata=self.tracer.metadata)
+
+    async def get_task_status(self, swarm_id: ID, task_id: ID) -> TaskStatus:
+        """
+        Get the status of a container
+
+        Parameters
+        ----------
+        swarm_id : ID
+            Swarm ID
+        task_id : ID
+            Task ID
+
+        Returns
+        -------
+        TaskStatus
+            Status of the container
+
+        Raises
+        ------
+        MantaDockerError
+            Container not found
+            Cannot reload container
+        """
+        container = self.containers.get((swarm_id, task_id))
+        if container is None:
+            self.tracer.error("Cannot found container")
+            raise MantaDockerError(
+                "Cannot found container", metadata=self.tracer.metadata
+            )
+
+        # self.reload_or_raise(container)
+        status = await self.get_status_or_raise(container)
+        self.tracer.debug(f"Task {task_id} container's status: {status}")
+
+        # Use domain service to map container status to task status
+        return TaskLifecycleService.calculate_status_from_container(status)
+
+    async def get_container_logs(
+        self, swarm_id: ID, task_id: ID, since: Optional[datetime] = None
+    ) -> bytes:
+        """
+        Get the logs of a container (async)
+
+        Parameters
+        ----------
+        swarm_id : ID
+            Swarm ID
+        task_id : ID
+            Task ID
+        since : Optional[datetime]
+            Show logs since a given datetime
+
+        Returns
+        -------
+        bytes
+            Logs from the container
+
+        Raises
+        ------
+        MantaDockerError
+            Cannot found container
+            Cannot get container's logs
+        """
+        container = self.containers.get((swarm_id, task_id))
+        if container is None:
+            message = f"Task {task_id} not found for swarm {swarm_id}."
+            self.tracer.error(message)
+            raise MantaDockerError(message, metadata=self.tracer.metadata)
+        return await self.get_container_logs_or_raise(container, since)
+
+    async def get_task_stats(self, swarm_id: ID, task_id: ID) -> dict:
+        """
+        Get the stats of a task
+
+        Parameters
+        ----------
+        swarm_id : ID
+            Swarm ID
+        task_id : ID
+            Task ID
+
+        Returns
+        -------
+        dict
+            Stats from the container
+
+        Raises
+        ------
+        MantaDockerError
+            Task not found
+        """
+        if container := self.containers.get((swarm_id, task_id)):
+            stats = cast(Dict[str, Any], await container.stats(stream=False))
+
+            cpu_usage = stats["cpu_stats"]["cpu_usage"]["total_usage"]
+            system_cpu_usage = stats["cpu_stats"]["system_cpu_usage"]
+            online_cpus = stats["cpu_stats"]["online_cpus"]
+            cpu_percent = (cpu_usage / system_cpu_usage) * online_cpus * 100
+
+            cpu = dict(
+                total=float(100.0),
+                available=100.0 - cpu_percent,
+                used=cpu_percent,
+            )
+
+            memory_usage = stats["memory_stats"]["usage"]
+            memory_limit = stats["memory_stats"]["limit"]
+
+            memory = dict(
+                total=float(memory_limit),
+                available=float(memory_usage / memory_limit * 100),
+                used=float(memory_usage),
+            )
+        else:
+            cpu = dict(
+                total=100.0,
+                available=nan,
+                used=nan,
+            )
+            memory = dict(
+                total=nan,
+                available=nan,
+                used=nan,
+            )
+
+        return {"cpu": cpu, "memory": memory}
+
+    # TODO : use (swarm_id, task_id) instead of only task_id
+    async def collect_system_summary(self) -> dict:
+        """
+        Collect the current system metrics
+
+        Returns
+        -------
+        dict
+            System informations
+        """
+        # Iterate over all containers (tasks) and get their stats
+        return {
+            task_id.xid: await self.get_task_stats(swarm_id, task_id)
+            for (swarm_id, task_id) in self.containers
+        }

manta_node/infrastructure/filesystem/__init__.py

@@ -0,0 +1,7 @@
+"""Filesystem infrastructure - Dataset and file management."""
+
+from .dataset_manager import DatasetManager
+
+__all__ = [
+    "DatasetManager",
+]