Spaces:
Running
Running
File size: 4,531 Bytes
a0c0093 bb2adf3 a0c0093 bb2adf3 a0c0093 bb2adf3 a0c0093 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | """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
|