Coverage for src/lilbee/providers/fleet/devices.py: 100%
214 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Enumerate and pin GPUs using the llama-server binary's own device view.
3The hazard this avoids: a device index from one API (Vulkan) is meaningless to
4another (CUDA); the same ordinal can be a different physical card. So both
5enumeration and pinning go through the binary's native backend index space,
6obtained from ``llama-server --list-devices``. The Vulkan VRAM probe is only a
7fallback when the binary can't enumerate. See docs/architecture.md.
8"""
10from __future__ import annotations
12import logging
13import re
14import subprocess
15from collections.abc import Mapping
16from dataclasses import dataclass
17from pathlib import Path
19from lilbee.providers.base import ProviderError, ProviderErrorKind
20from lilbee.providers.fleet.gpu_select import USABLE_VULKAN_TYPES, VkDeviceType
21from lilbee.providers.fleet.proc import run_bounded
22from lilbee.providers.roles import EngineBackend
24log = logging.getLogger(__name__)
26_PROVIDER = "llama-server"
27_LIST_DEVICES_TIMEOUT_S = 60.0
28# How long to wait for a killed probe to be reaped before abandoning it: a child
29# wedged in uninterruptible GPU-driver I/O ignores even SIGKILL.
30_PROBE_KILL_WAIT_S = 5.0
31# How much of the probe's own output to quote in a diagnostic. Enough to carry
32# the driver's error line, short enough to stay a readable message.
33_PROBE_TAIL_CHARS = 400
34_TOPO_TIMEOUT_S = 15.0
35_GPU_LABEL_RE = re.compile(r"^GPU(\d+)$")
36# llama-server prints this before the device loop, so a run that lists no GPUs
37# still prints it. Its absence means the binary never got as far as enumerating.
38_DEVICE_LIST_HEADER = "Available devices:"
39# nvidia-smi emits SGR escapes (e.g. an underlined header) even when stdout is
40# not a tty; strip them or the header's GPU labels never match.
41_ANSI_SGR_RE = re.compile(r"\x1b\[[0-9;]*m")
42# A topo-matrix header is 2+ leading GPU labels; a data row has exactly one. And
43# a link needs at least two GPUs to exist between.
44_TOPO_MIN_GPUS = 2
45MIB = 1024 * 1024
46# Per-backend visible-devices env vars (the probe inherits them; the children
47# re-emit them, composed through any parent restriction).
48_CUDA_VISIBLE_VAR = "CUDA_VISIBLE_DEVICES"
49_CUDA_ORDER_VAR = "CUDA_DEVICE_ORDER"
50_PCI_BUS_ID_ORDER = "PCI_BUS_ID"
51_ROCR_VISIBLE_VAR = "ROCR_VISIBLE_DEVICES"
52_HIP_VISIBLE_VAR = "HIP_VISIBLE_DEVICES"
53# ROCm's third numeric visibility variable, filtering exactly as the other two do.
54_GPU_DEVICE_ORDINAL_VAR = "GPU_DEVICE_ORDINAL"
55_VK_VISIBLE_VAR = "GGML_VK_VISIBLE_DEVICES"
56# " CUDA0: NVIDIA GeForce RTX 3090 (24268 MiB, 23500 MiB free)"
57_DEVICE_RE = re.compile(
58 r"^\s*([A-Za-z]+)(\d+):\s*(.+?)\s*\((\d+)\s*MiB(?:,\s*(\d+)\s*MiB\s*free)?\)\s*$"
59)
60# The engine's own name for the backend. Vendor-agnostic, so several rules key
61# on it: a Vulkan device's type has to be asked of the loader, and its util
62# source is chosen by the vendor in its device name rather than by the backend.
63VULKAN_BACKEND = "Vulkan"
64# One row per backend string the engine prints: its pin priority, and the name a
65# client reports for it. Several engine names mean one backend (HIP is ROCm, MTL
66# is Metal), so a surface that printed the raw name told two hosts apart that are
67# the same machine class.
68#
69# The two tables below are derived from this one rather than transcribed beside
70# it. Adding a backend is what makes a new GPU class usable at all, and with two
71# hand-written tables that edit ranked the new backend while leaving it nameless,
72# so a working GPU host reported "unknown".
73_BACKENDS: dict[str, tuple[int, EngineBackend]] = {
74 "CUDA": (3, EngineBackend.CUDA),
75 "ROCm": (3, EngineBackend.ROCM),
76 "HIP": (3, EngineBackend.ROCM),
77 "MTL": (3, EngineBackend.METAL),
78 "Metal": (3, EngineBackend.METAL),
79 "SYCL": (2, EngineBackend.SYCL),
80 VULKAN_BACKEND: (1, EngineBackend.VULKAN),
81}
82# Pin priority when a build reports more than one GPU backend: a real GPU backend
83# always wins over Vulkan, which wins over CPU.
84_BACKEND_RANK = {name: rank for name, (rank, _) in _BACKENDS.items()}
85# The engine's own backend names mapped to the name a client reports.
86_REPORTED_BACKEND = {name: reported for name, (_, reported) in _BACKENDS.items()}
87# Backends whose memory is always the host's: Apple Silicon reports a working-set
88# slice of system RAM, never a dedicated pool.
89_UNIFIED_BACKENDS = frozenset({"MTL", "Metal"})
90# Below this, a reported total is a BIOS carveout rather than a card's own pool.
91# An APU hands out a fixed slice of system RAM as "VRAM", often a few hundred
92# MiB, and planned as a dedicated device that size it refuses every role while
93# the machine has the whole system's memory to share. No real discrete GPU worth
94# serving from ships with less.
95_DEDICATED_VRAM_FLOOR = 2 * 1024 * MIB
98@dataclass(frozen=True)
99class FleetDevice:
100 """One GPU as the binary's backend enumerates it (native index space)."""
102 backend: str
103 index: int
104 name: str
105 total_bytes: int
106 free_bytes: int
107 # Whether this device's memory is the host's memory. An integrated GPU or an
108 # Apple Silicon Mac has no dedicated VRAM, so its reported total is a slice
109 # of the same RAM the OS and every other process is using, and placement
110 # must stay inside the system budget rather than treating it as headroom.
111 unified: bool = False
112 # Whether this device came from the host's Vulkan loader rather than from the
113 # engine's own listing. Its index is then a raw loader ordinal, which is a
114 # different space from the one the engine names its devices in, so it can be
115 # sized against but never pinned by.
116 from_loader: bool = False
119@dataclass(frozen=True)
120class DeviceProbe:
121 """The device probe's parsed devices plus its raw output for diagnostics."""
123 devices: list[FleetDevice]
124 output: str
125 # Whether the engine answered --list-devices at all: exited cleanly and
126 # printed the header it always prints. False means the binary does not speak
127 # this protocol (a build predating the flag prints usage text and exits
128 # non-zero), so its silence about devices is not a statement that there are
129 # none. Defaults False so a probe that never ran is never mistaken for one
130 # that ran and found nothing.
131 spoke_protocol: bool = False
132 # Whether the engine listed GPU devices and every one was rejected. Distinct
133 # from a host that simply has none: the engine will still pick one of those
134 # devices at launch unless it is told not to.
135 refused_all: bool = False
136 # Which backend the engine selected. Defaults to UNKNOWN so a probe that never
137 # ran reports no claim rather than CPU.
138 backend: EngineBackend = EngineBackend.UNKNOWN
141def _parse_topo_matrix(topo_text: str) -> tuple[set[int], set[frozenset[int]]]:
142 """GPU row indices and NVLink-joined pairs from ``nvidia-smi topo -m`` output.
144 The matrix header row labels the GPU columns; each ``GPU<r>`` row lists the
145 link type to each column (``NV#`` is NVLink; ``PIX``/``PHB``/``SYS`` are PCIe).
146 """
147 header_cols: list[int] = []
148 gpu_rows: set[int] = set()
149 pairs: set[frozenset[int]] = set()
150 for line in _ANSI_SGR_RE.sub("", topo_text).splitlines():
151 tokens = line.split()
152 # Leading run of GPU-label tokens: the header is all labels (>=2), a data
153 # row is one label ("GPU3") followed by link-type cells. split() strips the
154 # header's leading whitespace, so this run length is what tells them apart.
155 leading_labels: list[int] = []
156 for token in tokens:
157 match = _GPU_LABEL_RE.match(token)
158 if match is None:
159 break
160 leading_labels.append(int(match.group(1)))
161 if len(leading_labels) >= _TOPO_MIN_GPUS:
162 header_cols = leading_labels
163 elif len(leading_labels) == 1:
164 row_idx = leading_labels[0]
165 gpu_rows.add(row_idx)
166 for col_idx, cell in zip(header_cols, tokens[1:], strict=False):
167 if row_idx != col_idx and cell.startswith("NV"):
168 pairs.add(frozenset({row_idx, col_idx}))
169 return gpu_rows, pairs
172def host_lacks_nvlink() -> bool:
173 """Whether this host's GPUs are joined only by PCIe (no NVLink anywhere).
175 Tensor-splitting a large model across PCIe-only cards is all-reduce bound and
176 much slower than over NVLink. Deliberately a host-level claim: the fleet's
177 device indices live in the serving binary's backend index space, which does
178 not map onto ``nvidia-smi``'s physical numbering under a visible-devices
179 restriction (the very hazard this module exists to avoid), so per-pair
180 verdicts against plan indices would be unreliable. Returns False (no claim)
181 when the probe fails or reports fewer than two GPUs, so a non-NVIDIA or
182 single-card host stays silent rather than warning wrongly.
183 """
184 try:
185 stdout, _ = run_bounded(
186 ["nvidia-smi", "topo", "-m"],
187 timeout_s=_TOPO_TIMEOUT_S,
188 kill_wait_s=_PROBE_KILL_WAIT_S,
189 label="nvidia-smi topo",
190 )
191 except (OSError, subprocess.SubprocessError):
192 return False
193 gpu_rows, pairs = _parse_topo_matrix(stdout)
194 return len(gpu_rows) >= _TOPO_MIN_GPUS and not pairs
197def _probe_env(binary: Path) -> dict[str, str]:
198 """Env for the probe: stable PCI ordering so CUDA indices match what we pin.
200 A preset ``CUDA_DEVICE_ORDER`` is respected; ``visible_env`` re-emits the same
201 order var, so the probe and the spawned servers see one device ordering. The
202 probe runs under the engine's own environment (pin, loader safety) and with the
203 CUDA runtime libraries *binary* needs on its library path, so a zero-device
204 result reflects an unusable GPU rather than a probe that lacked them.
205 """
206 from lilbee.providers.fleet.cuda_runtime import cuda_runtime_env
207 from lilbee.providers.fleet.gpu_env import engine_environ
209 env = engine_environ()
210 env.update(cuda_runtime_env(binary))
211 env.setdefault(_CUDA_ORDER_VAR, _PCI_BUS_ID_ORDER)
212 return env
215def probe_devices(binary: Path, *, timeout_s: float = _LIST_DEVICES_TIMEOUT_S) -> DeviceProbe:
216 """Parse ``<binary> --list-devices``; empty devices when unavailable/unparseable.
218 Filtered to a single GPU backend (the highest-ranked one present) so device
219 indices are unambiguous when a build exposes several backends. A probe that
220 does not respond within *timeout_s* raises a ``ProviderError`` naming the
221 stuck probe: that is a wedged GPU driver, not a GPU-less host, and treating
222 it as "no devices" would silently plan a CPU fleet on a GPU box.
223 """
224 try:
225 output, returncode = _run_list_devices(binary, timeout_s)
226 except (OSError, subprocess.SubprocessError) as exc:
227 # Silently returning an empty probe here made an unrunnable binary look
228 # exactly like a host with no GPU, and the fleet planned for CPU with
229 # nothing said. The reason is the whole diagnosis: a wrong architecture,
230 # a missing loader, a permission denial.
231 log.warning(
232 "Could not run the GPU device probe (%s --list-devices): %s. Continuing "
233 "as though this host has no GPU; check that the engine binary is "
234 "executable and built for this machine.",
235 binary.name,
236 exc,
237 )
238 return DeviceProbe([], "")
239 parsed = _parse_devices(output)
240 selected = _select_backend(parsed)
241 offered = [d for d in parsed if d.backend in _BACKEND_RANK]
242 answered = _DEVICE_LIST_HEADER in output
243 spoke = returncode == 0 and answered
244 if not spoke and answered:
245 # It knew the flag and started answering, then died. Blaming the flag
246 # here sent the reader looking for the wrong engine build, when what
247 # they have is a crash partway through enumeration.
248 log.warning(
249 "%s --list-devices printed its device header then crashed (exit %d), so the "
250 "device list may be incomplete. This is usually a GPU driver or ICD fault "
251 "during enumeration. The probe reported: %s",
252 binary.name,
253 returncode,
254 _probe_tail(output),
255 )
256 elif not spoke:
257 log.warning(
258 "%s --list-devices exited %d without printing its device header, so it "
259 "does not appear to support the flag. Falling back to the host's Vulkan "
260 "loader to find GPUs; set %s if this is not the engine you meant to use.",
261 binary.name,
262 returncode,
263 "LILBEE_ENGINE_DIR",
264 )
265 return DeviceProbe(
266 selected,
267 output,
268 spoke_protocol=spoke,
269 refused_all=bool(offered) and not selected,
270 backend=_selected_backend(selected, spoke=spoke),
271 )
274def _selected_backend(selected: list[FleetDevice], *, spoke: bool) -> EngineBackend:
275 """Which backend the engine selected, or UNKNOWN when it did not answer.
277 An empty device list means CPU only when the engine answered the question. A
278 binary that never spoke the protocol said nothing about its backend, and
279 calling that CPU labels a GPU host as a CPU one in every diagnostic.
281 A selected device's backend is always a row of ``_BACKENDS``, because that is
282 what the selector filters on, so this indexes rather than defaulting. A
283 default here would turn a table the selector and the reporter disagree about
284 into a working GPU host that reports "unknown", which is this function's own
285 defect one backend along.
286 """
287 if selected:
288 return _REPORTED_BACKEND[selected[0].backend]
289 return EngineBackend.CPU if spoke else EngineBackend.UNKNOWN
292def _run_list_devices(binary: Path, timeout_s: float) -> tuple[str, int]:
293 """Run the probe with a bounded reap; raise on timeout.
295 A probe wedged in uninterruptible GPU-driver I/O would otherwise hang the
296 caller forever, since ``subprocess.run``'s timeout waits unbounded for the
297 reap; ``run_bounded`` abandons an unkillable child after a short wait.
299 The probe holds a device context and writes no state file, so nothing can reap
300 it later by record. It is the one caller that opts into the lifetime binding,
301 where the kernel offers one, and it is killed on the way out of every abort,
302 not just the timeout.
303 """
304 try:
305 return run_bounded(
306 [str(binary), "--list-devices"],
307 timeout_s=timeout_s,
308 kill_wait_s=_PROBE_KILL_WAIT_S,
309 env=_probe_env(binary),
310 merge_stderr=True,
311 label=f"{binary.name} --list-devices",
312 bind_lifetime=True,
313 )
314 except subprocess.TimeoutExpired as exc:
315 # Whatever the probe managed to print before it wedged says more than any
316 # fixed advice can, and the fixed advice named one vendor's tool at a host
317 # that may have neither that vendor nor that tool.
318 raise ProviderError(
319 f"The GPU device probe ({binary.name} --list-devices) did not respond "
320 f"within {timeout_s:.0f}s, so the engine cannot start. The GPU driver is "
321 "most likely wedged; check that your vendor's tool responds (nvidia-smi, "
322 "rocm-smi, xpu-smi) and reboot the host if it hangs.\n"
323 f"The probe reported: {_probe_tail(_decoded_output(exc.output))}",
324 provider=_PROVIDER,
325 kind=ProviderErrorKind.SERVER,
326 ) from None
329def _decoded_output(output: object) -> str:
330 """Partial child output from a timeout, which arrives as bytes even under text mode."""
331 if isinstance(output, bytes):
332 return output.decode(errors="replace")
333 return output if isinstance(output, str) else ""
336def _probe_tail(output: str) -> str:
337 """The tail of what the probe printed, for a message that has to stay readable."""
338 text = output.strip()
339 return text[-_PROBE_TAIL_CHARS:] if text else "(nothing)"
342def _parse_devices(text: str) -> list[FleetDevice]:
343 devices: list[FleetDevice] = []
344 # Sampled at most once per parse, and only when a line actually needs it:
345 # free memory is live, so it is read fresh here rather than cached, and the
346 # loader must not be opened once per device line to answer the same question.
347 loader_free: dict[str, int] | None = None
348 for line in text.splitlines():
349 match = _DEVICE_RE.match(line)
350 if match is None:
351 continue
352 backend, index, name, total_mib, free_mib = match.groups()
353 total = int(total_mib) * MIB
354 if total == 0:
355 # No memory is not a small GPU, it is one that cannot hold a model:
356 # a driver listing an adapter before its memory is queryable. Kept, it
357 # is the smallest card in the fleet and collapses every budget sized
358 # against the smallest, while the non-empty list switches off the
359 # shared-memory budget a host with no usable GPU depends on.
360 log.warning(
361 "Ignoring GPU %s%s (%s): it reports no memory, so nothing can be "
362 "placed on it. Check the GPU driver if this device should be usable.",
363 backend,
364 index,
365 name.strip(),
366 )
367 continue
368 if free_mib:
369 free = int(free_mib) * MIB
370 else:
371 if loader_free is None:
372 loader_free = _loader_free_bytes(backend)
373 free = loader_free.get(name.strip(), total)
374 devices.append(
375 FleetDevice(
376 backend,
377 int(index),
378 name.strip(),
379 total,
380 free,
381 unified=_is_unified(backend, name.strip()) or total < _DEDICATED_VRAM_FLOOR,
382 )
383 )
384 return devices
387# Mesa and friends expose CPU rasterizers through the Vulkan loader, and
388# llama.cpp's Vulkan backend enumerates them exactly like a GPU: same
389# "VulkanN: <name> (<total> MiB, <free> MiB free)" shape, with system RAM
390# reported as VRAM. Planning against one is worse than having no GPU at all,
391# because the "VRAM" looks enormous: a host with a real iGPU beside lavapipe
392# can be planned as a two-GPU machine and tensor-split across a real adapter
393# and a software renderer, which runs orders of magnitude slower than either
394# CPU inference or the iGPU alone.
395_SOFTWARE_RENDERER_MARKERS = ("llvmpipe", "lavapipe", "softpipe", "swiftshader")
398def _is_software_renderer(device: FleetDevice) -> bool:
399 """Whether *device* is a CPU rasterizer masquerading as a GPU.
401 A name test, so it only recognizes the rasterizers it already knows, and a
402 renamed or newly written one walks past it. It stays as the answer for hosts
403 where the Vulkan loader can't be opened from this process and the device
404 type is therefore unavailable; where the type is available,
405 ``_is_unusable_vulkan`` decides and this never gets the chance to be wrong.
406 """
407 name = device.name.casefold()
408 return any(marker in name for marker in _SOFTWARE_RENDERER_MARKERS)
411def _loader_free_bytes(backend: str) -> dict[str, int]:
412 """Live free memory per device name, for a listing that printed no free figure.
414 ggml omits the figure when the driver has no ``VK_EXT_memory_budget``, and
415 treating the omission as "all of it" is how a desktop holding gigabytes of
416 compositor and browser VRAM was planned as an empty card. The loader exposes
417 that extension to this process even when the engine build cannot use it, so
418 it is asked directly; a name it cannot speak for keeps the heap size.
420 Empty for any other backend: the Vulkan loader knows nothing about the
421 devices a CUDA or ROCm listing names.
422 """
423 if backend != VULKAN_BACKEND:
424 return {}
425 from lilbee.providers.fleet.gpu_select import vulkan_free_bytes_by_name
427 return vulkan_free_bytes_by_name()
430def _vulkan_device_type(name: str) -> VkDeviceType | None:
431 """The loader's type for the Vulkan adapter the engine printed as *name*.
433 ``None`` when the loader can't be reached or reports no adapter by that
434 name, which reads as "no opinion": the device is kept and assumed dedicated,
435 preserving the behaviour of hosts that never had a type to consult.
436 """
437 from lilbee.providers.fleet.gpu_select import vulkan_device_types_by_name
439 return vulkan_device_types_by_name().get(name)
442def _is_unusable_vulkan(device: FleetDevice) -> bool:
443 """Whether *device* is a Vulkan adapter ggml would not choose to run on.
445 ggml's Vulkan backend builds its device pool from discrete and integrated
446 adapters only, and falls back to the first non-CPU adapter when it finds
447 neither. In a VM that fallback is a paravirtual adapter (VMware SVGA,
448 VirtIO-GPU Venus, QXL), which reports guest RAM as VRAM and is typically
449 compute-incomplete or fails at allocation. Planning a fleet onto one costs
450 more than planning no GPU at all, since a non-empty device list also turns
451 off the shared-RAM budget.
453 Only a positive claim counts. VIRTUAL_GPU and CPU are the loader naming what
454 the adapter is; OTHER is it declining to, and refusing on a shrug took the
455 GPU away from real hardware whose driver simply does not classify itself.
456 """
457 if device.backend != VULKAN_BACKEND:
458 return False
459 device_type = _vulkan_device_type(device.name)
460 if device_type is None or device_type in USABLE_VULKAN_TYPES:
461 return False
462 # OTHER is the loader shrugging, not an accusation. The spec's own wording is
463 # "does not match any other available types", which a driver reaches for when
464 # it cannot classify itself, and some real adapters do. Refusing on it took a
465 # working GPU away from a machine the engine had already listed one for.
466 # VIRTUAL_GPU and CPU are positive claims and keep their veto.
467 return device_type is not VkDeviceType.OTHER
470def _is_unified(backend: str, name: str) -> bool:
471 """Whether the device *backend* printed as *name* shares its memory with the host.
473 Metal is unified by construction on Apple Silicon: the figure it reports is
474 ``recommendedMaxWorkingSetSize``, a slice of system RAM rather than a
475 separate pool. For Vulkan the loader knows the device type, so the type is
476 asked for rather than guessed; a size heuristic cannot work here, since a
477 24 GB discrete card in a 32 GB host and an Apple GPU reporting two thirds of
478 RAM are indistinguishable by proportion.
480 CUDA, ROCm and SYCL print no type at all, and an AMD APU or a Jetson looks
481 exactly like a discrete card there while reporting system RAM as its memory.
482 Those fall back to a question about the machine rather than the device: a
483 host whose Vulkan loader sees adapters but no discrete one has no discrete
484 GPU for another backend to be enumerating.
485 """
486 if backend in _UNIFIED_BACKENDS:
487 return True
488 if backend == VULKAN_BACKEND:
489 return _vulkan_device_type(name) is VkDeviceType.INTEGRATED_GPU
490 from lilbee.providers.fleet.gpu_select import host_has_no_discrete_gpu
492 return host_has_no_discrete_gpu()
495def _select_backend(devices: list[FleetDevice]) -> list[FleetDevice]:
496 """Keep one GPU backend's devices: highest rank, then most memory.
498 Returns a single backend so pinning is unambiguous: ``visible_env`` keys off
499 one backend, and mixing index spaces is the very hazard this module avoids.
501 CUDA, ROCm, HIP and Metal all rank alike, and a build that loads several
502 backends (``ggml_backend_load_all`` does) makes the tie real. Breaking it on
503 the backend's name meant a host with a 4090 beside an RX 6600 planned onto
504 the AMD card because "ROCm" sorts after "CUDA", and the NVIDIA card idled
505 with nothing said. Total memory decides instead; the name is only the last
506 resort that keeps the choice deterministic.
507 """
508 ranked = [
509 d
510 for d in devices
511 if d.backend in _BACKEND_RANK
512 and not _is_software_renderer(d)
513 and not _is_unusable_vulkan(d)
514 ]
515 if not ranked:
516 return []
517 by_backend: dict[str, list[FleetDevice]] = {}
518 for device in ranked:
519 by_backend.setdefault(device.backend, []).append(device)
520 backend, chosen = max(by_backend.items(), key=_backend_preference)
521 for other, group in by_backend.items():
522 if other != backend:
523 log.info(
524 "Engine reports %d %s device(s) beside %d %s device(s); planning onto %s, "
525 "which has more memory. Backends cannot be mixed: their device indexes "
526 "name different cards.",
527 len(group),
528 other,
529 len(chosen),
530 backend,
531 backend,
532 )
533 return chosen
536def _backend_preference(item: tuple[str, list[FleetDevice]]) -> tuple[int, int, int, str]:
537 """Sort key for choosing one backend's devices: rank, dedicated bytes, size.
539 Dedicated bytes come before raw size because the discrete backends all tie at
540 the same rank, and a shared-heap carveout reports a total that is host RAM
541 the host budget already counts. Left on raw size, an APU advertising a large
542 carveout beat a discrete card, which was then discarded and left idle while
543 the plan double-promised memory it did not have.
544 """
545 backend, group = item
546 dedicated = sum(d.total_bytes for d in group if not d.unified)
547 return _BACKEND_RANK[backend], dedicated, sum(d.total_bytes for d in group), backend
550def _compose_visible(indices: list[int], parent_value: str | None) -> str:
551 """Visible-devices value naming the same physical devices the probe saw.
553 When the parent env already restricts the var, the probe's indices are
554 relative to that comma-separated list (integer or UUID entries), so each
555 index maps through it; the child's value then names the same physical
556 devices instead of being re-interpreted as absolute.
557 """
558 if parent_value is None:
559 return ",".join(str(i) for i in indices)
560 entries = [entry.strip() for entry in parent_value.split(",") if entry.strip()]
561 out: list[str] = []
562 for i in indices:
563 if i >= len(entries):
564 # The probe enumerates devices under the parent restriction, so every
565 # index must map into it. An out-of-range index is an invariant
566 # violation; emitting a bare ``str(i)`` would pin an absolute integer
567 # into a possibly UUID-namespaced list, silently selecting the wrong
568 # GPU. Fail loudly instead.
569 raise ValueError(
570 f"device index {i} is outside the parent visible-devices list "
571 f"{parent_value!r}; cannot compose a child pin without selecting the wrong GPU"
572 )
573 out.append(entries[i])
574 return ",".join(out)
577def visible_env(devices: tuple[FleetDevice, ...]) -> dict[str, str]:
578 """Env that pins a child to *devices* via the right var for their backend.
580 Indices are the backend-native ones from ``probe_devices``, composed through
581 any parent visible-devices restriction so the child names the same physical
582 devices the probe enumerated; no cross-API index translation occurs.
583 """
584 if not devices:
585 return {}
586 from lilbee.providers.fleet.gpu_env import gpu_environ
588 env = gpu_environ()
589 backend = devices[0].backend
590 indices = [d.index for d in devices]
591 if backend == "CUDA":
592 return {
593 _CUDA_VISIBLE_VAR: _compose_visible(indices, env.get(_CUDA_VISIBLE_VAR)),
594 _CUDA_ORDER_VAR: env.get(_CUDA_ORDER_VAR, _PCI_BUS_ID_ORDER),
595 }
596 if backend in ("ROCm", "HIP"):
597 return _amd_visible_env(indices, env)
598 if backend == VULKAN_BACKEND:
599 # Deliberately not GGML_VK_VISIBLE_DEVICES. That variable indexes the raw
600 # loader enumeration, while these indices come from the engine's own
601 # filtered list, so the two disagree wherever ggml drops or merges a
602 # device -- two ICDs for one card being the clear case. Setting it also
603 # disables ggml's type filter, support check and dedup. Vulkan is pinned
604 # with --device instead, in the same space the names were parsed from.
605 return {}
606 if backend == "SYCL":
607 # Deliberately no ONEAPI_DEVICE_SELECTOR. It is a selector grammar over a
608 # backend runtime, not the index space --list-devices numbers, so a
609 # composed level_zero ordinal can name a different physical card than the
610 # one the probe enumerated. SYCL pins by --device instead, in the space
611 # the indices were read from. An inherited parent selector still applies:
612 # the engine enumerated behind it, so its names are already relative to it.
613 return {}
614 return {}
617def amd_visible_var(env: Mapping[str, str]) -> str:
618 """The one AMD visibility var in *env* an index list may be written to.
620 ``ROCR_VISIBLE_DEVICES`` and ``HIP_VISIBLE_DEVICES`` are applied sequentially:
621 ROCr filters first, then HIP re-indexes within the survivors. Writing the same
622 indices to both double-filters and selects the wrong cards, or none at all.
623 ``GPU_DEVICE_ORDINAL`` is the third and filters the same way, so writing HIP
624 over an ordinal mask both overrides it and re-exposes cards it had hidden.
626 So exactly one is ever written: whichever the environment already restricts,
627 in the runtime's precedence (HIP, then the ordinal, then ROCr), or HIP when
628 nothing restricts. An empty value means "no devices" rather than "this is the
629 variable in use", so it does not claim precedence. Every caller writing an AMD
630 pin asks here; two callers each picking their own would put the pair back.
631 """
632 for name in (_HIP_VISIBLE_VAR, _GPU_DEVICE_ORDINAL_VAR, _ROCR_VISIBLE_VAR):
633 if env.get(name, "").strip():
634 return name
635 return _HIP_VISIBLE_VAR
638def _amd_visible_env(indices: list[int], env: Mapping[str, str]) -> dict[str, str]:
639 """Pin an AMD ROCm/HIP child to the probe's *indices* with one visibility var.
641 The probe enumerated a single index space already filtered by whichever var
642 the parent set, so the chosen var is composed against that parent value and
643 the other is left inherited untouched. The child inherits the parent env, so
644 an unset override keeps any inherited sibling var in force.
645 """
646 var = amd_visible_var(env)
647 return {var: _compose_visible(indices, env.get(var))}