Coverage for src/lilbee/providers/model_cache.py: 100%
136 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +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(meta: dict[str, str] | None, kv_elem_bytes: int = _KV_ELEM_BYTES_F16) -> int:
45 """Estimate per-token KV cache size in bytes from GGUF metadata.
47 Formula: 2 (K + V) * n_layers * n_kv_heads * head_dim * elem_bytes.
48 Falls back to ``_KV_BYTES_PER_CTX_TOKEN`` when metadata is missing.
49 """
50 if not meta:
51 return _KV_BYTES_PER_CTX_TOKEN
52 try:
53 n_layers = int(meta["block_count"])
54 head_count_kv = int(meta.get("head_count_kv") or meta["head_count"])
55 if "key_length" in meta and "value_length" in meta:
56 kv_dim = int(meta["key_length"]) + int(meta["value_length"])
57 else:
58 embed = int(meta["embedding_length"])
59 head_count = int(meta.get("head_count") or head_count_kv)
60 head_dim = embed // head_count
61 kv_dim = 2 * head_dim
62 except (KeyError, ValueError, ZeroDivisionError):
63 return _KV_BYTES_PER_CTX_TOKEN
64 return n_layers * head_count_kv * kv_dim * kv_elem_bytes
67def estimate_model_memory(
68 model_path: Path,
69 n_ctx: int = _DEFAULT_CTX_LEN,
70 kv_bytes_per_tok: int = _KV_BYTES_PER_CTX_TOKEN,
71) -> int:
72 """Estimate memory consumption for a GGUF model.
73 Approximation: file_size (weights) + KV cache + 10% buffer overhead.
74 """
75 file_bytes = model_path.stat().st_size if model_path.exists() else 0
76 kv_bytes = n_ctx * kv_bytes_per_tok
77 overhead = int(file_bytes * _BUFFER_OVERHEAD_FRACTION)
78 return file_bytes + kv_bytes + overhead
81def compute_dynamic_ctx(
82 *,
83 model_bytes: int,
84 available_bytes: int,
85 training_ctx: int,
86 kv_bytes_per_tok: int,
87 ceiling: int,
88 target: int | None = None,
89 floor: int = _DYNAMIC_CTX_FLOOR,
90 quantum: int = _DYNAMIC_CTX_QUANTUM,
91) -> int:
92 """Pick the n_ctx that best fits target, ceiling, and ``available_bytes``.
94 Selection rule, in order:
96 1. ``upper = min(training_ctx, ceiling)`` is the hard upper bound; the
97 model cannot exceed its training window and the caller may cap below it.
98 2. If ``target`` is provided, prefer it (clamped to ``[floor, upper]``)
99 so a 40K-context model still loads at 8K when chat doesn't need more,
100 rather than maximising n_ctx just because the memory allows it.
101 3. ``raw_ctx = budget // kv_bytes_per_tok`` is the largest n_ctx the
102 available memory can physically back. The result is clamped to
103 ``raw_ctx`` so we never over-allocate on memory-constrained boxes.
104 4. Result is quantized down to ``quantum`` and floored at ``floor``.
105 """
106 upper = min(training_ctx, ceiling)
107 if kv_bytes_per_tok <= 0:
108 if target is not None:
109 return max(floor, min(target, upper))
110 return upper
111 overhead = int(model_bytes * _BUFFER_OVERHEAD_FRACTION)
112 budget = available_bytes - model_bytes - overhead
113 if budget <= 0:
114 return floor
115 raw_ctx = budget // kv_bytes_per_tok
116 # Aim for target when set, but never above what the memory or training_ctx permit.
117 desired = min(target, raw_ctx, upper) if target is not None else min(raw_ctx, upper)
118 bounded = max(floor, desired)
119 quantized = (bounded // quantum) * quantum
120 return max(floor, quantized)
123def get_available_memory(fraction: float, *, total: bool = False) -> int:
124 """Return usable GPU/unified memory in bytes, scaled by *fraction*.
125 - macOS (Apple Silicon): unified memory via psutil
126 - Linux with NVIDIA GPU: pynvml -> nvidia-smi -> psutil fallback
127 - Other: psutil system memory
129 With multiple NVIDIA GPUs, *total* sums every card's memory (whole-fleet
130 capacity, for deciding whether a model can run tensor-split across all of
131 them); the default sizes against the smallest single card.
133 A coarse figure for callers with no device list to hand. The fleet has one
134 and sizes against it instead
135 (:func:`lilbee.providers.fleet.planning.plan_sizing_budget`), because this
136 answers with system RAM on every host without an NVIDIA card. That system
137 figure is the process's, cgroup cap included, not the machine's.
138 """
139 system = platform.system()
141 if system == "Darwin":
142 return int(total_system_memory() * fraction)
144 if system in ("Linux", "Windows"):
145 nvidia_mem = _try_nvidia_memory(sum if total else min)
146 if nvidia_mem is not None:
147 return int(nvidia_mem * fraction)
149 return int(total_system_memory() * fraction)
152def free_system_memory() -> int:
153 """Live allocatable system RAM in bytes (free + reclaimable), right now.
155 The load-time counterpart to :func:`get_available_memory`, which scales total
156 capacity for sizing rather than reporting what is free this instant.
158 Bounded by what this process's cgroup still has, for the reason in
159 :func:`lilbee.core.system.cgroup_memory_limit`.
160 """
161 import psutil
163 from lilbee.core.system import cgroup_memory_limit, cgroup_memory_used
165 host_free = int(psutil.virtual_memory().available)
166 limit = cgroup_memory_limit()
167 if limit is None:
168 return host_free
169 used = cgroup_memory_used()
170 return min(host_free, limit if used is None else max(0, limit - used))
173def total_system_memory() -> int:
174 """Total system RAM in bytes this process may use, cgroup cap included.
176 Raises rather than answering zero when the host cannot be read: every caller
177 here is sizing a real placement, and a budget computed from zero refuses
178 every model without saying why.
179 """
180 from lilbee.core.system import capped_total_memory
182 return capped_total_memory()
185def has_nvidia_gpu() -> bool:
186 """Whether an NVIDIA GPU is physically present on this host (NVML or nvidia-smi).
188 Deliberately unmasked. ``CUDA_VISIBLE_DEVICES`` says what a CUDA process may
189 use, not what the machine has, and the callers of this ask the second
190 question: one of them exists to delete an empty mask that an orchestrator
191 left behind, which it could never do if the empty mask hid the card first.
192 """
193 return _nvidia_device_totals() is not None
196def _try_nvidia_memory(reducer: Callable[[list[int]], int] = min) -> int | None:
197 """NVIDIA GPU total memory the CUDA runtime can actually reach, or ``None``.
199 *reducer* combines the per-device totals. ``min`` (the default) sizes against
200 the smallest card, the safe budget for a single server that has not been told
201 which card it will run on. ``sum`` gives whole-fleet capacity, used only by
202 the catalog fit chip to decide whether a model can run split across every card.
204 Restricted to the devices ``CUDA_VISIBLE_DEVICES`` exposes. Neither NVML nor
205 nvidia-smi applies that mask on its own: it is read by the CUDA runtime, and
206 both tools report every card the driver knows about. Unmasked, a container
207 given one card of an eight-card host summed all eight and approved models
208 eight times too large for the card it had, and a fleet whose smallest card
209 was masked out sized every budget against a card the engine cannot see.
210 """
211 totals = _nvidia_device_totals()
212 if not totals:
213 return None
214 visible = _apply_cuda_visible_mask(totals)
215 return reducer([total for _uuid, total in visible]) if visible else None
218def _nvidia_device_totals() -> list[tuple[str, int]] | None:
219 """``[(uuid, total_bytes), ...]`` in driver enumeration order, or ``None``.
221 ``None`` means no NVIDIA GPU was detectable at all, which is the expected
222 outcome on every non-NVIDIA host.
223 """
224 try:
225 import pynvml # type: ignore[import-untyped]
227 pynvml.nvmlInit()
228 totals = []
229 for i in range(pynvml.nvmlDeviceGetCount()):
230 handle = pynvml.nvmlDeviceGetHandleByIndex(i)
231 totals.append(
232 (
233 _decoded(pynvml.nvmlDeviceGetUUID(handle)),
234 int(pynvml.nvmlDeviceGetMemoryInfo(handle).total),
235 )
236 )
237 pynvml.nvmlShutdown()
238 if totals:
239 return totals
240 except Exception: # noqa: S110 -- optional GPU detect; absence is expected on non-NVIDIA hosts
241 pass
243 try:
244 import subprocess
246 # nvidia-smi ships with the NVIDIA driver and is always on PATH when
247 # present; fully-qualifying it would break on every install layout.
248 result = subprocess.run(
249 ["nvidia-smi", "--query-gpu=memory.total,uuid", "--format=csv,noheader,nounits"], # noqa: S607
250 capture_output=True,
251 text=True,
252 encoding="utf-8",
253 errors="replace",
254 timeout=5,
255 )
256 if result.returncode == 0:
257 rows = [_parse_smi_row(line) for line in result.stdout.strip().splitlines()]
258 parsed = [row for row in rows if row is not None]
259 if parsed:
260 return parsed
261 except Exception: # noqa: S110 -- optional GPU detect; same rationale as above
262 pass
264 return None
267def _decoded(value: str | bytes) -> str:
268 """pynvml returns ``str`` on recent versions and ``bytes`` on older ones."""
269 return value.decode() if isinstance(value, bytes) else value
272def _parse_smi_row(line: str) -> tuple[str, int] | None:
273 """One ``memory.total,uuid`` CSV row as ``(uuid, total_bytes)``.
275 The UUID column is optional so an older nvidia-smi that only echoes the
276 memory still yields a device; only a UUID-keyed mask needs it.
277 """
278 fields = [field.strip() for field in line.split(",")]
279 if not fields or not fields[0]:
280 return None
281 try:
282 mib = int(fields[0])
283 except ValueError:
284 return None
285 return (fields[1] if len(fields) > 1 else "", mib * 1024 * 1024)
288def _apply_cuda_visible_mask(devices: list[tuple[str, int]]) -> list[tuple[str, int]]:
289 """The subset of *devices* ``CUDA_VISIBLE_DEVICES`` exposes, in its order.
291 Entries are driver indexes or ``GPU-``/``MIG-`` UUIDs. An unset variable
292 masks nothing; an empty one exposes nothing. CUDA stops enumerating at the
293 first entry that names no device, and so does this, which is what makes
294 ``0,9,1`` on a two-card host mean one card rather than two.
295 """
296 raw = os.environ.get("CUDA_VISIBLE_DEVICES")
297 if raw is None:
298 return devices
299 visible: list[tuple[str, int]] = []
300 for entry in (part.strip() for part in raw.split(",")):
301 matched = _resolve_cuda_entry(entry, devices)
302 if matched is None:
303 break
304 visible.append(matched)
305 return visible
308def _resolve_cuda_entry(entry: str, devices: list[tuple[str, int]]) -> tuple[str, int] | None:
309 """The device an entry of ``CUDA_VISIBLE_DEVICES`` names, ``None`` if it names none."""
310 if entry.isdigit():
311 index = int(entry)
312 return devices[index] if index < len(devices) else None
313 # UUIDs may be abbreviated to any unique prefix.
314 if entry.startswith(("GPU-", "MIG-")):
315 matches = [device for device in devices if device[0].startswith(entry)]
316 return matches[0] if len(matches) == 1 else None
317 return None