Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 56 additions & 1 deletion gpustack_runtime/deployer/kuberentes.py
Original file line number Diff line number Diff line change
Expand Up @@ -325,6 +325,61 @@ def _pin_pod_for_kueue(
}


def _is_device_plugin_resource(resource_key: str) -> bool:
"""
Report whether a resource key belongs to a device plugin resource family,
which is either a CDI kind ("nvidia.com/gpu") or one of its suffixed
variants ("nvidia.com/gpu.shared", "nvidia.com/gpu.sliced.units",
"nvidia.com/gpu.partitioned.mig-1g.20gb").
"""
return any(
resource_key == cdi or resource_key.startswith(f"{cdi}.")
for cdi in envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_CDI.values()
)


def _resolve_privileged(container: Container) -> bool:
"""
Resolve whether a container runs privileged.

Privilege is dropped when the container's devices are handed out by a
device plugin, which is the case for every device plugin resource family
and, under the KDP injection policy, for every mapped device request.

A privileged container receives all device nodes of the host, so it
enumerates -- and can use -- every accelerator on the node, no matter
which one the device plugin allocated to it. That silently undoes
slicing: a workload holding a single MIG device or a single memory slice
still sees the untouched cards next to it, and a soft-slicing limit
lands on whichever device comes first instead of the allocated one.
"""
if not container.execution or not container.execution.privileged:
return False
if not container.resources:
return True

kdp = get_resource_injection_policy() == "kdp"
for r_k in container.resources:
if r_k in ("cpu", "memory"):
continue
if _is_device_plugin_resource(r_k) or (
kdp
and (
r_k
in envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES
or r_k == envs.GPUSTACK_RUNTIME_DEPLOY_AUTOMAP_RESOURCE_KEY
)
):
Comment on lines +365 to +372

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES environment variable map can potentially be None according to its type definition in envs.py. If it is None, performing an in membership check will raise a TypeError.\n\nTo prevent potential runtime crashes, please use a fallback empty dictionary or {} when performing the membership check, similar to how envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_CDI is handled in _is_device_plugin_resource.

Suggested change
if _is_device_plugin_resource(r_k) or (
kdp
and (
r_k
in envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES
or r_k == envs.GPUSTACK_RUNTIME_DEPLOY_AUTOMAP_RESOURCE_KEY
)
):
if _is_device_plugin_resource(r_k) or (
kdp
and (
r_k
in (envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES or {})
or r_k == envs.GPUSTACK_RUNTIME_DEPLOY_AUTOMAP_RESOURCE_KEY
)
):

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not applicable — this map can never be None at runtime.

envs.__getattr__ always evaluates the lambda in variables, and the lambda is to_dict(getenv(...)); to_dict short-circuits to {} on an empty value:

>>> from gpustack_runtime import envs
>>> type(envs.GPUSTACK_RUNTIME_DEPLOY_RESOURCE_KEY_MAP_RUNTIME_VISIBLE_DEVICES)
<class 'dict'>
>>> from gpustack_runtime.envs import to_dict
>>> to_dict('')
{}

The | None in the module-level annotation is a placeholder for the lazy-attribute declaration, not a reachable value. Six existing call sites read this same map unguarded — deployer/__types__.py:1391, deployer/kuberentes.py:1137 and :1141, deployer/docker.py:1004, deployer/podman.py:1003, deployer/cdi/__types__.py:680 — so a guard here would be the only one of its kind.

You did point at a real inconsistency, though: I had or {} on the CDI map two functions up and not here. Resolved in the other direction — the stray or {} is gone, so both reads now match the idiom used everywhere else in the file. Leaving this thread open since it is the opposite of the suggested change.

clogger.info(
"Dropping privilege of container '%s', "
"as its device request '%s' is allocated by a device plugin",
container.name,
r_k,
)
return False
return True


class KubernetesDeployer(EndoscopicDeployer):
"""
Deployer implementation for Kubernetes.
Expand Down Expand Up @@ -1048,7 +1103,7 @@ def _create_pod(
run_as_user=c.execution.run_as_user,
run_as_group=c.execution.run_as_group,
read_only_root_filesystem=c.execution.readonly_rootfs,
privileged=c.execution.privileged,
privileged=_resolve_privileged(c),
capabilities=(
kubernetes.client.V1Capabilities(
add=c.execution.capabilities.add,
Expand Down
137 changes: 137 additions & 0 deletions tests/gpustack_runtime/deployer/test_privileged.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,137 @@
import pytest

from gpustack_runtime.deployer.__types__ import (
Container,
ContainerExecution,
ContainerResources,
)
from gpustack_runtime.deployer.kuberentes import _resolve_privileged


def _container(privileged: bool | None, resources: dict | None = None) -> Container:
container_resources = None
if resources is not None:
container_resources = ContainerResources()
container_resources.update(resources)
return Container(
name="default",
image="gpustack/runner:latest",
execution=(
ContainerExecution(privileged=privileged)
if privileged is not None
else None
),
resources=container_resources,
)


@pytest.mark.parametrize(
"name, privileged, resources, policy, expected",
[
(
"no execution",
None,
None,
"env",
False,
),
(
"not requested",
False,
{"nvidia.com/devices": "0"},
"env",
False,
),
(
"no resources",
True,
None,
"env",
True,
),
(
"non-device resources only",
True,
{"cpu": "2", "memory": "4Gi"},
"env",
True,
),
(
"specific whole cards, env injection",
True,
{"cpu": "2", "nvidia.com/devices": "0,1"},
"env",
True,
),
(
"all devices, env injection",
True,
{"nvidia.com/devices": "all"},
"env",
True,
),
(
"specific whole cards, kdp injection",
True,
{"nvidia.com/devices": "0,1"},
"kdp",
False,
),
(
"auto-mapped devices, kdp injection",
True,
{"gpustack.ai/devices": "0"},
"kdp",
False,
),
(
"exclusive whole card",
True,
{"nvidia.com/gpu": "1"},
"env",
False,
),
(
"soft slice",
True,
{
"nvidia.com/gpu.sliced": "1",
"nvidia.com/gpu.sliced.memory-percentage": "50",
"nvidia.com/gpu.sliced.cores-percentage": "50",
},
"env",
False,
),
(
"hard partition",
True,
{
"nvidia.com/gpu.partitioned": "1",
"nvidia.com/gpu.partitioned.mig-1g.20gb": "1",
},
"env",
False,
),
(
"non-NVIDIA soft slice",
True,
{"amd.com/gpu.sliced": "1"},
"env",
False,
),
(
"resource key merely prefixed by a CDI kind",
True,
{"nvidia.com/gpu-alike": "1"},
"env",
True,
),
],
)
def test_resolve_privileged(name, privileged, resources, policy, expected, monkeypatch):
monkeypatch.setattr(
"gpustack_runtime.deployer.kuberentes.get_resource_injection_policy",
lambda: policy,
)
actual = _resolve_privileged(_container(privileged, resources))
assert actual == expected, f"case {name} expected {expected}, but got {actual}"
Loading