Coverage for src/lilbee/providers/fleet/gpu_env.py: 100%

1"""The environment a fleet child runs in: GPU visibility and Vulkan-loader safety. 

2 

3Binding-free: composes the backend visible-device env vars (from 

4``cfg.gpu_devices``) and the dual-vendor Vulkan crash mitigations over a copy of 

5this process's environment. Nothing here writes ``os.environ``: ``setenv`` 

6reallocates the C ``environ`` array, and a native thread reading it at that moment 

7(lancedb's runtime) can fault. A caller passes the returned mapping as ``env=`` at 

8spawn, and a reader that needs the effective value reads it from that mapping. 

9""" 

10 

11from __future__ import annotations 

12 

13import logging 

14import os 

15import sys 

16from collections.abc import MutableMapping 

17 

18from lilbee.runtime.shard_process import shard_process 

19 

20log = logging.getLogger(__name__) 

21 

22# Backend env vars set from ``cfg.gpu_devices``. Vulkan, CUDA, and ROCm each read 

23# their own; a user-set ``cfg.gpu_devices`` is applied to all four because the 

24# user is specifying their own indexes and opting in to all wheel flavors. 

25_GPU_VISIBLE_ENV_VARS = ( 

26 "GGML_VK_VISIBLE_DEVICES", 

27 "CUDA_VISIBLE_DEVICES", 

28 "HIP_VISIBLE_DEVICES", 

29 "ROCR_VISIBLE_DEVICES", 

30) 

31# The AMD pair is deliberately absent: ROCr filters before HIP re-indexes within 

32# the survivors, so a pin may only ever be written to one of them. Which one is 

33# devices.amd_visible_var's decision. 

34_NON_AMD_VISIBLE_ENV_VARS = ("GGML_VK_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES") 

35_CUDA_VISIBLE_VAR = "CUDA_VISIBLE_DEVICES" 

36# What the NVIDIA container runtime sets to say which GPUs it gave this 

37# container. Its own words for "none" are below. 

38_NVIDIA_RUNTIME_VISIBLE_VAR = "NVIDIA_VISIBLE_DEVICES" 

39_NVIDIA_RUNTIME_NO_GPU_VALUES = frozenset({"void", "none"}) 

40 

41_VK_LOADER_LAYERS_DISABLE_ENV_VAR = "VK_LOADER_LAYERS_DISABLE" 

42 

43# Layers with documented crashes against multi-VkDevice apps: 

44# https://github.com/ggml-org/llama.cpp/issues/18109 (RTSS / OBS / HudSight) 

45# https://github.com/ValveSoftware/steam-for-linux/issues/9120 (Steam overlay) 

46# https://alegruz.github.io/graphics/2025/03/22/galaxyoverlayvklayer-issue.html (Galaxy) 

47# Vendor-dispatch layers (NV_optimus, AMD_switchable_graphics, MESA_device_select) 

48# and user-opt-in overlays (MangoHud) are intentionally absent so GPU routing 

49# stays identical to what every other Vulkan app on the host sees. 

50_VK_LOADER_LAYERS_DISABLE_GLOBS: tuple[str, ...] = ( 

51 "VK_LAYER_VALVE_steam_overlay*", 

52 "VK_LAYER_VALVE_steam_fossilize*", 

53 "VK_LAYER_RTSS*", 

54 "VK_LAYER_OBS_HOOK*", 

55 "VK_LAYER_HudSight*", 

56 "GalaxyOverlayVkLayer*", 

57 "VK_LAYER_GalaxyOverlay*", 

58 "VK_LAYER_DISCORD_overlay*", 

59 "VK_LAYER_EOS_Overlay*", 

60 "VK_LAYER_RESHADE*", 

61 "VK_LAYER_VKBASALT*", 

62) 

63_VK_LOADER_LAYERS_DISABLE_VALUE = ",".join(_VK_LOADER_LAYERS_DISABLE_GLOBS) 

64 

65 

66def _apply_vulkan_loader_safety(env: MutableMapping[str, str]) -> None: 

67 """Disable known-crashing overlay layers and conflicting dual-vendor ICDs in *env*. 

68 

69 Must reach every child that calls ``vkCreateInstance``: the loader loads every 

70 ICD and layer at instance creation, before ``GGML_VK_VISIBLE_DEVICES`` is 

71 consulted, so device pinning alone cannot stop a buggy second-vendor ICD from 

72 corrupting the heap. A value already in *env* is kept. 

73 """ 

74 from lilbee.providers.fleet.gpu_select import ( 

75 VulkanIcdEnvVar, 

76 disable_conflicting_vulkan_icds, 

77 ) 

78 

79 if sys.platform == "win32" or sys.platform.startswith("linux"): 

80 env.setdefault(_VK_LOADER_LAYERS_DISABLE_ENV_VAR, _VK_LOADER_LAYERS_DISABLE_VALUE) 

81 disable_glob = disable_conflicting_vulkan_icds() 

82 if disable_glob is not None: 

83 env.setdefault(VulkanIcdEnvVar.LOADER_DRIVERS_DISABLE, disable_glob) 

84 log.info("Disabling conflicting Vulkan ICDs: %s", disable_glob) 

85 

86 

87def _apply_gpu_devices_pin(env: MutableMapping[str, str]) -> bool: 

88 """Put the user's ``cfg.gpu_devices`` pin in *env* for every backend's visible-devices var. 

89 

90 Returns ``True`` when a pin was applied (so the caller can skip autodetect). 

91 The pin goes to all four backend vars because the user is naming indexes 

92 that match their own wheel's enumeration. A non-empty value *env* already 

93 carries is kept, since that is an equally explicit instruction arriving 

94 closer to the process. 

95 

96 An empty one is replaced. It carries no index to respect, and the pin is the 

97 more specific statement of the two: somebody wrote it into this lilbee's own 

98 configuration, where an empty mask is usually inherited from whatever 

99 launched the process. 

100 """ 

101 from lilbee.core.config import cfg 

102 from lilbee.providers.fleet.devices import amd_visible_var 

103 

104 if not cfg.gpu_devices: 

105 return False 

106 for name in (*_NON_AMD_VISIBLE_ENV_VARS, amd_visible_var(env)): 

107 if env.get(name, "").strip(): 

108 continue 

109 env[name] = cfg.gpu_devices 

110 return True 

111 

112 

113def shard_visible_devices(device: int) -> dict[str, str]: 

114 """The backend visible-device vars that pin a process to one card. 

115 

116 Composed against whatever mask this process already carries, so *device* names 

117 the card at that position in lilbee's own enumeration rather than the host's. 

118 """ 

119 from lilbee.providers.fleet.devices import amd_visible_var 

120 

121 env = gpu_environ() 

122 return { 

123 name: _mask_entry(env.get(name, ""), device) 

124 for name in (*_NON_AMD_VISIBLE_ENV_VARS, amd_visible_var(env)) 

125 } 

126 

127 

128def _mask_entry(mask: str, device: int) -> str: 

129 """The *device*-th entry of a visible-devices *mask*, or its index when unmasked.""" 

130 entries = [entry.strip() for entry in mask.split(",") if entry.strip()] 

131 if not entries: 

132 return str(device) 

133 return entries[device % len(entries)] 

134 

135 

136def _clear_empty_visible_device_vars(env: MutableMapping[str, str]) -> None: 

137 """Drop an empty ``CUDA_VISIBLE_DEVICES`` from *env* when the container runtime contradicts it. 

138 

139 An empty visibility variable is not a mistake to be corrected. It is the 

140 documented way to say "no devices", it is what SLURM and Kubernetes export on 

141 an allocation without a GPU, and it is what a user writes to force CPU. These 

142 variables are read-only filters; deleting one overrides a decision somebody 

143 made on purpose and can hand backend selection to a vendor that was fenced 

144 off deliberately. 

145 

146 The one exception is a genuine contradiction: the NVIDIA container runtime 

147 exposes a card through ``NVIDIA_VISIBLE_DEVICES`` while leaving 

148 ``CUDA_VISIBLE_DEVICES`` empty, so the two disagree and the runtime's own 

149 statement is the newer one. That marker speaks only for NVIDIA, so only the 

150 CUDA variable is touched; it says nothing about an AMD or Vulkan opt-out, and 

151 those are always left alone. 

152 """ 

153 if not _container_runtime_exposes_a_gpu(env): 

154 return 

155 if _CUDA_VISIBLE_VAR in env and not env[_CUDA_VISIBLE_VAR].strip(): 

156 del env[_CUDA_VISIBLE_VAR] 

157 log.debug( 

158 "%s was empty while %s exposes a GPU; clearing the empty mask so the engine " 

159 "can see the card the container runtime provided.", 

160 _CUDA_VISIBLE_VAR, 

161 _NVIDIA_RUNTIME_VISIBLE_VAR, 

162 ) 

163 

164 

165def _container_runtime_exposes_a_gpu(env: MutableMapping[str, str]) -> bool: 

166 """Whether the NVIDIA container runtime says this container was given a GPU. 

167 

168 ``void`` and ``none`` are its own words for "no GPU", so they confirm the 

169 empty mask rather than contradict it. 

170 """ 

171 value = env.get(_NVIDIA_RUNTIME_VISIBLE_VAR, "").strip().casefold() 

172 return bool(value) and value not in _NVIDIA_RUNTIME_NO_GPU_VALUES 

173 

174 

175def gpu_environ() -> dict[str, str]: 

176 """This process's environment with the ``cfg.gpu_devices`` pin and the empty-mask fix applied. 

177 

178 A copy: the process environment is not touched. Cheap and free of disk reads, 

179 so a sampler may call it every second. Anything that reads a visible-devices 

180 variable to learn which cards the engine sees reads it here, not from 

181 ``os.environ``, because the pin lives only in this mapping. 

182 

183 Nothing here chooses a device. The fleet selects through its own placement, 

184 and anything pinning ``GGML_VK_VISIBLE_DEVICES`` before ``probe_devices`` 

185 runs would hide every other GPU from it and switch off ggml's own device 

186 filtering besides. A ``cfg.gpu_devices`` pin is still honored, since there 

187 the user is naming their own indexes. An empty backend visible-devices var from 

188 the orchestrator is cleared first so it does not hide a present GPU, and so a 

189 pin can replace it rather than be blocked. 

190 """ 

191 env = dict(os.environ) 

192 env.update(shard_process.visible_devices) 

193 _clear_empty_visible_device_vars(env) 

194 _apply_gpu_devices_pin(env) 

195 return env 

196 

197 

198def engine_environ() -> dict[str, str]: 

199 """The environment for a child that loads the engine or the Vulkan loader. 

200 

201 :func:`gpu_environ` plus the Vulkan loader safety. The loader reads its 

202 variables when the child creates its first instance, so this must be the 

203 ``env=`` of every such spawn. The ICD discovery reads the host's device tree, 

204 which is why the sampling readers use :func:`gpu_environ` instead. 

205 """ 

206 env = gpu_environ() 

207 _apply_vulkan_loader_safety(env) 

208 return env