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:
@@ -1,7 +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, Union
|
||||
|
||||
@@ -32,10 +35,55 @@ app = FastAPI(title="ASR Faster-Whisper Worker")
|
||||
_MODEL_CACHE: Dict[str, WhisperModel] = {}
|
||||
_DIARIZATION_PIPELINE: Any = None
|
||||
|
||||
# Idle-unload / model-switch eviction: only one whisper model stays resident
|
||||
# at a time (switching sizes frees the old one), and everything (including
|
||||
# the diarization pipeline) 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 _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_whisper_models() -> None:
|
||||
if not _MODEL_CACHE:
|
||||
return
|
||||
print(f"[faster_whisper] unloading {list(_MODEL_CACHE.keys())}", flush=True)
|
||||
_MODEL_CACHE.clear()
|
||||
_free_gpu()
|
||||
|
||||
|
||||
def _unload_diarization() -> None:
|
||||
global _DIARIZATION_PIPELINE
|
||||
if _DIARIZATION_PIPELINE is None:
|
||||
return
|
||||
print("[faster_whisper] unloading diarization pipeline", flush=True)
|
||||
_DIARIZATION_PIPELINE = None
|
||||
_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_whisper_models()
|
||||
_unload_diarization()
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup() -> None:
|
||||
async def startup() -> None:
|
||||
ensure_runtime_dirs()
|
||||
asyncio.create_task(_idle_unload_loop())
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
@@ -75,6 +123,8 @@ async def transcribe(
|
||||
tmp.write(await file.read())
|
||||
tmp_path = Path(tmp.name)
|
||||
|
||||
global _active_requests, _last_used
|
||||
_active_requests += 1
|
||||
try:
|
||||
resolved_model = resolve_custom_model_path(custom_model_path) or model
|
||||
result = _transcribe(
|
||||
@@ -107,10 +157,16 @@ async def transcribe(
|
||||
raise HTTPException(status_code=500, detail=f"Worker failure: {exc}") from exc
|
||||
finally:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
_free_gpu()
|
||||
_active_requests -= 1
|
||||
_last_used = time.monotonic()
|
||||
|
||||
|
||||
def _load_model(model_name: str) -> WhisperModel:
|
||||
if model_name not in _MODEL_CACHE:
|
||||
if _MODEL_CACHE:
|
||||
# Only one whisper model resident at a time — switching sizes frees the old one.
|
||||
_unload_whisper_models()
|
||||
_MODEL_CACHE[model_name] = WhisperModel(
|
||||
model_name,
|
||||
device=DEVICE,
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import gc
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
@@ -29,6 +31,14 @@ app = FastAPI(title="ASR Qwen3 Worker")
|
||||
# and drive it through Qwen3ASRProcessor.apply_transcription_request().
|
||||
_MODEL_CACHE: Dict[str, Tuple[Any, 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
|
||||
|
||||
LANG_MAP = {
|
||||
"ko": "korean", "en": "english", "ja": "japanese", "zh": "chinese",
|
||||
"fr": "french", "de": "german", "es": "spanish", "ru": "russian",
|
||||
@@ -70,10 +80,29 @@ def _to_wav(src_path: Path) -> Path:
|
||||
return wav_path
|
||||
|
||||
|
||||
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) -> Tuple[Any, 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 transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
|
||||
|
||||
@@ -96,8 +125,9 @@ def _load_model(model_id: str) -> Tuple[Any, Any]:
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup() -> None:
|
||||
async def startup() -> None:
|
||||
ensure_runtime_dirs()
|
||||
asyncio.create_task(_idle_unload_loop())
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
@@ -118,6 +148,8 @@ async def transcribe(
|
||||
tmp.write(await file.read())
|
||||
tmp_path = Path(tmp.name)
|
||||
|
||||
global _active_requests, _last_used
|
||||
_active_requests += 1
|
||||
wav_path: Optional[Path] = None
|
||||
try:
|
||||
import torch
|
||||
@@ -162,6 +194,21 @@ async def transcribe(
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
if wav_path is not None:
|
||||
wav_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 generated
|
||||
except NameError:
|
||||
pass
|
||||
_free_gpu()
|
||||
_active_requests -= 1
|
||||
_last_used = time.monotonic()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user