GPU-resident model cache with LRU eviction. Key features: - Keeps models on GPU with LRU eviction by default - Evicts when free VRAM drops below MIN_FREE_VRAM_GB (if KEEP_ALL_ON_GPU=1) - Enforces MAX_CACHED_MODELS when set - LRU eviction order when eviction is necessary
| 43 | |
| 44 | |
| 45 | class ModelCache: |
| 46 | """ |
| 47 | GPU-resident model cache with LRU eviction. |
| 48 | |
| 49 | Key features: |
| 50 | - Keeps models on GPU with LRU eviction by default |
| 51 | - Evicts when free VRAM drops below MIN_FREE_VRAM_GB (if KEEP_ALL_ON_GPU=1) |
| 52 | - Enforces MAX_CACHED_MODELS when set |
| 53 | - LRU eviction order when eviction is necessary |
| 54 | - Fast model switching with no CPU offload |
| 55 | """ |
| 56 | |
| 57 | def __init__(self, max_size_gb: int = 28): |
| 58 | self.max_size_gb = max_size_gb |
| 59 | self.cache: OrderedDict[str, any] = OrderedDict() |
| 60 | self.most_recent_name: Optional[str] = None |
| 61 | self._log_memory_status() |
| 62 | |
| 63 | def _log_memory_status(self): |
| 64 | """Log current GPU memory status.""" |
| 65 | if DEVICE == "cuda": |
| 66 | mem = get_gpu_memory_info() |
| 67 | logger.info(f"GPU Memory: {mem['used']:.1f}GB used / {mem['total']:.1f}GB total ({mem['free']:.1f}GB free)") |
| 68 | |
| 69 | def __len__(self): |
| 70 | return len(self.cache) |
| 71 | |
| 72 | def _should_evict_for_memory(self) -> bool: |
| 73 | """Check if we need to evict models based on memory pressure.""" |
| 74 | if DEVICE != "cuda": |
| 75 | return False |
| 76 | mem = get_gpu_memory_info() |
| 77 | should_evict = mem["free"] < MIN_FREE_VRAM_GB |
| 78 | if should_evict: |
| 79 | logger.warning(f"Low VRAM: {mem['free']:.1f}GB free < {MIN_FREE_VRAM_GB}GB threshold") |
| 80 | return should_evict |
| 81 | |
| 82 | def _should_evict_for_size(self) -> bool: |
| 83 | """Check if we exceed the max cached models limit.""" |
| 84 | if MAX_CACHED_MODELS <= 0: |
| 85 | return False |
| 86 | return len(self.cache) >= MAX_CACHED_MODELS |
| 87 | |
| 88 | def _evict_lru_model(self, exclude_name: Optional[str]) -> bool: |
| 89 | """Evict the least recently used model (except the excluded one).""" |
| 90 | for model_name in list(self.cache.keys()): |
| 91 | if model_name == exclude_name: |
| 92 | continue |
| 93 | |
| 94 | logger.info(f"Evicting LRU model: {model_name}") |
| 95 | model = self.cache.pop(model_name) |
| 96 | |
| 97 | try: |
| 98 | if isinstance(model, (list, tuple)): |
| 99 | for m in model: |
| 100 | if hasattr(m, 'to'): |
| 101 | m.to("cpu") |
| 102 | del m |
no outgoing calls