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

1"""The Vulkan adapter enumeration, as a child process that prints JSON. 

2 

3``vkCreateInstance`` loads every vendor ICD installed on the host. A broken or 

4conflicting one faults inside the loader, and a fault is a signal, not an 

5exception: no ``except`` clause in the calling process can catch it. Running the 

6enumeration in the daemon therefore means any user's bad driver can kill lilbee, 

7and not only on the Vulkan path, since the unified-memory classification asks 

8this question for CUDA, ROCm and SYCL devices too. 

9 

10So it runs here instead, where dying is an answer. The parent reads the JSON on 

11stdout; a non-zero exit, a signal, or unparseable output all mean the same thing, 

12which is that the loader has no opinion and the host is treated as it was before 

13anything asked. 

14 

15The enumeration itself is unchanged and still lives in 

16:mod:`lilbee.providers.fleet.gpu_select`; this module is the process boundary 

17around it. 

18""" 

19 

20from __future__ import annotations 

21 

22import json 

23import sys 

24from typing import Any 

25 

26from lilbee.providers.fleet.gpu_select import VulkanDevice, enumerate_in_process 

27 

28# Exit code for "the loader could not be reached or failed", as distinct from a 

29# clean run that found no adapters, which exits 0 with an empty list. 

30NO_OPINION_EXIT = 2 

31 

32 

33def _as_json(device: VulkanDevice) -> dict[str, object]: 

34 """One adapter as JSON-safe fields; the UUID is bytes and travels as hex.""" 

35 return { 

36 "index": device.index, 

37 "device_type": device.device_type, 

38 "device_name": device.device_name, 

39 "vendor_id": device.vendor_id, 

40 "vram_bytes": device.vram_bytes, 

41 "device_uuid": device.device_uuid.hex(), 

42 "storage_buffer_16bit": device.storage_buffer_16bit, 

43 "free_bytes": device.free_bytes, 

44 } 

45 

46 

47def from_json(payload: list[dict[str, Any]]) -> list[VulkanDevice]: 

48 """Rebuild the adapters the child printed. 

49 

50 Raises on a payload that is not the shape :func:`_as_json` writes, which the 

51 parent reads as the loader having no opinion. 

52 """ 

53 return [ 

54 VulkanDevice( 

55 index=int(entry["index"]), 

56 device_type=int(entry["device_type"]), 

57 device_name=str(entry["device_name"]), 

58 vendor_id=int(entry["vendor_id"]), 

59 vram_bytes=int(entry["vram_bytes"]), 

60 device_uuid=bytes.fromhex(str(entry["device_uuid"])), 

61 storage_buffer_16bit=entry["storage_buffer_16bit"], 

62 free_bytes=entry["free_bytes"], 

63 ) 

64 for entry in payload 

65 ] 

66 

67 

68def main() -> None: 

69 """Print the enumeration as JSON, or exit non-zero when there is no answer.""" 

70 devices = enumerate_in_process() 

71 if devices is None: 

72 sys.exit(NO_OPINION_EXIT) 

73 json.dump([_as_json(device) for device in devices], sys.stdout) 

74 

75 

76if __name__ == "__main__": # pragma: no cover - process entry glue 

77 main()