dee/core/scoring.py

"""Shared ESM-2 scorer cache + concurrency guard.

Lives here — not in dee/server.py, where this cache originated — so both
the REST /api/run pipeline (dee/server.py) and the Turing agent's
design_variant_library tool (dee/core/agent_tools.py) can share ONE loaded
model instance per config. dee/server.py already imports dee/core/agent.py,
so dee/core/agent_tools.py importing FROM dee/server.py would be a circular
import; and having each caller load its own copy would multiply memory on
a box that already OOMs on a second full-size model (see effective_model
below) — the cache existed for exactly that reason before this move.

Concurrency: the deploy runs one process, ~2 vCPUs, gunicorn --workers 1
--threads 8 (see Dockerfile). Multiple threads calling score_all_
substitutions() concurrently don't corrupt anything (inference runs under
torch.no_grad(), no shared mutable state gets written) — but they DO
contend for the same handful of CPU cores, so two concurrent callers each
get proportionally slower instead of one finishing fast and the other
waiting cleanly. SCORING_SEMAPHORE serializes actual scoring calls so
that's a controlled, visible wait instead of silent cross-request
slowdown that's invisible until someone notices everything's crawling.
"""

from __future__ import annotations

import logging
import threading
from typing import Any, Dict, Optional

import pandas as pd

logger = logging.getLogger("dee.scoring")

_SCORER_CACHE: Dict[str, Any] = {}
_SCORER_CACHE_LOCK = threading.Lock()
_GPU_AVAILABLE: Optional[bool] = None

# One scoring call at a time, system-wide. Revisit only after confirming
# the deploy tier actually has CPU headroom for real concurrent inference
# — it doesn't today (free HF cpu-basic tier).
SCORING_SEMAPHORE = threading.Semaphore(1)


class ScoringBusyError(Exception):
    """score_guarded() couldn't acquire the semaphore within wait_timeout —
    the engine is already scoring another request."""


def gpu_available() -> bool:
    """True only if a CUDA device is actually present. Cached; treats a
    missing/broken torch as CPU-only."""
    global _GPU_AVAILABLE
    if _GPU_AVAILABLE is None:
        try:
            import torch
            _GPU_AVAILABLE = bool(torch.cuda.is_available())
        except Exception:  # noqa: BLE001
            _GPU_AVAILABLE = False
    return _GPU_AVAILABLE


def effective_model(requested: Optional[str]) -> str:
    """Downgrade a GPU-class model to 'small' on a CPU-only box so a single
    request can't OOM-kill the shared instance. The 650M/3B models are a
    ~12-24 GB fp32 load — that OOM-kills the whole container, every user,
    on the free CPU tier."""
    model = requested or "small"
    if model in ("medium", "large") and not gpu_available():
        logger.warning("Model %r requested on a CPU-only host — downgrading to 'small'.", model)
        return "small"
    return model


def get_scorer(model_name: Optional[str] = None,
               device: Optional[str] = None,
               quantization: Optional[str] = None):
    """Process-global, lock-guarded ESM-2 scorer cache. Weights load at most
    once per (effective) configuration and are reused by every caller —
    the REST pipeline and the agent tool both come through here."""
    from dee.models.scorer import ESM2Scorer, ScorerConfig
    model = effective_model(model_name)
    if not gpu_available():
        device = None          # ignore a user 'cuda'/'mps' on a CPU box
        quantization = None    # bitsandbytes quant is CUDA-only anyway
    key = f"{model}|{device or 'auto'}|{quantization or 'none'}"
    cached = _SCORER_CACHE.get(key)
    if cached is not None:
        return cached
    with _SCORER_CACHE_LOCK:
        cached = _SCORER_CACHE.get(key)   # double-checked under the lock
        if cached is None:
            cached = ESM2Scorer(ScorerConfig(
                model_name=model, device=device, quantization=quantization))
            _SCORER_CACHE[key] = cached
        return cached


def score_guarded(scorer, protein: str, *, wait_timeout: Optional[float] = None) -> pd.DataFrame:
    """scorer.score_all_substitutions(protein), serialized through
    SCORING_SEMAPHORE.

    wait_timeout=None blocks until the semaphore is free — used by the REST
    /api/run pipeline, which already has an async job-polling UI, so a
    longer wait there is invisible to the user, not a hang.

    A float bounds the wait and raises ScoringBusyError on timeout instead
    — used by the chat tool, which has one HTTP request's worth of budget
    and needs to fail fast with a clear message rather than hold the
    request open indefinitely."""
    if not SCORING_SEMAPHORE.acquire(timeout=wait_timeout):
        raise ScoringBusyError("ESM-2 scorer is busy with another request")
    try:
        return scorer.score_all_substitutions(protein)
    finally:
        SCORING_SEMAPHORE.release()