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