Auto-unload idle models and evict on model switch across all workers

Every ASR/TTS worker (faster-whisper, qwen3, vibevoice, xtts) kept every
model it ever loaded resident in GPU memory forever, and switching to a
different model (e.g. a different whisper size) just added another one
alongside it rather than freeing the old one. Combined with the shared
24GB GPU, this made memory pressure only ever go up.

Now: only one model stays resident per worker at a time (loading a
different model_id evicts the previous one first), and the whole cache
(plus, for faster-whisper, the diarization pipeline) is dropped after 2
minutes of no requests. Verified end-to-end: the idle timer actually fires
and reclaims memory, and switching qwen3 models evicts the old one.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
du5t
2026-07-23 23:29:13 +09:00
parent ce064d6894
commit 9472387d1b
4 changed files with 211 additions and 4 deletions

View File

@@ -1,8 +1,10 @@
from __future__ import annotations
import argparse
import asyncio
import gc
import tempfile
import time
from pathlib import Path
from typing import Any, Dict, List, Optional
@@ -24,15 +26,52 @@ app = FastAPI(title="ASR VibeVoice Worker")
_MODEL_CACHE: Dict[str, Any] = {}
# Idle-unload / model-switch eviction: only one model stays resident at a
# time, and the whole cache is dropped after IDLE_UNLOAD_SECONDS of no
# requests, so this backend doesn't permanently hog GPU memory shared with
# the other ASR workers.
IDLE_UNLOAD_SECONDS = 120
_last_used: float = 0.0
_active_requests: int = 0
def _log(msg: str) -> None:
print(f"[vibevoice] {msg}", flush=True)
def _free_gpu() -> None:
gc.collect()
try:
import torch
if DEVICE == "cuda" and torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
def _unload_models() -> None:
if not _MODEL_CACHE:
return
_log(f"unloading {list(_MODEL_CACHE.keys())}")
_MODEL_CACHE.clear()
_free_gpu()
async def _idle_unload_loop() -> None:
while True:
await asyncio.sleep(30)
if _active_requests == 0 and _last_used and (time.monotonic() - _last_used) >= IDLE_UNLOAD_SECONDS:
_unload_models()
def _load_model(model_id: str) -> Any:
if model_id in _MODEL_CACHE:
return _MODEL_CACHE[model_id]
if _MODEL_CACHE:
# Only one model resident at a time — switching models frees the old one.
_unload_models()
import torch
from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
@@ -59,8 +98,9 @@ def _load_model(model_id: str) -> Any:
@app.on_event("startup")
def startup() -> None:
async def startup() -> None:
ensure_runtime_dirs()
asyncio.create_task(_idle_unload_loop())
@app.get("/health")
@@ -81,6 +121,8 @@ async def transcribe(
tmp.write(await file.read())
tmp_path = Path(tmp.name)
global _active_requests, _last_used
_active_requests += 1
try:
import torch
@@ -153,6 +195,25 @@ async def transcribe(
raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
finally:
tmp_path.unlink(missing_ok=True)
# Long files blow up KV-cache/activation memory; without this, that
# memory stays reserved by this process and starves the other backends
# sharing the same GPU until the container restarts. (del locals()[...]
# does NOT work in CPython, hence the explicit names.)
try:
del inputs
except NameError:
pass
try:
del output_ids
except NameError:
pass
try:
del generated_ids
except NameError:
pass
_free_gpu()
_active_requests -= 1
_last_used = time.monotonic()
if __name__ == "__main__":