Coverage for src/lilbee/providers/fleet/rocm_runtime.py: 100%
96 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Guard a ROCm engine build against AMD hosts it cannot actually serve.
3Reads the kernel driver and the bundled artifacts rather than ROCm tooling,
4because ROCm being broken is the case these checks exist to catch.
5"""
7from __future__ import annotations
9import logging
10import os
11import sys
12from pathlib import Path
13from typing import TYPE_CHECKING
15from lilbee.providers.base import ProviderError
16from lilbee.providers.fleet.engine_diagnostics import device_probe_diagnostic, links_any
17from lilbee.providers.fleet.gpu_select import PCIVendorID, discrete_gpu_from_vendor
19if TYPE_CHECKING:
20 from lilbee.providers.fleet.devices import FleetDevice
22log = logging.getLogger(__name__)
24# A ROCm build links these and none of the CUDA sonames.
25_HIP_SONAMES: tuple[str, ...] = ("libamdhip64.so", "librocblas.so", "libhipblas.so")
26# PCI vendor id for AMD, as sysfs reports it.
27_AMD_PCI_VENDOR_ID = "0x1002"
28# Backends whose devices run rocBLAS; ggml prints either name depending on version.
29_AMD_BACKENDS = ("ROCm", "HIP")
30# AMD's escape hatch: the runtime treats every device as the given gfx version.
31_HSA_OVERRIDE_VAR = "HSA_OVERRIDE_GFX_VERSION"
34def _links_hip_runtime(binary: Path, env: dict[str, str]) -> bool:
35 """True when *binary* lists a HIP runtime soname (a ROCm build), resolved or not."""
36 return links_any(binary, env, _HIP_SONAMES)
39def _amd_gpu_present() -> bool:
40 """Whether the kernel exposes an AMD GPU, without needing ROCm to work.
42 Read from sysfs rather than from amd-smi or rocm-smi: those ship with ROCm,
43 and the failure this guards is precisely ROCm being installed wrong, so a
44 tool-based check would report "no GPU" for the case it exists to catch.
45 """
46 if not Path("/dev/kfd").exists():
47 return False
48 for vendor in Path("/sys/class/drm").glob("card*/device/vendor"):
49 try:
50 if vendor.read_text(encoding="utf-8").strip().lower() == _AMD_PCI_VENDOR_ID:
51 return True
52 except OSError:
53 continue
54 return False
57def _amd_discrete_gpu_proven() -> bool:
58 """Whether a discrete AMD card is positively known to be present.
60 Positive evidence only. The sysfs checks that find an AMD GPU cannot tell a
61 discrete card from an APU, and the difference decides between refusing to
62 start and merely running slower, so the question is put to the Vulkan loader,
63 which reports the device type. An unreachable loader proves nothing and
64 answers no, which keeps the softer path.
65 """
66 return discrete_gpu_from_vendor(PCIVendorID.AMD) is True
69def _gfx_name(target_version: int) -> str:
70 """The gfx name for a KFD ``gfx_target_version``; minor and step print as hex."""
71 major, rest = divmod(target_version, 10000)
72 minor, step = divmod(rest, 100)
73 return f"gfx{major}{minor:x}{step:x}"
76def _host_amd_gfx_targets() -> set[str]:
77 """The gfx targets of this host's AMD GPUs, from the driver's KFD topology.
79 CPU nodes report a target version of 0. Empty means "no claim".
80 """
81 targets: set[str] = set()
82 for props in Path("/sys/class/kfd/kfd/topology/nodes").glob("*/properties"):
83 try:
84 text = props.read_text(encoding="utf-8")
85 except OSError:
86 continue
87 for line in text.splitlines():
88 name, _, value = line.partition(" ")
89 if name == "gfx_target_version" and value.strip().isdigit() and int(value.strip()):
90 targets.add(_gfx_name(int(value.strip())))
91 return targets
94def _bundled_rocblas_gfx_targets(binary: Path) -> set[str] | None:
95 """The gfx targets covered by the rocBLAS Tensile masters bundled beside *binary*.
97 None when no bundle sits beside the binary (a system-ROCm engine): that is
98 "no claim", where an empty set would mean "supports nothing".
99 """
100 library = binary.parent / "rocblas" / "library"
101 if not library.is_dir():
102 return None
103 return {
104 f.name.removeprefix("TensileLibrary_lazy_").removesuffix(".dat")
105 for f in library.glob("TensileLibrary_lazy_gfx*.dat")
106 }
109def _hsa_override_gfx() -> str | None:
110 """The gfx target a user-set ``HSA_OVERRIDE_GFX_VERSION`` maps every device to."""
111 raw = os.environ.get(_HSA_OVERRIDE_VAR, "")
112 try:
113 major, minor, step = (int(part) for part in raw.split("."))
114 except ValueError:
115 return None
116 return _gfx_name(major * 10000 + minor * 100 + step)
119def _rocm_support_facts(binary: Path) -> str:
120 """What is known about shipped kernels and host gfx targets, as message text."""
121 shipped = _bundled_rocblas_gfx_targets(binary)
122 parts = []
123 if shipped:
124 parts.append(f"This build ships GPU kernels for: {', '.join(sorted(shipped))}.")
125 if host := _host_amd_gfx_targets():
126 parts.append(f"This host's AMD GPU targets: {', '.join(sorted(host))}.")
127 return " " + " ".join(parts) if parts else ""
130def _warn_if_override_uncovered(override: str, shipped: set[str]) -> None:
131 """The user overrode explicitly; respect it, but say what will happen."""
132 if override in shipped:
133 return
134 log.warning(
135 "%s maps every AMD device to %s, but this build ships GPU kernels only "
136 "for: %s. The engine will abort at the first matrix multiplication if a "
137 "model runs on the GPU.",
138 _HSA_OVERRIDE_VAR,
139 override,
140 ", ".join(sorted(shipped)),
141 )
144def _assert_rocblas_covers_enumerated_devices(binary: Path, devices: list[FleetDevice]) -> None:
145 """Refuse a card the bundle ships no rocBLAS kernels for, before rocBLAS aborts.
147 An enumerated device is no proof of support: a card with engine device code
148 but no rocBLAS kernels initializes fine and aborts at the first batched GEMM.
149 """
150 if not any(d.backend in _AMD_BACKENDS for d in devices):
151 return
152 shipped = _bundled_rocblas_gfx_targets(binary)
153 if shipped is None:
154 return
155 if (override := _hsa_override_gfx()) is not None:
156 _warn_if_override_uncovered(override, shipped)
157 return
158 host = _host_amd_gfx_targets()
159 if not host or host <= shipped:
160 return
161 if host & shipped:
162 log.warning(
163 "This host has AMD GPU(s) with target %s, which this build ships no GPU "
164 "kernels for; the engine will abort if a model is placed on one. Restrict "
165 "HIP_VISIBLE_DEVICES to the supported cards, or set %s if the card is a "
166 "near miss of a shipped target (gfx1031 runs gfx1030 kernels with "
167 "%s=10.3.0).",
168 ", ".join(sorted(host - shipped)),
169 _HSA_OVERRIDE_VAR,
170 _HSA_OVERRIDE_VAR,
171 )
172 return
173 raise ProviderError(
174 f"This host's AMD GPU is {', '.join(sorted(host))}, but this engine build ships "
175 f"GPU kernels only for: {', '.join(sorted(shipped))}. The engine would start and "
176 "then abort at the first matrix multiplication, so it is refused up front.\n"
177 f"If the card is a near miss of a shipped target, set {_HSA_OVERRIDE_VAR} to that "
178 f"target's version (a gfx1031 card runs the gfx1030 kernels with "
179 f"{_HSA_OVERRIDE_VAR}=10.3.0). Otherwise install lilbee's Vulkan build, which "
180 "supports AMD cards ROCm does not."
181 )
184def assert_rocm_devices_usable(binary: Path, devices: list[FleetDevice], probe_output: str) -> None:
185 """Fail loud when a ROCm build cannot serve the AMD hardware in front of it.
187 *devices* and *probe_output* are the engine's own ``--list-devices`` result.
188 An enumerated card is checked against the bundled rocBLAS kernels; an empty
189 list on an AMD host means the runtime loaded and enumerated no device, which
190 would otherwise silently fall back to CPU.
191 """
192 if not sys.platform.startswith("linux"):
193 return
194 _assert_rocblas_covers_enumerated_devices(binary, devices)
195 if devices:
196 return
197 if not (_links_hip_runtime(binary, dict(os.environ)) and _amd_gpu_present()):
198 return
199 if not _amd_discrete_gpu_proven():
200 # An APU is an AMD GPU by every check above: amdgpu exposes /dev/kfd for
201 # integrated parts too, and the iGPU carries vendor 0x1002. But AMD's
202 # population of GPUs ROCm legitimately does not support is large, and an
203 # unsupported gfx target is the normal case for an APU rather than a
204 # misconfiguration. Failing here would stop the engine on a laptop where
205 # CPU serving worked, which is worse than the slow fallback this guard
206 # exists to catch, so say so and let it start.
207 log.warning(
208 "The engine links the ROCm/HIP runtime and this host has an AMD GPU, but it "
209 "enumerated no device, so GPU work will fall back to CPU. No discrete AMD card "
210 "was found, so this is most likely an APU whose gfx target this ROCm build does "
211 "not support (check with 'rocminfo').%s The engine reported: %s",
212 _rocm_support_facts(binary),
213 device_probe_diagnostic(probe_output),
214 )
215 return
216 raise ProviderError(
217 "The engine links the ROCm/HIP runtime and this host has an AMD GPU, but it "
218 "enumerated no device, so GPU work would silently fall back to CPU.\n"
219 f"The engine reported: {device_probe_diagnostic(probe_output)}\n"
220 "Likely causes: the ROCm user-space version does not match the amdgpu kernel "
221 "driver; the GPU's gfx target is not supported by this ROCm build (check with "
222 "'rocminfo'); no read/write permission on /dev/kfd (the user is usually added "
223 "to the 'render' and 'video' groups); or a restrictive ROCR_VISIBLE_DEVICES or "
224 f"HIP_VISIBLE_DEVICES.{_rocm_support_facts(binary)}"
225 )