diff --git a/engine.py b/engine.py index a2be1e9..fc48086 100644 --- a/engine.py +++ b/engine.py @@ -5,6 +5,8 @@ import logging import os import random +import threading +from collections import OrderedDict import numpy as np import torch from typing import Optional, Tuple @@ -114,7 +116,78 @@ def _resolve_bf16_setting() -> bool: # Voice conditioning cache: avoids re-encoding the same voice file on every request. # Key: (resolved_path, file_mtime, exaggeration) — mtime invalidates if file changes. -_conds_cache: dict = {} +# +# Bounded, because each entry pins GPU tensors for the process lifetime and this +# was previously an unbounded dict cleared only by unload_model(). Measured on a +# P106-100 (5.93 GiB) with chatterbox-turbo: cycling 28 reference voices grew +# live tensors from 4.09 to 5.46 GiB and left 0.11 GiB of headroom, after which +# 406 of 453 requests failed. The same load with one voice ran clean. +# +# This is not only a small-card concern. /upload_reference stores each uploaded +# file under its own path, so every distinct voice a deployment ever serves mints +# a permanent key — a larger card buys a higher ceiling, not a different outcome. +# +# CONDS_CACHE_MAX bounds the entry count. Measured cost is ~175 MB of VRAM per +# entry (live tensors grew 1.37 GiB across exactly 8 cached voices), so this is a +# far more expensive cache than it looks — budget it against free VRAM, not +# against an entry count that sounds small. +# +# The default of 4 costs ~700 MB, which leaves a 6 GB card ~1.1 GiB for +# generation after the 4.09 GiB model. On a 12 GB card 16+ is comfortable. Set 0 +# to disable caching and re-encode the voice on every request. +def _resolve_conds_cache_max() -> int: + """Read CONDS_CACHE_MAX, tolerating garbage rather than failing to import.""" + raw = os.environ.get("CONDS_CACHE_MAX", "4").strip() + try: + return max(0, int(raw)) + except ValueError: + logger.warning( + f"CONDS_CACHE_MAX={raw!r} is not an integer; using the default of 4." + ) + return 4 + + +_CONDS_CACHE_MAX: int = _resolve_conds_cache_max() +_conds_cache: "OrderedDict[tuple, object]" = OrderedDict() + +# synthesize() runs concurrently: server.py dispatches it through +# loop.run_in_executor(), so several threads can touch this cache at once. +# Reads are two steps (look up, then refresh recency) and writes are three +# (insert, refresh, evict), none of which are atomic — without this lock a +# thread can evict a key between another thread's lookup and its use, raising +# KeyError. The unbounded version could not hit this because nothing was ever +# removed. +_conds_cache_lock = threading.Lock() + + +def _conds_cache_get(key: tuple): + """Return cached conds for `key` and mark it most-recently-used, else None.""" + if _CONDS_CACHE_MAX == 0: + return None + with _conds_cache_lock: + conds = _conds_cache.get(key) + if conds is not None: + _conds_cache.move_to_end(key) + return conds + + +def _conds_cache_store(key: tuple, conds) -> None: + """Insert under an LRU bound, evicting the least recently used entries. + + Eviction drops the last reference to that entry's tensors; the caching + allocator reuses the freed blocks, so VRAM is recovered without an explicit + empty_cache() on the request path. + """ + if _CONDS_CACHE_MAX == 0: + return + with _conds_cache_lock: + _conds_cache[key] = conds + _conds_cache.move_to_end(key) + while len(_conds_cache) > _CONDS_CACHE_MAX: + evicted_key, _ = _conds_cache.popitem(last=False) + logger.debug( + f"Voice cache evicted (LRU, max={_CONDS_CACHE_MAX}): {evicted_key[0]}" + ) def _conds_cache_key(path: str, exaggeration: float) -> tuple: @@ -479,8 +552,11 @@ def synthesize( if audio_prompt_path and hasattr(chatterbox_model, "conds"): ex_for_key = 0.0 if loaded_model_type == "turbo" else exaggeration conds_key = _conds_cache_key(audio_prompt_path, ex_for_key) - if conds_key in _conds_cache: - chatterbox_model.conds = _conds_cache[conds_key] + # Single locked lookup that also refreshes recency — checking + # membership and then indexing would race with eviction. + cached_conds = _conds_cache_get(conds_key) + if cached_conds is not None: + chatterbox_model.conds = cached_conds effective_prompt = None # conds already set, skip prepare_conditionals logger.debug(f"Voice cache hit: {audio_prompt_path}") @@ -509,7 +585,7 @@ def synthesize( # Store conds in cache after first compute for this voice. if conds_key is not None and effective_prompt is not None: if chatterbox_model.conds is not None: - _conds_cache[conds_key] = chatterbox_model.conds + _conds_cache_store(conds_key, chatterbox_model.conds) logger.debug(f"Cached voice conditionals for: {audio_prompt_path}") # The ChatterboxTTS.generate method already returns a CPU tensor. @@ -590,7 +666,8 @@ def reload_model() -> bool: MODEL_LOADED = False loaded_model_type = None loaded_model_class_name = None - _conds_cache.clear() + with _conds_cache_lock: + _conds_cache.clear() logger.info("Voice conditioning cache cleared.") # 3. Force Python Garbage Collection