Coverage for src/lilbee/providers/fleet/gpu_env.py: 100%
52 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Apply GPU-visibility and Vulkan-loader environment before the engine starts.
3Binding-free: sets the backend visible-device env vars (from ``cfg.gpu_devices``
4or Vulkan autodetect) and the dual-vendor Vulkan crash mitigations. These are
5process-wide ``setdefault`` writes, so a child llama-server inherits them.
6"""
8from __future__ import annotations
10import logging
11import os
12import sys
14log = logging.getLogger(__name__)
16# Backend env vars set from ``cfg.gpu_devices``. Vulkan, CUDA, and ROCm each read
17# their own; a user-set ``cfg.gpu_devices`` is applied to all four because the
18# user is specifying their own indexes and opting in to all wheel flavors.
19_GPU_VISIBLE_ENV_VARS = (
20 "GGML_VK_VISIBLE_DEVICES",
21 "CUDA_VISIBLE_DEVICES",
22 "HIP_VISIBLE_DEVICES",
23 "ROCR_VISIBLE_DEVICES",
24)
25# The AMD pair is deliberately absent: ROCr filters before HIP re-indexes within
26# the survivors, so a pin may only ever be written to one of them. Which one is
27# devices.amd_visible_var's decision.
28_NON_AMD_VISIBLE_ENV_VARS = ("GGML_VK_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES")
29_CUDA_VISIBLE_VAR = "CUDA_VISIBLE_DEVICES"
30# What the NVIDIA container runtime sets to say which GPUs it gave this
31# container. Its own words for "none" are below.
32_NVIDIA_RUNTIME_VISIBLE_VAR = "NVIDIA_VISIBLE_DEVICES"
33_NVIDIA_RUNTIME_NO_GPU_VALUES = frozenset({"void", "none"})
35_VK_LOADER_LAYERS_DISABLE_ENV_VAR = "VK_LOADER_LAYERS_DISABLE"
37# Layers with documented crashes against multi-VkDevice apps:
38# https://github.com/ggml-org/llama.cpp/issues/18109 (RTSS / OBS / HudSight)
39# https://github.com/ValveSoftware/steam-for-linux/issues/9120 (Steam overlay)
40# https://alegruz.github.io/graphics/2025/03/22/galaxyoverlayvklayer-issue.html (Galaxy)
41# Vendor-dispatch layers (NV_optimus, AMD_switchable_graphics, MESA_device_select)
42# and user-opt-in overlays (MangoHud) are intentionally absent so GPU routing
43# stays identical to what every other Vulkan app on the host sees.
44_VK_LOADER_LAYERS_DISABLE_GLOBS: tuple[str, ...] = (
45 "VK_LAYER_VALVE_steam_overlay*",
46 "VK_LAYER_VALVE_steam_fossilize*",
47 "VK_LAYER_RTSS*",
48 "VK_LAYER_OBS_HOOK*",
49 "VK_LAYER_HudSight*",
50 "GalaxyOverlayVkLayer*",
51 "VK_LAYER_GalaxyOverlay*",
52 "VK_LAYER_DISCORD_overlay*",
53 "VK_LAYER_EOS_Overlay*",
54 "VK_LAYER_RESHADE*",
55 "VK_LAYER_VKBASALT*",
56)
57_VK_LOADER_LAYERS_DISABLE_VALUE = ",".join(_VK_LOADER_LAYERS_DISABLE_GLOBS)
60def _apply_vulkan_loader_safety() -> None:
61 """Disable known-crashing overlay layers and conflicting dual-vendor ICDs.
63 Must precede every ``vkCreateInstance``: the loader loads every ICD and layer
64 at instance creation, before ``GGML_VK_VISIBLE_DEVICES`` is consulted, so
65 device pinning alone cannot stop a buggy second-vendor ICD from corrupting the
66 heap. ``setdefault`` preserves any user-set value.
67 """
68 from lilbee.providers.fleet.gpu_select import (
69 VulkanIcdEnvVar,
70 disable_conflicting_vulkan_icds,
71 )
73 if sys.platform == "win32" or sys.platform.startswith("linux"):
74 os.environ.setdefault(_VK_LOADER_LAYERS_DISABLE_ENV_VAR, _VK_LOADER_LAYERS_DISABLE_VALUE)
75 disable_glob = disable_conflicting_vulkan_icds()
76 if disable_glob is not None:
77 os.environ.setdefault(VulkanIcdEnvVar.LOADER_DRIVERS_DISABLE, disable_glob)
78 log.info("Disabling conflicting Vulkan ICDs: %s", disable_glob)
81def _apply_gpu_devices_pin() -> bool:
82 """Apply the user's ``cfg.gpu_devices`` pin to every backend's visible-devices var.
84 Returns ``True`` when a pin was applied (so the caller can skip autodetect).
85 The pin goes to all four backend vars because the user is naming indexes
86 that match their own wheel's enumeration. A non-empty env var the caller
87 already set is kept, since that is an equally explicit instruction arriving
88 closer to the process.
90 An empty one is replaced. It carries no index to respect, and the pin is the
91 more specific statement of the two: somebody wrote it into this lilbee's own
92 configuration, where an empty mask is usually inherited from whatever
93 launched the process.
94 """
95 from lilbee.core.config import cfg
96 from lilbee.providers.fleet.devices import amd_visible_var
98 if not cfg.gpu_devices:
99 return False
100 for name in (*_NON_AMD_VISIBLE_ENV_VARS, amd_visible_var()):
101 if os.environ.get(name, "").strip():
102 continue
103 os.environ[name] = cfg.gpu_devices
104 return True
107def shard_visible_devices(device: int) -> dict[str, str]:
108 """The backend visible-device vars that pin a process to one card.
110 Composed against whatever mask this process already carries, so *device* names
111 the card at that position in lilbee's own enumeration rather than the host's.
112 """
113 from lilbee.providers.fleet.devices import amd_visible_var
115 return {
116 name: _mask_entry(os.environ.get(name, ""), device)
117 for name in (*_NON_AMD_VISIBLE_ENV_VARS, amd_visible_var())
118 }
121def _mask_entry(mask: str, device: int) -> str:
122 """The *device*-th entry of a visible-devices *mask*, or its index when unmasked."""
123 entries = [entry.strip() for entry in mask.split(",") if entry.strip()]
124 if not entries:
125 return str(device)
126 return entries[device % len(entries)]
129def _clear_empty_visible_device_vars() -> None:
130 """Drop an empty ``CUDA_VISIBLE_DEVICES`` only when the container runtime contradicts it.
132 An empty visibility variable is not a mistake to be corrected. It is the
133 documented way to say "no devices", it is what SLURM and Kubernetes export on
134 an allocation without a GPU, and it is what a user writes to force CPU. These
135 variables are read-only filters; deleting one overrides a decision somebody
136 made on purpose and can hand backend selection to a vendor that was fenced
137 off deliberately.
139 The one exception is a genuine contradiction: the NVIDIA container runtime
140 exposes a card through ``NVIDIA_VISIBLE_DEVICES`` while leaving
141 ``CUDA_VISIBLE_DEVICES`` empty, so the two disagree and the runtime's own
142 statement is the newer one. That marker speaks only for NVIDIA, so only the
143 CUDA variable is touched; it says nothing about an AMD or Vulkan opt-out, and
144 those are always left alone.
145 """
146 if not _container_runtime_exposes_a_gpu():
147 return
148 if _CUDA_VISIBLE_VAR in os.environ and not os.environ[_CUDA_VISIBLE_VAR].strip():
149 del os.environ[_CUDA_VISIBLE_VAR]
150 log.info(
151 "%s was empty while %s exposes a GPU; clearing the empty mask so the engine "
152 "can see the card the container runtime provided.",
153 _CUDA_VISIBLE_VAR,
154 _NVIDIA_RUNTIME_VISIBLE_VAR,
155 )
158def _container_runtime_exposes_a_gpu() -> bool:
159 """Whether the NVIDIA container runtime says this container was given a GPU.
161 ``void`` and ``none`` are its own words for "no GPU", so they confirm the
162 empty mask rather than contradict it.
163 """
164 value = os.environ.get(_NVIDIA_RUNTIME_VISIBLE_VAR, "").strip().casefold()
165 return bool(value) and value not in _NVIDIA_RUNTIME_NO_GPU_VALUES
168def apply_fleet_gpu_env() -> None:
169 """Fleet engine bootstrap: loader safety plus the ``cfg.gpu_devices`` pin only.
171 Nothing here chooses a device. The fleet selects through its own placement,
172 and anything pinning ``GGML_VK_VISIBLE_DEVICES`` before ``probe_devices``
173 runs would hide every other GPU from it and switch off ggml's own device
174 filtering besides. A ``cfg.gpu_devices`` pin is still honored, since there
175 the user is naming their own indexes (the probe inherits this environment).
176 An empty backend visible-devices var from the orchestrator is cleared first so it
177 does not hide a present GPU, and so a pin can replace it rather than be blocked.
178 """
179 _apply_vulkan_loader_safety()
180 _clear_empty_visible_device_vars()
181 _apply_gpu_devices_pin()