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

1"""Cross-platform discovery of installed Vulkan ICD manifests. 

2 

3Mirrors the Vulkan loader's own discovery so callers can identify which 

4vendors are installed without calling ``vkCreateInstance`` (which would 

5pre-load every vendor's ICD into the process before any disable env 

6var can take effect). Windows reads the registry; Linux walks the XDG 

7``vulkan/icd.d`` hierarchy; macOS yields nothing. See 

8https://github.com/KhronosGroup/Vulkan-Loader/blob/main/docs/LoaderDriverInterface.md 

9for the loader-side spec. 

10""" 

11 

12from __future__ import annotations 

13 

14import logging 

15import os 

16import sys 

17from pathlib import Path 

18from typing import TYPE_CHECKING, Any 

19 

20if TYPE_CHECKING: 

21 from collections.abc import Iterator 

22 

23log = logging.getLogger(__name__) 

24 

25 

26def iter_vulkan_manifest_paths() -> Iterator[str]: 

27 """Yield absolute ``.json`` manifest paths the Vulkan loader would discover. 

28 

29 Returns an empty iterator on macOS (Metal-only build, no Vulkan loader). 

30 """ 

31 if sys.platform == "win32": 

32 yield from _iter_windows_vulkan_manifest_paths() 

33 elif sys.platform.startswith("linux"): 

34 yield from _iter_linux_vulkan_manifest_paths() 

35 else: 

36 yield from () 

37 

38 

39# Microsoft-defined PnP device-setup class GUIDs that host Vulkan ICD manifests. 

40# Both GUIDs and the Khronos software-driver key are documented in 

41# https://github.com/KhronosGroup/Vulkan-Loader/blob/main/docs/LoaderDriverInterface.md#driver-discovery-on-windows 

42# (the GUIDs themselves are the public Windows 

43# https://learn.microsoft.com/en-us/windows-hardware/drivers/install/system-defined-device-setup-classes-available-to-vendors). 

44PNP_DISPLAY_ADAPTER_CLASS_GUID = "{4d36e968-e325-11ce-bfc1-08002be10318}" 

45_PNP_SOFTWARE_COMPONENT_CLASS_GUID = "{5c4c3332-344d-483c-8739-259e934c9cc8}" 

46_PNP_CLASS_ROOT = r"SYSTEM\CurrentControlSet\Control\Class" 

47 

48# Legacy software-driver paths (HKLM + WOW6432Node mirror for 32-bit ICDs). 

49# Each value name is a manifest path; the DWORD value is 0=enabled. 

50_KHRONOS_DRIVERS_KEYS = ( 

51 r"SOFTWARE\Khronos\Vulkan\Drivers", 

52 r"SOFTWARE\WOW6432Node\Khronos\Vulkan\Drivers", 

53) 

54 

55_PNP_VULKAN_VALUE_NAMES = ("VulkanDriverName", "VulkanDriverNameWow") 

56 

57 

58def _iter_windows_vulkan_manifest_paths() -> Iterator[str]: 

59 """Yield manifest paths from the four Windows ICD-discovery locations.""" 

60 try: 

61 import winreg 

62 except ImportError: # pragma: no cover - winreg ships with CPython on Windows 

63 return 

64 yield from _iter_khronos_software_manifests(winreg) 

65 yield from _iter_pnp_class_manifests(winreg, PNP_DISPLAY_ADAPTER_CLASS_GUID) 

66 yield from _iter_pnp_class_manifests(winreg, _PNP_SOFTWARE_COMPONENT_CLASS_GUID) 

67 

68 

69def _iter_khronos_software_manifests(winreg: Any) -> Iterator[str]: 

70 """Yield enabled-flag (DWORD=0) manifest paths from the Khronos software keys.""" 

71 hklm = winreg.HKEY_LOCAL_MACHINE 

72 for sub_key in _KHRONOS_DRIVERS_KEYS: 

73 try: 

74 key = winreg.OpenKey(hklm, sub_key) 

75 except OSError: 

76 continue 

77 try: 

78 i = 0 

79 while True: 

80 try: 

81 name, value, _value_type = winreg.EnumValue(key, i) 

82 except OSError: 

83 break 

84 i += 1 

85 if value == 0 and name: 

86 yield name 

87 finally: 

88 winreg.CloseKey(key) 

89 

90 

91def iter_pnp_class_subkeys(winreg: Any, class_guid: str) -> Iterator[Any]: 

92 """Yield each open device subkey under one PnP device-class GUID. 

93 

94 Walks ``HKLM\\SYSTEM\\CurrentControlSet\\Control\\Class\\{GUID}\\NNNN``, 

95 closing every handle once the consumer has read it. A subkey that cannot be 

96 opened is skipped rather than ending the walk. 

97 """ 

98 hklm = winreg.HKEY_LOCAL_MACHINE 

99 try: 

100 class_root = winreg.OpenKey(hklm, rf"{_PNP_CLASS_ROOT}\{class_guid}") 

101 except OSError: 

102 return 

103 try: 

104 i = 0 

105 while True: 

106 try: 

107 subkey_name = winreg.EnumKey(class_root, i) 

108 except OSError: 

109 break 

110 i += 1 

111 try: 

112 subkey = winreg.OpenKey(class_root, subkey_name) 

113 except OSError: 

114 continue 

115 try: 

116 yield subkey 

117 finally: 

118 winreg.CloseKey(subkey) 

119 finally: 

120 winreg.CloseKey(class_root) 

121 

122 

123def _iter_pnp_class_manifests(winreg: Any, class_guid: str) -> Iterator[str]: 

124 """Yield manifest paths from PnP keys under one device-class GUID. 

125 

126 Reads each subkey's ``VulkanDriverName`` and ``VulkanDriverNameWow`` values. 

127 Both ``REG_SZ`` (single path) and ``REG_MULTI_SZ`` (path list) are honoured 

128 because the loader spec allows both. 

129 """ 

130 for subkey in iter_pnp_class_subkeys(winreg, class_guid): 

131 yield from _read_vulkan_driver_name_values(winreg, subkey) 

132 

133 

134def _read_vulkan_driver_name_values(winreg: Any, subkey: Any) -> Iterator[str]: 

135 """Yield non-empty paths from one PnP subkey's ``VulkanDriverName*`` values. 

136 

137 Handles both REG_SZ (single string) and REG_MULTI_SZ (list of strings) 

138 that the loader spec allows. 

139 """ 

140 for value_name in _PNP_VULKAN_VALUE_NAMES: 

141 try: 

142 value, _value_type = winreg.QueryValueEx(subkey, value_name) 

143 except OSError: 

144 continue 

145 if isinstance(value, str): 

146 if value: 

147 yield value 

148 elif isinstance(value, list): 

149 for entry in value: 

150 if isinstance(entry, str) and entry: 

151 yield entry 

152 

153 

154# Linux ICD-discovery search-path config. Defaults follow the XDG basedir 

155# spec (https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html); 

156# SYSCONFDIR / EXTRASYSCONFDIR are loader build-time constants that expand 

157# to /usr/local/etc and /etc on the distros lilbee ships against. The 

158# Flatpak export trees aren't in the Khronos spec but the loader picks them 

159# up via XDG_DATA_DIRS inside a Flatpak runtime; we walk them defensively 

160# in case lilbee runs outside the sandbox. 

161_VULKAN_ICD_SUBPATH = "vulkan/icd.d" 

162_LINUX_FIXED_ETC_ICD_DIRS: tuple[str, ...] = ( 

163 "/usr/local/etc/vulkan/icd.d", 

164 "/etc/vulkan/icd.d", 

165) 

166_LINUX_FLATPAK_ICD_DIRS: tuple[str, ...] = ( 

167 "~/.local/share/flatpak/exports/share/vulkan/icd.d", 

168 "/var/lib/flatpak/exports/share/vulkan/icd.d", 

169) 

170 

171 

172def _iter_linux_vulkan_manifest_paths() -> Iterator[str]: 

173 """Glob ``*.json`` across the Linux ICD-discovery directories, deduping.""" 

174 seen_dirs: set[Path] = set() 

175 seen_files: set[Path] = set() 

176 for directory in _linux_vulkan_icd_directories(): 

177 try: 

178 resolved = directory.expanduser() 

179 except RuntimeError: 

180 # PosixPath.expanduser raises when HOME is unset; skip. 

181 continue 

182 if resolved in seen_dirs: 

183 continue 

184 seen_dirs.add(resolved) 

185 if not resolved.is_dir(): 

186 continue 

187 try: 

188 entries = sorted(resolved.glob("*.json")) 

189 except OSError: 

190 log.debug("Vulkan ICD dir %s could not be read", resolved, exc_info=True) 

191 continue 

192 for entry in entries: 

193 if entry in seen_files or not entry.is_file(): 

194 continue 

195 seen_files.add(entry) 

196 yield str(entry) 

197 

198 

199def _linux_vulkan_icd_directories() -> Iterator[Path]: 

200 """Yield each Linux ICD search directory in loader-spec order.""" 

201 yield from _xdg_dirs("XDG_CONFIG_HOME", "~/.config", _VULKAN_ICD_SUBPATH) 

202 yield from _xdg_dirs("XDG_CONFIG_DIRS", "/etc/xdg", _VULKAN_ICD_SUBPATH) 

203 for fixed in _LINUX_FIXED_ETC_ICD_DIRS: 

204 yield Path(fixed) 

205 yield from _xdg_dirs("XDG_DATA_HOME", "~/.local/share", _VULKAN_ICD_SUBPATH) 

206 yield from _xdg_dirs("XDG_DATA_DIRS", "/usr/local/share:/usr/share", _VULKAN_ICD_SUBPATH) 

207 for flatpak in _LINUX_FLATPAK_ICD_DIRS: 

208 yield Path(flatpak) 

209 

210 

211def _xdg_dirs(env_var: str, default: str, subpath: str) -> Iterator[Path]: 

212 """Split *env_var* (or *default*) on ``:``, append *subpath* to each. 

213 

214 Empty components are dropped (the "extra slash in XDG_DATA_DIRS" loader 

215 quirk, Vulkan-Loader#2331) and appends *subpath* to each remaining 

216 entry. Falls back to *default* when the env var is unset. 

217 

218 The ":" separator is the XDG/Linux convention; this function is only 

219 called from Linux-gated discovery paths and must not be changed to ";" 

220 for cross-platform support. 

221 """ 

222 raw = os.environ.get(env_var) or default 

223 for component in raw.split(":"): 

224 stripped = component.strip() 

225 if not stripped: 

226 continue 

227 yield Path(stripped) / subpath