app/ml/runtime.py

"""Inference runtime for the email-risk model (PyTorch / transformers).

Strictly optional: `load()` returning False (or `predict()` returning None)
means "no ML this run" — the deterministic pipeline continues unchanged.
Imports are lazy so the API boots with or without torch installed, and the test
suite never requires a model download.

The deployed Space image DOES contain CPU torch (Dockerfile installs it from
the PyTorch CPU index) and ships the bundled model at `app/ml/model`. Do not
strip either: the loader fails open, so removing them silently disables the
advisory layer instead of raising.

CPU inference is offloaded to a worker thread (`predict_async`, `load_async`)
because running it inline stalls every concurrent request — measured at 42.6s
cold start before this.
"""
from __future__ import annotations

import asyncio
import threading
from typing import Callable, Optional

from ..config import settings

def default_pipeline_factory(repo: str, batch_size: int):
    """Production pipeline builder: transformers (PyTorch).

    The ONNX export produced confidently-wrong labels under two transformers
    versions while the PyTorch checkpoint scored 0.999 held-out - so inference
    runs on the trained weights directly. transformers is imported lazily here:
    a module-level import can fail under some import contexts (importlib_metadata
    'Version' bug on Python 3.14), which must degrade to "no ML" only when
    inference is actually attempted, not kill the capability at import time.

    The repo is a PUBLIC HF Hub download (free, unmetered); no token is used —
    Part B removed every HF token from the deployment surface.
    """
    import transformers

    tokenizer = transformers.AutoTokenizer.from_pretrained(repo)
    model = transformers.AutoModelForSequenceClassification.from_pretrained(repo)
    return transformers.pipeline(
        "text-classification",
        model=model,
        tokenizer=tokenizer,
        device="cpu",
        truncation=True,
        max_length=64,
    )


class EmailRiskModel:
    """Lazily-loaded ONNX text-classification wrapper. Fail-open by design."""

    def __init__(
        self,
        repo: str,
        batch_size: int = 64,
        pipeline_factory: Optional[Callable[[str, int], object]] = None,
    ):
        self.repo = repo
        self.batch_size = batch_size
        self._factory = pipeline_factory or default_pipeline_factory
        self._pipeline = None
        self._loaded = False
        self.last_error: Optional[str] = None
        self._lock = threading.Lock()

    @property
    def loaded(self) -> bool:
        return self._loaded

    def load(self) -> bool:
        """Load the model. Returns False on ANY failure (fail-open design)."""
        with self._lock:
            if self._loaded:
                return True
            try:
                self._pipeline = self._factory(
                    self.repo, self.batch_size
                )
                self._loaded = True
                return True
            except Exception as exc:
                self._pipeline = None
                self._loaded = False
                self.last_error = f"{type(exc).__name__}: {exc}"
                return False

    async def load_async(self) -> bool:
        """Load the model without blocking the event loop.

        Loading torch + a checkpoint takes seconds; doing it inline in an async
        handler stalls every concurrent request (measured: 42.6s cold start
        before this).
        """
        return await asyncio.to_thread(self.load)

    def predict(self, texts: list[str]) -> Optional[list[tuple[str, float]]]:
        """Classify texts. Returns (LABEL, confidence) tuples, or None on failure.

        The label is uppercased and guaranteed to be one of the training labels'
        names only if the model provides it; unknown labels pass through as-is
        for the caller to ignore.
        """
        if not self._loaded or self._pipeline is None or not texts:
            return None
        try:
            preds = self._pipeline(texts, batch_size=self.batch_size)
            out: list[tuple[str, float]] = []
            for batch in preds:
                # transformers may return a flat list or one list per element.
                if isinstance(batch, list):
                    preds_nested = batch
                else:
                    preds_nested = [batch]
                for p in preds_nested:
                    # An unexpected payload shape must fail open, not raise:
                    # the batch path calls predict() outside its own try.
                    if not isinstance(p, dict):
                        raise TypeError(f"unexpected prediction: {type(p).__name__}")
                    label = str(p.get("label", "")).upper()
                    try:
                        score = float(p.get("score", 0.0))
                    except (TypeError, ValueError):
                        score = 0.0
                    out.append((label, score))
            return out if len(out) == len(texts) else None
        except Exception as exc:
            self.last_error = f"predict: {type(exc).__name__}: {exc}"
            return None


_model: Optional[EmailRiskModel] = None
_model_lock = threading.Lock()


async def predict_async(model, texts: list[str]):
    """Classify texts without blocking the event loop.

    Works with either an `EmailRiskModel` (CPU inference in a worker thread) or
    any injected test double exposing a sync `predict` (called directly).
    """
    if model is None or not texts:
        return None
    predict = getattr(model, "predict", None)
    if predict is None:
        return None
    if isinstance(model, EmailRiskModel):
        return await asyncio.to_thread(predict, texts)
    return predict(texts)


async def get_model_async() -> Optional[EmailRiskModel]:
    """`get_model()` with the (slow) first load kept off the event loop."""
    global _model
    if not settings.ml_disposable_enabled or not settings.ml_model_repo:
        return None
    if _model is not None:
        return _model
    await asyncio.to_thread(get_model)
    return _model


def get_model() -> Optional[EmailRiskModel]:
    """Process-wide model, built AND loaded when flag + repo are configured."""
    global _model
    if not settings.ml_disposable_enabled or not settings.ml_model_repo:
        return None
    with _model_lock:
        if _model is None:
            _model = EmailRiskModel(
                repo=settings.ml_model_repo,
                batch_size=settings.ml_batch_size,
            )
            # Load at construction: without this the model fail-opens to None
            # on every predict and the feature is silently dead (found live).
            _model.load()
        return _model