Self-restart worker process when model loading fails

Discovered vibevoice sitting at 3.2GB resident GPU memory with
loaded_models: [] and no idle-unload ever firing for it again. Root cause:
a load attempt had OOM'd partway through (competing with an unrelated
ollama process on the same GPU), so _MODEL_CACHE was never populated —
the idle-unload loop only clears that cache, so it had nothing to act on,
even though the partially-constructed model had already left memory
allocated. gc.collect()+empty_cache() don't reliably reclaim memory from
an interrupted from_pretrained() call.

Confirmed a plain process restart does fully reclaim it, so each worker
now treats any load failure as fatal: log it and os._exit(1), letting
supervisord's autorestart=true respawn a clean process immediately.
Verified with a bogus model id — worker exits, respawns, and passes the
smoke test right after.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
du5t
2026-07-24 00:22:59 +09:00
parent 89f8643426
commit e3b2a53c54
4 changed files with 71 additions and 30 deletions

View File

@@ -167,12 +167,23 @@ def _load_model(model_name: str) -> WhisperModel:
if _MODEL_CACHE: if _MODEL_CACHE:
# Only one whisper model resident at a time — switching sizes frees the old one. # Only one whisper model resident at a time — switching sizes frees the old one.
_unload_whisper_models() _unload_whisper_models()
try:
_MODEL_CACHE[model_name] = WhisperModel( _MODEL_CACHE[model_name] = WhisperModel(
model_name, model_name,
device=DEVICE, device=DEVICE,
compute_type=COMPUTE_TYPE, compute_type=COMPUTE_TYPE,
download_root=str(MODEL_CACHE), download_root=str(MODEL_CACHE),
) )
except Exception as e:
# A load interrupted partway (e.g. OOM) can leave CUDA memory
# fragmented/leaked in ways gc.collect()+empty_cache() don't
# reliably reclaim, and since _MODEL_CACHE never got populated
# the idle-unload loop has nothing to clean up either.
# Restarting the whole process is the only guaranteed way to get
# that memory back — supervisord's autorestart=true respawns it
# immediately.
print(f"[faster_whisper] model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}", flush=True)
os._exit(1)
return _MODEL_CACHE[model_name] return _MODEL_CACHE[model_name]

View File

@@ -109,6 +109,7 @@ def _load_model(model_id: str) -> Tuple[Any, Any]:
_log(f"loading model {model_id}") _log(f"loading model {model_id}")
dtype = torch.float16 if DEVICE == "cuda" else torch.float32 dtype = torch.float16 if DEVICE == "cuda" else torch.float32
try:
processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE)) processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE))
model = AutoModelForSpeechSeq2Seq.from_pretrained( model = AutoModelForSpeechSeq2Seq.from_pretrained(
model_id, model_id,
@@ -118,6 +119,15 @@ def _load_model(model_id: str) -> Tuple[Any, Any]:
) )
model.to(DEVICE) model.to(DEVICE)
model.eval() model.eval()
except Exception as e:
# A load interrupted partway (e.g. OOM) can leave CUDA memory
# fragmented/leaked in ways gc.collect()+empty_cache() don't reliably
# reclaim, and since _MODEL_CACHE never got populated the idle-unload
# loop has nothing to clean up either. Restarting the whole process
# is the only guaranteed way to get that memory back — supervisord's
# autorestart=true respawns it immediately.
_log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}")
os._exit(1)
_MODEL_CACHE[model_id] = (model, processor) _MODEL_CACHE[model_id] = (model, processor)
_log(f"model {model_id} loaded") _log(f"model {model_id} loaded")

View File

@@ -77,6 +77,7 @@ def _load_model(model_id: str) -> Any:
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
_log(f"loading model {model_id}") _log(f"loading model {model_id}")
try:
processor = VibeVoiceASRProcessor.from_pretrained( processor = VibeVoiceASRProcessor.from_pretrained(
model_id, model_id,
language_model_pretrained_name="Qwen/Qwen2.5-1.5B", language_model_pretrained_name="Qwen/Qwen2.5-1.5B",
@@ -91,6 +92,15 @@ def _load_model(model_id: str) -> Any:
) )
model.to(DEVICE) model.to(DEVICE)
model.eval() model.eval()
except Exception as e:
# A load interrupted partway (e.g. OOM) can leave CUDA memory
# fragmented/leaked in ways gc.collect()+empty_cache() don't reliably
# reclaim, and since _MODEL_CACHE never got populated the idle-unload
# loop has nothing to clean up either. Restarting the whole process
# is the only guaranteed way to get that memory back — supervisord's
# autorestart=true respawns it immediately.
_log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}")
os._exit(1)
_MODEL_CACHE[model_id] = (model, processor) _MODEL_CACHE[model_id] = (model, processor)
_log(f"model {model_id} loaded") _log(f"model {model_id} loaded")

View File

@@ -83,7 +83,17 @@ def _load_model() -> Any:
from TTS.api import TTS from TTS.api import TTS
_log(f"loading model {MODEL_NAME}") _log(f"loading model {MODEL_NAME}")
try:
tts = TTS(MODEL_NAME).to(_device()) tts = TTS(MODEL_NAME).to(_device())
except Exception as e:
# A load interrupted partway (e.g. OOM) can leave CUDA memory
# fragmented/leaked in ways gc.collect()+empty_cache() don't reliably
# reclaim, and since _MODEL_CACHE never got populated the idle-unload
# loop has nothing to clean up either. Restarting the whole process
# is the only guaranteed way to get that memory back — supervisord's
# autorestart=true respawns it immediately.
_log(f"model load failed, restarting process to reclaim GPU memory: {type(e).__name__}: {e}")
os._exit(1)
_MODEL_CACHE[MODEL_NAME] = tts _MODEL_CACHE[MODEL_NAME] = tts
_log("model loaded") _log("model loaded")
return tts return tts