gpustack-runtime 0.1.41__py3-none-any.whl → 0.1.41.post2__py3-none-any.whl

gpustack_runtime/_version.py

@@ -27,8 +27,8 @@ version_tuple: VERSION_TUPLE
 __commit_id__: COMMIT_ID
 commit_id: COMMIT_ID
 
-__version__ = version = '0.1.41'
-__version_tuple__ = version_tuple = (0, 1, 41)
+__version__ = version = '0.1.41.post2'
+__version_tuple__ = version_tuple = (0, 1, 41, 'post2')
 try:
     from ._version_appendix import git_commit
     __commit_id__ = commit_id = git_commit

gpustack_runtime/_version_appendix.py

@@ -1 +1 @@
-git_commit = "f706cfd"
+git_commit = "ed331dc"

gpustack_runtime/deployer/cdi/__init__.py

@@ -151,6 +151,7 @@ def available_backends() -> list[str]:
 
 
 __all__ = [
+    "Config",
     "available_backends",
     "available_manufacturers",
     "dump_config",

gpustack_runtime/deployer/cdi/__utils__.py

@@ -68,6 +68,8 @@ def linux_device_from_path(path: Path | str | None) -> LinuxDevice | None:
         return None
 
     path_stat = path.lstat()
+    if not path_stat:
+        return None
 
     dev_mode = stat.S_IFMT(path_stat.st_mode)
     match dev_mode:

gpustack_runtime/deployer/k8s/deviceplugin/__init__.py

@@ -4,12 +4,16 @@ import asyncio
 import contextlib
 import logging
 import signal
+import stat
 import threading
-from typing import TYPE_CHECKING
+from functools import lru_cache
+from typing import TYPE_CHECKING, Literal
 
 from .... import envs
 from ....deployer.cdi import dump_config as cdi_dump_config
 from ....detector import ManufacturerEnum, detect_devices, supported_manufacturers
+from ...cdi import Config, manufacturer_to_cdi_kind
+from ..types.kubelet.deviceplugin.v1beta1 import KubeletSocket
 from .__types__ import GroupedError, PluginServer
 from .plugin import SharableDevicePlugin, cdi_kind_to_kdp_resource
 
@@ -61,9 +65,19 @@ async def serve_async(
         if not devices:
             continue
 
+        allocation_policy = _get_device_allocation_policy(manu)
+        logger.info(
+            "Using device allocation policy '%s' for manufacturer '%s'",
+            allocation_policy,
+            manu,
+        )
+
         # 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:
+        if (
+            allocation_policy == "cdi"
+            and envs.GPUSTACK_RUNTIME_KUBERNETES_KDP_CDI_SPECS_GENERATE
+        ):
             generated_content, generated_path = cdi_dump_config(
                 manufacturer=manu,
                 output=cdi_generation_output,
@@ -95,6 +109,8 @@ async def serve_async(
                 SharableDevicePlugin(
                     device=dev,
                     id_by="index" if manu == ManufacturerEnum.ASCEND else "uuid",
+                    allocation_policy=allocation_policy,
+                    max_allocations=envs.GPUSTACK_RUNTIME_KUBERNETES_KDP_PER_DEVICE_MAX_ALLOCATIONS,
                 ),
             )
 
@@ -104,6 +120,9 @@ async def serve_async(
         logger.info("Stop event triggered, shutting down...")
         return
 
+    if not kubelet_endpoint:
+        kubelet_endpoint = KubeletSocket
+
     # Create tasks to start all servers.
     serve_tasks = [
         asyncio.create_task(
@@ -233,8 +252,75 @@ def serve(
             pass
 
 
+def is_kubelet_socket_accessible(
+    kubelet_endpoint: Path | None = None,
+) -> bool:
+    """
+    Check if the kubelet socket is accessible.
+
+    Args:
+        kubelet_endpoint:
+            The path to the kubelet endpoint.
+
+    Returns:
+        True if the socket is accessible, False otherwise.
+
+    """
+    if not kubelet_endpoint:
+        kubelet_endpoint = KubeletSocket
+
+    if kubelet_endpoint.exists():
+        path_stat = kubelet_endpoint.lstat()
+        if path_stat and stat.S_ISSOCK(path_stat.st_mode):
+            return True
+    return False
+
+
+@lru_cache
+def _get_device_allocation_policy(
+    manufacturer: ManufacturerEnum,
+) -> Literal["env", "cdi", "opaque"]:
+    """
+    Get the device allocation policy (in lowercase) for the device plugin.
+
+    Args:
+        manufacturer:
+            The manufacturer of the device.
+
+    Returns:
+        The device allocation policy.
+
+    """
+    policy = envs.GPUSTACK_RUNTIME_KUBERNETES_KDP_DEVICE_ALLOCATION_POLICY.lower()
+    if policy != "auto":
+        return policy
+
+    cdi_kind = manufacturer_to_cdi_kind(manufacturer)
+
+    cdi_dir = envs.GPUSTACK_RUNTIME_DEPLOY_CDI_SPECS_DIRECTORY
+    for suffix in ["*.yaml", "*.yml", "*.json"]:
+        for file in cdi_dir.glob(suffix):
+            with contextlib.suppress(Exception):
+                config = Config.from_file(file)
+                if config and config.kind == cdi_kind:
+                    return "cdi"
+
+    if manufacturer in [
+        ManufacturerEnum.AMD,
+        # ManufacturerEnum.ASCEND, # Prioritize using Env policy for Ascend.
+        ManufacturerEnum.HYGON,
+        ManufacturerEnum.ILUVATAR,
+        ManufacturerEnum.METAX,
+        ManufacturerEnum.THEAD,
+    ]:
+        return "opaque"
+
+    return "env"
+
+
 __all__ = [
     "cdi_kind_to_kdp_resource",
+    "is_kubelet_socket_accessible",
     "serve",
     "serve_async",
 ]

gpustack_runtime/deployer/k8s/deviceplugin/__types__.py

@@ -107,7 +107,7 @@ class PluginServer(ABC):
     async def serve(
         self,
         stop_event: asyncio.Event,
-        kubelet_endpoint: Path | None = None,
+        kubelet_endpoint: Path,
         start_timeout: int = 5,
         register_timeout: int = 5,
     ):

gpustack_runtime/deployer/k8s/deviceplugin/plugin.py

@@ -11,7 +11,6 @@ import grpc
 from grpc_interceptor import AsyncServerInterceptor
 from grpc_interceptor.exceptions import GrpcException
 
-from .... import envs
 from ....detector import Device, str_range_to_list
 from ...cdi import (
     generate_config,
@@ -31,7 +30,6 @@ from ..types.kubelet.deviceplugin.v1beta1 import (
     DeviceSpec,
     Empty,
     Healthy,
-    KubeletSocket,
     ListAndWatchResponse,
     Mount,
     NUMANode,
@@ -110,6 +108,10 @@ class SharableDevicePlugin(PluginServer, DevicePluginServicer):
     """
     Controls how the device IDs of the Kubernetes Device Plugin are generated.
     """
+    _allocation_policy: Literal["env", "cdi", "opaque"]
+    """
+    Controls the device allocation policy.
+    """
     _max_allocations: int
     """
     Controls the maximum shards per underlying device.
@@ -131,7 +133,8 @@ class SharableDevicePlugin(PluginServer, DevicePluginServicer):
         self,
         device: Device,
         id_by: Literal["uuid", "index"] = "uuid",
-        max_allocations: int | None = None,
+        allocation_policy: Literal["env", "cdi", "opaque"] = "cdi",
+        max_allocations: int | None = 10,
     ):
         """
         Initializes the SharableDevicePlugin.
@@ -142,19 +145,16 @@ class SharableDevicePlugin(PluginServer, DevicePluginServicer):
             id_by:
                 Controls how the device IDs of the Kubernetes Device Plugin are generated.
                 Either "uuid" or "index". Default is "uuid".
+            allocation_policy:
+                Controls the device allocation policy.
             max_allocations:
                 Controls the maximum allocations per underlying device.
-                If None, uses the environment variable `GPUSTACK_RUNTIME_KUBERNETES_KDP_PER_DEVICE_MAX_ALLOCATIONS`.
 
         """
         self._device = device
         self._id_by = id_by
-        self._max_allocations = max_allocations
-        if not self._max_allocations:
-            self._max_allocations = (
-                envs.GPUSTACK_RUNTIME_KUBERNETES_KDP_PER_DEVICE_MAX_ALLOCATIONS
-            )
-        self._max_allocations = max(self._max_allocations, 1)
+        self._allocation_policy = allocation_policy
+        self._max_allocations = max(max_allocations, 1)
         self._cdi_kind = manufacturer_to_cdi_kind(device.manufacturer)
         self._runtime_env = manufacturer_to_runtime_env(device.manufacturer)
         self._kdp_resource = cdi_kind_to_kdp_resource(
@@ -229,7 +229,7 @@ class SharableDevicePlugin(PluginServer, DevicePluginServicer):
     async def serve(
         self,
         stop_event: asyncio.Event,
-        kubelet_endpoint: Path | None = None,
+        kubelet_endpoint: Path,
         start_timeout: int = 5,
         register_timeout: int = 5,
     ):
@@ -247,9 +247,6 @@ class SharableDevicePlugin(PluginServer, DevicePluginServicer):
                 The timeout in seconds for registering the device plugin.
 
         """
-        if not kubelet_endpoint:
-            kubelet_endpoint = KubeletSocket
-
         resource_name = self._kdp_resource
         endpoint = kubelet_endpoint.parent / f"{resource_name.replace('/', '.')}.sock"
 
@@ -421,8 +418,7 @@ class SharableDevicePlugin(PluginServer, DevicePluginServicer):
         self,
         req: ContainerAllocateRequest,
     ) -> ContainerAllocateResponse:
-        policy = envs.GPUSTACK_RUNTIME_KUBERNETES_KDP_DEVICE_ALLOCATION_POLICY.lower()
-
+        policy = self._allocation_policy
         request_dp_device_ids = req.devices_ids
 
         # CDI device allocation.