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:
@@ -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()
|
||||||
_MODEL_CACHE[model_name] = WhisperModel(
|
try:
|
||||||
model_name,
|
_MODEL_CACHE[model_name] = WhisperModel(
|
||||||
device=DEVICE,
|
model_name,
|
||||||
compute_type=COMPUTE_TYPE,
|
device=DEVICE,
|
||||||
download_root=str(MODEL_CACHE),
|
compute_type=COMPUTE_TYPE,
|
||||||
)
|
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]
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -109,15 +109,25 @@ 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
|
||||||
|
|
||||||
processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE))
|
try:
|
||||||
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
processor = AutoProcessor.from_pretrained(model_id, cache_dir=str(MODEL_CACHE))
|
||||||
model_id,
|
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
||||||
dtype=dtype,
|
model_id,
|
||||||
low_cpu_mem_usage=True,
|
dtype=dtype,
|
||||||
cache_dir=str(MODEL_CACHE),
|
low_cpu_mem_usage=True,
|
||||||
)
|
cache_dir=str(MODEL_CACHE),
|
||||||
model.to(DEVICE)
|
)
|
||||||
model.eval()
|
model.to(DEVICE)
|
||||||
|
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")
|
||||||
|
|||||||
@@ -77,20 +77,30 @@ 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}")
|
||||||
processor = VibeVoiceASRProcessor.from_pretrained(
|
try:
|
||||||
model_id,
|
processor = VibeVoiceASRProcessor.from_pretrained(
|
||||||
language_model_pretrained_name="Qwen/Qwen2.5-1.5B",
|
model_id,
|
||||||
cache_dir=str(MODEL_CACHE),
|
language_model_pretrained_name="Qwen/Qwen2.5-1.5B",
|
||||||
)
|
cache_dir=str(MODEL_CACHE),
|
||||||
model = VibeVoiceASRForConditionalGeneration.from_pretrained(
|
)
|
||||||
model_id,
|
model = VibeVoiceASRForConditionalGeneration.from_pretrained(
|
||||||
dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
|
model_id,
|
||||||
attn_implementation="sdpa",
|
dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32,
|
||||||
trust_remote_code=True,
|
attn_implementation="sdpa",
|
||||||
cache_dir=str(MODEL_CACHE),
|
trust_remote_code=True,
|
||||||
)
|
cache_dir=str(MODEL_CACHE),
|
||||||
model.to(DEVICE)
|
)
|
||||||
model.eval()
|
model.to(DEVICE)
|
||||||
|
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")
|
||||||
|
|||||||
@@ -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}")
|
||||||
tts = TTS(MODEL_NAME).to(_device())
|
try:
|
||||||
|
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
|
||||||
|
|||||||
Reference in New Issue
Block a user