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