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

1"""Put the CUDA 12 runtime wheels on the engine's library search path. 

2 

3Driver-only GPU images (common on RunPod) ship ``libcuda.so`` from the kernel driver 

4but not the CUDA 12 runtime that llama-server links (``libcudart.so.12``, 

5``libcublas.so.12``, ``libnvrtc.so.12``). The bundled engine now carries those beside 

6the binary and resolves them through its baked ``$ORIGIN`` rpath, so this module is 

7the fallback for an engine built elsewhere: installing lilbee with the ``cuda12`` 

8extra pulls the ``nvidia-cuda-runtime-cu12`` / ``nvidia-cublas-cu12`` / 

9``nvidia-cuda-nvrtc-cu12`` wheels, which carry those libraries under 

10``site-packages/nvidia``. :func:`cuda_runtime_env` adds their ``lib`` directories to 

11the spawned server's ``LD_LIBRARY_PATH`` -- the path can't be a baked rpath because 

12the wheels' location is only known at install time. 

13""" 

14 

15from __future__ import annotations 

16 

17import importlib.util 

18import os 

19import re 

20import sys 

21from pathlib import Path 

22from typing import TYPE_CHECKING 

23 

24from lilbee.providers import model_cache 

25from lilbee.providers.base import ProviderError 

26from lilbee.providers.fleet.engine_diagnostics import device_probe_diagnostic, ldd_output 

27 

28if TYPE_CHECKING: 

29 from lilbee.providers.fleet.devices import FleetDevice 

30 

31# Subpackages the NVIDIA runtime wheels install under the ``nvidia`` namespace. 

32# The distribution name carries the CUDA major (nvidia-cuda-runtime-cu12, 

33# -cu13) but the import path does not, so this resolves whichever major is 

34# installed and needs no version of its own. Only the packaging extra is 

35# major-specific; see the ``cuda12`` extra in pyproject.toml. 

36_CUDA_WHEEL_IMPORTS: tuple[str, ...] = ( 

37 "nvidia.cuda_runtime", 

38 "nvidia.cublas", 

39 "nvidia.cuda_nvrtc", 

40) 

41# The CUDA runtime sonames, matched by library name with the major read out of 

42# the version suffix rather than pinned into the string. A build linking 

43# libcudart.so.13 is as much a CUDA build as one linking .so.12, and pinning the 

44# major meant the whole guard returned early on the newer one: a cu13 engine that 

45# could not initialize a device fell to CPU in exactly the silence this exists to 

46# break. 

47_CUDA_SONAME_RE = re.compile(r"\blib(?:cudart|cublas|nvrtc)\.so\.(\d+)") 

48 

49 

50def _wheel_lib_dir(import_name: str) -> Path | None: 

51 """The ``lib/`` directory of an installed nvidia CUDA wheel subpackage, or None.""" 

52 try: 

53 spec = importlib.util.find_spec(import_name) 

54 except ModuleNotFoundError: 

55 # The ``nvidia`` namespace parent is not installed at all. 

56 return None 

57 if spec is None or not spec.submodule_search_locations: 

58 return None 

59 lib = Path(next(iter(spec.submodule_search_locations))) / "lib" 

60 return lib if lib.is_dir() else None 

61 

62 

63def _cuda_wheel_lib_dirs() -> list[Path]: 

64 """Lib directories of every installed CUDA-runtime wheel, in link order.""" 

65 return [lib for name in _CUDA_WHEEL_IMPORTS if (lib := _wheel_lib_dir(name)) is not None] 

66 

67 

68def _ships_its_own_cuda_runtime(binary: Path) -> bool: 

69 """Whether the CUDA runtime sits in the same directory as *binary*. 

70 

71 The bundled engine ships its libraries beside itself, and a wheel directory 

72 on the search path can only shadow them with a different build. 

73 """ 

74 parent = binary.parent 

75 return any(_CUDA_SONAME_RE.search(entry.name) for entry in _dir_entries(parent)) 

76 

77 

78def _dir_entries(directory: Path) -> list[Path]: 

79 """Entries of *directory*, empty when it cannot be read.""" 

80 try: 

81 return list(directory.iterdir()) 

82 except OSError: 

83 return [] 

84 

85 

86def cuda_runtime_env(binary: Path | None = None) -> dict[str, str]: 

87 """``LD_LIBRARY_PATH`` for running *binary*, or empty when there is nothing to add. 

88 

89 Ordering, which matters more than it looks: the binary's own directory, then 

90 any CUDA-runtime wheel directories, then whatever the caller already had. 

91 ``$ORIGIN`` lands in ``DT_RUNPATH``, which the loader searches *after* 

92 ``LD_LIBRARY_PATH``, so a wheel directory in front silently replaces the 

93 libraries the engine ships beside itself. On a host that merely has torch 

94 installed, that swapped the bundled engine's CUDA runtime for torch's. 

95 

96 Wheel directories are added only for a binary that actually links CUDA and 

97 does not already carry its own runtime. A Vulkan or CPU build has no use for 

98 them, and putting them on its path only gives an unrelated install a way to 

99 interfere. Without a *binary* to reason about, the wheel directories are 

100 returned as before. 

101 

102 Empty off Linux, where neither the wheels nor ``LD_LIBRARY_PATH`` apply. 

103 """ 

104 if not sys.platform.startswith("linux"): 

105 return {} 

106 parts: list[str] = [] 

107 if binary is not None: 

108 parts.append(str(binary.parent)) 

109 wants_wheels = _links_cuda_runtime(binary, dict(os.environ)) and not ( 

110 _ships_its_own_cuda_runtime(binary) 

111 ) 

112 else: 

113 wants_wheels = True 

114 dirs = _cuda_wheel_lib_dirs() if wants_wheels else [] 

115 if not parts and not dirs: 

116 return {} 

117 parts += [str(d) for d in dirs] 

118 # Drop existing entries that are already wheel dirs so an inherited path that 

119 # already carries them does not gain duplicate copies in each spawned server's 

120 # environment. 

121 wheel_dirs = set(parts) 

122 existing = os.environ.get("LD_LIBRARY_PATH", "") 

123 parts.extend(entry for entry in existing.split(os.pathsep) if entry and entry not in wheel_dirs) 

124 return {"LD_LIBRARY_PATH": os.pathsep.join(parts)} 

125 

126 

127def _linked_cuda_major(ldd_output: str) -> int | None: 

128 """The CUDA runtime major *ldd_output* links, or ``None`` when it links none.""" 

129 match = _CUDA_SONAME_RE.search(ldd_output) 

130 return int(match.group(1)) if match else None 

131 

132 

133def _links_cuda_runtime(binary: Path, env: dict[str, str]) -> bool: 

134 """True when *binary* lists a CUDA runtime soname (a CUDA build), resolved or not.""" 

135 out = ldd_output(binary, env) 

136 return out is not None and _linked_cuda_major(out) is not None 

137 

138 

139def assert_cuda_devices_usable(binary: Path, devices: list[FleetDevice], probe_output: str) -> None: 

140 """Fail loud when a CUDA build links a runtime it cannot initialize a GPU with. 

141 

142 *devices* and *probe_output* are the engine's own ``--list-devices`` result. 

143 When the list is empty yet *binary* is a CUDA build and the host has an NVIDIA 

144 GPU, the runtime loaded but enumerated no device. The probe's own diagnostic is 

145 surfaced and the likely causes are listed (rather than asserting one), so 

146 placement does not silently fall to CPU. 

147 """ 

148 if not sys.platform.startswith("linux"): 

149 return 

150 if devices: 

151 return 

152 env = {**os.environ, **cuda_runtime_env(binary)} 

153 if not _links_cuda_runtime(binary, env): 

154 return 

155 if not model_cache.has_nvidia_gpu(): 

156 return 

157 diagnostic = device_probe_diagnostic(probe_output) 

158 raise ProviderError( 

159 "The engine links the CUDA runtime and this host has an NVIDIA GPU, but it " 

160 "enumerated no CUDA-capable device, so GPU work would silently fall back to CPU.\n" 

161 f"The engine reported: {diagnostic}\n" 

162 "Likely causes: MIG is enabled on the card, whose parent device answers as an " 

163 "NVIDIA GPU while CUDA enumerates only its instances (list them with " 

164 "'nvidia-smi -L' and name one in CUDA_VISIBLE_DEVICES by its MIG- UUID); the " 

165 "installed CUDA runtime is newer than the GPU driver supports (check the driver's " 

166 "CUDA version with 'nvidia-smi' and match the nvidia-cuda-runtime / nvidia-cublas / " 

167 "nvidia-cuda-nvrtc wheels to the engine's CUDA build, or update the driver); a " 

168 "restrictive CUDA_VISIBLE_DEVICES; or the runtime libraries missing from the path." 

169 )