fix(system): report reserved VRAM alongside allocated in flush-memory

memory_allocated counts live tensors only, so after an unload it reads
near zero while nvidia-smi still shows gigabytes. That gap is the whole
substance of every "flush says it worked, the GPU says it didn't" report,
and the endpoint was reporting only the half that looks good.

memory_reserved is what the caching allocator holds from the driver; the
remainder between that and the driver's own figure is the CUDA context and
kernel workspaces, which nothing in-process can hand back.
This commit is contained in:
velixio
2026-08-12 02:58:42 +05:30
parent 090cc37144
commit 642513d205
+12 -1
View File
@@ -556,15 +556,25 @@ async def flush_memory(unload_model: bool = False):
free_vram()
# Snapshot after flush
# Snapshot after flush. Two numbers, because one of them is a lie by
# omission: `memory_allocated` counts live tensors only, so it reads ~0
# after an unload while nvidia-smi still shows gigabytes — which is exactly
# the report we keep getting ("flush says it worked, the GPU says it
# didn't"). `memory_reserved` is what the caching allocator holds from the
# driver, and the gap between reserved and the driver's own figure is the
# CUDA context plus kernel workspaces, which no in-process call can return.
vram_after = 0.0
vram_reserved = 0.0
try:
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
driver = getattr(torch.mps, "driver_allocated_memory", None)
if driver:
vram_after = driver() / (1024**3)
current = getattr(torch.mps, "current_allocated_memory", None)
vram_reserved = (current() / (1024**3)) if current else vram_after
elif torch.cuda.is_available():
vram_after = torch.cuda.memory_allocated() / (1024**3)
vram_reserved = torch.cuda.memory_reserved() / (1024**3)
except Exception:
pass
@@ -575,6 +585,7 @@ async def flush_memory(unload_model: bool = False):
"unloaded_model": freed_model,
"ram_after": round(ram_after, 2),
"vram_after": round(vram_after, 2),
"vram_reserved": round(vram_reserved, 2),
}