gpustack-runtime 0.1.39.post2__py3-none-any.whl → 0.1.40__py3-none-any.whl

gpustack_runtime/deployer/cdi/thead.py

@@ -0,0 +1,103 @@
+from __future__ import annotations as __future_annotations__
+
+from ...detector import (
+    Devices,
+    ManufacturerEnum,
+    detect_devices,
+    filter_devices_by_manufacturer,
+)
+from .__types__ import (
+    Config,
+    ConfigContainerEdits,
+    ConfigDevice,
+    Generator,
+    manufacturer_to_config_kind,
+)
+
+
+class THeadGenerator(Generator):
+    """
+    CDI generator for T-Head devices.
+    """
+
+    def __init__(self):
+        super().__init__(ManufacturerEnum.THEAD)
+
+    def generate(self, devices: Devices | None = None) -> Config | None:
+        """
+        Generate the CDI configuration for T-Head devices.
+
+        Args:
+            devices: The detected devices.
+            If None, all available devices are considered.
+
+        Returns:
+            The Config object, or None if not supported.
+
+        """
+        if devices is None:
+            devices = detect_devices(manufacturer=self.manufacturer)
+        else:
+            devices = filter_devices_by_manufacturer(
+                devices,
+                manufacturer=self.manufacturer,
+            )
+
+        if not devices:
+            return None
+
+        kind = manufacturer_to_config_kind(self.manufacturer)
+        if not kind:
+            return None
+
+        cdi_devices: list[ConfigDevice] = []
+
+        all_container_edits_device_nodes = [
+            "/dev/alixpu",
+            "/dev/alixpu_ctl",
+        ]
+        for dev in devices:
+            if not dev:
+                continue
+            all_container_edits_device_nodes.append(
+                f"/dev/alixpu_ppu{dev.index}",
+            )
+
+            # Add specific container edits for each device
+            cdi_container_edits = ConfigContainerEdits(
+                device_nodes=[
+                    "/dev/alixpu",
+                    "/dev/alixpu_ctl",
+                    f"/dev/alixpu_ppu{dev.index}",
+                ],
+            )
+            cdi_devices.append(
+                ConfigDevice(
+                    name=str(dev.index),
+                    container_edits=cdi_container_edits,
+                ),
+            )
+            cdi_devices.append(
+                ConfigDevice(
+                    name=dev.uuid,
+                    container_edits=cdi_container_edits,
+                ),
+            )
+
+        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,
+                ),
+            ),
+        )
+
+        return Config(
+            kind=kind,
+            devices=cdi_devices,
+        )

gpustack_runtime/deployer/docker.py

@@ -22,7 +22,7 @@ import docker.models.images
 import docker.models.volumes
 import docker.types
 from dataclasses_json import dataclass_json
-from docker.utils import parse_repository_tag
+from gpustack_runner import split_image
 from tqdm import tqdm
 
 from .. import envs
@@ -48,11 +48,13 @@ from .__types__ import (
 )
 from .__utils__ import (
     _MiB,
+    adjust_image_with_envs,
     bytes_to_human_readable,
-    replace_image_with,
+    isexception,
     safe_json,
     sensitive_env_var,
 )
+from .cdi import generate_config as cdi_generate_config
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
@@ -146,16 +148,11 @@ class DockerWorkloadPlan(WorkloadPlan):
         # Default and validate in the base class.
         super().validate_and_default()
 
-        # Adjust default image namespace if needed.
-        if namespace := envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_NAMESPACE:
-            self.pause_image = replace_image_with(
-                image=self.pause_image,
-                namespace=namespace,
-            )
-            self.unhealthy_restart_image = replace_image_with(
-                image=self.unhealthy_restart_image,
-                namespace=namespace,
-            )
+        # Adjust images.
+        self.pause_image = adjust_image_with_envs(self.pause_image)
+        self.unhealthy_restart_image = adjust_image_with_envs(
+            self.unhealthy_restart_image,
+        )
 
 
 @dataclass_json
@@ -330,7 +327,7 @@ class DockerDeployer(EndoscopicDeployer):
         if envs.GPUSTACK_RUNTIME_DEPLOY.lower() not in ("auto", _NAME):
             return supported
 
-        client = DockerDeployer._get_client()
+        client = DockerDeployer._get_client(timeout=3)
         if client:
             try:
                 supported = client.ping()
@@ -340,16 +337,24 @@ class DockerDeployer(EndoscopicDeployer):
                         "Connected to Docker API server: %s",
                         version_info,
                     )
-            except docker.errors.APIError:
-                debug_log_exception(logger, "Failed to connect to Docker API server")
+            except docker.errors.APIError as e:
+                if not isexception(e, FileNotFoundError):
+                    debug_log_exception(
+                        logger,
+                        "Failed to connect to Docker API server",
+                    )
 
         return supported
 
     @staticmethod
-    def _get_client() -> docker.DockerClient | None:
+    def _get_client(**kwargs) -> docker.DockerClient | None:
         """
         Return a Docker client.
 
+        Args:
+            **kwargs:
+                Additional arguments to pass to docker.from_env().
+
         Returns:
             A Docker client if available, None otherwise.
 
@@ -365,9 +370,9 @@ class DockerDeployer(EndoscopicDeployer):
                 os_env = os.environ.copy()
                 if envs.GPUSTACK_RUNTIME_DOCKER_HOST:
                     os_env["DOCKER_HOST"] = envs.GPUSTACK_RUNTIME_DOCKER_HOST
-                client = docker.from_env(environment=os_env)
+                client = docker.from_env(environment=os_env, **kwargs)
         except docker.errors.DockerException as e:
-            if "FileNotFoundError" not in str(e):
+            if not isexception(e, FileNotFoundError):
                 debug_log_exception(logger, "Failed to get Docker client")
 
         return client
@@ -428,8 +433,7 @@ class DockerDeployer(EndoscopicDeployer):
         try:
             logger.info("Pulling image %s", image)
 
-            repo, tag = parse_repository_tag(image)
-            tag = tag or "latest"
+            repo, tag = split_image(image, fill_blank_tag=True)
             auth_config = None
             if (
                 envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_REGISTRY_USERNAME
@@ -840,7 +844,7 @@ class DockerDeployer(EndoscopicDeployer):
             msg = f"Failed to upload ephemeral files to container {container.name}"
             raise OperationError(msg)
 
-    def _create_containers(
+    def _create_containers(  # noqa: C901
         self,
         workload: DockerWorkloadPlan,
         ephemeral_volume_name_mapping: dict[str, str],
@@ -953,7 +957,9 @@ class DockerDeployer(EndoscopicDeployer):
 
                 r_k_runtime_env = workload.resource_key_runtime_env_mapping or {}
                 r_k_backend_env = workload.resource_key_backend_env_mapping or {}
-                vd_env, vd_cdis, vd_values = self.get_visible_devices_values()
+                vd_manus, vd_env, vd_cdis, vd_values = (
+                    self.get_visible_devices_materials()
+                )
                 for r_k, r_v in c.resources.items():
                     match r_k:
                         case "cpu":
@@ -995,6 +1001,14 @@ 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:
+                                for re in runtime_env:
+                                    cdi_generate_config(
+                                        manufacturer=vd_manus[re],
+                                        output=envs.GPUSTACK_RUNTIME_DEPLOY_CDI_SPECS_DIRECTORY,
+                                    )
+
                             # Configure device access environment variable.
                             if r_v == "all" and backend_env:
                                 # Configure privileged if requested all devices.
@@ -1213,8 +1227,12 @@ class DockerDeployer(EndoscopicDeployer):
         self_container_envs: dict[str, str] = dict(
             item.split("=", 1) for item in self_container.attrs["Config"].get("Env", [])
         )
-        self_image_envs: dict[str, str] = dict(
-            item.split("=", 1) for item in self_image.attrs["Config"].get("Env", [])
+        self_image_envs: dict[str, str] = (
+            dict(
+                item.split("=", 1) for item in self_image.attrs["Config"].get("Env", [])
+            )
+            if self_image.attrs["Config"]
+            else {}
         )
         mirrored_envs: dict[str, str] = {
             # Filter out gpustack-internal envs and same-as-image envs.

gpustack_runtime/deployer/kuberentes.py

@@ -319,7 +319,7 @@ class KubernetesDeployer(EndoscopicDeployer):
         if client:
             try:
                 version_api = kubernetes.client.VersionApi(client)
-                version_info = version_api.get_code()
+                version_info = version_api.get_code(_request_timeout=3)
                 supported = version_info is not None
                 if envs.GPUSTACK_RUNTIME_LOG_EXCEPTION:
                     logger.debug(
@@ -337,10 +337,14 @@ class KubernetesDeployer(EndoscopicDeployer):
         return supported
 
     @staticmethod
-    def _get_client() -> kubernetes.client.ApiClient | None:
+    def _get_client(**kwargs) -> kubernetes.client.ApiClient | None:
         """
         Return a Kubernetes API client.
 
+        Args:
+            **kwargs:
+                Additional arguments to pass to the Kubernetes config loader.
+
         Returns:
             A Kubernetes API client if the configuration is valid, None otherwise.
 
@@ -353,7 +357,7 @@ class KubernetesDeployer(EndoscopicDeployer):
                 contextlib.redirect_stdout(dev_null),
                 contextlib.redirect_stderr(dev_null),
             ):
-                kubernetes.config.load_config()
+                kubernetes.config.load_config(**kwargs)
                 client = kubernetes.client.ApiClient()
                 client.user_agent = "gpustack/runtime"
         except kubernetes.config.config_exception.ConfigException:
@@ -989,7 +993,7 @@ class KubernetesDeployer(EndoscopicDeployer):
                 resources: dict[str, str] = {}
                 r_k_runtime_env = workload.resource_key_runtime_env_mapping or {}
                 r_k_backend_env = workload.resource_key_backend_env_mapping or {}
-                vd_env, _, vd_values = self.get_visible_devices_values()
+                _, vd_env, _, vd_values = self.get_visible_devices_materials()
                 for r_k, r_v in c.resources.items():
                     if r_k in ("cpu", "memory"):
                         resources[r_k] = str(r_v)

gpustack_runtime/deployer/podman.py

@@ -23,7 +23,7 @@ import podman.domain.images
 import podman.domain.volumes
 import podman.errors
 from dataclasses_json import dataclass_json
-from podman.api import parse_repository
+from gpustack_runner import split_image
 from podman.domain.containers_create import CreateMixin
 from tqdm import tqdm
 
@@ -51,11 +51,13 @@ from .__types__ import (
 )
 from .__utils__ import (
     _MiB,
+    adjust_image_with_envs,
     bytes_to_human_readable,
-    replace_image_with,
+    isexception,
     safe_json,
     sensitive_env_var,
 )
+from .cdi import generate_config as cdi_generate_config
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
@@ -149,16 +151,11 @@ class PodmanWorkloadPlan(WorkloadPlan):
         # Default and validate in the base class.
         super().validate_and_default()
 
-        # Adjust default image namespace if needed.
-        if namespace := envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_NAMESPACE:
-            self.pause_image = replace_image_with(
-                image=self.pause_image,
-                namespace=namespace,
-            )
-            self.unhealthy_restart_image = replace_image_with(
-                image=self.unhealthy_restart_image,
-                namespace=namespace,
-            )
+        # Adjust images.
+        self.pause_image = adjust_image_with_envs(self.pause_image)
+        self.unhealthy_restart_image = adjust_image_with_envs(
+            self.unhealthy_restart_image,
+        )
 
 
 @dataclass_json
@@ -333,7 +330,7 @@ class PodmanDeployer(EndoscopicDeployer):
         if envs.GPUSTACK_RUNTIME_DEPLOY.lower() not in ("auto", _NAME):
             return supported
 
-        client = PodmanDeployer._get_client()
+        client = PodmanDeployer._get_client(timeout=3)
         if client:
             try:
                 supported = client.ping()
@@ -343,16 +340,24 @@ class PodmanDeployer(EndoscopicDeployer):
                         "Connected to Podman API server: %s",
                         version_info,
                     )
-            except podman.errors.APIError:
-                debug_log_exception(logger, "Failed to connect to Podman API server")
+            except podman.errors.APIError as e:
+                if not isexception(e, FileNotFoundError):
+                    debug_log_exception(
+                        logger,
+                        "Failed to connect to Podman API server",
+                    )
 
         return supported
 
     @staticmethod
-    def _get_client() -> podman.PodmanClient | None:
+    def _get_client(**kwargs) -> podman.PodmanClient | None:
         """
         Return a Podman client.
 
+        Args:
+            **kwargs:
+                Additional arguments to pass to podman.from_env().
+
         Returns:
             A Podman client if available, None otherwise.
 
@@ -368,9 +373,9 @@ class PodmanDeployer(EndoscopicDeployer):
                 os_env = os.environ.copy()
                 if envs.GPUSTACK_RUNTIME_PODMAN_HOST:
                     os_env["CONTAINER_HOST"] = envs.GPUSTACK_RUNTIME_PODMAN_HOST
-                client = podman.from_env(environment=os_env)
+                client = podman.from_env(environment=os_env, **kwargs)
         except podman.errors.DockerException as e:
-            if "FileNotFoundError" not in str(e):
+            if not isexception(e, FileNotFoundError):
                 debug_log_exception(logger, "Failed to get Podman client")
 
         return client
@@ -431,8 +436,7 @@ class PodmanDeployer(EndoscopicDeployer):
         try:
             logger.info("Pulling image %s", image)
 
-            repo, tag = parse_repository(image)
-            tag = tag or "latest"
+            repo, tag = split_image(image, fill_blank_tag=True)
             auth_config = None
             if (
                 envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_REGISTRY_USERNAME
@@ -949,7 +953,9 @@ class PodmanDeployer(EndoscopicDeployer):
             if c.resources:
                 r_k_runtime_env = workload.resource_key_runtime_env_mapping or {}
                 r_k_backend_env = workload.resource_key_backend_env_mapping or {}
-                vd_env, vd_cdis, vd_values = self.get_visible_devices_values()
+                vd_manus, vd_env, vd_cdis, vd_values = (
+                    self.get_visible_devices_materials()
+                )
                 for r_k, r_v in c.resources.items():
                     match r_k:
                         case "cpu":
@@ -991,6 +997,14 @@ class PodmanDeployer(EndoscopicDeployer):
 
                             privileged = create_options.get("privileged", False)
 
+                            # Generate CDI config if not yet.
+                            if envs.GPUSTACK_RUNTIME_DEPLOY_CDI_SPECS_GENERATE:
+                                for re in runtime_env:
+                                    cdi_generate_config(
+                                        manufacturer=vd_manus[re],
+                                        output=envs.GPUSTACK_RUNTIME_DEPLOY_CDI_SPECS_DIRECTORY,
+                                    )
+
                             # Configure device access environment variable.
                             if r_v == "all" and backend_env:
                                 # Configure privileged if requested all devices.
@@ -1189,8 +1203,12 @@ class PodmanDeployer(EndoscopicDeployer):
         self_container_envs: dict[str, str] = dict(
             item.split("=", 1) for item in self_container.attrs["Config"].get("Env", [])
         )
-        self_image_envs: dict[str, str] = dict(
-            item.split("=", 1) for item in self_image.attrs["Config"].get("Env", [])
+        self_image_envs: dict[str, str] = (
+            dict(
+                item.split("=", 1) for item in self_image.attrs["Config"].get("Env", [])
+            )
+            if self_image.attrs["Config"]
+            else {}
         )
         mirrored_envs: dict[str, str] = {
             # Filter out gpustack-internal envs and same-as-image envs.

gpustack_runtime/detector/__init__.py

@@ -24,6 +24,7 @@ from .iluvatar import IluvatarDetector
 from .metax import MetaXDetector
 from .mthreads import MThreadsDetector
 from .nvidia import NVIDIADetector
+from .thead import THeadDetector
 
 logger = logging.getLogger(__package__)
 
@@ -36,6 +37,7 @@ _DETECTORS: list[Detector] = [
     MetaXDetector(),
     MThreadsDetector(),
     NVIDIADetector(),
+    THeadDetector(),
 ]
 """
 List of all detectors.
@@ -60,7 +62,10 @@ def supported_list() -> list[Detector]:
     return [det for det in _DETECTORS if det.is_supported()]
 
 
-def detect_backend(fast: bool = True) -> str | list[str]:
+def detect_backend(
+    fast: bool = True,
+    manufacturer: ManufacturerEnum = None,
+) -> str | list[str]:
     """
     Detect all supported backend.
 
@@ -68,12 +73,21 @@ def detect_backend(fast: bool = True) -> str | list[str]:
         fast:
             If True, return the first detected backend.
             Otherwise, return a list of all detected backends.
+        manufacturer:
+            Manufacturer to filter the detection, implies `fast=True`.
+            If None, detect all available manufacturers.
 
     Returns:
         A string of the detected backend if `fast` is True and a backend is found.
         A list of detected backends if `fast` is False.
 
     """
+    if manufacturer:
+        det = _DETECTORS_MAP.get(manufacturer)
+        if det and det.is_supported():
+            return det.backend
+        return ""
+
     backends: list[str] = []
 
     for det in _DETECTORS:
@@ -88,7 +102,10 @@ def detect_backend(fast: bool = True) -> str | list[str]:
     return backends
 
 
-def detect_devices(fast: bool = True) -> Devices:
+def detect_devices(
+    fast: bool = True,
+    manufacturer: ManufacturerEnum = None,
+) -> Devices:
     """
     Detect all available devices.
 
@@ -96,6 +113,9 @@ def detect_devices(fast: bool = True) -> Devices:
         fast:
             If True, return devices from the first supported detector.
             Otherwise, return devices from all supported detectors.
+        manufacturer:
+            Manufacturer to filter the detection, implies `fast=True`.
+            If None, detect all available manufacturers.
 
     Returns:
         A list of detected devices.
@@ -105,6 +125,18 @@ def detect_devices(fast: bool = True) -> Devices:
         If detection fails for the target detector specified by the `GPUSTACK_RUNTIME_DETECT` environment variable.
 
     """
+    if manufacturer:
+        det = _DETECTORS_MAP.get(manufacturer)
+        if det and det.is_supported():
+            try:
+                return det.detect()
+            except Exception:
+                detect_target = envs.GPUSTACK_RUNTIME_DETECT.lower()
+                if detect_target == det.name:
+                    raise
+                debug_log_exception(logger, "Failed to detect devices for %s", det.name)
+        return []
+
     devices: Devices = []
 
     for det in _DETECTORS:
@@ -128,6 +160,7 @@ def detect_devices(fast: bool = True) -> Devices:
 def get_devices_topologies(
     devices: Devices | None = None,
     fast: bool = True,
+    manufacturer: ManufacturerEnum = None,
 ) -> list[Topology]:
     """
     Get the topology information of the given devices.
@@ -140,6 +173,9 @@ def get_devices_topologies(
             If True, return topologies from the first supported detector.
             Otherwise, return topologies from all supported detectors.
             Only works when `devices` is None.
+        manufacturer:
+            Manufacturer to filter the detection.
+            If None, detect all available manufacturers.
 
     Returns:
         A list of Topology objects for each manufacturer group.
@@ -147,7 +183,7 @@ def get_devices_topologies(
     """
     group = False
     if not devices:
-        devices = detect_devices(fast=fast)
+        devices = detect_devices(fast=fast, manufacturer=manufacturer)
         if not devices:
             return []
         group = True and not fast
@@ -160,6 +196,7 @@ def get_devices_topologies(
 
     # Get topology for each group.
     topologies: list[Topology] = []
+
     for manu, devs in group_devices.items():
         det = _DETECTORS_MAP.get(manu)
         if det is not None:
@@ -172,6 +209,7 @@ def get_devices_topologies(
                 if detect_target == det.name:
                     raise
                 debug_log_exception(logger, "Failed to get topology for %s", det.name)
+
     return topologies
 
 
@@ -197,6 +235,26 @@ def group_devices_by_manufacturer(
     return group_devices
 
 
+def filter_devices_by_manufacturer(
+    devices: Devices | None,
+    manufacturer: ManufacturerEnum,
+) -> Devices:
+    """
+    Filter devices by their manufacturer.
+
+    Args:
+        devices:
+            A list of devices to be filtered.
+        manufacturer:
+            The manufacturer to filter by.
+
+    Returns:
+        A list of devices that match the specified manufacturer.
+
+    """
+    return [dev for dev in devices or [] if dev.manufacturer == manufacturer]
+
+
 __all__ = [
     "Device",
     "Devices",
@@ -205,6 +263,7 @@ __all__ = [
     "backend_to_manufacturer",
     "detect_backend",
     "detect_devices",
+    "filter_devices_by_manufacturer",
     "get_devices_topologies",
     "group_devices_by_manufacturer",
     "manufacturer_to_backend",

gpustack_runtime/detector/__types__.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
+from functools import lru_cache
 from typing import Any
 
 from dataclasses_json import dataclass_json
@@ -45,6 +46,10 @@ class ManufacturerEnum(str, Enum):
     """
     NVIDIA Corporation
     """
+    THEAD = "thead"
+    """
+    T-Head Semiconductor Co., Ltd.
+    """
     UNKNOWN = "unknown"
     """
     Unknown Manufacturer
@@ -63,6 +68,7 @@ _MANUFACTURER_BACKEND_MAPPING: dict[ManufacturerEnum, str] = {
     ManufacturerEnum.METAX: "maca",
     ManufacturerEnum.MTHREADS: "musa",
     ManufacturerEnum.NVIDIA: "cuda",
+    ManufacturerEnum.THEAD: "hggc",
 }
 """
 Mapping of manufacturer to runtime backend,
@@ -70,6 +76,7 @@ which should map to the gpustack-runner's backend names.
 """
 
 
+@lru_cache
 def manufacturer_to_backend(manufacturer: ManufacturerEnum) -> str:
     """
     Convert manufacturer to runtime backend,
@@ -92,6 +99,7 @@ def manufacturer_to_backend(manufacturer: ManufacturerEnum) -> str:
     return ManufacturerEnum.UNKNOWN.value
 
 
+@lru_cache
 def backend_to_manufacturer(backend: str) -> ManufacturerEnum:
     """
     Convert runtime backend to manufacturer,
@@ -449,6 +457,9 @@ class Detector(ABC):
     """
 
     manufacturer: ManufacturerEnum = ManufacturerEnum.UNKNOWN
+    """
+    Manufacturer of the detector.
+    """
 
     @staticmethod
     @abstractmethod

gpustack_runtime/detector/__utils__.py

@@ -951,3 +951,26 @@ def bitmask_to_str(bitmask_list: list) -> str:
         offset += get_bits_size()
 
     return list_to_range_str(sorted(bits_lists))
+
+
+def get_physical_function_by_bdf(bdf: str) -> str:
+    """
+    Get the physical function BDF for a given PCI device BDF address.
+
+    Args:
+        bdf:
+            The PCI device BDF address (e.g., "0000:00:1f.0").
+
+    Returns:
+        The physical function BDF if found, otherwise returns the original BDF.
+
+    """
+    if bdf:
+        with contextlib.suppress(Exception):
+            dev_path = Path(f"/sys/bus/pci/devices/{bdf}")
+            if dev_path.exists():
+                physfn_path = dev_path / "physfn"
+                if physfn_path.exists():
+                    physfn_realpath = physfn_path.resolve()
+                    return physfn_realpath.name
+    return bdf