Coverage for src/lilbee/providers/fleet/cuda_runtime.py: 100%
67 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Put the CUDA 12 runtime wheels on the engine's library search path.
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"""
15from __future__ import annotations
17import importlib.util
18import os
19import re
20import sys
21from pathlib import Path
22from typing import TYPE_CHECKING
24from lilbee.providers import model_cache
25from lilbee.providers.base import ProviderError
26from lilbee.providers.fleet.engine_diagnostics import device_probe_diagnostic, ldd_output
28if TYPE_CHECKING:
29 from lilbee.providers.fleet.devices import FleetDevice
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+)")
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
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]
68def _ships_its_own_cuda_runtime(binary: Path) -> bool:
69 """Whether the CUDA runtime sits in the same directory as *binary*.
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))
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 []
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.
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.
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.
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 calling this on every
119 # reload pass (apply_cuda_runtime_env) is idempotent instead of accumulating
120 # duplicate copies that get baked into each spawned server's 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)}
127def apply_cuda_runtime_env(binary: Path | None = None) -> None:
128 """Put the CUDA-runtime wheel libs on this process's ``LD_LIBRARY_PATH``.
130 The device probe and the child servers then resolve the same runtime, so a
131 zero-device probe reflects a genuinely unusable GPU rather than a probe that
132 merely ran without the wheel libraries on its search path.
133 """
134 os.environ.update(cuda_runtime_env(binary))
137def _linked_cuda_major(ldd_output: str) -> int | None:
138 """The CUDA runtime major *ldd_output* links, or ``None`` when it links none."""
139 match = _CUDA_SONAME_RE.search(ldd_output)
140 return int(match.group(1)) if match else None
143def _links_cuda_runtime(binary: Path, env: dict[str, str]) -> bool:
144 """True when *binary* lists a CUDA runtime soname (a CUDA build), resolved or not."""
145 out = ldd_output(binary, env)
146 return out is not None and _linked_cuda_major(out) is not None
149def assert_cuda_devices_usable(binary: Path, devices: list[FleetDevice], probe_output: str) -> None:
150 """Fail loud when a CUDA build links a runtime it cannot initialize a GPU with.
152 *devices* and *probe_output* are the engine's own ``--list-devices`` result.
153 When the list is empty yet *binary* is a CUDA build and the host has an NVIDIA
154 GPU, the runtime loaded but enumerated no device. The probe's own diagnostic is
155 surfaced and the likely causes are listed (rather than asserting one), so
156 placement does not silently fall to CPU.
157 """
158 if not sys.platform.startswith("linux"):
159 return
160 if devices:
161 return
162 env = {**os.environ, **cuda_runtime_env(binary)}
163 if not _links_cuda_runtime(binary, env):
164 return
165 if not model_cache.has_nvidia_gpu():
166 return
167 diagnostic = device_probe_diagnostic(probe_output)
168 raise ProviderError(
169 "The engine links the CUDA runtime and this host has an NVIDIA GPU, but it "
170 "enumerated no CUDA-capable device, so GPU work would silently fall back to CPU.\n"
171 f"The engine reported: {diagnostic}\n"
172 "Likely causes: MIG is enabled on the card, whose parent device answers as an "
173 "NVIDIA GPU while CUDA enumerates only its instances (list them with "
174 "'nvidia-smi -L' and name one in CUDA_VISIBLE_DEVICES by its MIG- UUID); the "
175 "installed CUDA runtime is newer than the GPU driver supports (check the driver's "
176 "CUDA version with 'nvidia-smi' and match the nvidia-cuda-runtime / nvidia-cublas / "
177 "nvidia-cuda-nvrtc wheels to the engine's CUDA build, or update the driver); a "
178 "restrictive CUDA_VISIBLE_DEVICES; or the runtime libraries missing from the path."
179 )