Coverage for src / lilbee / providers / llama_cpp / 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). 

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

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

46 

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

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

49_KHRONOS_DRIVERS_KEYS = ( 

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

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

52) 

53 

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

55 

56 

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

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

59 try: 

60 import winreg 

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

62 return 

63 yield from _iter_khronos_software_manifests(winreg) 

64 yield from _iter_pnp_class_manifests(winreg, _PNP_DISPLAY_ADAPTER_CLASS_GUID) 

65 yield from _iter_pnp_class_manifests(winreg, _PNP_SOFTWARE_COMPONENT_CLASS_GUID) 

66 

67 

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

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

70 hklm = winreg.HKEY_LOCAL_MACHINE 

71 for sub_key in _KHRONOS_DRIVERS_KEYS: 

72 try: 

73 key = winreg.OpenKey(hklm, sub_key) 

74 except OSError: 

75 continue 

76 try: 

77 i = 0 

78 while True: 

79 try: 

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

81 except OSError: 

82 break 

83 i += 1 

84 if value == 0 and name: 

85 yield name 

86 finally: 

87 winreg.CloseKey(key) 

88 

89 

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

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

92 

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

94 reading each subkey's ``VulkanDriverName`` and ``VulkanDriverNameWow`` 

95 values. Both ``REG_SZ`` (single path) and ``REG_MULTI_SZ`` (path list) 

96 are honoured because the loader spec allows both. 

97 """ 

98 hklm = winreg.HKEY_LOCAL_MACHINE 

99 class_root_path = rf"SYSTEM\CurrentControlSet\Control\Class\{class_guid}" 

100 try: 

101 class_root = winreg.OpenKey(hklm, class_root_path) 

102 except OSError: 

103 return 

104 try: 

105 i = 0 

106 while True: 

107 try: 

108 subkey_name = winreg.EnumKey(class_root, i) 

109 except OSError: 

110 break 

111 i += 1 

112 try: 

113 subkey = winreg.OpenKey(class_root, subkey_name) 

114 except OSError: 

115 continue 

116 try: 

117 yield from _read_vulkan_driver_name_values(winreg, subkey) 

118 finally: 

119 winreg.CloseKey(subkey) 

120 finally: 

121 winreg.CloseKey(class_root) 

122 

123 

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

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

126 

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

128 that the loader spec allows. 

129 """ 

130 for value_name in _PNP_VULKAN_VALUE_NAMES: 

131 try: 

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

133 except OSError: 

134 continue 

135 if isinstance(value, str): 

136 if value: 

137 yield value 

138 elif isinstance(value, list): 

139 for entry in value: 

140 if isinstance(entry, str) and entry: 

141 yield entry 

142 

143 

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

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

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

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

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

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

150# in case lilbee runs outside the sandbox. 

151_VULKAN_ICD_SUBPATH = "vulkan/icd.d" 

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

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

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

155) 

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

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

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

159) 

160 

161 

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

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

164 seen_dirs: set[Path] = set() 

165 seen_files: set[Path] = set() 

166 for directory in _linux_vulkan_icd_directories(): 

167 try: 

168 resolved = directory.expanduser() 

169 except RuntimeError: 

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

171 continue 

172 if resolved in seen_dirs: 

173 continue 

174 seen_dirs.add(resolved) 

175 if not resolved.is_dir(): 

176 continue 

177 try: 

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

179 except OSError: 

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

181 continue 

182 for entry in entries: 

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

184 continue 

185 seen_files.add(entry) 

186 yield str(entry) 

187 

188 

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

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

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

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

193 for fixed in _LINUX_FIXED_ETC_ICD_DIRS: 

194 yield Path(fixed) 

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

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

197 for flatpak in _LINUX_FLATPAK_ICD_DIRS: 

198 yield Path(flatpak) 

199 

200 

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

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

203 

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

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

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

207 """ 

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

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

210 stripped = component.strip() 

211 if not stripped: 

212 continue 

213 yield Path(stripped) / subpath