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

gpustack_runtime/__main__.py

@@ -7,10 +7,12 @@ import time
 from argparse import ArgumentParser
 
 import argcomplete
+from gpustack_runner.cmds import LoadImagesSubCommand
 
 from . import deployer, detector
 from ._version import commit_id, version
 from .cmds import (
+    CDIGenerateSubCommand,
     CopyImagesSubCommand,
     CreateWorkloadSubCommand,
     DeleteWorkloadsSubCommand,
@@ -71,12 +73,14 @@ def main():
     InspectWorkloadSubCommand.register(subcommand_parser)
     DetectDevicesSubCommand.register(subcommand_parser)
     GetDevicesTopologySubCommand.register(subcommand_parser)
-    ListImagesSubCommand.register(subcommand_parser)
-    SaveImagesSubCommand.register(subcommand_parser)
-    CopyImagesSubCommand.register(subcommand_parser)
     LogsSelfSubCommand.register(subcommand_parser)
     ExecSelfSubCommand.register(subcommand_parser)
     InspectSelfSubCommand.register(subcommand_parser)
+    CDIGenerateSubCommand.register(subcommand_parser)
+    ListImagesSubCommand.register(subcommand_parser)
+    SaveImagesSubCommand.register(subcommand_parser)
+    LoadImagesSubCommand.register(subcommand_parser)
+    CopyImagesSubCommand.register(subcommand_parser)
 
     # Autocomplete
     argcomplete.autocomplete(parser)

gpustack_runtime/_version.py

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

gpustack_runtime/_version_appendix.py

@@ -1 +1 @@
-git_commit = "e044bab"
+git_commit = "1f4627e"

gpustack_runtime/cmds/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from .deployer import (
+    CDIGenerateSubCommand,
     CreateWorkloadSubCommand,
     DeleteWorkloadsSubCommand,
     DeleteWorkloadSubCommand,
@@ -24,6 +25,7 @@ from .images import (
 )
 
 __all__ = [
+    "CDIGenerateSubCommand",
     "CopyImagesSubCommand",
     "CreateWorkloadSubCommand",
     "DeleteWorkloadSubCommand",

gpustack_runtime/cmds/deployer.py

@@ -5,7 +5,7 @@ import json
 import os
 import sys
 import time
-from argparse import REMAINDER
+from argparse import OPTIONAL, REMAINDER
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -25,6 +25,7 @@ from ..deployer import (
     WorkloadStatusStateEnum,
     async_logs_self,
     async_logs_workload,
+    cdi_generate_config,
     create_workload,
     delete_workload,
     exec_self,
@@ -34,7 +35,7 @@ from ..deployer import (
     inspect_workload,
     list_workloads,
 )
-from ..detector import supported_backends
+from ..detector import supported_backends, supported_manufacturers
 from .__types__ import SubCommand
 
 if TYPE_CHECKING:
@@ -92,6 +93,7 @@ class CreateWorkloadSubCommand(SubCommand):
     command_script: str | None
     port: int
     host_network: bool
+    privileged: bool
     check: bool
     namespace: str
     name: str
@@ -133,11 +135,19 @@ class CreateWorkloadSubCommand(SubCommand):
 
         deploy_parser.add_argument(
             "--host-network",
+            "--network-host",
             action="store_true",
             help="Use host network (default: False)",
             default=False,
         )
 
+        deploy_parser.add_argument(
+            "--privileged",
+            action="store_true",
+            help="Run the container in privileged mode (default: False)",
+            default=False,
+        )
+
         deploy_parser.add_argument(
             "--check",
             action="store_true",
@@ -183,6 +193,7 @@ class CreateWorkloadSubCommand(SubCommand):
         self.command_script = None
         self.port = args.port
         self.host_network = args.host_network
+        self.privileged = args.privileged
         self.check = args.check
         self.namespace = args.namespace
         self.name = args.name
@@ -237,6 +248,7 @@ class CreateWorkloadSubCommand(SubCommand):
         execution = ContainerExecution(
             command_script=self.command_script,
             args=self.extra_args,
+            privileged=self.privileged,
         )
         ports = (
             [
@@ -945,6 +957,76 @@ class InspectSelfSubCommand(SubCommand):
         print(inspect_self())
 
 
+class CDIGenerateSubCommand(SubCommand):
+    """
+    Command to generate CDI configurations.
+    """
+
+    format: str
+    output: Path | None
+
+    @staticmethod
+    def register(parser: _SubParsersAction):
+        cdi_parser = parser.add_parser(
+            "cdi-generate",
+            help="Generate CDI configurations according to the current environment",
+            aliases=["cdi-gen"],
+        )
+
+        cdi_parser.add_argument(
+            "--format",
+            type=str,
+            choices=["yaml", "json"],
+            default="yaml",
+            help="Format of the CDI configurations",
+        )
+
+        cdi_parser.add_argument(
+            "output",
+            nargs=OPTIONAL,
+            help="Output directory to save CDI configurations (default: current directory)",
+        )
+
+        cdi_parser.set_defaults(func=CDIGenerateSubCommand)
+
+    def __init__(self, args: Namespace):
+        self.format = args.format
+        self.output = Path(args.output) if args.output else None
+
+        if self.output:
+            try:
+                if not self.output.exists():
+                    self.output.mkdir(parents=True, exist_ok=True)
+            except OSError as e:
+                msg = f"Failed to prepare output directory '{self.output}' for CDI configurations"
+                raise RuntimeError(msg) from e
+
+            if not self.output.is_dir():
+                msg = f"The output path '{self.output}' is not a directory"
+                raise RuntimeError(msg)
+
+    def run(self):
+        print("\033[2J\033[H", end="")
+
+        generated = False
+        for manu in supported_manufacturers():
+            content, path = cdi_generate_config(
+                manufacturer=manu,
+                output=self.output,
+            )
+            if content:
+                generated = True
+                if path:
+                    print(f"Generated CDI configuration for '{manu}' at {path}:\n")
+                else:
+                    print(f"Generated CDI configuration for '{manu}':\n")
+                print(content)
+                print()
+
+        if not generated:
+            print("No CDI configurations were generated.")
+
+
 def format_workloads_json(sts: list[WorkloadStatus]) -> str:
     return json.dumps([st.to_dict() for st in sts], indent=2)
 

gpustack_runtime/cmds/images.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from gpustack_runner.cmds import (
     CopyImagesSubCommand,
     ListImagesSubCommand,
+    LoadImagesSubCommand,
     PlatformedImage,
     SaveImagesSubCommand,
     append_images,
@@ -20,6 +21,7 @@ append_images(
 __all__ = [
     "CopyImagesSubCommand",
     "ListImagesSubCommand",
+    "LoadImagesSubCommand",
     "PlatformedImage",
     "SaveImagesSubCommand",
     "append_images",

gpustack_runtime/deployer/__init__.py

@@ -31,6 +31,7 @@ from .__types__ import (
     WorkloadStatus,
     WorkloadStatusStateEnum,
 )
+from .cdi import generate_config as cdi_generate_config
 from .docker import (
     DockerDeployer,
     DockerWorkloadPlan,
@@ -602,6 +603,7 @@ __all__ = [
     "WorkloadStatusStateEnum",
     "async_logs_self",
     "async_logs_workload",
+    "cdi_generate_config",
     "create_workload",
     "delete_workload",
     "exec_self",

gpustack_runtime/deployer/__types__.py

@@ -14,6 +14,7 @@ from dataclasses_json import dataclass_json
 
 from .. import envs
 from ..detector import (
+    ManufacturerEnum,
     Topology,
     detect_devices,
     get_devices_topologies,
@@ -21,6 +22,7 @@ from ..detector import (
     manufacturer_to_backend,
 )
 from .__utils__ import (
+    adjust_image_with_envs,
     correct_runner_image,
     fnv1a_32_hex,
     fnv1a_64_hex,
@@ -1018,20 +1020,8 @@ class WorkloadPlan(WorkloadSecurity):
                 c.files.append(command_script)
                 c.execution.command = [command_script_name]  # Override command.
                 c.execution.command_script = None
-            # Add default registry if needed.
-            if (
-                envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_REGISTRY
-                and envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_REGISTRY
-                not in ["docker.io", "index.docker.io"]
-            ):
-                image_registry = (
-                    envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_NAMESPACE
-                )
-                image_split = c.image.split("/")
-                if len(image_split) == 1:
-                    c.image = f"{image_registry}/library/{c.image}"
-                elif len(image_split) == 2:
-                    c.image = f"{image_registry}/{c.image}"
+            # Adjust images.
+            c.image = adjust_image_with_envs(c.image)
             # Correct runner image if needed.
             if envs.GPUSTACK_RUNTIME_DEPLOY_CORRECT_RUNNER_IMAGE:
                 c.image, ok = correct_runner_image(c.image)
@@ -1279,6 +1269,17 @@ class Deployer(ABC):
     """
     Thread pool for the deployer.
     """
+    _visible_devices_manufacturers: dict[str, ManufacturerEnum] | None = None
+    """
+    Recorded visible devices manufacturers,
+    the key is the runtime visible devices env name,
+    the value is the corresponding manufacturer.
+    For example:
+    {
+        "NVIDIA_VISIBLE_DEVICES": ManufacturerEnum.NVIDIA,
+        "AMD_VISIBLE_DEVICES": ManufacturerEnum.AMD
+    }.
+    """
     _visible_devices_env: dict[str, list[str]] | None = None
     """
     Recorded visible devices envs,
@@ -1358,14 +1359,16 @@ class Deployer(ABC):
 
     def _prepare(self):
         """
-        Detect devices once, and construct critical elements for post processing, including:
+        Detect devices once, and construct critical elements for post-processing, including:
+        - Prepare visible devices manufacturers mapping.
         - Prepare visible devices environment variables mapping.
         - Prepare visible devices values mapping.
-        - Prepare topology.
+        - Prepare visible devices topologies mapping.
         """
-        if self._visible_devices_env:
+        if self._visible_devices_manufacturers is not None:
             return
 
+        self._visible_devices_manufacturers = {}
         self._visible_devices_env = {}
         self._visible_devices_cdis = {}
         self._visible_devices_values = {}
@@ -1379,15 +1382,19 @@ class Deployer(ABC):
         if group_devices:
             for manu, devs in group_devices.items():
                 backend = manufacturer_to_backend(manu)
-                rk = envs.GPUSTACK_RUNTIME_DETECT_BACKEND_MAP_RESOURCE_KEY.get(backend)
+                resource_key = (
+                    envs.GPUSTACK_RUNTIME_DETECT_BACKEND_MAP_RESOURCE_KEY.get(backend)
+                )
+                if resource_key is None:
+                    continue
                 ren = envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES.get(
-                    rk,
+                    resource_key,
                 )
                 ben_list = envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_BACKEND_VISIBLE_DEVICES.get(
-                    rk,
+                    resource_key,
                 )
-                cdi = envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_CONTAINER_DEVICE_INTERFACES.get(
-                    rk,
+                cdi = envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_CDI.get(
+                    resource_key,
                 )
                 if ren and ben_list:
                     valued_uuid = (
@@ -1401,6 +1408,8 @@ class Deployer(ABC):
                         dev_uuids.append(dev.uuid)
                         dev_indexes.append(str(dev.index))
                         dev_indexes_alignment[str(dev.index)] = str(dev_i)
+                    # Map runtime visible devices env <-> manufacturer.
+                    self._visible_devices_manufacturers[ren] = manu
                     # Map runtime visible devices env <-> backend visible devices env list.
                     self._visible_devices_env[ren] = ben_list
                     # Map runtime visible devices env <-> CDI key.
@@ -1433,16 +1442,28 @@ class Deployer(ABC):
                 return
 
         # Fallback to unknown backend
-        self._visible_devices_env["UNKNOWN_RUNTIME_VISIBLE_DEVICES"] = []
-        self._visible_devices_values["UNKNOWN_RUNTIME_VISIBLE_DEVICES"] = ["all"]
+        ren = "UNKNOWN_RUNTIME_VISIBLE_DEVICES"
+        self._visible_devices_manufacturers[ren] = ManufacturerEnum.UNKNOWN
+        self._visible_devices_env[ren] = []
+        self._visible_devices_cdis[ren] = "unknown/devices"
+        self._visible_devices_values[ren] = ["all"]
 
-    def get_visible_devices_values(
+    def get_visible_devices_materials(
         self,
-    ) -> (dict[str, list[str]], dict[str, str], dict[str, list[str]]):
+    ) -> (
+        dict[str, ManufacturerEnum],
+        dict[str, list[str]],
+        dict[str, str],
+        dict[str, list[str]],
+    ):
         """
         Return the visible devices environment variables, cdis and values mappings.
         For example:
         (
+            {
+                "NVIDIA_VISIBLE_DEVICES": ManufacturerEnum.NVIDIA,
+                "AMD_VISIBLE_DEVICES": ManufacturerEnum.AMD
+            },
             {
                 "NVIDIA_VISIBLE_DEVICES": ["CUDA_VISIBLE_DEVICES"],
                 "AMD_VISIBLE_DEVICES": ["HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES"]
@@ -1458,10 +1479,12 @@ class Deployer(ABC):
         ).
 
         Returns:
-            A tuple of two dictionaries:
+            A tuple of four dictionaries:
             - The first dictionary maps runtime visible devices environment variable names
-              to lists of backend visible devices environment variable names.
+              to corresponding manufacturers.
             - The second dictionary maps runtime visible devices environment variable names
+              to lists of backend visible devices environment variable names.
+            - The third dictionary maps runtime visible devices environment variable names
               to corresponding CDI keys.
             - The last dictionary maps runtime visible devices environment variable names
               to lists of device indexes or UUIDs.
@@ -1469,6 +1492,7 @@ class Deployer(ABC):
         """
         self._prepare()
         return (
+            self._visible_devices_manufacturers,
             self._visible_devices_env,
             self._visible_devices_cdis,
             self._visible_devices_values,

gpustack_runtime/deployer/__utils__.py

@@ -6,11 +6,18 @@ import json
 import platform
 import re
 from functools import lru_cache
+from pathlib import Path
 from typing import Any
 
 import yaml
-from gpustack_runner import DockerImage, list_backend_runners
+from gpustack_runner import (
+    DockerImage,
+    list_backend_runners,
+    parse_image,
+    replace_image_with,
+)
 
+from .. import envs
 from ..detector import backend_to_manufacturer, detect_backend, detect_devices
 from ..detector.ascend import get_ascend_cann_variant
 
@@ -338,6 +345,49 @@ def safe_yaml(obj: Any, **kwargs) -> str:
     return yaml.dump(dict_data, **kwargs)
 
 
+def load_yaml_or_json(path: str | Path) -> list[dict] | dict:
+    """
+    Load a YAML or JSON string into to a dict.
+
+    Args:
+        path:
+            The path to the CDI configuration file.
+
+    Returns:
+        The loaded dict.
+
+    """
+    if isinstance(path, str):
+        path = Path(path)
+
+    if not path.exists():
+        msg = f"File not found: {path}"
+        raise FileNotFoundError(msg)
+
+    content = path.read_text(encoding="utf-8")
+
+    if path.suffix in {".yaml", ".yml"}:
+        try:
+            ret = list(yaml.safe_load_all(content))
+        except yaml.YAMLError as e:
+            msg = f"Failed to parse YAML file: {path}"
+            raise RuntimeError(msg) from e
+        else:
+            if len(ret) == 1:
+                return ret[0]
+            return ret
+
+    if path.suffix == ".json":
+        try:
+            return json.loads(content)
+        except json.JSONDecodeError as e:
+            msg = f"Failed to parse JSON file: {path}"
+            raise RuntimeError(msg) from e
+
+    msg = f"Unsupported file format: {path.suffix}"
+    raise RuntimeError(msg)
+
+
 @lru_cache
 def compare_versions(v1: str | None, v2: str | None) -> int:
     """
@@ -672,140 +722,77 @@ def bytes_to_human_readable(size_in_bytes: int) -> str:
     return f"{size_in_bytes} B"
 
 
-def replace_image_with(
-    image: str,
-    registry: str | None = None,
-    namespace: str | None = None,
-    repository: str | None = None,
-) -> str:
+def sensitive_env_var(name: str) -> bool:
     """
-    Replace the registry, namespace, and repository of a Docker image string.
-
-    The given image string is parsed into its components (registry, namespace, repository, tag),
-    and the specified components are replaced with the provided values.
-
-    The format of a Docker image string is:
-    [registry/][namespace/]repository[:tag|@digest]
+    Check if the given environment variable name is considered sensitive.
 
     Args:
-        image:
-            The original Docker image string.
-        registry:
-            The new registry to use. If None, keep the original registry.
-        namespace:
-            The new namespace to use. If None, keep the original namespace.
-        repository:
-            The new repository to use. If None, keep the original repository.
+        name:
+            The environment variable name to check.
 
     Returns:
-        The modified Docker image string.
+        True if the name is considered sensitive, False otherwise.
 
     """
-    if not image or (not registry and not namespace and not repository):
-        return image
-
-    registry = registry.strip() if registry else None
-    namespace = namespace.strip() if namespace else None
-    repository = repository.strip() if repository else None
-
-    image_reg, image_ns, image_repo, image_tag = (
-        None,
-        None,
-        None,
-        None,
-    )
-    image_rest = image.strip()
-
-    # Get tag.
-    parts = image_rest.rsplit("@", maxsplit=1)
-    if len(parts) == 2:
-        image_rest, image_tag = parts
-    else:
-        parts = image_rest.rsplit(":", maxsplit=1)
-        if len(parts) == 2 and "/" not in parts[1]:
-            image_rest, image_tag = parts
-    if not image_rest:
-        return image
-
-    # Get repository.
-    parts = image_rest.rsplit("/", maxsplit=1)
-    if len(parts) == 2:
-        image_rest, image_repo = parts
-    else:
-        image_rest, image_repo = None, image_rest
-
-    # Get namespace.
-    if image_rest:
-        parts = image_rest.rsplit("/", maxsplit=1)
-        if len(parts) == 2:
-            image_reg, image_ns = parts
-        else:
-            image_reg, image_ns = None, image_rest
-
-    return make_image_with(
-        repository=repository or image_repo,
-        registry=registry or image_reg,
-        namespace=namespace or image_ns,
-        tag=image_tag,
-    )
+    return name.lower().endswith(_SENSITIVE_ENVS_SUFFIX)
 
 
-def make_image_with(
-    repository: str,
-    registry: str | None = None,
-    namespace: str | None = None,
-    tag: str | None = None,
-) -> str:
+def isexception(
+    e: Exception,
+    target_exception_types: type[Exception] | tuple[type[Exception], ...],
+) -> bool:
     """
-    Make a Docker image string from the given registry, namespace, repository, and tag.
-
-    The format of a Docker image string is:
-    [registry/][namespace/]repository[:tag|@digest]
+    Check if the given exception is caused by any of the target exception types.
 
     Args:
-        repository:
-            The repository name.
-        registry:
-            The registry to use. If None, no registry will be used.
-        namespace:
-            The namespace to use. If None, no namespace will be used.
-        tag:
-            The tag to use. If None, no tag will be used.
+        e:
+            The exception to check.
+        target_exception_types:
+            A tuple of target exception types.
 
     Returns:
-        The Docker image string.
+        True if the exception is caused by any of the target exception types, False otherwise.
 
     """
-    if not repository or (not registry and not namespace and not tag):
-        return repository
+    if isinstance(e, target_exception_types):
+        return True
 
-    image = ""
-    if registry:
-        image += f"{registry}/"
-    if namespace:
-        image += f"{namespace}/"
-    elif registry:
-        image += "library/"
-    image += repository
-    if not tag:
-        return image
-    if tag.startswith("sha256:"):
-        image += f"@{tag}"
-    else:
-        image += f":{tag}"
-    return image
+    cause = getattr(e, "__cause__", None)
+    if cause and isexception(cause, target_exception_types):
+        return True
 
+    context = getattr(e, "__context__", None)
+    return bool(context and isexception(context, target_exception_types))
 
-def sensitive_env_var(name: str) -> bool:
+
+@lru_cache
+def adjust_image_with_envs(image: str) -> str:
     """
-    Check if the given environment variable name is considered sensitive.
+    Replace the registry and namespace of the given image
+    with the default ones from environment variables if applicable.
 
     Args:
-        name:
-            The environment variable name to check.
+        image:
+            The image to replace.
 
     Returns:
-        True if the name is considered sensitive, False otherwise.
+        The replaced image.
 
     """
-    return name.lower().endswith(_SENSITIVE_ENVS_SUFFIX)
+    original_reg, original_ns, _, _ = parse_image(image)
+
+    target_reg = original_reg or envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_REGISTRY
+    target_ns = (
+        original_ns
+        if original_ns != "gpustack"
+        else envs.GPUSTACK_RUNTIME_DEPLOY_DEFAULT_CONTAINER_NAMESPACE
+    )
+
+    if original_reg == target_reg and original_ns == target_ns:
+        return image
+
+    return replace_image_with(
+        image,
+        registry=target_reg,
+        namespace=target_ns,
+    )