Coverage for src/lilbee/providers/model_cache.py: 100%
138 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Loader-mode constants and dynamic-context / GPU-memory helpers for llama-server."""
3from __future__ import annotations
5import logging
6import platform
7from collections.abc import Callable
8from enum import StrEnum
9from pathlib import Path
11from lilbee.providers.fleet.gpu_env import gpu_environ
13log = logging.getLogger(__name__)
16class LoaderMode(StrEnum):
17 """Which task to configure llama.cpp for at load time."""
19 CHAT = "chat"
20 EMBED = "embed"
21 RERANK = "rerank"
24# Fallback KV cache estimate when GGUF metadata can't be read.
25# 2048 bytes/token undershoots real KV size for modern models (Gemma3-4B is
26# ~640 KB/token f16) but is fine as a coarse pre-load eviction signal.
27_KV_BYTES_PER_CTX_TOKEN = 2048
29# Metal/CUDA buffer overhead as fraction of model weight memory
30_BUFFER_OVERHEAD_FRACTION = 0.10
32# Default context length for estimation when metadata unavailable
33_DEFAULT_CTX_LEN = 2048
35# Floor for the dynamic n_ctx computation (smaller is unusable for chat)
36_DYNAMIC_CTX_FLOOR = 512
38# Round dynamic n_ctx down to a multiple of this (clean batch sizes)
39_DYNAMIC_CTX_QUANTUM = 256
41# KV cache element size for f16 (bytes). Quantized KV reduces this.
42_KV_ELEM_BYTES_F16 = 2
45def kv_bytes_per_token(
46 meta: dict[str, str] | None,
47 k_elem_bytes: float = _KV_ELEM_BYTES_F16,
48 v_elem_bytes: float | None = None,
49) -> int:
50 """Estimate per-token KV cache size in bytes from GGUF metadata.
52 Formula: n_layers * n_kv_heads * (k_dim * k_elem_bytes + v_dim * v_elem_bytes).
53 K and V are always one type in a launch, so callers pass the same cost for
54 both. ``v_elem_bytes`` defaults to the K cost. Falls back to
55 ``_KV_BYTES_PER_CTX_TOKEN`` when metadata is missing.
56 """
57 if v_elem_bytes is None:
58 v_elem_bytes = k_elem_bytes
59 if not meta:
60 return _KV_BYTES_PER_CTX_TOKEN
61 try:
62 n_layers = int(meta["block_count"])
63 head_count_kv = int(meta.get("head_count_kv") or meta["head_count"])
64 if "key_length" in meta and "value_length" in meta:
65 k_dim = int(meta["key_length"])
66 v_dim = int(meta["value_length"])
67 else:
68 embed = int(meta["embedding_length"])
69 head_count = int(meta.get("head_count") or head_count_kv)
70 k_dim = v_dim = embed // head_count
71 except (KeyError, ValueError, ZeroDivisionError):
72 return _KV_BYTES_PER_CTX_TOKEN
73 return int(n_layers * head_count_kv * (k_dim * k_elem_bytes + v_dim * v_elem_bytes))
76def estimate_model_memory(
77 model_path: Path,
78 n_ctx: int = _DEFAULT_CTX_LEN,
79 kv_bytes_per_tok: int = _KV_BYTES_PER_CTX_TOKEN,
80) -> int:
81 """Estimate memory consumption for a GGUF model.
82 Approximation: file_size (weights) + KV cache + 10% buffer overhead.
83 """
84 file_bytes = model_path.stat().st_size if model_path.exists() else 0
85 kv_bytes = n_ctx * kv_bytes_per_tok
86 overhead = int(file_bytes * _BUFFER_OVERHEAD_FRACTION)
87 return file_bytes + kv_bytes + overhead
90def compute_dynamic_ctx(
91 *,
92 model_bytes: int,
93 available_bytes: int,
94 training_ctx: int,
95 kv_bytes_per_tok: int,
96 ceiling: int,
97 target: int | None = None,
98 floor: int = _DYNAMIC_CTX_FLOOR,
99 quantum: int = _DYNAMIC_CTX_QUANTUM,
100) -> int:
101 """Pick the n_ctx that best fits target, ceiling, and ``available_bytes``.
103 Selection rule, in order:
105 1. ``upper = min(training_ctx, ceiling)`` is the hard upper bound; the
106 model cannot exceed its training window and the caller may cap below it.
107 2. If ``target`` is provided, prefer it (clamped to ``[floor, upper]``)
108 so a 40K-context model still loads at 8K when chat doesn't need more,
109 rather than maximising n_ctx just because the memory allows it.
110 3. ``raw_ctx = budget // kv_bytes_per_tok`` is the largest n_ctx the
111 available memory can physically back. The result is clamped to
112 ``raw_ctx`` so we never over-allocate on memory-constrained boxes.
113 4. Result is quantized down to ``quantum`` and floored at ``floor``.
114 """
115 upper = min(training_ctx, ceiling)
116 if kv_bytes_per_tok <= 0:
117 if target is not None:
118 return max(floor, min(target, upper))
119 return upper
120 overhead = int(model_bytes * _BUFFER_OVERHEAD_FRACTION)
121 budget = available_bytes - model_bytes - overhead
122 if budget <= 0:
123 return floor
124 raw_ctx = budget // kv_bytes_per_tok
125 # Aim for target when set, but never above what the memory or training_ctx permit.
126 desired = min(target, raw_ctx, upper) if target is not None else min(raw_ctx, upper)
127 bounded = max(floor, desired)
128 quantized = (bounded // quantum) * quantum
129 return max(floor, quantized)
132def get_available_memory(fraction: float, *, total: bool = False) -> int:
133 """Return usable GPU/unified memory in bytes, scaled by *fraction*.
134 - macOS (Apple Silicon): unified memory via psutil
135 - Linux with NVIDIA GPU: pynvml -> nvidia-smi -> psutil fallback
136 - Other: psutil system memory
138 With multiple NVIDIA GPUs, *total* sums every card's memory (whole-fleet
139 capacity, for deciding whether a model can run tensor-split across all of
140 them); the default sizes against the smallest single card.
142 A coarse figure for callers with no device list to hand. The fleet has one
143 and sizes against it instead
144 (:func:`lilbee.providers.fleet.planning.plan_sizing_budget`), because this
145 answers with system RAM on every host without an NVIDIA card. That system
146 figure is the process's, cgroup cap included, not the machine's.
147 """
148 system = platform.system()
150 if system == "Darwin":
151 return int(total_system_memory() * fraction)
153 if system in ("Linux", "Windows"):
154 nvidia_mem = _try_nvidia_memory(sum if total else min)
155 if nvidia_mem is not None:
156 return int(nvidia_mem * fraction)
158 return int(total_system_memory() * fraction)
161def free_system_memory() -> int:
162 """Live allocatable system RAM in bytes (free + reclaimable), right now.
164 The load-time counterpart to :func:`get_available_memory`, which scales total
165 capacity for sizing rather than reporting what is free this instant.
167 Bounded by what this process's cgroup still has, for the reason in
168 :func:`lilbee.core.system.cgroup_memory_limit`.
169 """
170 import psutil
172 from lilbee.core.system import cgroup_memory_limit, cgroup_memory_used
174 host_free = int(psutil.virtual_memory().available)
175 limit = cgroup_memory_limit()
176 if limit is None:
177 return host_free
178 used = cgroup_memory_used()
179 return min(host_free, limit if used is None else max(0, limit - used))
182def total_system_memory() -> int:
183 """Total system RAM in bytes this process may use, cgroup cap included.
185 Raises rather than answering zero when the host cannot be read: every caller
186 here is sizing a real placement, and a budget computed from zero refuses
187 every model without saying why.
188 """
189 from lilbee.core.system import capped_total_memory
191 return capped_total_memory()
194def has_nvidia_gpu() -> bool:
195 """Whether an NVIDIA GPU is physically present on this host (NVML or nvidia-smi).
197 Deliberately unmasked. ``CUDA_VISIBLE_DEVICES`` says what a CUDA process may
198 use, not what the machine has, and the callers of this ask the second
199 question: one of them exists to delete an empty mask that an orchestrator
200 left behind, which it could never do if the empty mask hid the card first.
201 """
202 return _nvidia_device_totals() is not None
205def _try_nvidia_memory(reducer: Callable[[list[int]], int] = min) -> int | None:
206 """NVIDIA GPU total memory the CUDA runtime can actually reach, or ``None``.
208 *reducer* combines the per-device totals. ``min`` (the default) sizes against
209 the smallest card, the safe budget for a single server that has not been told
210 which card it will run on. ``sum`` gives whole-fleet capacity, used only by
211 the catalog fit chip to decide whether a model can run split across every card.
213 Restricted to the devices ``CUDA_VISIBLE_DEVICES`` exposes. Neither NVML nor
214 nvidia-smi applies that mask on its own: it is read by the CUDA runtime, and
215 both tools report every card the driver knows about. Unmasked, a container
216 given one card of an eight-card host summed all eight and approved models
217 eight times too large for the card it had, and a fleet whose smallest card
218 was masked out sized every budget against a card the engine cannot see.
219 """
220 totals = _nvidia_device_totals()
221 if not totals:
222 return None
223 visible = _apply_cuda_visible_mask(totals)
224 return reducer([total for _uuid, total in visible]) if visible else None
227def _nvidia_device_totals() -> list[tuple[str, int]] | None:
228 """``[(uuid, total_bytes), ...]`` in driver enumeration order, or ``None``.
230 ``None`` means no NVIDIA GPU was detectable at all, which is the expected
231 outcome on every non-NVIDIA host.
232 """
233 try:
234 import pynvml # type: ignore[import-untyped]
236 pynvml.nvmlInit()
237 totals = []
238 for i in range(pynvml.nvmlDeviceGetCount()):
239 handle = pynvml.nvmlDeviceGetHandleByIndex(i)
240 totals.append(
241 (
242 _decoded(pynvml.nvmlDeviceGetUUID(handle)),
243 int(pynvml.nvmlDeviceGetMemoryInfo(handle).total),
244 )
245 )
246 pynvml.nvmlShutdown()
247 if totals:
248 return totals
249 except Exception: # noqa: S110 -- optional GPU detect; absence is expected on non-NVIDIA hosts
250 pass
252 try:
253 import subprocess
255 # nvidia-smi ships with the NVIDIA driver and is always on PATH when
256 # present; fully-qualifying it would break on every install layout.
257 result = subprocess.run(
258 ["nvidia-smi", "--query-gpu=memory.total,uuid", "--format=csv,noheader,nounits"], # noqa: S607
259 capture_output=True,
260 text=True,
261 encoding="utf-8",
262 errors="replace",
263 timeout=5,
264 )
265 if result.returncode == 0:
266 rows = [_parse_smi_row(line) for line in result.stdout.strip().splitlines()]
267 parsed = [row for row in rows if row is not None]
268 if parsed:
269 return parsed
270 except Exception: # noqa: S110 -- optional GPU detect; same rationale as above
271 pass
273 return None
276def _decoded(value: str | bytes) -> str:
277 """pynvml returns ``str`` on recent versions and ``bytes`` on older ones."""
278 return value.decode() if isinstance(value, bytes) else value
281def _parse_smi_row(line: str) -> tuple[str, int] | None:
282 """One ``memory.total,uuid`` CSV row as ``(uuid, total_bytes)``.
284 The UUID column is optional so an older nvidia-smi that only echoes the
285 memory still yields a device; only a UUID-keyed mask needs it.
286 """
287 fields = [field.strip() for field in line.split(",")]
288 if not fields or not fields[0]:
289 return None
290 try:
291 mib = int(fields[0])
292 except ValueError:
293 return None
294 return (fields[1] if len(fields) > 1 else "", mib * 1024 * 1024)
297def _apply_cuda_visible_mask(devices: list[tuple[str, int]]) -> list[tuple[str, int]]:
298 """The subset of *devices* ``CUDA_VISIBLE_DEVICES`` exposes, in its order.
300 Entries are driver indexes or ``GPU-``/``MIG-`` UUIDs. An unset variable
301 masks nothing; an empty one exposes nothing. CUDA stops enumerating at the
302 first entry that names no device, and so does this, which is what makes
303 ``0,9,1`` on a two-card host mean one card rather than two.
304 """
305 raw = gpu_environ().get("CUDA_VISIBLE_DEVICES")
306 if raw is None:
307 return devices
308 visible: list[tuple[str, int]] = []
309 for entry in (part.strip() for part in raw.split(",")):
310 matched = _resolve_cuda_entry(entry, devices)
311 if matched is None:
312 break
313 visible.append(matched)
314 return visible
317def _resolve_cuda_entry(entry: str, devices: list[tuple[str, int]]) -> tuple[str, int] | None:
318 """The device an entry of ``CUDA_VISIBLE_DEVICES`` names, ``None`` if it names none."""
319 if entry.isdigit():
320 index = int(entry)
321 return devices[index] if index < len(devices) else None
322 # UUIDs may be abbreviated to any unique prefix.
323 if entry.startswith(("GPU-", "MIG-")):
324 matches = [device for device in devices if device[0].startswith(entry)]
325 return matches[0] if len(matches) == 1 else None
326 return None