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,7 +1,10 @@
from __future__ import annotations from __future__ import annotations
import argparse import argparse
import asyncio
import gc
import tempfile import tempfile
import time
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional, Union
@@ -32,10 +35,55 @@ app = FastAPI(title="ASR Faster-Whisper Worker")
_MODEL_CACHE: Dict[str, WhisperModel] = {} _MODEL_CACHE: Dict[str, WhisperModel] = {}
_DIARIZATION_PIPELINE: Any = None _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") @app.on_event("startup")
def startup() -> None: async def startup() -> None:
ensure_runtime_dirs() ensure_runtime_dirs()
asyncio.create_task(_idle_unload_loop())
@app.get("/health") @app.get("/health")
@@ -75,6 +123,8 @@ async def transcribe(
tmp.write(await file.read()) tmp.write(await file.read())
tmp_path = Path(tmp.name) tmp_path = Path(tmp.name)
global _active_requests, _last_used
_active_requests += 1
try: try:
resolved_model = resolve_custom_model_path(custom_model_path) or model resolved_model = resolve_custom_model_path(custom_model_path) or model
result = _transcribe( result = _transcribe(
@@ -107,10 +157,16 @@ async def transcribe(
raise HTTPException(status_code=500, detail=f"Worker failure: {exc}") from exc raise HTTPException(status_code=500, detail=f"Worker failure: {exc}") from exc
finally: finally:
tmp_path.unlink(missing_ok=True) tmp_path.unlink(missing_ok=True)
_free_gpu()
_active_requests -= 1
_last_used = time.monotonic()
def _load_model(model_name: str) -> WhisperModel: def _load_model(model_name: str) -> WhisperModel:
if model_name not in _MODEL_CACHE: 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_CACHE[model_name] = WhisperModel(
model_name, model_name,
device=DEVICE, device=DEVICE,

View File

@@ -1,9 +1,11 @@
from __future__ import annotations from __future__ import annotations
import argparse import argparse
import asyncio
import gc import gc
import subprocess import subprocess
import tempfile import tempfile
import time
from pathlib import Path from pathlib import Path
from typing import Any, Dict, Optional, Tuple 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(). # and drive it through Qwen3ASRProcessor.apply_transcription_request().
_MODEL_CACHE: Dict[str, Tuple[Any, Any]] = {} _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 = { LANG_MAP = {
"ko": "korean", "en": "english", "ja": "japanese", "zh": "chinese", "ko": "korean", "en": "english", "ja": "japanese", "zh": "chinese",
"fr": "french", "de": "german", "es": "spanish", "ru": "russian", "fr": "french", "de": "german", "es": "spanish", "ru": "russian",
@@ -70,10 +80,29 @@ def _to_wav(src_path: Path) -> Path:
return wav_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]: def _load_model(model_id: str) -> Tuple[Any, Any]:
if model_id in _MODEL_CACHE: if model_id in _MODEL_CACHE:
return _MODEL_CACHE[model_id] 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 import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
@@ -96,8 +125,9 @@ def _load_model(model_id: str) -> Tuple[Any, Any]:
@app.on_event("startup") @app.on_event("startup")
def startup() -> None: async def startup() -> None:
ensure_runtime_dirs() ensure_runtime_dirs()
asyncio.create_task(_idle_unload_loop())
@app.get("/health") @app.get("/health")
@@ -118,6 +148,8 @@ async def transcribe(
tmp.write(await file.read()) tmp.write(await file.read())
tmp_path = Path(tmp.name) tmp_path = Path(tmp.name)
global _active_requests, _last_used
_active_requests += 1
wav_path: Optional[Path] = None wav_path: Optional[Path] = None
try: try:
import torch import torch
@@ -162,6 +194,21 @@ async def transcribe(
tmp_path.unlink(missing_ok=True) tmp_path.unlink(missing_ok=True)
if wav_path is not None: if wav_path is not None:
wav_path.unlink(missing_ok=True) 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__": if __name__ == "__main__":

View File

@@ -1,8 +1,10 @@
from __future__ import annotations from __future__ import annotations
import argparse import argparse
import asyncio
import gc import gc
import tempfile import tempfile
import time
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
@@ -24,15 +26,52 @@ app = FastAPI(title="ASR VibeVoice Worker")
_MODEL_CACHE: Dict[str, Any] = {} _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: def _log(msg: str) -> None:
print(f"[vibevoice] {msg}", flush=True) 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: def _load_model(model_id: str) -> Any:
if model_id in _MODEL_CACHE: if model_id in _MODEL_CACHE:
return _MODEL_CACHE[model_id] 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 import torch
from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
@@ -59,8 +98,9 @@ def _load_model(model_id: str) -> Any:
@app.on_event("startup") @app.on_event("startup")
def startup() -> None: async def startup() -> None:
ensure_runtime_dirs() ensure_runtime_dirs()
asyncio.create_task(_idle_unload_loop())
@app.get("/health") @app.get("/health")
@@ -81,6 +121,8 @@ async def transcribe(
tmp.write(await file.read()) tmp.write(await file.read())
tmp_path = Path(tmp.name) tmp_path = Path(tmp.name)
global _active_requests, _last_used
_active_requests += 1
try: try:
import torch import torch
@@ -153,6 +195,25 @@ async def transcribe(
raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
finally: finally:
tmp_path.unlink(missing_ok=True) 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__": if __name__ == "__main__":

View File

@@ -1,8 +1,11 @@
from __future__ import annotations from __future__ import annotations
import argparse import argparse
import asyncio
import gc
import os import os
import tempfile import tempfile
import time
from pathlib import Path from pathlib import Path
from typing import Any, Dict 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_NAME = "tts_models/multilingual/multi-dataset/xtts_v2"
_MODEL_CACHE: Dict[str, Any] = {} _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: def _log(msg: str) -> None:
print(f"[xtts] {msg}", flush=True) print(f"[xtts] {msg}", flush=True)
@@ -41,6 +51,31 @@ def _device() -> str:
return "cpu" 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: def _load_model() -> Any:
if MODEL_NAME in _MODEL_CACHE: if MODEL_NAME in _MODEL_CACHE:
return _MODEL_CACHE[MODEL_NAME] return _MODEL_CACHE[MODEL_NAME]
@@ -55,8 +90,9 @@ def _load_model() -> Any:
@app.on_event("startup") @app.on_event("startup")
def startup() -> None: async def startup() -> None:
ensure_runtime_dirs() ensure_runtime_dirs()
asyncio.create_task(_idle_unload_loop())
@app.get("/health") @app.get("/health")
@@ -74,10 +110,14 @@ def synthesize(
if not speaker_path.exists(): if not speaker_path.exists():
raise HTTPException(status_code=400, detail=f"speaker_wav not found: {speaker_wav}") raise HTTPException(status_code=400, detail=f"speaker_wav not found: {speaker_wav}")
global _active_requests, _last_used
_active_requests += 1
try: try:
tts = _load_model() tts = _load_model()
except Exception as e: except Exception as e:
_log(f"model load error: {type(e).__name__}: {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}") raise HTTPException(status_code=500, detail=f"모델 로드 실패: {type(e).__name__}: {e}")
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: 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}") raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
finally: finally:
out_path.unlink(missing_ok=True) out_path.unlink(missing_ok=True)
_free_gpu()
_active_requests -= 1
_last_used = time.monotonic()
if __name__ == "__main__": if __name__ == "__main__":