Spaces:
Running
Running
| """Models are loaded when first asked for and kept in a bounded cache. | |
| The Space has 16 GB of memory and the registry lists more models each | |
| release, so start-up loads only the default model and the others are loaded | |
| on first use. Loaded models are kept, least recently used first out, while | |
| their estimated memory stays within a budget (``TAXONOVA_WEB_MODEL_BUDGET_MB``, | |
| default ``DEFAULT_BUDGET_MB``). The estimate is the checkpoint file size times | |
| ``MEMORY_PER_FILE_BYTE``. The model being asked for is always loaded, even if | |
| it alone exceeds the budget. | |
| Measured 2026-11-04 for the 26.11 models (taxonova 26.11.0, torch 2.5.1 CPU, | |
| 2 threads, one model per process): resident memory grew by 263 to 640 MB per | |
| loaded model (1.2 to 1.3 times the file size; v4 334 MB, v4v5 640 MB), and | |
| loading peaked about 0.35 GB above that. A 2,000-read job (rRNA screen and | |
| classification) took 25 to 58 s and raised the process peak to 0.8 to 1.2 GB. | |
| Jobs run one at a time in one worker thread, so a model is never evicted | |
| while it is in use; the lock only protects the bookkeeping against the page | |
| and health requests that read it. | |
| """ | |
| from __future__ import annotations | |
| import ctypes | |
| import gc | |
| import logging | |
| import os | |
| import threading | |
| from collections import OrderedDict | |
| from pathlib import Path | |
| from typing import Callable, Dict, List, Optional | |
| logger = logging.getLogger(__name__) | |
| #: Estimated memory per byte of checkpoint file: the measured 1.2 to 1.3 once | |
| #: loaded, with margin for the load peak and allocator overhead. | |
| MEMORY_PER_FILE_BYTE = 2.0 | |
| #: About 4 GB resident at the measured ratio, which leaves most of the | |
| #: Space's 16 GB for a 20,000-sequence job, the rRNA screen and the server. | |
| DEFAULT_BUDGET_MB = 6000 | |
| BUDGET_ENV = "TAXONOVA_WEB_MODEL_BUDGET_MB" | |
| def budget_mb() -> int: | |
| try: | |
| return max(0, int(os.environ.get(BUDGET_ENV, DEFAULT_BUDGET_MB))) | |
| except ValueError: | |
| return DEFAULT_BUDGET_MB | |
| def estimate_mb(path) -> float: | |
| return Path(path).stat().st_size * MEMORY_PER_FILE_BYTE / 2 ** 20 | |
| def _release_memory() -> None: | |
| gc.collect() | |
| try: # hand freed heap pages back to the system (glibc only) | |
| ctypes.CDLL("libc.so.6").malloc_trim(0) | |
| except (OSError, AttributeError): | |
| pass | |
| class ModelCache: | |
| """Loaded models by id, least recently used first, within a memory budget.""" | |
| def __init__(self, loader: Callable[[str], object], budget: Optional[float] = None, | |
| estimate: Callable[[str], float] = estimate_mb): | |
| self._loader = loader | |
| self._estimate = estimate | |
| self.budget = budget_mb() if budget is None else budget | |
| self._items: "OrderedDict[str, object]" = OrderedDict() | |
| self._sizes: Dict[str, float] = {} | |
| self._lock = threading.Lock() | |
| self.loads = 0 | |
| self.evictions = 0 | |
| def loaded_ids(self) -> List[str]: | |
| with self._lock: | |
| return list(self._items) | |
| def used_mb(self) -> float: | |
| with self._lock: | |
| return sum(self._sizes.values()) | |
| def is_loaded(self, model_id: str) -> bool: | |
| with self._lock: | |
| return model_id in self._items | |
| def put(self, model_id: str, loaded, size_mb: float = 0.0) -> None: | |
| """Insert an already loaded model (tests, or a model loaded elsewhere).""" | |
| with self._lock: | |
| self._items[model_id] = loaded | |
| self._items.move_to_end(model_id) | |
| self._sizes[model_id] = size_mb | |
| def get(self, model_id: str, path) -> object: | |
| """The loaded model, loading it (and evicting others) if needed.""" | |
| with self._lock: | |
| if model_id in self._items: | |
| self._items.move_to_end(model_id) | |
| return self._items[model_id] | |
| need = self._estimate(str(path)) | |
| evicted = [] | |
| with self._lock: | |
| while self._items and sum(self._sizes.values()) + need > self.budget: | |
| old, _ = self._items.popitem(last=False) | |
| self._sizes.pop(old, None) | |
| evicted.append(old) | |
| if evicted: | |
| self.evictions += len(evicted) | |
| logger.info("model cache: unloaded %s to make room for %s", ", ".join(evicted), model_id) | |
| _release_memory() | |
| logger.info("loading %s from %s (about %.0f MB)", model_id, path, need) | |
| loaded = self._loader(str(path)) | |
| self.loads += 1 | |
| with self._lock: | |
| self._items[model_id] = loaded | |
| self._sizes[model_id] = need | |
| return loaded | |