rarelens/api/app/services/scoring.py

"""Score a job's variants with the registered MLflow model.

The registered model is a pyfunc that owns its feature engineering (rarelens_ml.features travels
with it as model code) and returns P(pathogenic). Serving therefore only sends the raw columns
below and cannot drift from training.
"""
import asyncio
import uuid
from collections.abc import Sequence
from typing import Any

import mlflow
import pandas as pd
from mlflow import MlflowClient
from sqlalchemy import func, select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.models import Prediction, Variant

# Must match rarelens_ml.features.RAW_COLUMNS.
RAW_COLUMNS = ["impact", "consequence", "gnomad_af", "cadd_phred", "am_pathogenicity"]
CHUNK_SIZE = 5000

_models: dict[str, Any] = {}  # model version -> loaded pyfunc


def load_model() -> tuple[Any, str]:
    """Return (model, version), from MODEL_URI if set, else from the registry alias."""
    if settings.model_uri:
        return _load_uri(settings.model_uri)
    client = MlflowClient(tracking_uri=settings.mlflow_tracking_uri)
    version = client.get_model_version_by_alias(settings.model_name, settings.model_alias).version
    if version not in _models:
        mlflow.set_tracking_uri(settings.mlflow_tracking_uri)
        _models.clear()
        _models[version] = mlflow.pyfunc.load_model(f"models:/{settings.model_name}/{version}")
    return _models[version], version


def _load_uri(uri: str) -> tuple[Any, str]:
    """Load a model artifact directly (gs://...): no tracking server, nothing running when idle."""
    version = uri.rstrip("/").rsplit("/", 1)[-1][:40] or "uri"
    if version not in _models:
        _models.clear()
        _models[version] = mlflow.pyfunc.load_model(uri)
    return _models[version], version


def raw_frame(variants: Sequence[Variant]) -> pd.DataFrame:
    return pd.DataFrame(
        {
            "impact": [v.impact for v in variants],
            "consequence": [v.consequence for v in variants],
            "gnomad_af": [v.gnomad_af if v.gnomad_af is not None else float("nan") for v in variants],
            "cadd_phred": [v.annotations.get("CADD_PHRED") for v in variants],
            "am_pathogenicity": [v.annotations.get("am_pathogenicity") for v in variants],
        },
        columns=RAW_COLUMNS,
    )


async def score_job(job_id: uuid.UUID, session: AsyncSession) -> tuple[int, str]:
    model, version = await asyncio.to_thread(load_model)
    scored = 0
    last_id = 0
    # Keyset pagination keeps memory flat for whole-genome jobs.
    while True:
        variants = (
            await session.scalars(
                select(Variant)
                .where(Variant.job_id == job_id, Variant.id > last_id)
                .order_by(Variant.id)
                .limit(CHUNK_SIZE)
            )
        ).all()
        if not variants:
            break
        scores = await asyncio.to_thread(model.predict, raw_frame(variants))
        stmt = insert(Prediction).values(
            [
                {
                    "variant_id": v.id,
                    "model_name": settings.model_name,
                    "model_version": version,
                    "score": float(s),
                }
                for v, s in zip(variants, scores, strict=True)
            ]
        )
        # Re-scoring (e.g. after a new model version) replaces the previous prediction.
        await session.execute(
            stmt.on_conflict_do_update(
                index_elements=[Prediction.variant_id],
                set_={
                    "model_name": stmt.excluded.model_name,
                    "model_version": stmt.excluded.model_version,
                    "score": stmt.excluded.score,
                    "created_at": func.now(),
                },
            )
        )
        scored += len(variants)
        last_id = variants[-1].id
    await session.commit()
    return scored, version