gpustack-runtime 0.1.40.post1__py3-none-any.whl → 0.1.41__py3-none-any.whl

gpustack_runtime/deployer/cdi/thead.py

@@ -11,8 +11,10 @@ from .__types__ import (
     ConfigContainerEdits,
     ConfigDevice,
     Generator,
-    manufacturer_to_config_kind,
+    manufacturer_to_cdi_kind,
+    manufacturer_to_runtime_env,
 )
+from .__utils__ import device_to_cdi_device_node
 
 
 class THeadGenerator(Generator):
@@ -23,13 +25,20 @@ class THeadGenerator(Generator):
     def __init__(self):
         super().__init__(ManufacturerEnum.THEAD)
 
-    def generate(self, devices: Devices | None = None) -> Config | None:
+    def generate(
+        self,
+        devices: Devices | None = None,
+        include_all_devices: bool = True,
+    ) -> Config | None:
         """
         Generate the CDI configuration for T-Head devices.
 
         Args:
-            devices: The detected devices.
-            If None, all available devices are considered.
+            devices:
+                The detected devices.
+                If None, all available devices are considered.
+            include_all_devices:
+                Whether to include a device entry that represents all T-Head devices.
 
         Returns:
             The Config object, or None if not supported.
@@ -46,30 +55,44 @@ class THeadGenerator(Generator):
         if not devices:
             return None
 
-        kind = manufacturer_to_config_kind(self.manufacturer)
+        kind = manufacturer_to_cdi_kind(self.manufacturer)
         if not kind:
             return None
 
         cdi_devices: list[ConfigDevice] = []
 
-        all_container_edits_device_nodes = [
+        common_device_nodes = []
+        for p in [
             "/dev/alixpu",
             "/dev/alixpu_ctl",
-        ]
+        ]:
+            cdn = device_to_cdi_device_node(
+                path=p,
+            )
+            if cdn:
+                common_device_nodes.append(cdn)
+        if not common_device_nodes:
+            return None
+
+        all_device_nodes = []
+
         for dev in devices:
             if not dev:
                 continue
-            all_container_edits_device_nodes.append(
-                f"/dev/alixpu_ppu{dev.index}",
+
+            container_device_nodes = []
+
+            cdn = device_to_cdi_device_node(
+                path=f"/dev/alixpu_ppu{dev.index}",
             )
+            if not cdn:
+                continue
+            all_device_nodes.append(cdn)
+            container_device_nodes.append(cdn)
 
-            # Add specific container edits for each device
+            # Add specific container edits for each device.
             cdi_container_edits = ConfigContainerEdits(
-                device_nodes=[
-                    "/dev/alixpu",
-                    "/dev/alixpu_ctl",
-                    f"/dev/alixpu_ppu{dev.index}",
-                ],
+                device_nodes=container_device_nodes,
             )
             cdi_devices.append(
                 ConfigDevice(
@@ -87,17 +110,28 @@ class THeadGenerator(Generator):
         if not cdi_devices:
             return None
 
-        # Add common container edits for all devices
-        cdi_devices.append(
-            ConfigDevice(
-                name="all",
-                container_edits=ConfigContainerEdits(
-                    device_nodes=all_container_edits_device_nodes,
+        # Add common container edits for all devices.
+        if include_all_devices:
+            cdi_devices.append(
+                ConfigDevice(
+                    name="all",
+                    container_edits=ConfigContainerEdits(
+                        device_nodes=all_device_nodes,
+                    ),
                 ),
-            ),
-        )
+            )
+
+        runtime_env = manufacturer_to_runtime_env(self.manufacturer)
 
         return Config(
             kind=kind,
             devices=cdi_devices,
+            container_edits=[
+                ConfigContainerEdits(
+                    env=[
+                        f"{runtime_env}=void",
+                    ],
+                    device_nodes=common_device_nodes,
+                ),
+            ],
         )

gpustack_runtime/deployer/docker.py

@@ -1,4 +1,4 @@
-from __future__ import annotations
+from __future__ import annotations as __future_annotations__
 
 import contextlib
 import io
@@ -10,7 +10,7 @@ import socket
 import sys
 import tarfile
 from dataclasses import dataclass, field
-from functools import lru_cache, reduce
+from functools import reduce
 from math import ceil
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -21,6 +21,7 @@ import docker.models.containers
 import docker.models.images
 import docker.models.volumes
 import docker.types
+from cachetools.func import ttl_cache
 from dataclasses_json import dataclass_json
 from gpustack_runner import split_image
 from tqdm import tqdm
@@ -54,7 +55,7 @@ from .__utils__ import (
     safe_json,
     sensitive_env_var,
 )
-from .cdi import generate_config as cdi_generate_config
+from .cdi import dump_config as cdi_dump_config
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
@@ -314,7 +315,7 @@ class DockerDeployer(EndoscopicDeployer):
     """
 
     @staticmethod
-    @lru_cache
+    @ttl_cache(maxsize=1, ttl=60)
     def is_supported() -> bool:
         """
         Check if Docker is supported in the current environment.
@@ -1002,9 +1003,9 @@ class DockerDeployer(EndoscopicDeployer):
                             privileged = create_options.get("privileged", False)
 
                             # Generate CDI config if not yet.
-                            if cdi and envs.GPUSTACK_RUNTIME_DEPLOY_CDI_SPECS_GENERATE:
+                            if cdi and envs.GPUSTACK_RUNTIME_DOCKER_CDI_SPECS_GENERATE:
                                 for re in runtime_env:
-                                    cdi_generate_config(
+                                    cdi_dump_config(
                                         manufacturer=vd_manus[re],
                                         output=envs.GPUSTACK_RUNTIME_DEPLOY_CDI_SPECS_DIRECTORY,
                                     )
@@ -1548,7 +1549,7 @@ class DockerDeployer(EndoscopicDeployer):
     def _get(
         self,
         name: WorkloadName,
-        namespace: WorkloadNamespace | None = None,  # noqa: ARG002
+        namespace: WorkloadNamespace | None = None,
     ) -> WorkloadStatus | None:
         """
         Get the status of a Docker workload.
@@ -1670,7 +1671,7 @@ class DockerDeployer(EndoscopicDeployer):
     @_supported
     def _list(
         self,
-        namespace: WorkloadNamespace | None = None,  # noqa: ARG002
+        namespace: WorkloadNamespace | None = None,
         labels: dict[str, str] | None = None,
     ) -> list[WorkloadStatus]:
         """

gpustack_runtime/deployer/k8s/deviceplugin/__init__.py

@@ -0,0 +1,240 @@
+from __future__ import annotations as __future_annotations__
+
+import asyncio
+import contextlib
+import logging
+import signal
+import threading
+from typing import TYPE_CHECKING
+
+from .... import envs
+from ....deployer.cdi import dump_config as cdi_dump_config
+from ....detector import ManufacturerEnum, detect_devices, supported_manufacturers
+from .__types__ import GroupedError, PluginServer
+from .plugin import SharableDevicePlugin, cdi_kind_to_kdp_resource
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+logger = logging.getLogger(__package__)
+
+
+async def serve_async(
+    stop_event: asyncio.Event,
+    manufacturer: ManufacturerEnum | None = None,
+    cdi_generation_output: Path | None = None,
+    kubelet_endpoint: Path | None = None,
+    start_timeout: int = 5,
+    register_timeout: int = 5,
+):
+    """
+    Serve device plugins for all detected devices asynchronously.
+
+    Args:
+        stop_event:
+            An asyncio event to signal stopping the server.
+        manufacturer:
+            Manufacturer for serving.
+            If None, detected from current environment.
+        cdi_generation_output:
+            A directory to store generated CDI configuration files.
+            If None, CDI configuration is not stored.
+        kubelet_endpoint:
+            The path to the kubelet endpoint.
+        start_timeout:
+            The timeout in seconds for starting the device plugin.
+        register_timeout:
+            The timeout in seconds for registering the device plugin.
+
+    Raises:
+        If any device plugin fails during serving, the exception is raised.
+
+    """
+    manufacturers = [manufacturer] if manufacturer else supported_manufacturers()
+
+    # Create servers for all detected devices.
+    servers: list[PluginServer] = []
+    for manu in manufacturers:
+        devices = detect_devices(
+            manufacturer=manu,
+        )
+        if not devices:
+            continue
+
+        # Also works if the manufacturer does not have a CDI generator,
+        # which means we are relying on other tools to generate CDI specs.
+        if envs.GPUSTACK_RUNTIME_KUBERNETES_KDP_CDI_SPECS_GENERATE:
+            generated_content, generated_path = cdi_dump_config(
+                manufacturer=manu,
+                output=cdi_generation_output,
+            )
+            if generated_content:
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(
+                        "Generated CDI configuration for '%s' at '%s':\n%s",
+                        manu,
+                        generated_path,
+                        generated_content,
+                    )
+                else:
+                    logger.info(
+                        "Generated CDI configuration for '%s' at '%s'",
+                        manu,
+                        generated_path,
+                    )
+            else:
+                logger.warning(
+                    "Delegated CDI configuration by other tools for manufacturer '%s', "
+                    "e.g. NVIDIA Container Toolkit Manual CDI Specification Generation, "
+                    "see https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/cdi-support.html#manual-cdi-specification-generation",
+                    manu,
+                )
+
+        for dev in devices:
+            servers.append(
+                SharableDevicePlugin(
+                    device=dev,
+                    id_by="index" if manu == ManufacturerEnum.ASCEND else "uuid",
+                ),
+            )
+
+    if not servers:
+        logger.warning("No supported devices found. Waiting for stop event...")
+        await stop_event.wait()
+        logger.info("Stop event triggered, shutting down...")
+        return
+
+    # Create tasks to start all servers.
+    serve_tasks = [
+        asyncio.create_task(
+            server.serve(
+                stop_event=stop_event,
+                kubelet_endpoint=kubelet_endpoint,
+                start_timeout=start_timeout,
+                register_timeout=register_timeout,
+            ),
+            name=f"serve-{server.name}",
+        )
+        for server in servers
+    ]
+
+    # Create a task to wait for the stop event.
+    stop_task = asyncio.create_task(
+        stop_event.wait(),
+        name="stop",
+    )
+
+    try:
+        # Wait for stop event or any server to fail.
+        done, pending = await asyncio.wait(
+            [*serve_tasks, stop_task],
+            return_when=asyncio.FIRST_COMPLETED,
+        )
+
+        # Check if stop event was triggered.
+        if stop_task in done:
+            logger.info("Stop event triggered, shutting down servers...")
+            for task in serve_tasks:
+                if not task.done():
+                    with contextlib.suppress(asyncio.CancelledError):
+                        task.cancel()
+            await asyncio.gather(*serve_tasks, return_exceptions=True)
+            logger.info("All servers shut down gracefully")
+            return
+
+        # Otherwise, one or more servers have failed.
+        errs = []
+        for task in pending:
+            with contextlib.suppress(asyncio.CancelledError):
+                task.cancel()
+        for task in done:
+            if err := task.exception():
+                errs.append(err)
+        if errs:
+            logger.error(
+                "One or more servers have failed, shutting down remaining servers...",
+            )
+            raise GroupedError(errs)
+    finally:
+        # Ensure all tasks are cleaned up.
+        for task in serve_tasks:
+            if not task.done():
+                with contextlib.suppress(asyncio.CancelledError):
+                    task.cancel()
+
+
+def serve(
+    manufacturer: ManufacturerEnum | None = None,
+    cdi_generation_output: Path | None = None,
+    kubelet_endpoint: Path | None = None,
+    start_timeout: int = 5,
+    register_timeout: int = 5,
+):
+    """
+    Serve device plugins for all detected devices.
+
+    Args:
+        manufacturer:
+            Manufacturer for serving.
+            If None, detected from current environment.
+        cdi_generation_output:
+            A directory to store generated CDI configuration files.
+            If None, CDI configuration is not stored.
+        kubelet_endpoint:
+            The path to the kubelet endpoint.
+        start_timeout:
+            The timeout in seconds for starting the device plugin.
+        register_timeout:
+            The timeout in seconds for registering the device plugin.
+
+    Raises:
+        If any device plugin fails during serving, the exception is raised.
+
+    """
+    # Ensure we're running in the main thread.
+    if threading.current_thread() != threading.main_thread():
+        logger.warning("Serve should be called from main thread")
+
+    # Create a stop event and register signal handlers.
+    loop = asyncio.get_event_loop()
+    stop_event = asyncio.Event()
+
+    def handle_signal(_s, _f):
+        logger.info("Received termination signal, stopping servers...")
+        stop_event.set()
+
+    # Register signal handlers for graceful shutdown.
+    signals = [signal.SIGINT, signal.SIGTERM]
+    try:
+        for sig in signals:
+            loop.add_signal_handler(sig, handle_signal, sig, None)
+    except (NotImplementedError, RuntimeError):
+        for sig in signals:
+            signal.signal(sig, handle_signal)
+
+    # Run the asynchronous server.
+    try:
+        loop.run_until_complete(
+            serve_async(
+                stop_event=stop_event,
+                manufacturer=manufacturer,
+                cdi_generation_output=cdi_generation_output,
+                kubelet_endpoint=kubelet_endpoint,
+                start_timeout=start_timeout,
+                register_timeout=register_timeout,
+            ),
+        )
+    finally:
+        # Remove signal handlers to avoid side effects.
+        try:
+            for sig in signals:
+                loop.remove_signal_handler(sig)
+        except (NotImplementedError, RuntimeError):
+            pass
+
+
+__all__ = [
+    "cdi_kind_to_kdp_resource",
+    "serve",
+    "serve_async",
+]

gpustack_runtime/deployer/k8s/deviceplugin/__types__.py

@@ -0,0 +1,131 @@
+from __future__ import annotations as __future_annotations__
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import asyncio
+    from pathlib import Path
+
+
+class GroupedError(Exception):
+    """
+    Exception to encapsulate multiple errors.
+
+    Args:
+        errors:
+            A list of exceptions that occurred.
+
+    """
+
+    _errors: list[BaseException | Exception]
+    """
+    Errors that occurred.
+    """
+
+    def __init__(self, errors: list[BaseException | Exception] | None = None):
+        """
+        Initialize the GroupedError exception.
+
+        Args:
+            errors:
+                A list of exceptions that occurred.
+
+        """
+        self._errors = errors or []
+
+    def append(self, error: BaseException | Exception):
+        """
+        Append an error to the GroupedError.
+
+        Args:
+            error:
+                The exception to append.
+
+        """
+        self._errors.append(error)
+
+    def extend(self, errors: list[BaseException | Exception]):
+        """
+        Extend the GroupedError with multiple errors.
+
+        Args:
+            errors:
+                A list of exceptions to extend with.
+
+        """
+        self._errors.extend(errors)
+
+    def __str__(self) -> str:
+        """
+        Get the string representation of the GroupedError exception.
+
+        Returns:
+            A string representation of the GroupedError exception.
+
+        """
+        error_messages = "\n".join(
+            f"{i + 1}. {error!s}" for i, error in enumerate(self._errors)
+        )
+        return f"{len(self._errors)} errors occurred:\n{error_messages}"
+
+
+class PluginServer(ABC):
+    """
+    Base class for Kubernetes device plugins.
+
+    """
+
+    _name: str
+    """
+    Name of the device plugin.
+    """
+
+    def __init__(self, name: str):
+        """
+        Initialize the device plugin.
+
+        Args:
+            name:
+                The name of the device plugin.
+
+        """
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        """
+        Get the name of the device plugin.
+
+        Returns:
+            The name of the device plugin.
+
+        """
+        return self._name
+
+    @abstractmethod
+    async def serve(
+        self,
+        stop_event: asyncio.Event,
+        kubelet_endpoint: Path | None = None,
+        start_timeout: int = 5,
+        register_timeout: int = 5,
+    ):
+        """
+        Serve the device plugin asynchronously.
+
+        Args:
+            stop_event:
+                An asyncio event to signal stopping the server.
+            kubelet_endpoint:
+                The path to the kubelet endpoint.
+                Default is None, which uses the default kubelet socket path.
+            start_timeout:
+                The timeout in seconds for starting the device plugin server.
+                Default is 5 seconds.
+            register_timeout:
+                The timeout in seconds for registering the device plugin.
+                Default is 5 seconds.
+
+        """
+        raise NotImplementedError