diff --git a/gpustack_runtime/deployer/__init__.py b/gpustack_runtime/deployer/__init__.py index d8c6c41..6665ccd 100644 --- a/gpustack_runtime/deployer/__init__.py +++ b/gpustack_runtime/deployer/__init__.py @@ -36,6 +36,7 @@ DockerWorkloadPlan, DockerWorkloadStatus, ) +from .k8s.devicemanager import KubernetesResourceInjectionPolicyEnum from .kubernetes import ( KubernetesDeployer, KubernetesWorkloadPlan, @@ -595,6 +596,7 @@ def inspect_self() -> str: "ContainerSecurity", "DockerWorkloadPlan", "DockerWorkloadStatus", + "KubernetesResourceInjectionPolicyEnum", "KubernetesWorkloadPlan", "KubernetesWorkloadStatus", "OperationError", diff --git a/gpustack_runtime/deployer/k8s/devicemanager/__init__.py b/gpustack_runtime/deployer/k8s/devicemanager/__init__.py index 9265080..d6b9772 100644 --- a/gpustack_runtime/deployer/k8s/devicemanager/__init__.py +++ b/gpustack_runtime/deployer/k8s/devicemanager/__init__.py @@ -1,6 +1,7 @@ from __future__ import annotations as __future_annotations__ import stat +from enum import Enum from functools import lru_cache from pathlib import Path from typing import TYPE_CHECKING, Any, Literal @@ -11,6 +12,36 @@ from collections.abc import Callable, Mapping +class KubernetesResourceInjectionPolicyEnum(str, Enum): + """ + Device injection policies for Kubernetes workloads. + """ + + AUTO = "Auto" + """ + Detect the policy from the target node's allocatable resources. + """ + ENV = "Env" + """ + Inject runtime-visible device environment variables without device-plugin requests. + """ + KDP = "KDP" + """ + Request device-plugin resources. + """ + + @classmethod + def _missing_(cls, value: object) -> KubernetesResourceInjectionPolicyEnum | None: + if isinstance(value, str): + for member in cls: + if member.value.lower() == value.lower(): + return member + return None + + def __str__(self) -> str: + return self.value + + def is_kubelet_socket_accessible( kubelet_endpoint: Path | None = None, ) -> bool: @@ -79,13 +110,16 @@ def node_has_device_plugin_resources( def get_resource_injection_policy( probe_node_allocatable: Callable[[], Mapping[str, Any] | None] | None = None, + *, + policy: KubernetesResourceInjectionPolicyEnum | None = None, ) -> Literal["env", "kdp"]: """ Get the resource injection policy (in lowercase) for the deployer. - An explicit policy always wins. Under "auto" the decision belongs to the - cluster, not to the process doing the deploying: the Kubernetes deployer - orchestrates remotely, so whether *it* can reach a kubelet socket says + A workload policy overrides the environment setting. Under "auto" the + decision belongs to the cluster, not to the process doing the deploying: + the Kubernetes deployer orchestrates remotely, so whether *it* can reach + a kubelet socket says nothing about whether the *target* node runs a device plugin. So the probe reads that node's allocatable resources and looks for a device-plugin resource family there. @@ -99,14 +133,29 @@ def get_resource_injection_policy( probe_node_allocatable: Called only under the "auto" policy, to read the target node's allocatable resources. Returns None when the node cannot be read. + policy: + Workload policy override. None inherits the environment setting; + Auto probes the node even if the environment selects Env or KDP. + + Raises: + ValueError: + If the policy is not Auto, Env or KDP. Returns: The resource injection policy. """ - policy = envs.GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY.lower() - if policy != "auto": - return policy + if policy is None: + policy = envs.GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY + try: + policy = KubernetesResourceInjectionPolicyEnum(policy) + except ValueError as e: + msg = f"Invalid resource injection policy '{policy}'" + raise ValueError(msg) from e + if policy == KubernetesResourceInjectionPolicyEnum.ENV: + return "env" + if policy == KubernetesResourceInjectionPolicyEnum.KDP: + return "kdp" if probe_node_allocatable is None: return "kdp" @@ -142,6 +191,7 @@ def cdi_kind_to_kdp_resource( __all__ = [ + "KubernetesResourceInjectionPolicyEnum", "cdi_kind_to_kdp_resource", "get_resource_injection_policy", "is_kubelet_socket_accessible", diff --git a/gpustack_runtime/deployer/kubernetes.py b/gpustack_runtime/deployer/kubernetes.py index e84d206..ffdedc7 100644 --- a/gpustack_runtime/deployer/kubernetes.py +++ b/gpustack_runtime/deployer/kubernetes.py @@ -7,12 +7,14 @@ import os import re import time +from copy import copy from dataclasses import dataclass, field from datetime import datetime, timezone from decimal import Decimal from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING +from threading import Lock +from typing import TYPE_CHECKING, Literal import kubernetes import kubernetes.stream.ws_client @@ -49,11 +51,16 @@ sensitive_env_var, validate_rfc1123_domain_name, ) -from .k8s.devicemanager import get_resource_injection_policy +from .k8s.devicemanager import ( + KubernetesResourceInjectionPolicyEnum, + get_resource_injection_policy, +) if TYPE_CHECKING: from collections.abc import Callable, Generator + from .__types__ import DevicesMaterial + logger = logging.getLogger(__name__) clogger = logger.getChild("conversion") @@ -117,6 +124,9 @@ class KubernetesWorkloadPlan(WorkloadPlan): Domain suffix for the cluster. Default is "cluster.local". service_type (KubernetesWorkloadServiceTypeEnum): Service type for the workload. Default is CLUSTER_IP. + resource_injection_policy (KubernetesResourceInjectionPolicyEnum | None): + Workload override for GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY. + None inherits the environment setting. Auto probes the target node. namespace (str | None): Namespace of the workload. name (str): @@ -159,6 +169,16 @@ class KubernetesWorkloadPlan(WorkloadPlan): Service type for the workload. """ + resource_injection_policy: KubernetesResourceInjectionPolicyEnum | None = None + """ + Device injection policy for this workload. None inherits + GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY (default Auto). + Env uses runtime-visible device environment variables; KDP requests + device-plugin resources; Auto detects the policy from the target node. + This setting does not change the environment or other workloads. + Env injection of all devices uses the runtime's privileged container path. + """ + def validate_and_default(self): """ Validate and set defaults for the workload plan. @@ -168,6 +188,15 @@ def validate_and_default(self): If the workload plan is invalid. """ + if self.resource_injection_policy is not None: + try: + self.resource_injection_policy = KubernetesResourceInjectionPolicyEnum( + self.resource_injection_policy, + ) + except ValueError as e: + msg = f"Invalid resource injection policy '{self.resource_injection_policy}'" + raise ValueError(msg) from e + if self.labels is None: self.labels = {} if self.containers is None: @@ -1153,6 +1182,9 @@ def _probe_node_allocatable(self) -> dict[str, str] | None: self._node_name = node.metadata.name return node.status.allocatable or {} + _resource_injection_policy: Literal["env", "kdp"] | None = None + """Resolved policy on the deployment instance dedicated to one workload.""" + def _resolve_resource_injection_policy(self) -> str: """ Resolve the resource injection policy for this deployer, probing the @@ -1162,6 +1194,8 @@ def _resolve_resource_injection_policy(self) -> str: The resource injection policy. """ + if self._resource_injection_policy is not None: + return self._resource_injection_policy return get_resource_injection_policy(self._probe_node_allocatable) def _get_default_node_name(self) -> str: @@ -1806,14 +1840,33 @@ def __init__(self): self._client = self._get_client() self._node_name = envs.GPUSTACK_RUNTIME_KUBERNETES_NODE_NAME self._runtime_uuid_values_allowed: bool | None = None + self._deployment_lock = Lock() + self._materials_cache: dict[str, dict[str, DevicesMaterial]] = {} + self._materials_locks = {policy: Lock() for policy in ("env", "kdp")} + + def _prepare(self): + """ + Reuse device materials built for the same resolved injection policy. + """ + if self._materials is not None: + return + policy = self._resolve_resource_injection_policy() + with self._materials_locks[policy]: + if policy not in self._materials_cache: + self._runtime_uuid_values_allowed = policy != "kdp" + try: + super()._prepare() + except Exception: + # Incomplete detection must not become a reusable device map. + self._materials = None + raise + self._materials_cache[policy] = self._materials + self._materials = self._materials_cache[policy] @property def allowed_runtime_uuid_values(self) -> bool: - # Resolved once per deployer, unlike the per-Pod resolution the - # creation path wants: this gates how `_prepare` builds the device - # materials, which are themselves built once, and it is read there once - # per manufacturer -- so probing on every read would spend one API call - # per manufacturer to answer a question already settled. + # Device materials and UUID eligibility belong to the same resolved + # policy. Workloads with the same policy reuse the device materials. if self._runtime_uuid_values_allowed is None: self._runtime_uuid_values_allowed = ( self._resolve_resource_injection_policy() != "kdp" @@ -2143,25 +2196,40 @@ def _create(self, workload: WorkloadPlan): msg = f"Invalid workload plan type: {type(workload)}" raise TypeError(msg) - self._prepare_mirrored_deployment() - - if isinstance(workload, WorkloadPlan): + if not isinstance(workload, KubernetesWorkloadPlan): workload = KubernetesWorkloadPlan(**workload.__dict__) workload.validate_and_default() + + # Node selection and mirrored options belong to the registered deployer, + # including the node used by later workload listing operations. + with self._deployment_lock: + self._prepare_mirrored_deployment() + deployer: KubernetesDeployer = copy(self) + + # Each workload resolves its policy independently while sharing cached + # device materials only with workloads using the same policy. + deployer._materials = None + deployer._resource_injection_policy = get_resource_injection_policy( + deployer._probe_node_allocatable, + policy=workload.resource_injection_policy, + ) + deployer._runtime_uuid_values_allowed = ( + deployer._resource_injection_policy != "kdp" + ) if logger.isEnabledFor(logging.DEBUG): logger.debug("Creating workload:\n%s", workload.to_yaml()) # Create ephemeral file if needed, # (container index, configured path): ephemeral_filename_mapping: dict[tuple[int, str], str] = ( - self._create_ephemeral_configmaps(workload) + deployer._create_ephemeral_configmaps(workload) ) # Create Service if needed. - self._create_service(workload) + deployer._create_service(workload) # Create Pod. - self._create_pod( + deployer._create_pod( workload, ephemeral_filename_mapping, ) diff --git a/tests/gpustack_runtime/deployer/test_workload_resource_injection_policy.py b/tests/gpustack_runtime/deployer/test_workload_resource_injection_policy.py new file mode 100644 index 0000000..417a9ae --- /dev/null +++ b/tests/gpustack_runtime/deployer/test_workload_resource_injection_policy.py @@ -0,0 +1,347 @@ +# The generated Pod and device mapping exercise the deployer's private state. +# ruff: noqa: SLF001 + +from __future__ import annotations as __future_annotations__ + +import json +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from threading import Barrier +from types import SimpleNamespace +from unittest.mock import MagicMock + +import kubernetes.client +import pytest + +from gpustack_runtime import envs +from gpustack_runtime.deployer import ( + Container, + ContainerExecution, + ContainerResources, + KubernetesDeployer, + KubernetesResourceInjectionPolicyEnum, + KubernetesWorkloadPlan, + OperationError, + WorkloadPlan, + create_workload, + list_workloads, +) +from gpustack_runtime.deployer.k8s.devicemanager import get_resource_injection_policy +from gpustack_runtime.detector import Device, ManufacturerEnum + + +@pytest.mark.parametrize( + "configured, override, allocatable, expected", + [ + ("KDP", "Env", {}, "env"), + ("Env", "KDP", {}, "kdp"), + ("KDP", "Auto", {}, "env"), + ("Env", "Auto", {"nvidia.com/gpu.shared": "2"}, "kdp"), + ("Env", "Auto", None, "kdp"), + ("Env", None, {}, "env"), + ("KDP", None, {}, "kdp"), + ("Auto", None, {}, "env"), + ("KDP", KubernetesResourceInjectionPolicyEnum.ENV, {}, "env"), + ("Env", KubernetesResourceInjectionPolicyEnum.KDP, {}, "kdp"), + ("KDP", KubernetesResourceInjectionPolicyEnum.AUTO, {}, "env"), + ], +) +def test_workload_policy_precedence( + monkeypatch, + configured, + override, + allocatable, + expected, +): + monkeypatch.setattr( + envs, + "GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY", + configured, + ) + probe = MagicMock(return_value=allocatable) + + assert get_resource_injection_policy(probe, policy=override) == expected + assert probe.call_count == int((override or configured) == "Auto") + assert configured == envs.GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY + + +def _plan(policy=None, name="test", devices="all"): + return KubernetesWorkloadPlan( + name=name, + namespace="default", + resource_injection_policy=policy, + host_ipc=True, + containers=[ + Container( + name="default", + image="example.com/cache:test", + execution=ContainerExecution(privileged=False), + resources=ContainerResources( + **{"nvidia.com/devices": devices, "cpu": 2, "memory": "20Gi"}, + ), + ), + ], + ) + + +@pytest.fixture +def cluster(monkeypatch): + monkeypatch.setattr(KubernetesDeployer, "is_supported", staticmethod(lambda: True)) + monkeypatch.setattr(KubernetesDeployer, "_get_client", staticmethod(lambda: None)) + monkeypatch.setattr(envs, "GPUSTACK_RUNTIME_KUBERNETES_NODE_NAME", "node-a") + monkeypatch.setattr(envs, "GPUSTACK_RUNTIME_DEPLOY_ASYNC", False) + monkeypatch.setattr(envs, "GPUSTACK_RUNTIME_DEPLOY_CORRECT_RUNNER_IMAGE", False) + monkeypatch.setattr( + envs, + "GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY", + "KDP", + ) + monkeypatch.setattr( + envs, + "GPUSTACK_RUNTIME_DEPLOY_RUNTIME_VISIBLE_DEVICES_VALUE_UUID", + {"NVIDIA_VISIBLE_DEVICES"}, + ) + devices = [ + Device( + manufacturer=ManufacturerEnum.NVIDIA, + index=i, + uuid=f"GPU-{i}", + appendix={}, + ) + for i in range(2) + ] + detect = MagicMock(return_value=devices) + monkeypatch.setattr("gpustack_runtime.deployer.__types__.detect_devices", detect) + worker = kubernetes.client.V1Pod( + metadata=kubernetes.client.V1ObjectMeta(name="worker", namespace="default"), + spec=kubernetes.client.V1PodSpec( + runtime_class_name="nvidia", + containers=[kubernetes.client.V1Container(name="default")], + ), + ) + monkeypatch.setattr(KubernetesDeployer, "_find_self_pod", lambda _self: worker) + core = MagicMock() + core.list_node.return_value = kubernetes.client.V1NodeList( + items=[ + kubernetes.client.V1Node( + metadata=kubernetes.client.V1ObjectMeta(name="node-a"), + status=kubernetes.client.V1NodeStatus( + allocatable={"nvidia.com/gpu.shared": "2"}, + ), + ), + ], + ) + core.read_namespaced_pod.side_effect = kubernetes.client.ApiException(status=404) + pods = {} + + def create_pod(namespace, body): + assert namespace == "default" + pods[body.metadata.name] = body + return body + + core.create_namespaced_pod.side_effect = create_pod + monkeypatch.setattr(kubernetes.client, "CoreV1Api", lambda _client: core) + monkeypatch.setattr(kubernetes.client, "NodeV1Api", lambda _client: MagicMock()) + deployer = KubernetesDeployer() + monkeypatch.setattr("gpustack_runtime.deployer._DEPLOYERS", [deployer]) + return SimpleNamespace( + deployer=deployer, + core=core, + pods=pods, + detect=detect, + devices=devices, + ) + + +def _assert_pod(pod, policy, devices="all"): + assert pod.spec.runtime_class_name == "nvidia" + assert pod.spec.host_ipc is True + assert pod.spec.node_name == "node-a" + container = pod.spec.containers[0] + expected = {"cpu": "2", "memory": "20Gi"} + variables = {e.name: e.value for e in container.env} + if policy == "KDP": + expected["nvidia.com/gpu.shared"] = "2" if devices == "all" else "1" + assert "NVIDIA_VISIBLE_DEVICES" not in variables + assert container.security_context.privileged is False + else: + assert variables["NVIDIA_VISIBLE_DEVICES"] == ( + "GPU-0,GPU-1" if devices == "all" else "GPU-0" + ) + assert container.security_context.privileged is (devices == "all") + assert container.resources.requests == expected + assert container.resources.limits == expected + + +@pytest.mark.parametrize( + "policy", + [None, "Auto", "Env", "KDP", *KubernetesResourceInjectionPolicyEnum], +) +def test_policy_survives_serialization(policy): + plan = _plan(policy) + plan.containers[0].resources = None + serialized = plan.to_json() + assert json.loads(serialized)["resource_injection_policy"] == ( + str(policy) if policy is not None else None + ) + restored = KubernetesWorkloadPlan.from_json(serialized) + if policy is not None: + assert ( + restored.resource_injection_policy + is KubernetesResourceInjectionPolicyEnum(policy) + ) + restored.validate_and_default() + assert restored.resource_injection_policy == policy + + +@pytest.mark.parametrize("policy", list(KubernetesResourceInjectionPolicyEnum)) +def test_enum_and_string_policies_create_equivalent_pods(cluster, policy): + expected = "Env" if policy == KubernetesResourceInjectionPolicyEnum.ENV else "KDP" + for name, value in (("enum", policy), ("string", policy.value)): + plan = _plan(value, name) + create_workload(plan) + assert plan.resource_injection_policy is policy + _assert_pod(cluster.pods[name], expected) + + +@pytest.mark.parametrize("policy", ["", "invalid", "CDI", 1]) +def test_invalid_policy_is_rejected_before_cluster_access(cluster, policy): + with pytest.raises(ValueError, match="Invalid resource injection policy"): + create_workload(_plan(policy)) + cluster.core.list_node.assert_not_called() + cluster.core.create_namespaced_pod.assert_not_called() + + +@pytest.mark.parametrize("devices", ["all", "0"]) +def test_sequential_workloads_keep_independent_policies_and_device_maps( + cluster, + devices, +): + for i, policy in enumerate(["KDP", "Env", None, "Env", "Auto"]): + name = f"workload-{i}" + create_workload(_plan(policy, name, devices)) + _assert_pod(cluster.pods[name], "Env" if policy == "Env" else "KDP", devices) + assert cluster.core.list_node.call_count == 1 + assert envs.GPUSTACK_RUNTIME_KUBERNETES_RESOURCE_INJECTION_POLICY == "KDP" + assert cluster.detect.call_count == 2 + + +def test_auto_probes_each_workload_once(cluster): + create_workload(_plan("Auto", "operator")) + _assert_pod(cluster.pods["operator"], "KDP") + cluster.core.list_node.return_value.items[0].status.allocatable = { + "nvidia.com/gpu": "2", + } + create_workload(_plan("Auto", "stock")) + _assert_pod(cluster.pods["stock"], "Env") + assert cluster.core.list_node.call_count == 2 + + +def test_generic_plan_inherits_worker_policy(cluster): + data = vars(_plan()).copy() + for key in ("resource_injection_policy", "domain_suffix", "service_type"): + data.pop(key) + create_workload(WorkloadPlan(**data)) + _assert_pod(cluster.pods["test"], "KDP") + + +def test_concurrent_workloads_keep_independent_device_maps(cluster): + barrier = Barrier(2) + + def detect(*, fast): + assert fast is False + barrier.wait(timeout=10) + return cluster.devices + + cluster.detect.side_effect = detect + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(create_workload, _plan(policy, policy.lower())) + for policy in ("Env", "KDP") + ] + for future in futures: + future.result(timeout=15) + _assert_pod(cluster.pods["env"], "Env") + _assert_pod(cluster.pods["kdp"], "KDP") + + +def test_failed_workload_does_not_change_next_workload_policy(cluster): + create_pod = cluster.core.create_namespaced_pod.side_effect + cluster.core.create_namespaced_pod.side_effect = kubernetes.client.ApiException( + status=500, + ) + with pytest.raises(OperationError): + create_workload(_plan("Env", "failed")) + cluster.core.create_namespaced_pod.side_effect = create_pod + create_workload(_plan(None, "model")) + _assert_pod(cluster.pods["model"], "KDP") + + +@pytest.mark.parametrize("policy", list(KubernetesResourceInjectionPolicyEnum)) +@pytest.mark.parametrize("case", [str.lower, str.upper, str.swapcase]) +def test_policy_accepts_case_insensitive_strings(cluster, policy, case): + plan = _plan(case(policy.value)) + create_workload(plan) + assert plan.resource_injection_policy is policy + _assert_pod(cluster.pods["test"], "Env" if policy.value == "Env" else "KDP") + + +@pytest.mark.parametrize("policy", ["Env", "KDP"]) +def test_concurrent_workloads_reuse_completed_device_maps(cluster, policy): + barrier = Barrier(2) + create_pod = cluster.core.create_namespaced_pod.side_effect + + def create(namespace, body): + barrier.wait(timeout=10) + return create_pod(namespace, body) + + cluster.core.create_namespaced_pod.side_effect = create + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(create_workload, _plan(policy, name)) + for name in ("first", "second") + ] + for future in futures: + future.result(timeout=15) + for pod in cluster.pods.values(): + _assert_pod(pod, policy) + cluster.detect.assert_called_once_with(fast=False) + + +def test_failed_detection_is_not_cached(cluster): + cluster.detect.side_effect = [RuntimeError("detection failed"), cluster.devices] + with pytest.raises(RuntimeError, match="detection failed"): + create_workload(_plan("Env", "failed")) + create_workload(_plan("Env", "retry")) + _assert_pod(cluster.pods["retry"], "Env") + assert cluster.detect.call_count == 2 + + +def test_workloads_share_default_node_and_remain_listable(cluster, monkeypatch): + cluster.deployer._node_name = None + select_node = MagicMock(side_effect=["node-a", "node-b"]) + monkeypatch.setattr(KubernetesDeployer, "_get_default_node_name", select_node) + barrier = Barrier(2) + create_pod = cluster.core.create_namespaced_pod.side_effect + + def create(namespace, body): + barrier.wait(timeout=10) + return create_pod(namespace, body) + + cluster.core.create_namespaced_pod.side_effect = create + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit(create_workload, _plan(policy, policy.lower())) + for policy in ("Env", "KDP") + ] + for future in futures: + future.result(timeout=15) + select_node.assert_called_once() + for pod in cluster.pods.values(): + assert pod.spec.node_name == "node-a" + pod.metadata.creation_timestamp = datetime(2026, 1, 1, tzinfo=timezone.utc) + pod.status = kubernetes.client.V1PodStatus(phase="Running") + cluster.core.list_namespaced_pod.return_value = kubernetes.client.V1PodList( + items=list(cluster.pods.values()), + ) + assert {status.name for status in list_workloads()} == {"env", "kdp"}