From 9472387d1b9fa8fa528fe763a230ab51746bbd99 Mon Sep 17 00:00:00 2001 From: du5t Date: Thu, 23 Jul 2026 23:29:13 +0900 Subject: [PATCH] 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 --- app/asr/workers/faster_whisper_worker.py | 58 +++++++++++++++++++++- app/asr/workers/qwen3_worker.py | 49 +++++++++++++++++- app/asr/workers/vibevoice_worker.py | 63 +++++++++++++++++++++++- app/tts/workers/xtts_worker.py | 45 ++++++++++++++++- 4 files changed, 211 insertions(+), 4 deletions(-) diff --git a/app/asr/workers/faster_whisper_worker.py b/app/asr/workers/faster_whisper_worker.py index a5c7284..4f3f8f6 100644 --- a/app/asr/workers/faster_whisper_worker.py +++ b/app/asr/workers/faster_whisper_worker.py @@ -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, diff --git a/app/asr/workers/qwen3_worker.py b/app/asr/workers/qwen3_worker.py index 3af8dc1..f36e04b 100644 --- a/app/asr/workers/qwen3_worker.py +++ b/app/asr/workers/qwen3_worker.py @@ -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__": diff --git a/app/asr/workers/vibevoice_worker.py b/app/asr/workers/vibevoice_worker.py index 5a1707e..d46e312 100644 --- a/app/asr/workers/vibevoice_worker.py +++ b/app/asr/workers/vibevoice_worker.py @@ -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__": diff --git a/app/tts/workers/xtts_worker.py b/app/tts/workers/xtts_worker.py index 21f7497..43cce17 100644 --- a/app/tts/workers/xtts_worker.py +++ b/app/tts/workers/xtts_worker.py @@ -1,8 +1,11 @@ from __future__ import annotations import argparse +import asyncio +import gc import os import tempfile +import time from pathlib import Path from typing import Any, Dict @@ -28,6 +31,13 @@ app = FastAPI(title="TTS XTTS Worker") MODEL_NAME = "tts_models/multilingual/multi-dataset/xtts_v2" _MODEL_CACHE: Dict[str, Any] = {} +# Idle-unload: the model is dropped after IDLE_UNLOAD_SECONDS of no synthesis +# requests, so this backend doesn't permanently hog GPU memory shared with +# the ASR workers. +IDLE_UNLOAD_SECONDS = 120 +_last_used: float = 0.0 +_active_requests: int = 0 + def _log(msg: str) -> None: print(f"[xtts] {msg}", flush=True) @@ -41,6 +51,31 @@ def _device() -> str: return "cpu" +def _free_gpu() -> None: + gc.collect() + try: + import torch + if torch.cuda.is_available(): + torch.cuda.empty_cache() + except Exception: + pass + + +def _unload_model() -> None: + if not _MODEL_CACHE: + return + _log("unloading model (idle)") + _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_model() + + def _load_model() -> Any: if MODEL_NAME in _MODEL_CACHE: return _MODEL_CACHE[MODEL_NAME] @@ -55,8 +90,9 @@ def _load_model() -> Any: @app.on_event("startup") -def startup() -> None: +async def startup() -> None: ensure_runtime_dirs() + asyncio.create_task(_idle_unload_loop()) @app.get("/health") @@ -74,10 +110,14 @@ def synthesize( if not speaker_path.exists(): raise HTTPException(status_code=400, detail=f"speaker_wav not found: {speaker_wav}") + global _active_requests, _last_used + _active_requests += 1 try: tts = _load_model() except Exception as e: _log(f"model load error: {type(e).__name__}: {e}") + _active_requests -= 1 + _last_used = time.monotonic() raise HTTPException(status_code=500, detail=f"모델 로드 실패: {type(e).__name__}: {e}") with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: @@ -105,6 +145,9 @@ def synthesize( raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") finally: out_path.unlink(missing_ok=True) + _free_gpu() + _active_requests -= 1 + _last_used = time.monotonic() if __name__ == "__main__":