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

1"""Per-vendor GPU utilization backends. 

2 

3Each vendor module exposes a class implementing the UtilBackend Protocol. Adding 

4a new vendor requires one new file and one line in _REGISTRY below; gpu_stats.py 

5stays untouched. 

6 

7resolve_backend(device_backend) returns the backend for a given llama-server 

8backend string, or None when no backend covers that vendor. 

9""" 

10 

11from __future__ import annotations 

12 

13from lilbee.providers.fleet.devices import VULKAN_BACKEND as _VULKAN 

14from lilbee.providers.fleet.gpu_backends.amd import AmdBackend 

15from lilbee.providers.fleet.gpu_backends.apple import BACKEND_KEY as _APPLE_KEY 

16from lilbee.providers.fleet.gpu_backends.apple import AppleBackend 

17from lilbee.providers.fleet.gpu_backends.base import UtilBackend, UtilSample 

18from lilbee.providers.fleet.gpu_backends.intel import ( 

19 IntelBackend, 

20 IntelHintKind, 

21 IntelUtilHint, 

22 intel_util_hint, 

23) 

24from lilbee.providers.fleet.gpu_backends.nvidia import NvidiaBackend 

25 

26# Maps the backend string that llama-server --list-devices emits to the backend 

27# instance. One entry per backend string (HIP and ROCm share an instance). 

28_apple = AppleBackend() 

29 

30_REGISTRY: dict[str, UtilBackend] = { 

31 "CUDA": NvidiaBackend(), 

32 "ROCm": AmdBackend(), 

33 "HIP": AmdBackend(), 

34 "SYCL": IntelBackend(), 

35 # Apple Metal: register both strings seen in the wild (build-dependent). 

36 _APPLE_KEY: _apple, 

37 "Metal": _apple, 

38} 

39 

40 

41def resolve_backend(device_backend: str) -> UtilBackend | None: 

42 """Return the UtilBackend for device_backend, or None if unregistered.""" 

43 return _REGISTRY.get(device_backend) 

44 

45 

46# Vulkan is vendor-agnostic, and a consumer GPU is often only exposed to the 

47# engine via Vulkan. Map a Vulkan device to a vendor's util backend by the vendor 

48# named in its device string so its utilization still reads. 

49_VENDOR_KEYS: tuple[tuple[str, str], ...] = ( 

50 ("intel", "SYCL"), 

51 ("nvidia", "CUDA"), 

52 ("radeon", "ROCm"), 

53 ("amd", "ROCm"), 

54) 

55 

56 

57def util_backend_name(backend: str, name: str) -> str: 

58 """Registry key for a device's util backend. 

59 

60 A recognized inference backend already implies the vendor. Vulkan does not, so 

61 a Vulkan device is mapped by the vendor named in *name*; anything unrecognized 

62 is returned unchanged (and resolves to no backend, i.e. structural fallback). 

63 """ 

64 if backend in _REGISTRY: 

65 return backend 

66 if backend == _VULKAN: 

67 lowered = name.lower() 

68 for hint, key in _VENDOR_KEYS: 

69 if hint in lowered: 

70 return key 

71 return backend 

72 

73 

74__all__ = [ 

75 "IntelHintKind", 

76 "IntelUtilHint", 

77 "UtilBackend", 

78 "UtilSample", 

79 "intel_util_hint", 

80 "resolve_backend", 

81 "util_backend_name", 

82]