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

1"""Which GPU vendors are physically installed, read from the OS device tree. 

2 

3The Vulkan loader cannot answer this. Creating an instance loads every installed 

4ICD, which is the crash the ICD disable exists to prevent, so the answer has to 

5come from the OS: Linux reads the PCI display controllers out of sysfs, Windows 

6reads the PnP display-adapter class out of the registry. 

7 

8An installed ICD manifest is not evidence of hardware. Mesa ships every vendor's 

9driver together on Linux (a Flatpak runtime always carries all of them), and a 

10Windows manifest outlives the card it arrived with. 

11 

12Read directly rather than through pyudev or WMI: this is one attribute read per 

13device on either platform, and neither dependency earns its place for that. 

14""" 

15 

16from __future__ import annotations 

17 

18import re 

19import sys 

20from pathlib import Path 

21from typing import TYPE_CHECKING, Any 

22 

23from lilbee.providers.fleet.vulkan_icd_discovery import ( 

24 PNP_DISPLAY_ADAPTER_CLASS_GUID, 

25 iter_pnp_class_subkeys, 

26) 

27 

28if TYPE_CHECKING: 

29 from collections.abc import Iterator 

30 

31# PCI base class 0x03 is "display controller"; sysfs writes the full 24-bit 

32# class code (0x030000 for VGA), so the base class is its high byte. 

33_PCI_DISPLAY_CONTROLLER_CLASS = 0x03 

34_PCI_CLASS_CODE_BASE_SHIFT = 16 

35_SYSFS_PCI_DEVICE_DIR = Path("/sys/bus/pci/devices") 

36_SYSFS_CLASS_FILE = "class" 

37_SYSFS_VENDOR_FILE = "vendor" 

38 

39# Windows PnP device-instance IDs carry the PCI vendor in their hardware IDs 

40# ("PCI\\VEN_10DE&DEV_1F95&..."). Both value names are set by the class 

41# installer; REG_SZ and REG_MULTI_SZ are both allowed. 

42_PNP_HARDWARE_ID_VALUE_NAMES = ("MatchingDeviceId", "HardwareID") 

43_PCI_VENDOR_ID_PATTERN = re.compile(r"ven_([0-9a-f]{4})", re.IGNORECASE) 

44 

45 

46def installed_gpu_vendor_ids() -> frozenset[int]: 

47 """PCI vendor IDs of this host's display controllers. 

48 

49 Empty means the device tree holds no PCI display controller or could not be 

50 read at all (macOS, an ARM SoC whose GPU is not on the PCI bus, a container 

51 with no ``/sys``). That is "cannot tell", not "there is no GPU", and callers 

52 must not read it as proof that a vendor is absent. 

53 """ 

54 if sys.platform == "win32": 

55 return frozenset(_windows_gpu_vendor_ids()) 

56 if sys.platform.startswith("linux"): 

57 return frozenset(_linux_gpu_vendor_ids()) 

58 return frozenset() 

59 

60 

61def _linux_gpu_vendor_ids() -> Iterator[int]: 

62 """Yield the vendor ID of every PCI display controller in sysfs.""" 

63 try: 

64 devices = sorted(_SYSFS_PCI_DEVICE_DIR.iterdir()) 

65 except OSError: 

66 return 

67 for device in devices: 

68 class_code = _read_sysfs_hex(device / _SYSFS_CLASS_FILE) 

69 if class_code is None: 

70 continue 

71 if class_code >> _PCI_CLASS_CODE_BASE_SHIFT != _PCI_DISPLAY_CONTROLLER_CLASS: 

72 continue 

73 vendor_id = _read_sysfs_hex(device / _SYSFS_VENDOR_FILE) 

74 if vendor_id is not None: 

75 yield vendor_id 

76 

77 

78def _read_sysfs_hex(path: Path) -> int | None: 

79 """The ``0x``-prefixed integer in a sysfs attribute, ``None`` when unreadable.""" 

80 try: 

81 return int(path.read_text(encoding="utf-8").strip(), 16) 

82 except (OSError, ValueError): 

83 return None 

84 

85 

86def _windows_gpu_vendor_ids() -> Iterator[int]: 

87 """Yield the vendor ID of every registered display adapter.""" 

88 try: 

89 import winreg 

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

91 return 

92 yield from _iter_windows_gpu_vendor_ids(winreg) 

93 

94 

95def _iter_windows_gpu_vendor_ids(winreg: Any) -> Iterator[int]: 

96 """Yield vendor IDs from the hardware IDs under the display-adapter class.""" 

97 for subkey in iter_pnp_class_subkeys(winreg, PNP_DISPLAY_ADAPTER_CLASS_GUID): 

98 for value_name in _PNP_HARDWARE_ID_VALUE_NAMES: 

99 try: 

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

101 except OSError: 

102 continue 

103 yield from _vendor_ids_in(value) 

104 

105 

106def _vendor_ids_in(value: object) -> Iterator[int]: 

107 """Yield the PCI vendor IDs named in one registry hardware-ID value.""" 

108 # Registry values are untyped: REG_SZ arrives as str, REG_MULTI_SZ as list[str]. 

109 entries = value if isinstance(value, list) else [value] 

110 for entry in entries: 

111 if not isinstance(entry, str): 

112 continue 

113 match = _PCI_VENDOR_ID_PATTERN.search(entry) 

114 if match is not None: 

115 yield int(match.group(1), 16)