Spaces:
Running on Zero
Running on Zero
| """LoRA cache quota management. | |
| Two helpers: | |
| - `_lora_cache_usage_bytes()` — best-effort total size of files in the cache dir. | |
| - `_evict_lora_cache_for_quota(needed)` — LRU-evict cache files (skipping any | |
| referenced by an in-memory LoRA store) until `needed` more bytes fit, or | |
| raise `gr.Error` if the quota still cannot accommodate the request. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import re | |
| import gradio as gr | |
| from state import ( | |
| _CUSTOM_LoRA, | |
| _H3_MAX_LORA_CACHE_BYTES, | |
| _LORA_CACHE_DIR, | |
| _log, | |
| ) | |
| def _lora_cache_usage_bytes() -> int: | |
| """Sum sizes of files under the LoRA cache dir (best-effort).""" | |
| total = 0 | |
| try: | |
| if not os.path.isdir(_LORA_CACHE_DIR): | |
| return 0 | |
| for name in os.listdir(_LORA_CACHE_DIR): | |
| p = os.path.join(_LORA_CACHE_DIR, name) | |
| if os.path.isfile(p) and not name.endswith(".tmp"): | |
| try: | |
| total += os.path.getsize(p) | |
| except OSError: | |
| pass | |
| except OSError: | |
| pass | |
| return total | |
| def _evict_lora_cache_for_quota(needed: int) -> None: | |
| """Evict oldest unreferenced cache files until *needed* bytes fit, or raise. | |
| Only deletes files that are not currently referenced by in-memory LoRA stores. | |
| Never crashes the process solely because the cache is full — rejects the new | |
| artifact instead. | |
| """ | |
| if needed <= 0: | |
| return | |
| usage = _lora_cache_usage_bytes() | |
| if usage + needed <= _H3_MAX_LORA_CACHE_BYTES: | |
| return | |
| # Build set of keys currently held in RAM (do not evict those). | |
| live = set() | |
| for k in list(_CUSTOM_LoRA.keys()): | |
| if isinstance(k, tuple) and len(k) >= 2: | |
| live.add(re.sub(r"[^A-Za-z0-9._-]", "_", f"{k[0].replace('/', '__')}__{k[1]}")) | |
| candidates = [] | |
| try: | |
| for name in os.listdir(_LORA_CACHE_DIR): | |
| if name.endswith(".tmp") or name.startswith("."): | |
| continue | |
| p = os.path.join(_LORA_CACHE_DIR, name) | |
| if not os.path.isfile(p): | |
| continue | |
| # Skip if name is a prefix match of a live key | |
| if any(name.startswith(lk) or lk.startswith(name.split("__rev_")[0]) for lk in live): | |
| continue | |
| try: | |
| candidates.append((os.path.getmtime(p), os.path.getsize(p), p, name)) | |
| except OSError: | |
| pass | |
| except OSError: | |
| pass | |
| candidates.sort() # oldest first | |
| freed = 0 | |
| for _mtime, size, p, name in candidates: | |
| if usage + needed - freed <= _H3_MAX_LORA_CACHE_BYTES: | |
| break | |
| try: | |
| os.remove(p) | |
| freed += size | |
| _log(f"LoRA cache evict (quota): {name} ({size / 1e6:.1f} MB)") | |
| except OSError as _e: | |
| _log(f"LoRA cache evict failed for {name}: {_e}", "WARN") | |
| if usage + needed - freed > _H3_MAX_LORA_CACHE_BYTES: | |
| raise gr.Error( | |
| f"LoRA cache quota exceeded (need {needed / 1e6:.0f} MB, " | |
| f"limit {_H3_MAX_LORA_CACHE_BYTES / 1e9:.1f} GB). " | |
| "Try again later or use a smaller LoRA." | |
| ) | |