Skip to content

C
coderai
Commit 8274ff3e authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

measure a model's memory from the model, and release what it leaves behind 

Six findings from one 21h30m production log (548,687 lines, 51 loads of 12
models). They share a cause: nothing measured a model directly, so every
decision downstream of a measurement was made on a number that was wrong.

1. The footprint was a process-RSS delta
---------------------------------------
`measured_ram = self._get_own_ram_gb() - ram_before`, where _get_own_ram_gb
sums RSS over the engine and every child. RSS counts mmap'd weight-file pages
the loader just walked; summed over a tree it counts pages SHARED between those
processes once per process; and as a delta it absorbs anything else that moved
in the window. Against ground truth:

    gemma-2-9b-it.Q4_K_M   5.37 GB on disk  ->  reported 30.81 GB   (5.7x)
    LTX-2.5, six loads, one unchanged config:
        22.11, 22.11, 23.28, 4.50, 19.02, 29.80 GB                  (6.6x)

Every other model sat at 1.0-1.2x, so this was not noise. Wan2.2's 3.4x spread
was the real quantization fix and is the one that belonged there.

codai/models/footprint.py measures the object instead. Quantization needs no
correction factor, which is the point: a bitsandbytes Params4bit IS a uint8
tensor, so numel x element_size already gives packed bytes. What those formats
hide is counted explicitly -- bnb's quant_state (absmax, code book, nested
double-quant state) lives outside parameters(), and GPTQ/AWQ keep
qweight/qzeros/scales in BUFFERS, so a parameters-only walk reports a 4-bit
model as nearly empty. Tied weights are deduped by storage pointer; accelerate
meta placeholders are sized from their dtype and flagged as declared, keeping
the placement-independence the field already documents; a diffusers pipeline is
walked through its components, since it is not an nn.Module.

Context is costed properly: KV from the KEY/VALUE head count, not the attention
head count (4x on a GQA model), with sliding-window layers sized to their
window. But a computed KV is only a fallback and says so -- this codebase
already learned that "predicting KV from GGUF metadata is architecture-specific
... and was off by 5x in BOTH directions", so measured_kv_gb for the context in
question always wins. A GGUF, which no walk can reach, falls back to its file
size; it is loaded as stored, so the file IS the weights.

2. A failed load overwrote the good measurement
-----------------------------------------------
    Device map: cpu: 1 layers
    Measured footprint: 4.50 GB total
    Saved measured_vram_gb=4.505 (force_vram_update)
    generation FAILED: StopIteration interacts badly with generators
    Full-GPU load looks viable (4.5 GB peak need)   <- 4 seconds later

A load that placed one layer and raised still reached the recorder, and
force_vram_update wrote it to disk as truth. Now a measurement that is
implausible on its own (~0 GB) or that collapses to under half of what the SAME
quantization signature previously measured is treated as a partial load and
discarded. Wan2.2's genuine 40.04 -> 11.73 GB drop is let through because it
arrived WITH a signature change, which is exactly what tells the two apart.
The rejection names which of the two checks fired.

The ratio guard alone was not enough. Walking a pipeline and finding that NO
component holds any weights is a failed load outright, and there is nothing to
fall back to -- a size taken from process memory measures everything except
the model in question. That case is now refused before any estimate is
considered, and the refusal is recorded in the history with its reason rather
than left looking like a successful load. The hole showed up as one of the new
tests passing alone and failing in the full suite: the fallback delta grows
with whatever else the process has loaded, so the figure it invented depended
on test order.

3. The teardown-leak breaker never looked at host RAM
-----------------------------------------------------
It registers under `if _tens.is_cuda` and sweeps under `_obj.is_cuda`. A
pipeline loaded with enable_model_cpu_offload keeps its weights in HOST RAM,
and 17 of 21 video loads took that path, so none of those weights were tracked
and nothing looked for them. With nothing loaded and 482 MiB of VRAM in use:

    coderai-nvidia   33.14 GB PSS_Anon
    30.85 GB anon in 6588 mappings; 6565 of them under 100 MB (20.25 GB)

Thousands of mid-sized allocations is retained tensors, not allocator slack --
which is why gc and malloc_trim ran every 15 s for hours with the figure never
moving, and why the ladder then evicted healthy models to chase it:

    RAM eviction: unloading 'video:Lightricks/LTX-2.5-Diffusers' to free host
    RAM (RSS 40.0 GB > cap target 36.5 GB)

codai/models/teardown.py adds the host arm, by weakref rather than by storage
pointer: a freed host address is readily handed back out by malloc, so a
pointer match could null a live tensor that by then belongs to another model. A
referent still alive after teardown IS this pipeline's leak, identity-exact. It
runs before the allocator trim, since trimming while the weights are still
referenced returns nothing.

4. The RAM cap counted shared pages and page cache
--------------------------------------------------
codai/models/proc_mem.py reads Pss_Anon + Pss_Shmem from smaps_rollup, which
has already divided shared pages by their sharers and excludes reclaimable page
cache. Measured correction on a 5-engine host: 40.06 GB (RSS) -> 37.91 GB
(PSS). Modest -- the leak above was the real problem -- but a cap should chase
memory that freeing an allocation can actually release. Falls back to the old
reading rather than to 0.0, which would read as "no memory in use".

5. A model nobody used stayed loaded
------------------------------------
    Ondemand mode - keeping 'audio_gen:ACE-Step' in VRAM alongside new model

It held 7.9 GB of a 3090 and 2.1 GB of host RAM for hours after its last
request. offload.idle_unload_seconds (default 1800) unloads a model untouched
that long. Deliberately generous: reloads are the expensive event -- that day
spent 2h04m of 21h30m inside loads -- so a short timeout would trade a VRAM
saving nobody asked for against a reload everybody waits on. A key that was
never stamped is left alone rather than assumed idle. It runs before the
leak-watch gate, so turning that off does not silently disable this.

6. The log left no room for the history anyone needed
-----------------------------------------------------
32,633 of 33,009 requests were embeddings, and one engine produced 426,682 of
the 548,687 lines -- 78% -- filling the 2,000-line view in about four minutes.
The per-item endpoints are now summarised to one line (model, count, shape)
instead of dumped, which also stops a caller's content being written to a
shared log. Prompts are deliberately NOT truncated: a 4,000-character script
prompt is what identified the music defect below.

GET /v1/models/history replaces inferring swaps from a log that rotates: time,
event, model, engine, plus seconds, vram_gb, offloaded_gb, strategy for a load
and reason, freed_gb for an eviction. Eviction reasons are passed at the call
sites. The front merges across engines, because a swap is a cross-engine event
-- one engine evicts so a co-located one can load -- so one engine's copy is
half the story.

Also: ACE-Step was asked for a song with no words
-------------------------------------------------
Motus sends a music_prompt and no lyrics, and empty lyrics do not mean
instrumental: the pipeline then conditions on a single padding token
(lyric_token_idx = [0], lyric_mask = [0]) and the model still tries to sing,
which comes out as babbling over the backing track. ACE-Step's own UI documents
the way to ask -- "Use [instrumental] or [inst] to generate instrumental
music" -- and every genre preset it ships ends in either "male vocals, raw
vocals" or "instrumental". Both are now sent when no lyrics are given.
Co-Authored-By: 's avatarClaude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014T3rZozi1Ehk96NG1WesgC
parent 1b83dc27
Show whitespace changes
Inline Side-by-side
Showing with 2447 additions and 19 deletions
codai/__init__.py
View file @ 8274ff3e
......@@ -16,7 +16,7 @@
# Canonical product version for CoderAI — single source of truth. Both the API
# metadata and the admin web UI read from here.
__version__ = "0.5.15"
__version__ = "0.5.16"
__tagline__ = "Complete orchestration, distribution and escalation of remotizable advanced inference."
# Configure the CUDA caching allocator BEFORE torch is imported anywhere.
......
codai/api/app.py
View file @ 8274ff3e
......@@ -937,6 +937,24 @@ async def list_models():
return ModelList(data=models)
@app.get("/v1/models/history", summary="Model load / eviction history", tags=["Core"])
async def models_history(limit: int = 200, event: str = None):
"""When models were loaded and evicted on this engine, and why.
Answers the question a slow render actually raises -- which swaps happened,
what they cost, what triggered them -- without having to reconstruct it from
the log, which rotates: on one production host 78% of the log was
per-request debug and the 2,000-line view covered about four minutes.
Each entry carries ``time``, ``event`` ('load' | 'evict'), ``model`` and
``engine``, plus whatever the call site knew: ``seconds`` (how long the load
took), ``vram_gb`` (measured footprint), ``offloaded_gb``, ``strategy``, and
for an eviction ``reason`` and ``freed_gb``. Newest first, in-memory and
bounded, so it is empty after a restart.
"""
return JSONResponse({"history": multi_model_manager.get_history(limit, event)})
def _admin_scoped_token(request) -> str:
"""The caller's bearer token when it carries the admin scope, else "".
......
codai/api/audio_gen.py
View file @ 8274ff3e
......@@ -407,9 +407,23 @@ def _generate_audio(pipe, model_name: str, request: AudioGenerationRequest, task
# for. A field left out falls back to ACE-Step's own default; a field
# explicitly set to 10 still means 10.
_sent = getattr(request, 'model_fields_set', set())
# ACE-Step is a SONG model: the prompt carries comma-separated tags and
# the lyrics carry the words. Handing it empty lyrics does not mean
# "instrumental" -- the pipeline then conditions on a single padding
# token (lyric_token_idx = [0], lyric_mask = [0]) and the model still
# tries to sing, which comes out as babbling over the backing track.
# Its own UI documents the way to ask: "Use [instrumental] or [inst] to
# generate instrumental music", and every genre preset it ships ends in
# either "male vocals, raw vocals" or "instrumental".
_lyrics = (getattr(request, 'lyrics', None) or "").strip()
_prompt = request.prompt or ""
if not _lyrics:
_lyrics = "[instrumental]"
if "instrumental" not in _prompt.lower():
_prompt = (_prompt + ", instrumental") if _prompt else "instrumental"
wav = pipe.generate(
prompt=request.prompt,
lyrics=getattr(request, 'lyrics', None) or "",
prompt=_prompt,
lyrics=_lyrics,
duration=(float(request.duration) if 'duration' in _sent
and request.duration else 60.0),
steps=int(getattr(request, 'steps', None) or 60),
......
codai/api/log.py
View file @ 8274ff3e
......@@ -54,6 +54,32 @@ def _redact_blobs(obj):
return obj
def _summarise(parsed) -> str:
"""One line describing a request, for the high-volume paths.
Says what is needed to follow what the server was asked to do -- the model,
how many items, the shape -- without reproducing the caller's content.
"""
if not isinstance(parsed, dict):
return "(no JSON body)"
bits = []
model = parsed.get("model")
if model:
bits.append(f"model={model}")
inp = parsed.get("input")
if isinstance(inp, list):
bits.append(f"inputs={len(inp)}")
chars = sum(len(x) for x in inp if isinstance(x, str))
if chars:
bits.append(f"chars={chars}")
elif isinstance(inp, str):
bits.append(f"inputs=1 chars={len(inp)}")
for key in ("dimensions", "encoding_format"):
if parsed.get(key) is not None:
bits.append(f"{key}={parsed[key]}")
return " ".join(bits) or "(empty body)"
def get_recent_activity():
return list(_activity)
......@@ -119,6 +145,25 @@ def _durable_request_debug(method, url, parsed):
pass
# NOTE: long free text is deliberately NOT truncated here. A prompt is the most
# useful thing in the log -- a 4,000-character script prompt was what identified
# a music defect -- and capping it would trade the diagnosis for bytes. The
# volume problem was one endpoint called per item rather than per user action,
# which _SUMMARY_ONLY_PATHS addresses without touching anything else.
#: Paths whose bodies are SUMMARISED rather than dumped, even under debug.
#:
#: These are called per item, not per user action, so a full dump is the log
#: rather than a part of it. One production day: 32,633 embedding requests out
#: of 33,009 total, and the radeon engine alone produced 426,682 of the log's
#: 548,687 lines -- about 78% -- which filled the 2,000-line view in roughly
#: four minutes and left no record of the model loads and evictions anyone was
#: trying to diagnose. The activity tracking below is unaffected: it is one
#: dict entry per request and is what the admin page reads.
_SUMMARY_ONLY_PATHS = {
"/v1/embeddings",
}
_TRACKED_PATHS = {
"/v1/chat/completions": "chat",
"/v1/completions": "completion",
......@@ -175,7 +220,9 @@ async def log_requests(request: Request, call_next):
parsed = json.loads(body_str)
model = parsed.get("model", "—")
if global_debug:
if global_debug and path in _SUMMARY_ONLY_PATHS:
print(f"[req] {request.method} {path} {_summarise(parsed)}")
elif global_debug:
print(f"\n{'='*80}")
print(f"=== FULL REQUEST DEBUG ===")
print(f"Method: {request.method} URL: {request.url}")
......
codai/api/video.py
View file @ 8274ff3e
......@@ -940,6 +940,21 @@ def _free_pipeline_vram(pipe) -> None:
continue
except Exception:
pass
# The pointer pass above is CUDA-only, by design and by necessity: a freed
# HOST address is readily handed back out by malloc, so matching host tensors
# by storage pointer could null a live tensor that by then belongs to another
# model. Track the host-resident ones by weak reference instead, so a referent
# still alive after teardown IS this pipeline's leaked tensor. Without this,
# a 'model CPU offload' pipeline -- 17 of 21 video loads on one production day
# -- had NONE of its weights tracked, and the engine sat on 33 GB of anonymous
# host RAM with an empty GPU. See codai/models/teardown.py.
_host_refs = []
try:
from codai.models import teardown as _td
_host_refs = _td.track_host_tensors(pipe)
except Exception:
_host_refs = []
try:
if pipe is not None:
_comps = getattr(pipe, 'components', {}) or {}
......@@ -1083,6 +1098,22 @@ def _free_pipeline_vram(pipe) -> None:
f"{_vram_used0:.1f} → {_vram_used1:.1f} GB")
except Exception as _lbe:
print(f" [video] teardown-leak breaker failed: {_lbe}")
# Host-side arm of the same breaker. Runs BEFORE the allocator trim: trimming
# while the weights are still referenced returns nothing, which is what the
# RAM mitigation ladder was doing every 15 s for hours without the figure
# moving.
if _host_refs:
try:
from codai.models import teardown as _td
_hr = _td.break_host_leaks(_host_refs)
if _hr["broken"] or _hr["survived_gb"] > 0.5:
print(f" [video] host teardown-leak breaker: "
f"{_hr['survived_gb']:.1f} GB of this pipeline's host-resident "
f"tensors survived teardown; nulled {_hr['broken']} external "
f"reference(s)"
+ (f"; holders: {_hr['holders']}" if _hr["holders"] else ""))
except Exception as _he:
print(f" [video] host teardown-leak breaker failed: {_he}")
try:
from codai.models.manager import _trim_cpu_ram
_trim_cpu_ram()
......
codai/config.py
View file @ 8274ff3e
......@@ -292,6 +292,17 @@ class OffloadConfig:
ram_watch_poll_seconds: float = 15.0 # how often the watcher samples RSS
ram_watch_soft_fraction: float = 0.90 # mitigate at/above this fraction of the cap
ram_watch_cuda: bool = True # allow mitigation to call CUDA empty_cache()
# Unload a model that has gone untouched for this long, in seconds (0 = never).
#
# Eviction already frees a model when something else needs the memory, so this
# is for the case where nothing does: a 3.5B music model sat on 7.9 GB of a
# 3090 and 2.1 GB of host RAM for hours after its last request, because the
# on-demand path keeps a model resident whenever there is room for it.
#
# Kept generous on purpose. Reloads are the expensive event -- one production
# day spent 2h04m of 21h30m inside model loads -- so a short timeout trades a
# VRAM saving nobody asked for against a reload everybody waits on.
idle_unload_seconds: float = 1800.0
# Cross-backend GPU pooling. OFF by default: each engine uses only its own
# backend's GPUs (CUDA for an NVIDIA engine, Vulkan for a Radeon engine) and a
# model still splits across multiple SAME-backend cards (e.g. 2× 3090). When ON,
......
codai/frontproxy/app.py
View file @ 8274ff3e
......@@ -3100,6 +3100,33 @@ def build_app(config, config_dir=None) -> FastAPI:
media_type=val.headers.get("content-type"))
return JSONResponse(val)
# /v1/models/history is the MERGE across engines, newest first: a swap is a
# cross-engine event (one engine evicts so a co-located one can load), so
# reading a single engine's copy tells only half the story. Registered before
# the catch-all so it is aggregated rather than proxied to one engine.
@app.get("/v1/models/history", include_in_schema=False)
async def _models_history(request: Request):
headers = front._filter_headers(request.headers, _DROP_REQ)
merged = []
for e in front.registry.all():
try:
r = await front._short.get(e.url + "/v1/models/history",
headers=headers,
params=request.query_params)
if r.status_code == 200:
rows = (r.json() or {}).get("history") or []
for row in rows:
row.setdefault("engine", getattr(e, "name", "") or "")
merged.extend(rows)
except Exception:
continue # an engine that is busy or down just contributes nothing
merged.sort(key=lambda r: r.get("time") or 0, reverse=True)
try:
limit = max(1, min(500, int(request.query_params.get("limit") or 200)))
except (TypeError, ValueError):
limit = 200
return JSONResponse({"history": merged[:limit]})
# Engine management (front-owned: it runs the supervisor). Registered before
# the catch-all so they aren't proxied to an engine.
@app.get("/admin/api/engines", include_in_schema=False)
......
codai/main.py
View file @ 8274ff3e
......@@ -1216,6 +1216,7 @@ def main():
# Global host-RAM cap + leak watch (read live by hf_loading, manager, ram_monitor).
global_args.max_ram_gb = config.offload.max_ram_gb
global_args.evict_idle_on_ram = config.offload.evict_idle_on_ram
global_args.idle_unload_seconds = config.offload.idle_unload_seconds
global_args.ram_leak_watch = config.offload.ram_leak_watch
global_args.ram_watch_poll_seconds = config.offload.ram_watch_poll_seconds
global_args.ram_watch_soft_fraction = config.offload.ram_watch_soft_fraction
......
codai/models/footprint.py 0 → 100644
View file @ 8274ff3e
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""What a loaded model actually occupies, measured from the model itself.
The previous measurement was a process-RSS delta taken across the load:
measured_ram = self._get_own_ram_gb() - ram_before
RSS is the wrong instrument for this. It counts file-backed pages that the
kernel maps on a read -- the mmap'd safetensors/GGUF the loader just walked --
and ``_get_own_ram_gb`` sums RSS over the engine *and every child*, so pages
shared between those processes (torch, the CUDA runtime, libc, the interpreter)
are added once per process. It is also a delta over a window, so anything else
loading, evicting or faulting pages in that window lands in the figure.
Measured against ground truth on a production host, 2026-10-29:
gemma-2-9b-it.Q4_K_M 5.37 GB on disk -> reported 30.81 GB (5.7x)
LTX-2.5, six loads, one config: 22.11, 22.11, 23.28, 4.50, 19.02, 29.80 GB
One of those readings (4.50 GB) came from a load that placed a single layer and
then raised; another (29.80 GB) was taken after an abandoned resident attempt
had moved a transformer and a text encoder to host RAM, so it counted that
attempt as well. Both were persisted as authoritative under
``force_vram_update`` and planned against, which is how a model that fits ends
up on the CPU.
This module answers the question directly instead: walk the loaded object and
add up the bytes its tensors really occupy, then add the parts that are not
weights but are still required -- the KV cache for the configured context, and
a runtime reserve for activations and the decode spike.
It is deliberately placement-independent, matching the contract
``measured_vram_gb`` already documents: offload moves weights between GPU and
host, it does not change how much the model needs. The split is reported
separately (:attr:`Footprint.weights_gpu_gb` / :attr:`weights_cpu_gb`) for
callers that care where the bytes currently sit.
Quantization needs no correction factor here, which is the point of measuring
the object: a bitsandbytes ``Params4bit`` *is* a uint8 tensor, so
``numel() * element_size()`` already yields the packed size, and GPTQ/AWQ keep
their ``qweight``/``qzeros``/``scales`` as ordinary tensors that the walk sees.
What those formats hide is accounted for explicitly: bnb's ``quant_state``
(absmax, the code book, nested second-level state) lives outside
``parameters()`` and is real memory.
"""
from __future__ import annotations
import os
from dataclasses import dataclass, field
# Bytes per element of a llama.cpp KV cache type, relative to the raw dtype.
# Mirrors codai/backends/vulkan.py's table; kept here so this module has no
# import-time dependency on a backend.
_KV_TYPE_BITS = {
"f32": 32.0, "f16": 16.0, "bf16": 16.0,
"q8_0": 8.5, "q8_1": 9.0,
"q5_0": 5.5, "q5_1": 6.0,
"q4_0": 4.5, "q4_1": 5.0,
}
# Fraction of the weight footprint to reserve for activations and the decode
# spike, by model type. Diffusion caps its spike with VAE tiling/slicing, so it
# does not scale with context the way an autoregressive model does.
_ACTIVATION_RESERVE = {
"video": 0.08,
"image": 0.08,
"audio_gen": 0.08,
"tts": 0.05,
"embedding": 0.05,
}
_ACTIVATION_RESERVE_DEFAULT = 0.08
# Autoregressive models get a flat activation share; their context cost is the
# KV cache, which is computed exactly rather than guessed as a fraction.
_ACTIVATION_RESERVE_AUTOREGRESSIVE = 0.10
@dataclass
class Footprint:
"""A model's memory requirement, broken into the parts it is made of."""
weights_gpu_gb: float = 0.0
weights_cpu_gb: float = 0.0
kv_gb: float = 0.0
activations_gb: float = 0.0
#: False when something could not be accounted for -- the caller must then
#: NOT persist this as authoritative. Every gap is named in :attr:`notes`.
complete: bool = True
#: True when :attr:`kv_gb` came from a measurement rather than the estimate.
kv_measured: bool = False
#: True when the object was walkable and held NO weights -- a failed or
#: partial load. Distinct from an object we simply cannot introspect, and
#: the difference decides whether a fallback estimate is legitimate.
empty_load: bool = False
notes: list = field(default_factory=list)
@property
def weights_gb(self) -> float:
return self.weights_gpu_gb + self.weights_cpu_gb
@property
def total_gb(self) -> float:
"""Everything the model needs, wherever its weights currently sit."""
return round(self.weights_gb + self.kv_gb + self.activations_gb, 3)
def describe(self) -> str:
bits = [f"weights {self.weights_gb:.2f} GB"]
if self.weights_cpu_gb > 0.005:
bits[-1] += (f" ({self.weights_gpu_gb:.2f} on GPU, "
f"{self.weights_cpu_gb:.2f} on host)")
if self.kv_gb > 0.005:
bits.append(f"KV {self.kv_gb:.2f} GB"
+ ("" if self.kv_measured else " (est)"))
if self.activations_gb > 0.005:
bits.append(f"activations {self.activations_gb:.2f} GB")
out = f"{self.total_gb:.2f} GB total ({', '.join(bits)})"
if not self.complete:
out += " [INCOMPLETE: " + "; ".join(self.notes) + "]"
return out
# --------------------------------------------------------------------------
# tensor-level accounting
# --------------------------------------------------------------------------
def _dtype_itemsize(dtype) -> int:
"""Bytes per element of a torch dtype, without needing a real tensor.
A meta tensor has no storage to ask, so size it from the dtype it declares.
"""
try:
return int(dtype.itemsize)
except (AttributeError, TypeError):
pass
try: # older torch
import torch
return int(torch.empty((), dtype=dtype).element_size())
except Exception:
return 2 # fp16/bf16 is the common case
def _tensor_bytes(t) -> tuple:
"""``(bytes, device_kind, exact)`` for one tensor.
``device_kind`` is 'gpu' or 'cpu'. ``exact`` is False for a meta tensor,
whose bytes are DECLARED (numel x dtype) rather than measured -- accelerate
leaves a meta placeholder where it has moved the real weights to host RAM
or to disk, and the model still needs them.
"""
try:
is_meta = bool(getattr(t, "is_meta", False)) or str(t.device) == "meta"
except Exception:
is_meta = False
if is_meta:
return _dtype_itemsize(t.dtype) * t.numel(), "cpu", False
try:
nbytes = t.numel() * t.element_size()
except Exception:
return 0, "cpu", False
try:
kind = "gpu" if t.device.type in ("cuda", "xpu", "hpu", "mps") else "cpu"
except Exception:
kind = "cpu"
return nbytes, kind, True
def _storage_key(t):
"""Identity of the tensor's backing storage, or None when it has none.
Tied weights (an embedding shared with the LM head) and shared buffers must
be counted once. Two tensors are the same memory when their storage pointer
matches, which is not the same thing as the tensors being equal.
"""
try:
return t.untyped_storage().data_ptr()
except Exception:
try:
return t.storage().data_ptr()
except Exception:
return None
def _quant_state_tensors(param):
"""Tensors a bitsandbytes ``Params4bit`` keeps outside ``parameters()``.
A 4-bit parameter is a uint8 blob plus a ``quant_state`` holding the
per-block absmax, the code book and -- when double quantization is on -- a
nested second-level state. That is real memory and is invisible to a plain
parameter walk.
"""
qs = getattr(param, "quant_state", None)
if qs is None:
return
seen = set()
stack = [qs]
while stack:
node = stack.pop()
if id(node) in seen:
continue
seen.add(id(node))
for attr in ("absmax", "code", "offset", "shape_offset"):
t = getattr(node, attr, None)
if t is not None and hasattr(t, "numel"):
yield t
nested = getattr(node, "state2", None)
if nested is not None:
stack.append(nested)
def walk_module_bytes(module) -> dict:
"""Bytes held by one module's parameters and buffers, split GPU/host.
Returns ``{'gpu': int, 'cpu': int, 'declared': int, 'tensors': int}``,
where ``declared`` is the subset sized from a meta placeholder rather than
from real storage. Buffers are walked as well as parameters because the
pre-quantized formats keep their payload there: GPTQ and AWQ register
``qweight``/``qzeros``/``scales`` as buffers, so a parameters-only walk
reports a 4-bit model as nearly empty.
"""
out = {"gpu": 0, "cpu": 0, "declared": 0, "tensors": 0}
if module is None or not hasattr(module, "parameters"):
return out
seen = set()
def _add(t):
key = _storage_key(t)
if key is not None:
if key in seen:
return
seen.add(key)
nbytes, kind, exact = _tensor_bytes(t)
if nbytes <= 0:
return
out[kind] += nbytes
out["tensors"] += 1
if not exact:
out["declared"] += nbytes
try:
params = list(module.parameters(recurse=True))
except Exception:
params = []
for p in params:
_add(p)
for extra in _quant_state_tensors(p):
_add(extra)
try:
for b in module.buffers(recurse=True):
_add(b)
except Exception:
pass
return out
def _submodules(obj):
"""``(name, module)`` pairs to account for, for a module or a pipeline.
A diffusers pipeline is a container: its weights live in components
(transformer, text_encoder, vae, ...) alongside things that hold none
(schedulers, tokenizers, feature extractors). Walking the pipeline itself
would miss the components entirely, since it is not an ``nn.Module``.
"""
comps = getattr(obj, "components", None)
if isinstance(comps, dict) and comps:
for name, comp in comps.items():
if comp is not None and hasattr(comp, "parameters"):
yield name, comp
return
if hasattr(obj, "parameters"):
yield "", obj
return
# A text/vision model is registered as a ModelManager wrapper, so the real
# module hangs off .backend.model. Without this an HF LLM falls through to
# the process-delta fallback even though it is perfectly walkable.
inner = getattr(getattr(obj, "backend", None), "model", None)
if inner is not None and hasattr(inner, "parameters"):
yield "backend.model", inner
# A vision projector is a separate module beside the LM.
for attr in ("mmproj", "vision_model", "vision_tower", "multi_modal_projector"):
extra = getattr(inner, attr, None) or getattr(obj, attr, None)
if extra is not None and hasattr(extra, "parameters") and extra is not inner:
yield f"backend.{attr}", extra
def measure_weights(obj) -> tuple:
"""``(gpu_gb, cpu_gb, notes, ok)`` for a loaded torch model or pipeline.
A note of ``EMPTY_PIPELINE`` means the object IS a pipeline and we walked
it and it holds no weights at all. That is a failed or partial load, not an
object we cannot introspect, and the caller must not substitute a guess for
it.
"""
notes, gpu, cpu, declared, tensors = [], 0, 0, 0, 0
parts = list(_submodules(obj))
if not parts:
if isinstance(getattr(obj, "components", None), dict):
return 0.0, 0.0, ["EMPTY_PIPELINE: walked every component and none "
"holds any weights — the load did not complete"], False
return 0.0, 0.0, ["object exposes no torch modules"], False
for name, mod in parts:
r = walk_module_bytes(mod)
gpu += r["gpu"]
cpu += r["cpu"]
declared += r["declared"]
tensors += r["tensors"]
if tensors == 0:
return 0.0, 0.0, ["EMPTY_PIPELINE: walked the model and found no "
"tensors — the load did not complete"], False
if declared:
notes.append(f"{declared / 1e9:.2f} GB sized from meta placeholders "
f"(accelerate moved it off the GPU)")
return gpu / 1e9, cpu / 1e9, notes, True
# --------------------------------------------------------------------------
# KV cache
# --------------------------------------------------------------------------
def _kv_bits(name) -> float:
return _KV_TYPE_BITS.get(str(name or "f16").strip().lower().replace("-", "_"), 16.0)
def kv_cache_gb(geom: dict, n_ctx: int, type_k=None, type_v=None,
batch: int = 1, swa_full: bool = True) -> float:
"""KV cache for ``n_ctx`` tokens: layers x kv_heads x (k+v dims) x bits.
Uses the number of KEY/VALUE heads, not attention heads: a GQA model such
as Qwen3 shares one KV head across several query heads, so sizing by
attention heads over-states the cache several-fold.
Sliding-window layers are sized to their window rather than the full
context when ``swa_full`` is False, which is what llama.cpp's ``swa_full``
controls: Gemma alternates a windowed layer with a global one, so a
full-context figure over-states those layers by ``n_ctx / window``.
This is an ESTIMATE and callers must treat it as one. Per-layer KV-head
arrays, shared KV and architecture-specific cache layouts are not
modelled, and on this codebase's own history such a prediction has been
wrong by 5x in both directions. Prefer a measured figure
(``measured_kv_gb`` for the context in question) whenever one exists.
Returns 0.0 when the geometry is unknown, so the caller can say so rather
than guess.
"""
layers = int(geom.get("block_count") or geom.get("num_hidden_layers") or 0)
kv_heads = int(geom.get("head_count_kv") or geom.get("num_key_value_heads")
or geom.get("head_count") or geom.get("num_attention_heads") or 0)
if not layers or not kv_heads or n_ctx <= 0:
return 0.0
k_dim = int(geom.get("key_length") or geom.get("head_dim") or 0)
v_dim = int(geom.get("value_length") or k_dim or 0)
if not k_dim:
heads = int(geom.get("head_count") or geom.get("num_attention_heads") or 0)
emb = int(geom.get("embedding_length") or geom.get("hidden_size") or 0)
if heads and emb:
k_dim = v_dim = emb // heads
if not k_dim:
return 0.0
per_layer = kv_heads * ((k_dim * _kv_bits(type_k)
+ v_dim * _kv_bits(type_v)) / 8.0)
window = int(geom.get("sliding_window") or geom.get("attention.sliding_window") or 0)
swa_layers = int(geom.get("swa_layer_count") or 0)
if window and not swa_full and 0 < window < n_ctx:
# Alternating pattern unless the geometry states otherwise: llama.cpp
# reports is_swa per layer, so a caller that knows the real count
# should pass swa_layer_count.
if not swa_layers:
swa_layers = layers // 2
swa_layers = min(swa_layers, layers)
tokens = swa_layers * window + (layers - swa_layers) * n_ctx
else:
tokens = layers * n_ctx
return max(0, batch) * per_layer * tokens / 1e9
def hf_geometry(obj) -> dict:
"""Geometry for KV sizing, read from a transformers config when present."""
cfg = getattr(obj, "config", None)
if cfg is None:
cfg = getattr(getattr(obj, "model", None), "config", None)
if cfg is None:
# A ModelManager wrapper keeps the real module at .backend.model, which
# is where its config lives too. Missing this silently drops the KV term
# for every HF text model -- the one kind that definitely has one.
_inner = getattr(getattr(obj, "backend", None), "model", None)
cfg = getattr(_inner, "config", None)
if cfg is None:
return {}
def _g(*names):
for n in names:
v = getattr(cfg, n, None)
if isinstance(v, int) and v > 0:
return v
return 0
geom = {
"num_hidden_layers": _g("num_hidden_layers", "n_layer", "num_layers"),
"num_attention_heads": _g("num_attention_heads", "n_head"),
"num_key_value_heads": _g("num_key_value_heads", "num_kv_heads"),
"hidden_size": _g("hidden_size", "n_embd", "d_model"),
"head_dim": _g("head_dim"),
}
return {k: v for k, v in geom.items() if v}
# --------------------------------------------------------------------------
# top level
# --------------------------------------------------------------------------
def _is_autoregressive(model_type: str) -> bool:
return str(model_type or "").lower() in ("text", "vision", "")
def _configured_ctx(cfg: dict) -> int:
for key in ("n_ctx", "context_size", "max_model_len", "max_position_embeddings"):
try:
v = int((cfg or {}).get(key) or 0)
except (TypeError, ValueError):
continue
if v > 0:
return v
return 0
def _measured_kv_for(measured, n_ctx: int) -> float:
"""A previously measured KV size for this context, or 0.0.
Recorded per n_ctx because the cache scales with it, so a figure taken at
4k says nothing about 128k. Accepts a bare number (treated as applying to
the context asked for) or a ``{n_ctx: gb}`` mapping, which is how
``measured_kv_gb`` is persisted.
"""
if measured is None:
return 0.0
if isinstance(measured, dict):
for key in (str(n_ctx), n_ctx):
if key in measured:
try:
return float(measured[key] or 0.0)
except (TypeError, ValueError):
return 0.0
return 0.0
try:
return float(measured or 0.0)
except (TypeError, ValueError):
return 0.0
def measure(obj, cfg: dict = None, model_type: str = "", *,
geom: dict = None, file_size_gb: float = 0.0,
batch: int = 1, measured_kv_gb=None, swa_full: bool = True) -> Footprint:
"""Measure what ``obj`` occupies, plus the context and runtime it needs.
``obj`` is a torch module, a diffusers pipeline, or None when the model is
not a torch object at all (a GGUF served by llama.cpp, or an isolated
subprocess worker). In that case pass ``file_size_gb``: the weights are
exactly the file, since a GGUF is loaded as-is with no dtype conversion.
The KV term prefers ``measured_kv_gb`` -- llama.cpp reports its own KV
buffer size exactly, and predicting it from metadata has been wrong by 5x
in both directions on this codebase's own models. It is computed from
``geom``/``cfg`` only when no measurement is available, and the footprint
then says so in :attr:`Footprint.notes`.
A footprint whose weights could not be accounted for comes back with
``complete=False`` and must not be persisted.
"""
cfg = cfg or {}
fp = Footprint()
if obj is not None:
gpu, cpu, notes, ok = measure_weights(obj)
fp.weights_gpu_gb, fp.weights_cpu_gb = gpu, cpu
fp.notes.extend(notes)
fp.empty_load = any(n.startswith("EMPTY_PIPELINE") for n in notes)
if not ok:
if fp.empty_load:
# Nothing to fall back TO: we looked, and the model has no
# weights. A size invented from elsewhere is how a load that
# placed one layer came to be recorded as the model's footprint.
fp.complete = False
elif file_size_gb > 0:
fp.weights_cpu_gb = float(file_size_gb)
fp.notes.append("fell back to the on-disk size")
else:
fp.complete = False
elif file_size_gb > 0:
fp.weights_cpu_gb = float(file_size_gb)
else:
fp.complete = False
fp.notes.append("nothing to measure: no model object and no file size")
# KV cache -- exact where the geometry is known, and only for the model
# kinds that keep one. A diffusion pipeline has no KV cache.
if _is_autoregressive(model_type):
n_ctx = _configured_ctx(cfg)
meas = _measured_kv_for(measured_kv_gb, n_ctx)
if meas > 0:
fp.kv_gb = round(meas, 3)
fp.kv_measured = True
else:
g = dict(geom or {})
if not g and obj is not None:
g = hf_geometry(obj)
if n_ctx and g:
fp.kv_gb = round(kv_cache_gb(g, n_ctx, cfg.get("cache_type_k"),
cfg.get("cache_type_v"), batch=batch,
swa_full=swa_full), 3)
if fp.kv_gb:
fp.notes.append(f"KV is an ESTIMATE ({fp.kv_gb:.2f} GB at "
f"n_ctx={n_ctx}); no measured_kv_gb recorded yet")
else:
fp.notes.append(f"KV not sized: geometry incomplete ({sorted(g)})")
elif n_ctx:
fp.notes.append("KV not sized: model geometry unknown")
base = fp.weights_gb + fp.kv_gb
share = (_ACTIVATION_RESERVE_AUTOREGRESSIVE if _is_autoregressive(model_type)
else _ACTIVATION_RESERVE.get(str(model_type or "").lower(),
_ACTIVATION_RESERVE_DEFAULT))
fp.activations_gb = round(base * share, 3)
fp.weights_gpu_gb = round(fp.weights_gpu_gb, 3)
fp.weights_cpu_gb = round(fp.weights_cpu_gb, 3)
return fp
def extra_file_gb(*paths) -> float:
"""Total size of companion weight files (an mmproj projector, a LoRA).
These load on top of the model and are not reachable by walking it when
they are held by a separate handler.
"""
total = 0
for p in paths:
if not p:
continue
try:
total += os.path.getsize(os.path.expanduser(str(p)))
except OSError:
continue
return total / 1e9
codai/models/manager.py
View file @ 8274ff3e
......@@ -26,6 +26,7 @@ import time
import contextlib
import threading
import gc
from collections import deque
# Import needed classes from other codai modules
from codai.models.parser import ModelParserAdapter
......@@ -1595,6 +1596,14 @@ class MultiModelManager:
self._last_used: Dict[str, float] = {} # model_key -> last-served monotonic ts (LRU)
self._active_in_vram: Optional[str] = None # backing field for active_in_vram property
self.models_in_vram: set = set() # all models currently in VRAM
# Why models came and went, newest last. Load and eviction decisions were
# previously only inferable from the log, and the log rotates: on one
# production host 78% of it was per-request debug and the 2,000-line view
# covered about four minutes, so by the time anyone looked at a slow render
# the swaps behind it were gone. Bounded, in-memory, and cheap enough to
# record unconditionally.
self._history: deque = deque(maxlen=500)
self._load_started: Dict[str, float] = {}
self.model_aliases: Dict[str, str] = {}
self.whisper_server: Optional[WhisperServerManager] = None # legacy single-instance compat
self.whisper_servers: Dict[str, WhisperServerManager] = {} # id -> manager
......@@ -3525,6 +3534,10 @@ class MultiModelManager:
gb = self._get_model_used_vram_gb(model_key)
if gb and gb > 0:
self._loading_reservations[model_key] = float(gb)
# Remember when this load began. Load time is the number everyone
# actually waits on -- one production day spent 2h04m of 21h30m
# inside loads -- and it cannot be recovered after the fact.
self._load_started[model_key] = time.monotonic()
except Exception:
pass
......@@ -3711,17 +3724,36 @@ class MultiModelManager:
"""Call immediately before loading a model; returns a snapshot for delta measurement."""
return self._free_vram_snapshot()
@staticmethod
def _fp_extra_gb(ref) -> float:
"""Size of a companion weight file (an mmproj projector), 0.0 if absent."""
try:
from codai.models.footprint import extra_file_gb
return extra_file_gb(ref)
except Exception:
return 0.0
def record_vram_delta(self, model_key: str, free_before: float,
offloaded: bool = False, ram_before: float = -1.0,
n_gpu_layers=None) -> None:
n_gpu_layers=None, model_obj=None) -> None:
"""Call immediately after a model finishes loading to record its REAL memory
footprint on this hardware, so the next load reuses it instead of repeating
the OOM→move-layers-to-RAM retry dance.
Records up to three SYSTEM-OWNED fields (never the user's ``used_vram_gb``):
* ``measured_vram_gb`` — GPU VRAM delta + runtime reserve (the GPU
portion; eviction frees against this).
* ``measured_ram_gb`` — host-RAM delta (the CPU-offloaded layers).
* ``measured_vram_gb`` — the model's FULL footprint: the bytes its
tensors occupy (wherever they sit) plus
the KV cache for the configured context
plus a runtime reserve. Measured by
walking the loaded object, which needs no
quantization factor because a 4-bit
parameter is literally a uint8 tensor.
Falls back to the old process-memory
delta when the object cannot be walked
(a GGUF, an isolated subprocess worker)
and no file size is available. Eviction
frees against this.
* ``measured_ram_gb`` — the part of those weights held on the host.
* ``measured_n_gpu_layers`` — the layer split llama.cpp actually settled
on (after any progressive CPU-offload retry).
The TOTAL memory the model needs is ``measured_vram_gb + measured_ram_gb``.
......@@ -3753,9 +3785,71 @@ class MultiModelManager:
force_update = bool(cfgd.get("force_vram_update"))
user_pinned = cfgd.get("used_vram_gb") is not None
# --- Measure the MODEL, not the process ---------------------------------
# Preferred path: walk the loaded object and add up the bytes its tensors
# really occupy, then add the KV cache for the configured context and a
# runtime reserve. See codai/models/footprint.py for why the RSS delta
# below is kept only as a fallback: it counts mmap'd weight-file pages and,
# summed over a process tree, pages shared between those processes, and as
# a delta it absorbs anything else that moved in the window. On one
# production host the same unchanged LTX-2.5 config measured 22.11, 22.11,
# 23.28, 4.50, 19.02 and 29.80 GB across six loads, and a 5.37 GB GGUF
# measured 30.81 GB.
fp = None
_obj = model_obj if model_obj is not None else self.models.get(model_key)
# A GGUF is not a torch object -- llama.cpp owns its memory and the
# registry holds a ModelManager wrapper with no parameters() to walk. Its
# weights ARE the file, byte for byte: it is loaded as stored, with no
# dtype conversion. Pass that size so the walk has something to fall back
# to; without it a 5.37 GB Q4_K_M measured 30.81 GB from the RSS delta.
_file_gb = 0.0
try:
for _cand in (cfgd.get("model_path"), cfgd.get("path"),
model_key.split(":", 1)[-1]):
if _cand and os.path.isfile(os.path.expanduser(str(_cand))):
_file_gb = os.path.getsize(os.path.expanduser(str(_cand))) / 1e9
break
if _file_gb:
_mmproj = cfgd.get("mmproj")
if _mmproj:
_file_gb += self._fp_extra_gb(_mmproj)
except Exception:
_file_gb = 0.0
try:
from codai.models import footprint as _fpmod
fp = _fpmod.measure(
_obj, cfgd, self._model_type_of(model_key),
file_size_gb=_file_gb,
measured_kv_gb=cfgd.get("measured_kv_gb"),
swa_full=bool(cfgd.get("swa_full", True)))
if fp is not None and fp.empty_load:
# We walked the model and it holds no weights. There is nothing
# to fall back TO: a size taken from process memory would be
# measuring everything except this model, which is how a load
# that placed one layer came to be recorded as its footprint.
print(f" footprint: {'; '.join(fp.notes)} — recording NOTHING "
f"for this load")
try:
self.record_history("load", model_key, reason="failed: "
"the loaded model holds no weights")
self._load_started.pop(model_key, None)
except Exception:
pass
return
if not fp.complete:
print(f" footprint: not accounted for ({'; '.join(fp.notes)}) — "
f"falling back to the process delta, and NOT persisting it")
fp = None
except Exception as e:
print(f" footprint: direct measurement unavailable ({e}); "
f"falling back to the process delta")
fp = None
# --- Host-RAM footprint (the portion of weights offloaded to CPU) -------
measured_ram = None
if ram_before is not None and ram_before >= 0:
if fp is not None:
measured_ram = round(fp.weights_cpu_gb, 3) or None
elif ram_before is not None and ram_before >= 0:
try:
rd = round(self._get_own_ram_gb() - ram_before, 3)
if rd > 0:
......@@ -3774,12 +3868,55 @@ class MultiModelManager:
delta_gb = (free_before - free_after) / 1e9
gpu_gb = max(0.0, delta_gb)
measured_vram = None
if gpu_gb > 0 or (measured_ram and measured_ram > 0):
if fp is not None:
measured_vram = fp.total_gb
gpu_gb = fp.weights_gpu_gb
self._measured_vram_gb[model_key] = measured_vram # in-memory truth
print(f" footprint: {fp.describe()}")
elif gpu_gb > 0 or (measured_ram and measured_ram > 0):
reserve_gb = self._runtime_reserve_gb(
cfgd, model_key, gpu_gb + (measured_ram or 0.0))
measured_vram = round(gpu_gb + (measured_ram or 0.0) + reserve_gb, 3)
self._measured_vram_gb[model_key] = measured_vram # in-memory truth
# --- Reject a partial load before it becomes authoritative --------------
# A load that places a handful of modules and then raises still reaches
# here, and force_vram_update would write its figure to disk as truth.
# That is how 'Full-GPU load looks viable (4.5 GB peak need)' came to be
# planned four seconds after a load that mapped ONE layer and crashed.
# Guard by self-comparison rather than a disk walk: a collapse to under
# half of what the SAME quantization previously measured is a partial
# load, whereas a real drop (Wan2.2's 40.04 -> 11.73 GB) arrives together
# with a quantization-signature change and is let through.
if measured_vram is not None:
try:
_prev = float(cfgd.get("measured_vram_gb") or 0.0)
except (TypeError, ValueError):
_prev = 0.0
_same_quant = (cfgd.get("measured_quant_sig")
and cfgd.get("measured_quant_sig") == self._quant_signature(cfgd))
# A real model is never ~0 GB, so treat that as a failed measurement
# whatever came before -- this is the case where the guard matters
# most, and gating it on "> 0" skipped exactly it.
_why = ""
if measured_vram < 0.01:
_why = "not a plausible size for a loaded model"
elif _prev > 0 and _same_quant and measured_vram < 0.5 * _prev:
_why = (f"less than half the {_prev:.2f} GB previously measured "
f"under the same quantization "
f"({cfgd.get('measured_quant_sig')})")
_collapsed = bool(_why)
if _collapsed:
print(f" footprint: {measured_vram:.3f} GB is {_why} — treating "
f"this as a PARTIAL load; keeping the previous figure and "
f"NOT persisting this one")
if _prev > 0:
self._measured_vram_gb[model_key] = _prev
else:
self._measured_vram_gb.pop(model_key, None)
measured_vram = None
measured_ram = None
# --- Settled GPU-layer split (the value the retry loop landed on) -------
measured_layers = None
try:
......@@ -3836,6 +3973,18 @@ class MultiModelManager:
except Exception as e:
print(f" Warning: could not persist {field}: {e}")
try:
_t0 = self._load_started.pop(model_key, None)
self.record_history(
"load", model_key,
seconds=(time.monotonic() - _t0) if _t0 else None,
vram_gb=measured_vram or self._measured_vram_gb.get(model_key),
offloaded_gb=measured_ram,
strategy=self.model_devices.get(model_key),
measured=("model" if fp is not None else "process-delta"))
except Exception:
pass
_persist("measured_vram_gb", measured_vram, vram_guard=True)
_persist("measured_ram_gb", measured_ram)
# Stamp WHICH quantization this measurement describes, so a later
......@@ -5295,11 +5444,14 @@ class MultiModelManager:
_trim_cpu_ram()
return True
def _evict_one(self, key):
def _evict_one(self, key, reason: str = ""):
"""Fully unload ONE model by key and free its VRAM + host RAM.
Shared by VRAM- and RAM-driven eviction. Never pulls a model that is still
mid-request; waits briefly first and skips it if it stays busy.
``reason`` is recorded in the history so a swap can be explained after
the fact rather than reconstructed from the log.
"""
# Make absolutely sure nothing is mid-request on this model before we
# pull its tensors off the GPU.
......@@ -5307,6 +5459,8 @@ class MultiModelManager:
if not self._wait_until_idle(key):
print(f" Skipping eviction of '{key}' — still busy after wait")
return
_free0 = self._get_free_vram_gb()
_known = self._measured_vram_gb.get(key)
model_obj = self.models.pop(key, None)
self.models_in_vram.discard(key)
self._last_used.pop(key, None)
......@@ -5382,19 +5536,45 @@ class MultiModelManager:
# otherwise RSS stays high across evict/load cycles and the machine
# slowly fills RAM + swap.
_trim_cpu_ram()
try:
_freed = self._get_free_vram_gb() - _free0
self.record_history("evict", key, reason=reason or "unspecified",
freed_gb=max(0.0, _freed),
expected_gb=_known)
except Exception:
pass
@staticmethod
def _get_process_ram_gb() -> float:
"""Resident-set size of the server process TREE, in GB (0.0 on failure).
"""Host memory the server process TREE holds, in GB (0.0 on failure).
Offloaded weights and worker subprocesses count against the global cap, so
sum the root process plus all children (mirrors thermal.read_process_tree_cpu).
this covers the root process plus all children.
Under the front/engine split the host-RAM cap is SHARED, not split: when the
front spawned this engine it set CODERAI_FRONT_PID, so the root is the
*front* — every engine then measures the same fleet-wide total (front + all
engines + their workers) and enforces the single cap against it. In
single-process mode the root is just this process, as before."""
single-process mode the root is just this process, as before.
Measured as proportional set size (``Pss_Anon + Pss_Shmem``), not summed
RSS. Summing RSS over a tree counts pages SHARED between those processes
once per process -- the interpreter, torch and the CUDA runtime are mapped
into every engine -- and counts file-backed pages, which are page cache the
kernel reclaims by itself and which evicting a model cannot free. PSS has
already divided each shared page by its sharers, so summing it across the
tree counts every page exactly once. On a 5-engine production host the two
readings were 40.06 GB (RSS) against 37.91 GB (PSS), so the correction is
modest -- but a cap should chase memory that freeing an allocation can
actually release. See codai/models/proc_mem.py."""
try:
import os
from codai.models import proc_mem
return proc_mem.tree_held_gb(os.environ.get("CODERAI_FRONT_PID"))
except Exception:
pass
# Fall back to the old summed-RSS reading rather than returning 0.0, which
# would read as "no memory in use" and disable the cap entirely.
try:
import os
import psutil
......@@ -5506,6 +5686,83 @@ class MultiModelManager:
"""Loaded model keys, least-recently-used first (for eviction)."""
return sorted(self.models.keys(), key=lambda k: self._last_used.get(k, 0.0))
def record_history(self, event: str, key: str, **fields) -> None:
"""Note that a model was loaded, evicted or swapped, and why.
``event`` is 'load', 'evict' or 'swap'. ``fields`` carries whatever the
caller knows -- ``reason``, ``seconds``, ``vram_gb``, ``freed_gb``,
``strategy`` -- and is passed through as-is so a new call site can add
detail without a schema change here.
"""
try:
rec = {"time": time.time(), "event": str(event), "model": key,
"engine": os.environ.get("CODERAI_ENGINE_NAME", "")}
for k, v in fields.items():
if v is not None:
rec[k] = round(v, 3) if isinstance(v, float) else v
self._history.append(rec)
except Exception:
pass
def get_history(self, limit: int = 200, event: str = None) -> list:
"""Most recent history entries first."""
try:
items = list(self._history)
except Exception:
return []
if event:
items = [r for r in items if r.get("event") == event]
items.reverse()
try:
limit = max(1, min(500, int(limit)))
except (TypeError, ValueError):
limit = 200
return items[:limit]
def unload_idle_models(self, idle_seconds: float) -> list:
"""Unload models untouched for ``idle_seconds``; returns the keys freed.
Eviction already frees a model when something else needs the memory.
This covers the case where nothing does: a 3.5B music model held 7.9 GB
of a 3090 and 2.1 GB of host RAM for hours after its last request,
because the on-demand path keeps a model resident whenever there is
room ("Ondemand mode - keeping 'audio_gen:ACE-Step' in VRAM alongside
new model").
Recency comes from ``_last_used``, which is stamped by the
``active_in_vram`` setter. A key that has never been stamped is left
alone rather than assumed idle -- that is indistinguishable from a
model loaded seconds ago by a path that forgot to stamp it, and
unloading it would be the opposite of what is wanted.
"""
try:
idle_seconds = float(idle_seconds or 0)
except (TypeError, ValueError):
return []
if idle_seconds <= 0:
return []
now = time.monotonic()
freed = []
for key in list(self._lru_order()):
if key in self._keep_resident:
continue
last = self._last_used.get(key)
if last is None or last <= 0:
continue # never stamped — don't guess
age = now - last
if age < idle_seconds:
continue
if self._is_key_busy(key):
continue
print(f"Idle unload: '{key}' has been untouched for {age / 60.0:.0f} min "
f"(limit {idle_seconds / 60.0:.0f} min) — freeing it")
try:
self._evict_one(key, reason=f"idle for {age / 60.0:.0f} min")
freed.append(key)
except Exception as e:
print(f" Idle unload of '{key}' failed: {e}")
return freed
def _evict_models_for_vram(self, needed_gb: float):
"""Unload loaded models (LRU first) until we have at least needed_gb free VRAM.
......@@ -5543,7 +5800,7 @@ class MultiModelManager:
print(f" '{key}' is busy serving a request — not evicting it")
continue
print(f"On-request VRAM eviction: unloading '{key}' ({_size_label(key)}) to free VRAM")
_evict_key(key)
_evict_key(key, reason=f"vram: another load needed {needed_gb:.1f} GB")
# Second pass: evict the active model if still not enough AND it is idle
# (unless it's keep-resident — that's handled by the last-resort pass).
......@@ -5557,7 +5814,8 @@ class MultiModelManager:
if self._wait_until_idle(_active):
print(f"On-request VRAM eviction: unloading active model '{_active}' "
f"({_size_label(_active)}) to free VRAM")
_evict_key(_active)
_evict_key(_active, reason=f"vram: another load needed {needed_gb:.1f} GB "
f"(was the active model)")
self.active_in_vram = None
else:
print(f" Active model '{_active}' still busy after wait — leaving it loaded")
......@@ -5662,7 +5920,7 @@ class MultiModelManager:
continue
try:
print(f" [cosite-evict] releasing '{key}' for a co-located engine")
self._evict_one(key)
self._evict_one(key, reason="cosite: a co-located engine needed the GPU")
if key == self.active_in_vram:
self.active_in_vram = None
except Exception as e:
......
codai/models/proc_mem.py 0 → 100644
View file @ 8274ff3e
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""How much host memory the server fleet actually holds.
The host-RAM cap used to be enforced against the sum of RSS over the front, the
engines and every worker. That number cannot be driven down by evicting models,
for two reasons:
* **Shared pages are counted once per process.** The interpreter, torch, the
CUDA runtime and libc are mapped into every engine. Summing RSS adds each of
those pages as many times as there are processes.
* **File-backed pages count as RSS.** Walking a mmap'd GGUF or safetensors file
maps its pages into the process; they are page cache, reclaimed by the kernel
on demand, and dropping a model does not release them.
Measured on a production host, 2026-10-29. The mitigation ladder ran every
~15 s for hours with the figure never moving:
RAM watch: RSS 40.0/43.0 GB (93%) -- mitigation [gc+trim+evict_idle] -> 39.1 GB
RAM watch: RSS 39.1/43.0 GB (91%) -- mitigation [gc+trim+evict_idle] -> 39.1 GB
RAM watch: RSS 39.1/43.0 GB (91%) -- mitigation [gc+trim+evict_idle] -> 39.1 GB
Before and after are the same value, 291 identical repetitions for one figure
alone, while the host itself had 14 GB available and 11 GB of reclaimable
cache. It was not merely noisy -- it evicted working models to chase it:
RAM eviction: unloading 'video:Lightricks/LTX-2.5-Diffusers' to free host
RAM (RSS 40.0 GB > cap target 36.5 GB)
The fix is to measure proportional set size instead. ``Pss`` divides each
shared page by the number of processes sharing it, so summing PSS across a tree
counts every page exactly once, and the ``Pss_Anon``/``Pss_File`` split
separates anonymous memory (which only freeing an allocation releases) from
page cache (which the kernel reclaims by itself). A cap should chase the
former.
"""
from __future__ import annotations
import os
#: Fields of ``/proc/<pid>/smaps_rollup`` we care about, in kB.
_WANTED = ("Pss", "Pss_Anon", "Pss_File", "Pss_Shmem", "Rss")
def read_rollup(pid) -> dict:
"""Parse ``/proc/<pid>/smaps_rollup`` into a dict of bytes, or ``{}``.
Cheap: the kernel aggregates the whole address space into one record, so
this is a single short read regardless of how many mappings a process has
(unlike ``smaps``, which emits a stanza per mapping and is what makes
naive memory accounting expensive).
"""
try:
with open(f"/proc/{int(pid)}/smaps_rollup") as fh:
out = {}
for line in fh:
name, _, rest = line.partition(":")
if name in _WANTED:
try:
out[name] = int(rest.split()[0]) * 1024
except (IndexError, ValueError):
continue
return out
except (OSError, ValueError, TypeError):
return {}
def process_held_bytes(pid) -> int:
"""Memory this one process holds that only it can release.
``Pss_Anon`` is anonymous memory -- allocations, model weights that were
copied rather than mapped -- shared pages already divided by their sharers.
``Pss_Shmem`` is shared-memory segments, which are not backed by a file on
disk and so are not reclaimable either. ``Pss_File`` is deliberately
excluded: those are mapped file pages, i.e. page cache.
Falls back to ``Pss`` and then ``Rss`` on a kernel or permission setup that
does not expose the breakdown, so the caller always gets a number -- a
coarse one, which is still better than none.
"""
r = read_rollup(pid)
if not r:
return 0
if "Pss_Anon" in r:
return r.get("Pss_Anon", 0) + r.get("Pss_Shmem", 0)
if "Pss" in r:
return r["Pss"]
return r.get("Rss", 0)
def tree_pids(root_pid=None) -> list:
"""The root process and all its descendants."""
try:
import psutil
except Exception:
return [int(root_pid)] if root_pid else [os.getpid()]
try:
proc = psutil.Process(int(root_pid)) if root_pid else psutil.Process()
except Exception:
try:
proc = psutil.Process()
except Exception:
return []
pids = [proc.pid]
try:
pids += [c.pid for c in proc.children(recursive=True)]
except Exception:
pass
return pids
def tree_held_gb(root_pid=None) -> float:
"""Unreclaimable host memory held by a process tree, in GB.
Summing PSS across the tree is sound precisely because PSS has already
divided shared pages by their sharers: a page mapped into five engines
contributes a fifth from each, totalling one.
"""
total = 0
for pid in tree_pids(root_pid):
total += process_held_bytes(pid)
return total / 1e9
def tree_rss_gb(root_pid=None) -> float:
"""Summed RSS of the tree, in GB -- the old reading, kept for comparison.
Reported alongside the PSS figure so the gap between them is visible in
diagnostics rather than inferred: on a multi-engine host it is large, and
it is entirely double-counted shared pages plus page cache.
"""
total = 0
for pid in tree_pids(root_pid):
r = read_rollup(pid)
total += r.get("Rss", 0)
return total / 1e9
def host_available_gb() -> float:
"""What the kernel says is actually available, including reclaimable cache.
``MemAvailable`` is the kernel's own estimate of what can be allocated
without swapping, and it already accounts for the page cache it is willing
to drop. A cap that disagrees with this is measuring the wrong thing.
"""
try:
with open("/proc/meminfo") as fh:
for line in fh:
if line.startswith("MemAvailable:"):
return int(line.split()[1]) / (1024 ** 2)
except OSError:
pass
try:
import psutil
return psutil.virtual_memory().available / (1024 ** 3)
except Exception:
return 0.0
codai/models/ram_monitor.py
View file @ 8274ff3e
......@@ -48,6 +48,17 @@ def get_status() -> Dict[str, Any]:
return dict(_state)
def _idle_unload_seconds() -> float:
"""``offload.idle_unload_seconds`` (0 = never), via the global args."""
try:
from codai.api.state import get_global_args
ga = get_global_args()
v = getattr(ga, "idle_unload_seconds", None) if ga else None
return float(v) if v else 0.0
except Exception:
return 0.0
def _cap_gb() -> Optional[float]:
try:
from codai.api.state import get_global_args
......@@ -211,6 +222,19 @@ def _loop():
while True:
try:
time.sleep(_poll_seconds())
# Idle unload is independent of the leak watch and of the RAM cap: a
# model nobody has used for half an hour is not holding its memory
# for any reason, and waiting for pressure is how a 3.5B music model
# came to sit on 7.9 GB of a 3090 with the card otherwise idle. Run
# it BEFORE the _watch_enabled() gate so turning the leak watch off
# does not silently disable it too.
try:
_idle_limit = _idle_unload_seconds()
if _idle_limit > 0:
from codai.models.manager import multi_model_manager as _mmm
_mmm.unload_idle_models(_idle_limit)
except Exception as _iue:
_log.debug("idle unload sweep failed: %s", _iue)
if not _watch_enabled():
continue
cap = _cap_gb()
......
codai/models/teardown.py 0 → 100644
View file @ 8274ff3e
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Finding a pipeline's weights that outlived its teardown -- on the HOST side.
``codai/api/video.py`` already carries a teardown-leak breaker, and it works:
on a production host it went from reporting "16.9 GB of this pipeline's tensors
survived teardown; nulled 7581 external reference(s)" to 0.0 GB once the
injected components were released in the right order. But both halves of it are
CUDA-only -- it records ``_orig_cuda_ptrs`` under ``if _tens.is_cuda`` and
sweeps under ``if not (isinstance(_obj, Tensor) and _obj.is_cuda): continue``.
A pipeline loaded with ``enable_model_cpu_offload`` keeps its weights in HOST
RAM and streams modules to the GPU for each forward. Seventeen of twenty-one
video loads in one production day took that path. None of those weights were
ever tracked, so nothing looked for them after teardown. Measured on the same
host, 2026-10-29, with nothing loaded and 482 MiB of VRAM in use:
coderai-nvidia 33.14 GB PSS_Anon
total anon: 30.85 GB in 6588 mappings
>= 1 GB: 2
100 MB - 1 GB: 21 (3.49 GB)
< 100 MB: 6,565 (20.25 GB)
Thousands of separate mid-sized allocations is retained tensors, not allocator
slack -- which is why the mitigation ladder's gc + ``malloc_trim`` could run
every 15 s for hours without the figure moving, and why it then evicted healthy
models to chase it.
Identity, not addresses
-----------------------
The CUDA breaker matches surviving tensors by storage pointer. That is sound
enough on the GPU, but a freed *host* address is readily handed back out by
malloc, so a pointer that once belonged to this pipeline may by then belong to
another model's tensor -- and nulling its referrers would corrupt a live model.
So this tracks :mod:`weakref` references to the tensor objects themselves. A
referent that is still alive after teardown IS the leaked tensor, with no
address-reuse ambiguity, and a referent that has gone is proof the teardown
worked.
"""
from __future__ import annotations
import gc
import weakref
def _component_tensors(pipe):
"""Every parameter and buffer of every component of a pipeline."""
comps = getattr(pipe, "components", None)
if not isinstance(comps, dict):
comps = {}
for comp in comps.values():
if comp is None or not hasattr(comp, "parameters"):
continue
try:
params = list(comp.parameters(recurse=True))
except Exception:
params = []
try:
bufs = list(comp.buffers(recurse=True))
except Exception:
bufs = []
for t in params + bufs:
yield t
# A bnb 4-bit parameter keeps its quantized payload on `.data`,
# which can be a different tensor object from the parameter.
d = getattr(t, "data", None)
if d is not None and d is not t:
yield d
def track_host_tensors(pipe) -> list:
"""Weak references to this pipeline's HOST-resident tensors.
Call before teardown. GPU-resident tensors are left to the existing CUDA
breaker, so the two do not fight over the same objects.
"""
refs = []
seen = set()
for t in _component_tensors(pipe):
try:
if getattr(t, "is_cuda", False) or t.device.type != "cpu":
continue
if t.numel() == 0:
continue
except Exception:
continue
if id(t) in seen:
continue
seen.add(id(t))
try:
refs.append(weakref.ref(t))
except TypeError:
continue # not weak-referenceable; nothing we can do
return refs
def _null_referrers(obj) -> int:
"""Clear the slots in lists/dicts that still point at ``obj``.
The same approach the CUDA breaker uses: a weight pinned by a stray entry
in some cache dict or list cannot be freed by dropping the pipeline, and
these two container kinds are what actually holds them in practice
(``_modules`` dicts, hook state, our own caches).
"""
broken = 0
for ref in gc.get_referrers(obj):
try:
if isinstance(ref, list):
for i, item in enumerate(ref):
if item is obj:
ref[i] = None
broken += 1
elif isinstance(ref, dict):
for k, v in list(ref.items()):
if v is obj:
ref[k] = None
broken += 1
except Exception:
continue
return broken
def break_host_leaks(refs, *, min_report_gb: float = 0.05) -> dict:
"""Null external references to host tensors that survived teardown.
``refs`` is what :func:`track_host_tensors` returned. Returns a report with
the surviving size, how many references were broken, and the biggest
holders, mirroring what the CUDA breaker prints so the two read alike.
Only tensors we tracked are touched, and each by identity, so another
model's memory is never at risk.
"""
out = {"survived_gb": 0.0, "tensors": 0, "broken": 0, "holders": []}
if not refs:
return out
gc.collect()
for ref in refs:
t = ref()
if t is None:
continue # freed as it should have been
try:
gb = t.numel() * t.element_size() / 1e9
except Exception:
continue
out["survived_gb"] += gb
out["tensors"] += 1
out["broken"] += _null_referrers(t)
if gb >= min_report_gb and len(out["holders"]) < 6:
try:
out["holders"].append(f"{gb:.2f} GB {tuple(t.shape)} {t.dtype}")
except Exception:
pass
del t
gc.collect()
out["survived_gb"] = round(out["survived_gb"], 3)
return out
def trim_host_allocator() -> bool:
"""Ask glibc to return freed heap to the OS.
Worth doing only AFTER the references are broken: ``malloc_trim`` can
return nothing while the memory is still referenced, which is exactly what
the mitigation ladder was doing every 15 s for hours.
"""
try:
import ctypes
return bool(ctypes.CDLL("libc.so.6").malloc_trim(0))
except Exception:
return False
docs/README.md
View file @ 8274ff3e
......@@ -33,6 +33,7 @@ reference; these are the deep dives.
| [zimage-lora-training.md](zimage-lora-training.md) | LoRA training for image models |
| [text-to-music.md](text-to-music.md) | `POST /v1/audio/generate` with ACE-Step: style prompt + sung lyrics → a 48 kHz stereo track, the python-3.12 worker it needs, and the two image variants (`--with-acestep`) |
| [dtype-auto-selection.md](dtype-auto-selection.md) | How model dtype is chosen automatically |
| [memory-accounting.md](memory-accounting.md) | Where `measured_vram_gb` comes from: walking the loaded model (quantization needs no factor), costing the KV cache from the real GQA/sliding-window geometry, refusing a measurement from a failed load, the PSS-based host-RAM cap, idle unload, and `GET /v1/models/history` |
| [model-defaults-and-sources.md](model-defaults-and-sources.md) | A model's own `defaults` (steps/guidance/sampler/modes/licence), the sampler names and the flow-matching boundary, and where weights come from: repo, single file, `org/repo:file` or `civitai:<versionId>` |
## Deployment
......
docs/memory-accounting.md 0 → 100644
View file @ 8274ff3e
# How coderai works out what a model needs
Every placement decision — resident on the GPU, accelerate CPU offload, how
much to evict before a load, whether to spill to disk — is made against one
number per model. This is where that number comes from, and what it does and
does not include.
## The model is measured, not the process
`measured_vram_gb` is a model's **full footprint**, and it is deliberately
placement-independent: offload moves weights between GPU and host, it does not
change how much the model needs. The split is reported separately.
It is obtained by walking the loaded object and adding up the bytes its tensors
really occupy (`codai/models/footprint.py`), plus the KV cache for the
configured context, plus a runtime reserve for activations and the decode
spike.
Quantization needs no correction factor, and that is the whole reason for
measuring the object rather than estimating from config:
* a bitsandbytes `Params4bit` **is** a `uint8` tensor, so `numel ×
element_size` already yields the packed size;
* GPTQ and AWQ keep `qweight` / `qzeros` / `scales` in **buffers**, so buffers
are walked as well as parameters — a parameters-only walk reports a 4-bit
model as very nearly empty;
* bitsandbytes keeps `absmax`, the code book and (under double quantization) a
nested second-level state in `quant_state`, outside `parameters()`. That is
real memory and is counted explicitly.
Tied weights — an embedding shared with the LM head — are one allocation and
are deduplicated by storage pointer. An accelerate `meta` placeholder, left
where the real weights were moved to host RAM or disk, is sized from its
declared dtype and flagged, because the model still needs those bytes.
A diffusers pipeline is not an `nn.Module`; it is walked through its
components, skipping the schedulers and tokenizers that hold nothing. A
text/vision model is registered as a `ModelManager` wrapper, so its real module
and its config are reached through `.backend.model`.
## Context is costed from the geometry
The KV cache is sized from the number of **key/value** heads, not attention
heads. Under GQA many query heads share one KV head, so sizing by attention
heads overstates the cache by the grouping factor — fourfold on a 32/8 model.
Sliding-window layers are sized to their window rather than the full context
when `swa_full` is off: a model that alternates a windowed layer with a global
one would otherwise be overstated on half its layers.
**A computed KV figure is an estimate and is labelled as one.** Per-layer
KV-head arrays, shared KV and architecture-specific cache layouts are not
modelled, and predicting this from GGUF metadata has been wrong by 5× in both
directions on these very models. A measured `measured_kv_gb` for the context in
question — llama.cpp reports its own KV buffer size exactly — always wins.
## What a GGUF needs
A GGUF is loaded as stored, with no dtype conversion, so its weights **are** the
file, byte for byte. llama.cpp owns that memory and there is no torch object to
walk, so the file size is used, plus the KV cache and the reserve. A companion
vision projector (`mmproj`) loads on top and is added.
## A measurement that cannot be trusted is refused
`force_vram_update` refreshes these fields on every load, so a bad reading
becomes authoritative and is planned against. Three checks stand in the way:
1. **An empty load.** If the object is walkable and holds no weights at all,
the load did not complete. There is nothing to fall back to — a size taken
from process memory measures everything *except* this model — so nothing is
recorded, and the refusal is written to the history with its reason.
2. **An implausible size.** No real model is under ~10 MB.
3. **A collapse against the previous reading.** Less than half of what the
*same quantization signature* measured before is a partial load. A genuine
drop arrives together with a signature change — adding 4-bit quantization
to a model's transformers took one from 40.04 GB to 11.73 GB — and is let
through. The signature covers precision, `load_in_4bit`, `load_in_8bit` and
`component_quantization`.
The log names which check fired, because the three call for opposite responses.
## Why not process memory
The footprint used to be a process-RSS delta across the load. RSS is the wrong
instrument: it counts file-backed pages that the kernel maps when the loader
reads a mmap'd safetensors or GGUF, and summed over a process tree it counts
pages **shared** between those processes once per process — the interpreter,
torch and the CUDA runtime are mapped into every engine. As a delta it also
absorbs anything else that loaded, evicted or faulted pages in the same window.
Measured against ground truth on a 5-engine host:
```
gemma-2-9b-it.Q4_K_M 5.37 GB on disk -> reported 30.81 GB (5.7x)
LTX-2.5, six loads, one unchanged config:
22.11, 22.11, 23.28, 4.50, 19.02, 29.80 GB (6.6x)
```
Every other model on that host sat at 1.0–1.2×. The 4.50 GB reading came from a
load that placed a single layer and then raised; `Full-GPU load looks viable
(4.5 GB peak need)` was planned off it four seconds later.
## The host-RAM cap
`offload.max_ram_gb` is enforced against proportional set size
(`Pss_Anon + Pss_Shmem`, from `/proc/<pid>/smaps_rollup`), summed over the
front, the engines and their workers. PSS has already divided each shared page
by the number of processes sharing it, so summing it across a tree counts every
page exactly once; `Pss_File` is excluded because mapped file pages are page
cache that the kernel reclaims by itself and that evicting a model cannot free.
The cap can only be met by releasing memory something actually holds, so the
metric has to be one that releasing memory moves.
## Idle unload
Eviction frees a model when something else needs the memory.
`offload.idle_unload_seconds` (default 1800) covers the case where nothing
does: the on-demand path keeps a model resident whenever there is room, which
left a 3.5B music model on 7.9 GB of a 24 GB card for hours after its last
request.
It is deliberately generous. Reloads are the expensive event — one production
day spent 2h04m of 21h30m inside model loads — so a short timeout trades a VRAM
saving nobody asked for against a reload everybody waits on. A model whose LRU
clock was never stamped is left alone rather than assumed idle.
## Seeing it happen
`GET /v1/models/history` returns what was loaded and evicted and why, newest
first, merged across engines — a swap is a cross-engine event, since one engine
evicts so a co-located one can load.
```
time, event ('load' | 'evict'), model, engine
load: seconds, vram_gb, offloaded_gb, strategy, measured
evict: reason, freed_gb, expected_gb
```
It is in-memory and bounded (500 entries), so it is empty after a restart. It
exists because the alternative was reconstructing swaps from the log, and the
log rotates: on one host 78% of it was per-request debug, which filled the
2,000-line view in about four minutes.
docs/text-to-music.md
View file @ 8274ff3e
......@@ -17,7 +17,7 @@ curl -s http://localhost:8000/v1/audio/generate \
| field | ACE-Step default | notes |
|---|---|---|
| `prompt` | — | the *style*, not the words: genre, instruments, tempo, mood |
| `lyrics` | `""` | supports its `[verse]` / `[chorus]` tags; `[inst]` for an instrumental |
| `lyrics` | `[instrumental]` | supports its `[verse]` / `[chorus]` tags; send words to get vocals |
| `duration` | 60 s | the whole track is generated at once, not stitched |
| `steps` | 60 | diffusion steps |
| `cfg_coef` | 15.0 | guidance; ACE-Step wants far more of it than MusicGen |
......@@ -31,6 +31,23 @@ them out gets **ACE-Step's** defaults instead, because a 10-second song at a
guidance it was not tuned for is not a useful reading of "unspecified". A value
you actually send is always used, including `duration: 10`.
## No lyrics means instrumental, and it has to be said
ACE-Step is a *song* model: the prompt carries the style and the lyrics carry
the words. Leaving `lyrics` empty does **not** ask for an instrumental — the
pipeline then conditions on a single padding token (`lyric_token_idx = [0]`,
`lyric_mask = [0]`) and the model still tries to sing, which comes out as
babbling over an otherwise reasonable backing track. ACE-Step's own interface
says how to ask: *"Use `[instrumental]` or `[inst]` to generate instrumental
music"*, and every genre preset it ships ends in either `male vocals, raw
vocals` or `instrumental`.
So a request with no `lyrics` is sent as `[instrumental]`, and `instrumental`
is appended to the style prompt when it is not already there. Background music
for a video — the common case, where the caller has a mood to describe and no
words — therefore comes out instrumental by default. To get vocals, send
lyrics.
## Why it runs in its own interpreter
ACE-Step ships its own `acestep` package, which pins `transformers==4.50` and
......
tests/test_footprint_is_measured_not_guessed.py 0 → 100644
View file @ 8274ff3e
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""The footprint must come from the model, and must survive quantization.
Replaces a process-RSS delta that was wrong by construction. Ground truth from
a production host, 2026-10-29:
gemma-2-9b-it.Q4_K_M 5.37 GB on disk -> reported 30.81 GB (5.7x)
LTX-2.5, six loads, one unchanged config:
22.11, 22.11, 23.28, 4.50, 19.02, 29.80 GB (6.6x)
RSS counts mmap'd weight-file pages and, summed over a process tree, counts
pages shared between those processes once per process; as a delta it also
absorbs whatever else moved in the window. The 4.50 GB reading came from a load
that placed one layer and then raised, and the 29.80 GB one was taken after an
abandoned resident attempt had pushed a transformer and text encoder to host
RAM. Both were persisted under force_vram_update and planned against.
"""
import io
import os
import sys
import pytest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from codai.models import footprint as F
torch = pytest.importorskip("torch")
nn = torch.nn
# ---------------------------------------------------------------- weights ---
def test_a_plain_module_measures_its_real_bytes():
m = nn.Linear(1000, 1000, bias=False).to(torch.float32)
r = F.walk_module_bytes(m)
assert r["gpu"] == 0
assert r["cpu"] == 1000 * 1000 * 4 # fp32, exactly
assert r["declared"] == 0
def test_dtype_is_read_from_the_tensor_not_assumed():
"""Half precision must halve the figure -- no factor table involved."""
big = F.walk_module_bytes(nn.Linear(512, 512, bias=False).to(torch.float32))
small = F.walk_module_bytes(nn.Linear(512, 512, bias=False).to(torch.bfloat16))
assert big["cpu"] == 2 * small["cpu"]
def test_a_four_bit_parameter_counts_its_packed_bytes_not_its_shape():
"""A bitsandbytes Params4bit IS a uint8 tensor, so numel x element_size
already yields the packed size. This is why measuring the object needs no
quantization correction factor, unlike an estimate from the config."""
m = nn.Module()
packed = torch.empty(512 * 512 // 2, dtype=torch.uint8) # 4 bits/weight
m.register_parameter("weight", nn.Parameter(packed, requires_grad=False))
r = F.walk_module_bytes(m)
assert r["cpu"] == 512 * 512 // 2
# ... and that is a quarter of the fp16 equivalent, not the same.
assert r["cpu"] == (512 * 512 * 2) // 4
def test_the_bnb_quant_state_is_counted():
"""absmax / the code book / the nested second-level state are real memory
and live OUTSIDE parameters(), so a plain walk misses them."""
class _QS:
def __init__(self, nested=None):
self.absmax = torch.empty(4096, dtype=torch.float32) # 16 KiB
self.code = torch.empty(256, dtype=torch.float32) # 1 KiB
self.state2 = nested
m = nn.Module()
p = nn.Parameter(torch.empty(1024, dtype=torch.uint8), requires_grad=False)
p.quant_state = _QS(nested=_QS()) # double quantization
m.register_parameter("weight", p)
r = F.walk_module_bytes(m)
expected = 1024 + 2 * (4096 * 4 + 256 * 4)
assert r["cpu"] == expected, "quant_state tensors were not counted"
def test_prequantized_buffers_are_counted():
"""GPTQ/AWQ keep the payload in BUFFERS (qweight/qzeros/scales), so a
parameters-only walk reports a 4-bit model as very nearly empty."""
m = nn.Module()
m.register_buffer("qweight", torch.empty((512, 64), dtype=torch.int32))
m.register_buffer("qzeros", torch.empty((4, 64), dtype=torch.int32))
m.register_buffer("scales", torch.empty((4, 512), dtype=torch.float16))
r = F.walk_module_bytes(m)
expected = 512 * 64 * 4 + 4 * 64 * 4 + 4 * 512 * 2
assert r["cpu"] == expected
assert r["tensors"] == 3
def test_tied_weights_are_counted_once():
"""An embedding shared with the LM head is ONE allocation. Counting the
same storage twice inflates every model that ties them."""
emb = nn.Embedding(5000, 256)
head = nn.Linear(256, 5000, bias=False)
head.weight = emb.weight # tied, as in most LLMs
m = nn.Module()
m.emb, m.head = emb, head
r = F.walk_module_bytes(m)
assert r["cpu"] == 5000 * 256 * 4
assert r["tensors"] == 1
def test_a_meta_placeholder_is_sized_from_its_dtype_and_flagged():
"""accelerate leaves a meta tensor where it moved the real weights to host
RAM or disk. The model still needs those bytes (the footprint is
placement-independent), but they are DECLARED, not measured."""
m = nn.Module()
with torch.device("meta"):
w = torch.empty((1000, 1000), dtype=torch.bfloat16)
m.register_parameter("weight", nn.Parameter(w, requires_grad=False))
r = F.walk_module_bytes(m)
assert r["cpu"] == 1000 * 1000 * 2
assert r["declared"] == r["cpu"], "meta bytes were not flagged as declared"
def test_a_pipeline_is_walked_through_its_components():
"""A diffusers pipeline is not an nn.Module; its weights are in components,
beside schedulers and tokenizers that hold none. Walking the container
alone would report ~0 for every video and image model."""
class _Pipe:
pass
p = _Pipe()
p.components = {
"transformer": nn.Linear(1000, 1000, bias=False).to(torch.bfloat16),
"vae": nn.Linear(500, 500, bias=False).to(torch.bfloat16),
"scheduler": object(), # no parameters()
"tokenizer": None,
}
gpu, cpu, notes, ok = F.measure_weights(p)
assert ok
assert cpu == pytest.approx((1000 * 1000 * 2 + 500 * 500 * 2) / 1e9, rel=1e-6)
# --------------------------------------------------------------------- KV ---
def test_kv_uses_key_value_heads_not_attention_heads():
"""Under GQA many query heads share one KV head. Sizing by attention heads
over-states the cache by the grouping factor -- 4x on this geometry."""
gqa = {"num_hidden_layers": 32, "num_attention_heads": 32,
"num_key_value_heads": 8, "head_dim": 128}
mha = dict(gqa, num_key_value_heads=32)
assert F.kv_cache_gb(mha, 4096) == pytest.approx(4 * F.kv_cache_gb(gqa, 4096))
def test_kv_scales_with_the_configured_context():
g = {"num_hidden_layers": 32, "num_key_value_heads": 8, "head_dim": 128}
assert F.kv_cache_gb(g, 32768) == pytest.approx(8 * F.kv_cache_gb(g, 4096))
def test_kv_is_exact_for_a_known_geometry():
"""2 (K+V) x layers x kv_heads x head_dim x 2 bytes x tokens."""
g = {"num_hidden_layers": 32, "num_key_value_heads": 8, "head_dim": 128}
expect = 32 * 8 * (128 * 16 + 128 * 16) / 8.0 * 32000 / 1e9
assert F.kv_cache_gb(g, 32000) == pytest.approx(expect, rel=1e-9)
def test_a_quantized_kv_cache_is_smaller():
"""cache_type_k/v=q8_0 roughly halves the cache; reserving an f16-sized one
forces needless CPU offload at long context."""
g = {"num_hidden_layers": 32, "num_key_value_heads": 8, "head_dim": 128}
f16 = F.kv_cache_gb(g, 32000)
q8 = F.kv_cache_gb(g, 32000, "q8_0", "q8_0")
assert q8 < f16
assert q8 == pytest.approx(f16 * 8.5 / 16.0, rel=1e-6)
def test_unknown_geometry_reports_zero_rather_than_guessing():
assert F.kv_cache_gb({}, 32000) == 0.0
assert F.kv_cache_gb({"num_hidden_layers": 32}, 32000) == 0.0
def test_a_diffusion_model_gets_no_kv_term():
pipe = type("P", (), {})()
pipe.components = {"transformer": nn.Linear(100, 100).to(torch.bfloat16)}
fp = F.measure(pipe, {"n_ctx": 32768}, "video")
assert fp.kv_gb == 0.0, "a diffusion pipeline has no KV cache"
# ------------------------------------------------------------ integration ---
def test_the_gemma_weights_are_the_file_not_six_times_it():
"""The production reading: a 5.37 GB Q4_K_M GGUF reported 30.81 GB total.
A GGUF is loaded as-is with no dtype conversion, so its WEIGHTS are exactly
the file -- that 5.7x was the RSS delta, not the model."""
g = {"num_hidden_layers": 42, "num_key_value_heads": 8, "head_dim": 256}
fp = F.measure(None, {"n_ctx": 32000}, "text", geom=g, file_size_gb=5.37)
assert fp.complete
assert fp.weights_gb == pytest.approx(5.37), (
f"the weights are no longer the file size: {fp.describe()}")
def test_the_kv_estimate_is_labelled_as_one():
"""Predicting KV from metadata has been wrong by 5x in both directions on
this codebase's own models, so a computed figure must announce itself."""
g = {"num_hidden_layers": 42, "num_key_value_heads": 8, "head_dim": 256}
fp = F.measure(None, {"n_ctx": 32000}, "text", geom=g, file_size_gb=5.37)
assert fp.kv_gb > 0
assert not fp.kv_measured
assert any("ESTIMATE" in n for n in fp.notes), fp.notes
assert "(est)" in fp.describe()
def test_a_measured_kv_size_beats_the_estimate():
"""llama.cpp reports its own KV buffer size exactly. When that has been
recorded for this context it must win over any computation."""
g = {"num_hidden_layers": 42, "num_key_value_heads": 8, "head_dim": 256}
est = F.measure(None, {"n_ctx": 32000}, "text", geom=g, file_size_gb=5.37)
meas = F.measure(None, {"n_ctx": 32000}, "text", geom=g, file_size_gb=5.37,
measured_kv_gb={"32000": 2.2})
assert meas.kv_gb == pytest.approx(2.2)
assert meas.kv_measured
assert meas.kv_gb < est.kv_gb
assert "(est)" not in meas.describe()
def test_a_measured_kv_size_is_per_context():
"""A figure taken at 4k says nothing about 128k, so it must not be reused
for a context it was not measured at."""
g = {"num_hidden_layers": 42, "num_key_value_heads": 8, "head_dim": 256}
fp = F.measure(None, {"n_ctx": 32000}, "text", geom=g, file_size_gb=5.37,
measured_kv_gb={"4096": 0.3})
assert not fp.kv_measured, "a 4k measurement was reused for a 32k context"
def test_sliding_window_layers_are_sized_to_their_window():
"""Gemma alternates a windowed layer with a global one. Sizing every layer
at the full context over-states the windowed half by n_ctx/window -- which
is why llama.cpp's own figure was far below the naive calculation."""
g = {"num_hidden_layers": 42, "num_key_value_heads": 8, "head_dim": 256,
"sliding_window": 4096}
full = F.kv_cache_gb(g, 32000, swa_full=True)
windowed = F.kv_cache_gb(g, 32000, swa_full=False)
assert windowed < full
# 21 global layers at 32000 tokens + 21 windowed layers at 4096.
expect_ratio = (21 * 32000 + 21 * 4096) / (42 * 32000)
assert windowed / full == pytest.approx(expect_ratio, rel=1e-6)
def test_a_known_sliding_layer_count_is_honoured():
"""llama.cpp prints is_swa per layer, so a caller that counted them should
not be overridden by the alternating assumption."""
g = {"num_hidden_layers": 42, "num_key_value_heads": 8, "head_dim": 256,
"sliding_window": 4096, "swa_layer_count": 36}
assumed = F.kv_cache_gb(dict(g, swa_layer_count=0), 32000, swa_full=False)
stated = F.kv_cache_gb(g, 32000, swa_full=False)
assert stated < assumed, "swa_layer_count was ignored"
def test_a_window_wider_than_the_context_changes_nothing():
g = {"num_hidden_layers": 32, "num_key_value_heads": 8, "head_dim": 128,
"sliding_window": 8192}
assert (F.kv_cache_gb(g, 4096, swa_full=False)
== pytest.approx(F.kv_cache_gb(g, 4096, swa_full=True)))
def test_an_unaccountable_model_is_not_reported_as_complete():
"""The 4.50 GB production reading came from a load that placed one layer
and then raised. A footprint that could not be accounted for must announce
that, so the caller declines to persist it as authoritative."""
fp = F.measure(object(), {}, "video")
assert not fp.complete
assert fp.notes
def test_the_total_is_placement_independent():
"""Moving weights GPU<->host must not change what the model needs -- the
contract measured_vram_gb already documents."""
m = nn.Linear(1000, 1000, bias=False).to(torch.bfloat16)
pipe = type("P", (), {})()
pipe.components = {"transformer": m}
on_host = F.measure(pipe, {}, "video").total_gb
if torch.cuda.is_available():
pipe.components = {"transformer": m.cuda()}
assert F.measure(pipe, {}, "video").total_gb == pytest.approx(on_host)
m.cpu()
else:
assert on_host > 0
def test_companion_files_are_added(tmp_path):
"""An mmproj projector loads ON TOP of the model and is held by a separate
handler, so walking the model cannot see it."""
p = tmp_path / "mmproj.gguf"
p.write_bytes(b"\0" * 1_000_000)
assert F.extra_file_gb(str(p)) == pytest.approx(0.001)
assert F.extra_file_gb(None, "", "/nonexistent") == 0.0
def test_a_gguf_wrapper_falls_back_to_the_file_and_stays_complete():
"""A GGUF is not a torch object: llama.cpp owns its memory and the registry
holds a ModelManager wrapper with no parameters() to walk. The weights are
the file byte for byte -- loaded as stored, no dtype conversion -- so with
that size and llama.cpp's own measured KV the production 30.81 GB reading
for a 5.37 GB Q4_K_M becomes explainable."""
class _ManagerLike:
def __init__(self):
self.backend = object()
fp = F.measure(_ManagerLike(), {"n_ctx": 32000}, "text",
geom={"num_hidden_layers": 42, "num_key_value_heads": 8,
"head_dim": 256},
file_size_gb=5.37, measured_kv_gb={"32000": 2.2})
assert fp.complete, fp.notes
assert fp.weights_gb == pytest.approx(5.37)
assert fp.kv_gb == pytest.approx(2.2)
assert fp.kv_measured
assert fp.total_gb < 10.0, f"still inflated: {fp.describe()}"
assert any("on-disk size" in n for n in fp.notes)
def test_without_a_file_size_an_unwalkable_object_is_incomplete():
"""Better to decline than to persist a guess: this is the path that wrote
4.50 GB for a 22 GB model and planned against it four seconds later."""
class _ManagerLike:
pass
fp = F.measure(_ManagerLike(), {"n_ctx": 4096}, "text")
assert not fp.complete
# ------------------------------------------------- the partial-load guard ---
def _manager():
from codai.models.manager import multi_model_manager
return multi_model_manager
def _load(m, key, width):
class Pipe:
pass
p = Pipe()
p.components = {"transformer": nn.Linear(width, width, bias=False).to(torch.bfloat16)}
m.models[key] = p
m.active_in_vram = key
m.record_vram_delta(key, free_before=10.0, ram_before=1.0)
def test_a_partial_load_cannot_overwrite_a_good_measurement():
"""The production failure: a load placed ONE layer, raised, and still wrote
4.505 GB to disk under force_vram_update -- and 'Full-GPU load looks viable
(4.5 GB peak need)' was planned off it four seconds later."""
m = _manager()
key = "video:__guard_partial"
m.config[key] = {"force_vram_update": True}
try:
_load(m, key, 20000)
good = m.config[key].get("measured_vram_gb")
assert good and good > 0.5
_load(m, key, 6000) # a fraction of the real model
assert m.config[key].get("measured_vram_gb") == good, (
"a partial load overwrote the good measurement")
finally:
m.models.pop(key, None)
m.config.pop(key, None)
def test_a_real_drop_with_a_quantization_change_is_accepted():
"""Wan2.2 genuinely went 40.04 -> 11.73 GB once its transformers were
quantized. That must not be mistaken for a partial load -- and it is told
apart by the quantization signature changing with it."""
m = _manager()
key = "video:__guard_quant"
m.config[key] = {"force_vram_update": True}
try:
_load(m, key, 20000)
before = m.config[key].get("measured_vram_gb")
m.config[key]["load_in_4bit"] = True # signature changes with it
_load(m, key, 6000)
after = m.config[key].get("measured_vram_gb")
assert after < before, "the quantization escape hatch did not open"
finally:
m.models.pop(key, None)
m.config.pop(key, None)
def test_an_empty_pipeline_is_reported_as_a_failed_load():
"""Walked every component and none holds weights: that is a failed load,
not an object we cannot introspect. The difference decides whether a
fallback estimate is legitimate -- and it is not, because a size taken from
process memory measures everything EXCEPT this model."""
class Pipe:
pass
p = Pipe()
p.components = {"scheduler": object(), "tokenizer": None}
fp = F.measure(p, {}, "video")
assert not fp.complete
assert fp.empty_load
assert any(n.startswith("EMPTY_PIPELINE") for n in fp.notes), fp.notes
def test_an_unintrospectable_object_is_not_called_an_empty_load():
"""A GGUF wrapper is not a failed load -- it is a model we cannot walk, and
its file size IS a legitimate answer."""
class _ManagerLike:
pass
fp = F.measure(_ManagerLike(), {}, "text", file_size_gb=5.37)
assert not fp.empty_load
assert fp.complete
assert fp.weights_gb == pytest.approx(5.37)
def test_a_failed_load_persists_nothing_and_says_why():
"""Deterministic, unlike a check on the resulting number: the previous
figure must survive and the history must record the refusal. Asserting on
process memory made this pass alone and fail in the full suite, because the
fallback delta grows with whatever else the process has loaded."""
m = _manager()
key = "video:__guard_zero"
m.config[key] = {"force_vram_update": True, "measured_vram_gb": 22.11}
try:
class Pipe:
pass
p = Pipe()
p.components = {"scheduler": object()} # nothing with weights
m.models[key] = p
m.active_in_vram = key
m.record_vram_delta(key, free_before=10.0, ram_before=1.0)
assert m.config[key].get("measured_vram_gb") == 22.11, (
"a failed load overwrote the previous good measurement")
last = m.get_history(limit=1)[0]
assert last["model"] == key and "no weights" in last.get("reason", "")
finally:
m.models.pop(key, None)
m.config.pop(key, None)
def test_the_rejection_names_which_check_fired():
"""Two different conditions reject a measurement -- an implausible absolute
size and a collapse against the previous one -- and the log must say which,
or it sends the reader after the wrong thing."""
src = io.open(os.path.join(os.path.dirname(os.path.dirname(
os.path.abspath(__file__))), "codai/models/manager.py"), encoding="utf-8").read()
blk = src.split("Reject a partial load", 1)[1][:2200]
assert "not a plausible size for a loaded model" in blk
assert "less than half the" in blk
def test_a_text_model_wrapper_is_walked_through_its_backend():
"""A text/vision model is registered as a ModelManager wrapper, so its real
module hangs off .backend.model. Without reaching it, an HF LLM falls
through to the process-delta fallback despite being perfectly walkable."""
class _Backend:
def __init__(self, m):
self.model = m
class _Wrapper:
def __init__(self, m):
self.backend = _Backend(m)
inner = nn.Linear(1000, 1000, bias=False).to(torch.bfloat16)
gpu, cpu, notes, ok = F.measure_weights(_Wrapper(inner))
assert ok, notes
assert cpu == pytest.approx(1000 * 1000 * 2 / 1e9, rel=1e-6)
def test_a_vision_projector_beside_the_lm_is_counted():
"""An mmproj loads ON TOP of the language model as a separate module."""
class _Backend:
def __init__(self, m):
self.model = m
class _Wrapper:
def __init__(self, m, proj):
self.backend = _Backend(m)
self.mmproj = proj
lm = nn.Linear(500, 500, bias=False).to(torch.bfloat16)
proj = nn.Linear(200, 200, bias=False).to(torch.bfloat16)
gpu, cpu, notes, ok = F.measure_weights(_Wrapper(lm, proj))
assert ok
assert cpu == pytest.approx((500 * 500 * 2 + 200 * 200 * 2) / 1e9, rel=1e-6)
def test_geometry_is_found_through_the_backend_wrapper_too():
"""Walking a wrapper's weights but missing its config would drop the KV
term for the one model kind that definitely has one."""
class _Cfg:
num_hidden_layers = 32
num_key_value_heads = 8
head_dim = 128
num_attention_heads = 32
hidden_size = 4096
class _Inner(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(100, 100)
self.config = _Cfg()
class _Backend:
def __init__(self, m):
self.model = m
class _Wrapper:
def __init__(self, m):
self.backend = _Backend(m)
w = _Wrapper(_Inner())
geom = F.hf_geometry(w)
assert geom.get("num_key_value_heads") == 8, geom
fp = F.measure(w, {"n_ctx": 8192}, "text")
assert fp.kv_gb > 0, f"the KV term was dropped: {fp.describe()}"
tests/test_host_teardown_leak.py 0 → 100644
View file @ 8274ff3e
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Host-resident pipeline weights must be released on teardown too.
The existing teardown-leak breaker in video.py is CUDA-only on both sides: it
registers under ``if _tens.is_cuda`` and sweeps under ``if not (... and
_obj.is_cuda): continue``. A pipeline loaded with ``enable_model_cpu_offload``
keeps its weights in HOST RAM, and 17 of 21 video loads in one production day
took that path, so none of those weights were tracked and nothing looked for
them afterwards.
Measured 2026-10-29 with nothing loaded and 482 MiB of VRAM in use:
coderai-nvidia 33.14 GB PSS_Anon
30.85 GB anon in 6588 mappings; 6565 of them under 100 MB (20.25 GB)
Thousands of mid-sized allocations is retained tensors, not allocator slack --
so gc + malloc_trim every 15 s never moved it, and the ladder evicted healthy
models instead.
"""
import gc
import io
import os
import sys
import pytest
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from codai.models import teardown as T
torch = pytest.importorskip("torch")
nn = torch.nn
def _pipe(**comps):
p = type("Pipe", (), {})()
p.components = comps
return p
def test_host_tensors_are_tracked():
pipe = _pipe(transformer=nn.Linear(64, 64), scheduler=object(), tokenizer=None)
refs = T.track_host_tensors(pipe)
assert refs, "no host tensors tracked for a CPU-resident pipeline"
def test_gpu_tensors_are_left_to_the_cuda_breaker():
"""The two arms must not both chase the same objects."""
if not torch.cuda.is_available():
pytest.skip("needs CUDA")
pipe = _pipe(transformer=nn.Linear(64, 64).cuda())
assert T.track_host_tensors(pipe) == []
def test_a_clean_teardown_reports_nothing_surviving():
pipe = _pipe(transformer=nn.Linear(256, 256))
refs = T.track_host_tensors(pipe)
del pipe
gc.collect()
r = T.break_host_leaks(refs)
assert r["survived_gb"] == 0.0
assert r["tensors"] == 0
assert r["broken"] == 0
def test_a_tensor_pinned_by_a_dict_is_found_and_released():
"""The real shape of the leak: a stray cache entry outlives the pipeline
and pins the weights, so dropping the pipeline frees nothing."""
pipe = _pipe(transformer=nn.Linear(512, 512))
stray_cache = {"weight": pipe.components["transformer"].weight}
refs = T.track_host_tensors(pipe)
del pipe
gc.collect()
r = T.break_host_leaks(refs)
assert r["tensors"] >= 1, "the surviving host tensor was not detected"
assert r["broken"] >= 1, "the reference pinning it was not broken"
assert stray_cache["weight"] is None, "the cache still pins the weight"
def test_a_tensor_pinned_by_a_list_is_released():
pipe = _pipe(transformer=nn.Linear(256, 256))
stray = [pipe.components["transformer"].weight]
refs = T.track_host_tensors(pipe)
del pipe
gc.collect()
T.break_host_leaks(refs)
assert stray[0] is None
def test_the_surviving_size_is_reported():
pipe = _pipe(transformer=nn.Linear(1024, 1024, bias=False).to(torch.float32))
held = {"w": pipe.components["transformer"].weight}
refs = T.track_host_tensors(pipe)
del pipe
gc.collect()
r = T.break_host_leaks(refs)
# reported to 3 dp, i.e. 1 MB granularity
assert r["survived_gb"] == pytest.approx(1024 * 1024 * 4 / 1e9, abs=0.001)
assert held["w"] is None
def test_only_tracked_tensors_are_touched():
"""Another model's memory must never be at risk. This is why the host arm
tracks weakrefs to the objects instead of matching storage addresses: a
freed host address is readily reused by malloc, and a pointer match would
then null a live tensor belonging to something else."""
victim = nn.Linear(128, 128)
bystander = {"w": victim.weight}
other = _pipe(transformer=nn.Linear(128, 128))
refs = T.track_host_tensors(other)
del other
gc.collect()
T.break_host_leaks(refs)
assert bystander["w"] is victim.weight, "an untracked tensor was nulled"
def test_buffers_are_tracked_not_just_parameters():
"""Pre-quantized formats keep their payload in buffers, so a
parameters-only sweep would leave a 4-bit model's weights pinned."""
m = nn.Module()
m.register_buffer("qweight", torch.empty((256, 64), dtype=torch.int32))
pipe = _pipe(transformer=m)
held = {"q": m.qweight}
refs = T.track_host_tensors(pipe)
del pipe, m
gc.collect()
r = T.break_host_leaks(refs)
assert r["tensors"] >= 1
assert held["q"] is None
def test_the_cuda_arm_still_only_tracks_cuda():
"""Guard the division of labour in video.py: if the CUDA registration ever
starts taking host tensors too, the two arms would double-handle them."""
src = io.open(os.path.join(os.path.dirname(os.path.dirname(
os.path.abspath(__file__))), "codai/api/video.py"), encoding="utf-8").read()
assert "_orig_cuda_ptrs.add" in src
reg = src.split("_orig_cuda_ptrs.add", 1)[0][-400:]
assert "is_cuda" in reg, (
"the CUDA breaker no longer gates registration on is_cuda")
def test_trim_is_available_and_safe_to_call():
assert T.trim_host_allocator() in (True, False)
tests/test_idle_unload_history_and_log_volume.py 0 → 100644
View file @ 8274ff3e
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""Four things a production day asked for.
Measured on one host over 21h30m, 2026-10-29:
* a 3.5B music model held 7.9 GB of a 3090 and 2.1 GB of host RAM for hours
after its last request, because the on-demand path keeps a model resident
whenever there is room ("Ondemand mode - keeping 'audio_gen:ACE-Step' in VRAM
alongside new model");
* 78% of the log (426,682 of 548,687 lines) was per-request debug from one
engine, filling the 2,000-line view in about four minutes and leaving no
record of the loads and evictions behind a slow render;
* those loads and evictions could only be inferred from that log;
* ACE-Step was asked for a song with no lyrics instead of an instrumental.
"""
import io
import os
import sys
import time
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# ------------------------------------------------------------ idle unload ---
class _FakeManager:
"""Only what unload_idle_models touches."""
from codai.models.manager import MultiModelManager as _MM
unload_idle_models = _MM.unload_idle_models
record_history = _MM.record_history
get_history = _MM.get_history
def __init__(self):
from collections import deque
self.models = {}
self._last_used = {}
self._keep_resident = set()
self._history = deque(maxlen=500)
self.evicted = []
self.busy = set()
def _lru_order(self):
return sorted(self.models, key=lambda k: self._last_used.get(k, 0.0))
def _is_key_busy(self, key):
return key in self.busy
def _evict_one(self, key, reason=""):
self.evicted.append((key, reason))
self.models.pop(key, None)
self._last_used.pop(key, None)
def _mgr(**ages):
m = _FakeManager()
now = time.monotonic()
for key, age in ages.items():
m.models[key] = object()
m._last_used[key] = now - age
return m
def test_an_idle_model_is_unloaded():
m = _mgr(**{"audio_gen:ACE-Step": 3600.0})
freed = m.unload_idle_models(1800.0)
assert freed == ["audio_gen:ACE-Step"]
assert m.evicted and "idle" in m.evicted[0][1]
def test_a_recently_used_model_is_kept():
m = _mgr(**{"video:X": 60.0})
assert m.unload_idle_models(1800.0) == []
assert m.models
def test_a_busy_model_is_never_pulled():
m = _mgr(**{"video:X": 9999.0})
m.busy.add("video:X")
assert m.unload_idle_models(1800.0) == []
def test_keep_resident_is_respected():
m = _mgr(**{"tts:Kokoro": 9999.0})
m._keep_resident.add("tts:Kokoro")
assert m.unload_idle_models(1800.0) == []
def test_an_unstamped_model_is_left_alone():
"""A key with no _last_used is indistinguishable from one loaded a second
ago by a path that forgot to stamp it, so unloading it would be the
opposite of what is wanted."""
m = _FakeManager()
m.models["video:X"] = object() # never stamped
assert m.unload_idle_models(1800.0) == []
def test_zero_disables_it():
m = _mgr(**{"video:X": 99999.0})
assert m.unload_idle_models(0) == []
assert m.unload_idle_models(None) == []
def test_the_setting_exists_and_is_plumbed():
from codai.config import OffloadConfig
assert hasattr(OffloadConfig(), "idle_unload_seconds")
main = io.open(os.path.join(_ROOT, "codai/main.py"), encoding="utf-8").read()
assert "global_args.idle_unload_seconds" in main, (
"the setting never reaches the running server")
rm = io.open(os.path.join(_ROOT, "codai/models/ram_monitor.py"), encoding="utf-8").read()
assert "unload_idle_models" in rm, "nothing ever runs the idle sweep"
# ---------------------------------------------------------------- history ---
def test_history_records_a_load_with_what_it_cost():
m = _FakeManager()
m.record_history("load", "video:X", seconds=44.2, vram_gb=15.68,
strategy="resident")
row = m.get_history()[0]
assert row["event"] == "load" and row["model"] == "video:X"
assert row["seconds"] == 44.2 and row["vram_gb"] == 15.68
assert row["time"] > 0
def test_history_records_why_a_model_was_evicted():
m = _FakeManager()
m.record_history("evict", "video:X", reason="cosite: a co-located engine "
"needed the GPU", freed_gb=11.1)
row = m.get_history()[0]
assert "co-located" in row["reason"]
assert row["freed_gb"] == 11.1
def test_history_is_newest_first_and_filterable():
m = _FakeManager()
m.record_history("load", "a")
m.record_history("evict", "b")
assert [r["model"] for r in m.get_history()] == ["b", "a"]
assert [r["model"] for r in m.get_history(event="load")] == ["a"]
def test_history_is_bounded():
"""It must not become the leak it is there to help diagnose."""
m = _FakeManager()
for i in range(1200):
m.record_history("load", f"m{i}")
assert len(m._history) == 500
def test_eviction_reasons_are_passed_at_the_call_sites():
src = io.open(os.path.join(_ROOT, "codai/models/manager.py"), encoding="utf-8").read()
for marker in ("reason=f\"vram:", "reason=\"cosite:", "reason=f\"idle for"):
assert marker in src, f"no reason given at the {marker!r} call site"
def test_the_endpoints_exist():
eng = io.open(os.path.join(_ROOT, "codai/api/app.py"), encoding="utf-8").read()
assert '"/v1/models/history"' in eng
front = io.open(os.path.join(_ROOT, "codai/frontproxy/app.py"), encoding="utf-8").read()
assert '"/v1/models/history"' in front, (
"the front does not aggregate it, so a caller sees one engine's half of "
"a cross-engine swap")
# It must be registered BEFORE the catch-all, or it is proxied to one engine.
assert front.index('"/v1/models/history"') < front.index("/admin/api/engines")
# ------------------------------------------------------------- log volume ---
def test_embeddings_are_summarised_not_dumped():
from codai.api import log as L
assert "/v1/embeddings" in L._SUMMARY_ONLY_PATHS
assert "/v1/embeddings" in L._TRACKED_PATHS, (
"summarising must not stop the admin page counting them")
def test_the_summary_says_what_was_asked_without_the_content():
from codai.api import log as L
secret = "a client's private property listing " * 20
line = L._summarise({"model": "Qwen3-Embedding-4B", "input": [secret, secret],
"dimensions": 1024})
assert "model=Qwen3-Embedding-4B" in line
assert "inputs=2" in line
assert "dimensions=1024" in line
assert "property listing" not in line, "the caller's content is still in the log"
def test_a_prompt_is_left_intact():
"""Deliberately NOT truncated: a prompt is the most useful thing in the log
-- a 4,000-character script prompt was what identified a music defect --
and the volume problem was one per-item endpoint, not prompts."""
from codai.api import log as L
# Real prose, not a run of base64-legal characters: a long unbroken token
# IS treated as a blob, and correctly so.
prompt = ("Write a 81-second video script about the rally. " * 100)
out = L._redact_blobs({"prompt": prompt})
assert out["prompt"] == prompt, "a script prompt is being truncated"
# A base64 payload is still redacted.
blob = L._redact_blobs({"init_image": "A" * 5000})["init_image"]
assert "truncated" in blob
def test_a_short_value_is_untouched():
from codai.api import log as L
assert L._redact_blobs({"model": "Krea-2-Turbo"}) == {"model": "Krea-2-Turbo"}
# ------------------------------------------------------- acestep lyrics ----
def test_an_instrumental_request_says_so():
"""ACE-Step is a SONG model. Empty lyrics leave it conditioning on a single
padding token (lyric_token_idx = [0], lyric_mask = [0]) and it still tries
to sing. Its own UI documents the way to ask: "Use [instrumental] or [inst]
to generate instrumental music"."""
src = io.open(os.path.join(_ROOT, "codai/api/audio_gen.py"), encoding="utf-8").read()
block = src.split("elif model_type == 'ace_step':", 1)[1][:2000]
assert "[instrumental]" in block, (
"an empty-lyrics request is still sent as a song with no words")
assert "lyrics=_lyrics" in block
assert "instrumental" in block.split("_prompt")[1][:400], (
"the tag prompt is not told it is instrumental either")
tests/test_ltx2_audio_and_engines.py
View file @ 8274ff3e
......@@ -80,7 +80,11 @@ def test_stable_audio_and_ace_step_each_have_their_own_loader():
assert "_load_musicgen" not in load
# ...and the generate branch returns the worker's wav rather than falling
# through to the array-to-int16 writer, which would see no audio_np at all.
gen = src[src.index("elif model_type == 'ace_step':"):][:1400]
# Slice to the NEXT branch rather than a character count: a fixed window
# silently stops covering the branch as soon as anything is added to it.
_start = src.index("elif model_type == 'ace_step':")
_end = src.index("elif model_type ==", _start + 10)
gen = src[_start:_end]
assert "return wav, 'wav'" in gen
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment