Serving¶

class FeatureLookupService:
    """Loads precomputed features for all matches from the batch inference output.

    The parquet file is produced by the ``batch_inference`` DVC stage and has
    the match ``id`` as its index.  It contains both upcoming matches and
    finished matches (with ``outcome_1x2``, ``homeScore``, ``awayScore``).
    The service is loaded lazily on the first call and cached in-process.

    Args:
        features_path: Absolute path to ``match_features.parquet``.
    """

    def __init__(self, features_path: Path | None = None) -> None:
        self._path = features_path or (
            DATA_PREDICTIONS_PATH / _FUTURE_FEATURES_FILENAME
        )
        self._df: pd.DataFrame | None = None
        # UTC timestamp of the loaded data: local st_mtime or MinIO LastModified.
        self._mtime: float | None = None
        # monotonic time of the last MinIO head_object check (avoids per-request calls).
        self._last_minio_check: float = 0.0

    # ------------------------------------------------------------------
    # MinIO helpers
    # ------------------------------------------------------------------

    def _minio_last_modified(self) -> float | None:
        """Return the MinIO LastModified timestamp for the features object, or None."""
        bucket = get_minio_settings().bucket_predictions
        if not bucket:
            return None
        try:
            client = boto3.client(
                "s3",
                endpoint_url=get_minio_settings().endpoint_url,
                aws_access_key_id=get_minio_settings().access_key,
                aws_secret_access_key=get_minio_settings().secret_key,
            )
            resp = client.head_object(Bucket=bucket, Key=_MINIO_PREDICTIONS_KEY)
            return resp["LastModified"].timestamp()
        except Exception as exc:  # noqa: BLE001
            logger.debug("MinIO head_object failed: %s", exc)
            return None

    def _load_from_minio(self) -> pd.DataFrame:
        """Read future_features.parquet directly from MinIO into memory."""
        bucket = get_minio_settings().bucket_predictions
        s3_url = f"s3://{bucket}/{_MINIO_PREDICTIONS_KEY}"
        try:
            df = pd.read_parquet(
                s3_url, storage_options=get_minio_settings().storage_options
            )
            logger.info(
                "Loaded future features from MinIO (%s): %d matches", s3_url, len(df)
            )
            return df
        except Exception:  # noqa: BLE001
            logger.exception(
                "Failed to load future_features.parquet from MinIO (%s)", s3_url
            )
            return pd.DataFrame()

    # ------------------------------------------------------------------
    # Core load logic
    # ------------------------------------------------------------------

    def _load(self) -> pd.DataFrame:
        # ── Local file (dev / Airflow server) ────────────────────────────
        if self._path.exists():
            current_mtime = self._path.stat().st_mtime
            if self._df is not None and current_mtime == self._mtime:
                return self._df
            self._df = pd.read_parquet(self._path)
            self._mtime = current_mtime
            logger.info(
                "Loaded future features from local: %d matches (mtime=%s)",
                len(self._df),
                datetime.fromtimestamp(current_mtime, tz=timezone.utc).isoformat(),
            )
            return self._df

        # ── MinIO (production K8s — no local data volume) ─────────────────
        now = time.monotonic()
        should_check = (now - self._last_minio_check) >= _FEATURE_CACHE_CHECK_INTERVAL
        if self._df is not None and not should_check:
            return self._df  # return cached; not time to re-check yet

        self._last_minio_check = now
        minio_mtime = self._minio_last_modified()

        if minio_mtime is None:
            if self._df is not None:
                logger.warning("MinIO unreachable — serving stale feature cache")
                return self._df
            logger.warning(
                "future_features.parquet not found locally or in MinIO — "
                "run the batch_inference DVC stage and ensure "
                "MINIO_BUCKET_PREDICTIONS is configured."
            )
            self._df = pd.DataFrame()
            self._mtime = None
            return self._df

        if self._df is not None and minio_mtime == self._mtime:
            return self._df  # MinIO file unchanged

        self._df = self._load_from_minio()
        self._mtime = minio_mtime
        return self._df

    def get_features(self, match_id: int) -> dict | None:
        """Return the feature dict for *match_id*, or ``None`` if not found."""
        df = self._load()
        if df.empty or match_id not in df.index:
            return None
        row = df.loc[match_id]
        # Convert to plain Python dict; drop NaN values to avoid serialisation issues
        return {k: v for k, v in row.to_dict().items() if pd.notna(v)}

    def list_matches(self) -> list[dict]:
        """Return a lightweight list of upcoming matches for UI display."""
        df = self._load()
        if df.empty:
            return []
        # Keep only matches that have not been played yet (no recorded outcome).
        if "outcome_1x2" in df.columns:
            df = df[df["outcome_1x2"].isna()]
        display_cols = [
            c
            for c in (
                "startTimeUtc",
                "homeTeamName",
                "awayTeamName",
                "homeTeamId",
                "awayTeamId",
            )
            if c in df.columns
        ]
        df_display = df[display_cols].copy()
        df_display.index.name = "match_id"
        return df_display.reset_index().to_dict(orient="records")

    @property
    def features_computed_at(self) -> datetime | None:
        """UTC datetime when the feature file was last written (= last batch_inference run)."""
        self._load()
        if self._mtime is None:
            return None
        return datetime.fromtimestamp(self._mtime, tz=timezone.utc)

class PredictionLookupService:
    """Serves precomputed batch predictions from the batch_inference DVC stage output.

    Mirrors the ``FeatureLookupService`` caching pattern:
    - Checks local file first (dev / CI).
    - Falls back to MinIO with a configurable re-check interval.

    The parquet file is indexed by match ``id`` and must contain columns:
    ``proba_home``, ``proba_draw``, ``proba_away``, ``predicted_class``,
    ``predicted_label``, optionally ``is_future``, ``startTimeUtc``,
    ``homeTeamName``, ``awayTeamName``, ``model_run_id``, ``model_stage``.
    """

    def __init__(self, predictions_path: Path | None = None) -> None:
        self._path = predictions_path or (DATA_PREDICTIONS_PATH / _PREDICTIONS_FILENAME)
        self._df: pd.DataFrame | None = None
        self._mtime: float | None = None
        self._last_minio_check: float = 0.0

    # ------------------------------------------------------------------
    # MinIO helpers
    # ------------------------------------------------------------------

    def _minio_last_modified(self) -> float | None:
        """Return the MinIO LastModified for the predictions object, or None."""
        bucket = get_minio_settings().bucket_predictions
        if not bucket:
            return None
        try:
            client = boto3.client(
                "s3",
                endpoint_url=get_minio_settings().endpoint_url,
                aws_access_key_id=get_minio_settings().access_key,
                aws_secret_access_key=get_minio_settings().secret_key,
            )
            resp = client.head_object(Bucket=bucket, Key=_MINIO_BATCH_PREDICTIONS_KEY)
            return resp["LastModified"].timestamp()
        except Exception as exc:  # noqa: BLE001
            logger.debug("MinIO head_object (predictions) failed: %s", exc)
            return None

    def _load_from_minio(self) -> pd.DataFrame:
        """Read predictions.parquet directly from MinIO into memory."""
        bucket = get_minio_settings().bucket_predictions
        s3_url = f"s3://{bucket}/{_MINIO_BATCH_PREDICTIONS_KEY}"
        try:
            df = pd.read_parquet(
                s3_url, storage_options=get_minio_settings().storage_options
            )
            logger.info(
                "Loaded batch predictions from MinIO (%s): %d rows", s3_url, len(df)
            )
            return df
        except Exception:  # noqa: BLE001
            logger.exception(
                "Failed to load predictions.parquet from MinIO (%s)", s3_url
            )
            return pd.DataFrame()

    # ------------------------------------------------------------------
    # Core load logic
    # ------------------------------------------------------------------

    def _load(self) -> pd.DataFrame:
        # ── Local file (dev / CI) ─────────────────────────────────────────
        if self._path.exists():
            current_mtime = self._path.stat().st_mtime
            if self._df is not None and current_mtime == self._mtime:
                return self._df
            self._df = pd.read_parquet(self._path)
            self._mtime = current_mtime
            logger.info(
                "Loaded batch predictions from local: %d rows (mtime=%s)",
                len(self._df),
                datetime.fromtimestamp(current_mtime, tz=timezone.utc).isoformat(),
            )
            return self._df

        # ── MinIO (production K8s) ────────────────────────────────────────
        now = time.monotonic()
        should_check = (
            now - self._last_minio_check
        ) >= _PREDICTIONS_CACHE_CHECK_INTERVAL
        if self._df is not None and not should_check:
            return self._df

        self._last_minio_check = now
        minio_mtime = self._minio_last_modified()

        if minio_mtime is None:
            if self._df is not None:
                logger.warning("MinIO unreachable — serving stale predictions cache")
                return self._df
            logger.warning(
                "predictions.parquet not found locally or in MinIO — "
                "run the batch_inference DVC stage and ensure "
                "MINIO_BUCKET_PREDICTIONS is configured."
            )
            self._df = pd.DataFrame()
            self._mtime = None
            return self._df

        if self._df is not None and minio_mtime == self._mtime:
            return self._df

        self._df = self._load_from_minio()
        self._mtime = minio_mtime
        return self._df

    def get_prediction(self, match_id: int) -> dict | None:
        """Return the prediction dict for *match_id*, or ``None`` if not found."""
        df = self._load()
        if df.empty or match_id not in df.index:
            return None
        row = df.loc[match_id]
        return {k: v for k, v in row.to_dict().items() if pd.notna(v)}

    def list_matches(self) -> list[dict]:
        """Return all rows as a list of dicts (for diagnostics/admin endpoints)."""
        df = self._load()
        if df.empty:
            return []
        return df.reset_index().to_dict(orient="records")

    @property
    def predictions_computed_at(self) -> datetime | None:
        """UTC datetime when predictions.parquet was last written."""
        self._load()
        if self._mtime is None:
            return None
        return datetime.fromtimestamp(self._mtime, tz=timezone.utc)

class PredictionService:
    """Loads and serves a model from the MLflow Model Registry.

    The model is loaded lazily on the first call to ``predict`` and then
    cached in-process for the lifetime of the worker.
    """

    def __init__(
        self, tracking_uri: str, model_name: str, model_stage: str = "Staging"
    ) -> None:
        self._tracking_uri = tracking_uri
        self._model_name = model_name
        self._model_stage = model_stage
        self._model: Any = None  # mlflow.pyfunc.PyFuncModel, imported lazily
        # Protects the lazy-load critical section from concurrent reads
        # (e.g. gunicorn workers sharing the same process via threads).
        self._lock: threading.Lock = threading.Lock()
        # Redis client initialised lazily; None means "not yet checked".
        self._redis: Any = None
        self._redis_checked: bool = False

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _load_model(self) -> Any:
        # Fast path: model already loaded (no lock needed after init).
        if self._model is not None:
            return self._model
        # Slow path: double-checked locking prevents duplicate load
        # when multiple threads/workers call predict() simultaneously on startup.
        with self._lock:
            if self._model is None:
                # Deferred import: mlflow is only installed in the ml environment.
                # Must never execute at module-load time (api env has no mlflow).
                import mlflow
                import mlflow.pyfunc

                # MLflow's boto3 client reads these from os.environ at download time.
                # pydantic-settings populates Python objects but does NOT write back
                # to os.environ, so we must do it explicitly here.
                os.environ.setdefault(
                    "MLFLOW_S3_ENDPOINT_URL", get_minio_settings().endpoint_url
                )
                os.environ.setdefault(
                    "AWS_ACCESS_KEY_ID", get_minio_settings().access_key
                )
                os.environ.setdefault(
                    "AWS_SECRET_ACCESS_KEY", get_minio_settings().secret_key
                )
                mlflow.set_tracking_uri(self._tracking_uri)
                # MLflow 3.x: use alias URI (models:/name@alias)
                # Falls back gracefully to stage URI for older MLflow servers
                model_uri = f"models:/{self._model_name}@{self._model_stage}"
                logger.info("Loading model from MLflow: %s", model_uri)
                self._model = mlflow.pyfunc.load_model(model_uri)
                logger.info(
                    "Model loaded. run_id=%s",
                    getattr(self._model.metadata, "run_id", "unknown"),
                )
        return self._model

    def load(self) -> Any:
        """Eagerly load the model. Safe to call multiple times (idempotent).

        Call this during application startup (e.g. FastAPI lifespan or
        Celery ``worker_process_init``) to avoid paying the cold-start
        penalty on the first user request.
        """
        return self._load_model()

    # ------------------------------------------------------------------
    # Redis prediction cache
    # ------------------------------------------------------------------

    def _get_redis(self) -> Any:
        """Return a connected Redis client or *None* if unavailable.

        Tried once per process; subsequent calls return the cached result
        without retrying so a missing Redis never blocks inference.
        """
        if self._redis_checked:
            return self._redis
        self._redis_checked = True
        try:
            import redis as redis_lib

            url = os.getenv(
                "REDIS_CACHE_URL",
                os.getenv("CELERY_RESULT_BACKEND", "redis://redis:6379/0"),
            )
            client = redis_lib.from_url(
                url, decode_responses=True, socket_connect_timeout=2
            )
            client.ping()
            self._redis = client
            logger.info("Prediction cache: connected to Redis at %s", url)
        except Exception as exc:
            logger.warning("Redis unavailable; prediction caching disabled: %s", exc)
        return self._redis

    def _cache_get(self, key: str) -> dict | None:
        r = self._get_redis()
        if r is None:
            return None
        try:
            value = r.get(key)
            return json.loads(value) if value else None
        except Exception as exc:
            logger.debug("Cache get failed for key=%s: %s", key, exc)
            return None

    def _cache_set(self, key: str, value: dict, ttl: int) -> None:
        r = self._get_redis()
        if r is None:
            return
        try:
            r.setex(key, ttl, json.dumps(value))
        except Exception as exc:
            logger.debug("Cache set failed for key=%s: %s", key, exc)

    def _predict_proba(self, df: pd.DataFrame) -> "Any":  # np.ndarray at runtime
        """Return (N, 3) probability array.

        Handles two common MLflow pyfunc flavours:
        - sklearn Pipeline logged with pyfunc_predict_fn='predict_proba'  → model.predict() returns (N, 3)
        - sklearn Pipeline logged with default pyfunc                     → model.predict() returns (N,) labels;
          in that case we fall back to the underlying python_model.
        """
        import numpy as np

        model = self._load_model()
        raw = model.predict(df)

        # Already a 2-D probability array
        if hasattr(raw, "ndim") and raw.ndim == 2:
            return raw

        # 1-D label output → try to access the underlying sklearn estimator
        logger.warning(
            "pyfunc.predict() returned 1-D output; attempting predict_proba fallback."
        )
        try:
            sklearn_model = model._model_impl.python_model
            proba = sklearn_model.predict_proba(df)
            return np.asarray(proba)
        except AttributeError:
            # Last resort: one-hot encode the hard labels
            labels = raw.astype(int)
            proba = np.zeros((len(labels), len(_LABEL_ORDER)), dtype=float)
            for i, lbl in enumerate(labels):
                if lbl in _LABEL_ORDER:
                    proba[i, _LABEL_ORDER.index(lbl)] = 1.0
            return proba

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def get_model_info(self) -> dict:
        """Return model metadata from the MLflow Model Registry.

        Queries the registered model for the configured stage, then fetches
        run metrics and params.  Does NOT require the model to be loaded.
        """
        import mlflow
        from mlflow.tracking import MlflowClient

        os.environ.setdefault(
            "MLFLOW_S3_ENDPOINT_URL", get_minio_settings().endpoint_url
        )
        os.environ.setdefault("AWS_ACCESS_KEY_ID", get_minio_settings().access_key)
        os.environ.setdefault("AWS_SECRET_ACCESS_KEY", get_minio_settings().secret_key)
        mlflow.set_tracking_uri(self._tracking_uri)

        client = MlflowClient()
        try:
            # MLflow 3.x: look up by alias (e.g. "Staging")
            mv = client.get_model_version_by_alias(self._model_name, self._model_stage)
        except Exception:
            mv = None
        if mv is None:
            return {
                "model_name": self._model_name,
                "stage": self._model_stage,
                "version": "unknown",
                "run_id": "unknown",
                "metrics": {},
                "params": {},
                "feature_names": [],
                "created_at": None,
            }

        run_id = mv.run_id
        if run_id is None:
            return {
                "params": {},
                "feature_names": [],
                "created_at": None,
            }
        run = client.get_run(run_id)
        metrics = dict(run.data.metrics)
        params = dict(run.data.params)

        # Try to retrieve feature names from the input schema of the loaded model
        feature_names: list[str] = []
        try:
            model = self._load_model()
            schema = model.metadata.get_input_schema()
            if schema is not None:
                feature_names = [c.name for c in schema.inputs]
        except Exception:
            logger.debug("Could not retrieve feature names from model schema.")

        return {
            "model_name": self._model_name,
            "stage": self._model_stage,
            "version": mv.version,
            "run_id": run_id,
            "metrics": metrics,
            "params": params,
            "feature_names": feature_names,
            "created_at": str(mv.creation_timestamp) if mv.creation_timestamp else None,
        }

    def predict(
        self,
        features: dict,
        match_id: int | None = None,
        features_computed_at: datetime | None = None,
    ) -> dict:
        """Run inference for a single match.

        Args:
            features: Feature dict matching model input schema.
            match_id: Optional identifier for downstream tracing.
            features_computed_at: UTC timestamp when features were produced
                (batch_inference stage).  Stored in the response for traceability.

        Returns:
            Dict compatible with ``PredictResponse`` schema.
        """
        model = self._load_model()
        run_id = getattr(model.metadata, "run_id", "unknown")
        cache_ttl = int(os.getenv("PREDICTION_CACHE_TTL", "3600"))

        # Cache key is scoped to (match_id, run_id) so it auto-invalidates
        # when the model is updated in the registry.
        cache_key: str | None = (
            f"predict:{match_id}:{run_id}" if match_id is not None else None
        )
        if cache_key:
            cached = self._cache_get(cache_key)
            if cached is not None:
                cached["cached"] = True
                return cached

        df = pd.DataFrame([features])

        input_schema = model.metadata.get_input_schema()

        # Restore columns that were dropped during JSON serialisation (NaN is not
        # JSON-safe so get_features() omits them).  The model's sklearn pipeline
        # handles NaN via its own imputer/transformer; we just need the column to
        # exist so ColumnTransformer does not raise "columns are missing".
        if input_schema is not None:
            for col_spec in input_schema.inputs:
                if col_spec.name not in df.columns:
                    df[col_spec.name] = float("nan")

        # Align column dtypes to the model's MLflow input schema.
        # Parquet round-trips can promote int → float (e.g. when NaNs are present);
        # MLflow's schema enforcement refuses unsafe float→int conversions.
        _MLFLOW_TO_NUMPY: dict[str, str] = {
            "integer": "int32",
            "long": "int64",
            "float": "float32",
            "double": "float64",
        }
        if input_schema is not None:
            for col_spec in input_schema.inputs:
                if col_spec.name in df.columns:
                    target = _MLFLOW_TO_NUMPY.get(col_spec.type.name)
                    if target is not None and str(df[col_spec.name].dtype) != target:
                        df[col_spec.name] = df[col_spec.name].astype(target)
            # Drop columns not in the model signature to avoid MLflow warnings
            # about extra inputs (metadata cols like homeTeamName, startTimeUtc…).
            schema_cols = {c.name for c in input_schema.inputs}
            extra_cols = [c for c in df.columns if c not in schema_cols]
            if extra_cols:
                df = df.drop(columns=extra_cols)

        proba = self._predict_proba(df)

        predicted_idx = int(proba.argmax(axis=1)[0])
        predicted_class = _LABEL_ORDER[predicted_idx]

        result: dict = {
            "match_id": match_id,
            "cached": False,
            "features_computed_at": (
                features_computed_at.isoformat() if features_computed_at else None
            ),
            "prediction": {
                "predicted_class": predicted_class,
                "probabilities": {
                    str(lbl): float(proba[0, i]) for i, lbl in enumerate(_LABEL_ORDER)
                },
                "model_version": self._model_stage,
                "model_run_id": run_id,
            },
        }
        if cache_key:
            self._cache_set(cache_key, result, ttl=cache_ttl)
        return result