monitoring.py

"""
monitoring.py — Data drift detection for the Churn Prediction API.

Fixes applied:
    - Reference stats loaded from models/telco_linear/reference_stats.json
      (computed and saved by train.py) — no more hardcoded values
    - Drift log persisted to SQLite (same DB as MLflow) instead of /tmp/
      so it survives Render free-tier container restarts
"""

import json
import logging
import math
import os
import sqlite3
from typing import Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MODEL_NAME = os.getenv("MODEL_NAME", "telco_linear")
REFERENCE_STATS_PATH = os.path.join("models", MODEL_NAME, "reference_stats.json")

# Use same SQLite DB as MLflow — already exists and is writable
_MLFLOW_URI = os.getenv("MLFLOW_TRACKING_URI", "sqlite:////tmp/mlruns.db")
if _MLFLOW_URI.startswith("sqlite:///"):
    DB_PATH = _MLFLOW_URI[len("sqlite:///"):]
    if DB_PATH.startswith("/"):
        pass  # absolute path like /tmp/mlruns.db
    # else relative path
else:
    DB_PATH = "/tmp/drift.db"

# Drift thresholds
PSI_WARNING   = 0.1
PSI_ALERT     = 0.2
CHI2_WARNING  = 0.2
CHI2_ALERT    = 0.5
MIN_SAMPLES   = 50


# ---------------------------------------------------------------------------
# PSI + Chi-Squared
# ---------------------------------------------------------------------------

def _psi(expected: list, actual: list) -> float:
    psi = 0.0
    for e, a in zip(expected, actual):
        e = max(e, 1e-6)
        a = max(a, 1e-6)
        psi += (a - e) * math.log(a / e)
    return round(psi, 4)


def _chi2(expected: dict, actual: dict) -> float:
    all_keys = set(expected) | set(actual)
    dist = 0.0
    for k in all_keys:
        e = max(expected.get(k, 0), 1e-6)
        a = actual.get(k, 0)
        dist += ((a - e) ** 2) / e
    return round(dist, 4)


# ---------------------------------------------------------------------------
# DriftMonitor
# ---------------------------------------------------------------------------

class DriftMonitor:
    def __init__(self):
        self.reference_stats = self._load_reference_stats()
        self._init_db()

    # ── Reference stats ─────────────────────────────────────────────────────

    def _load_reference_stats(self) -> dict:
        """
        Load reference stats saved by train.py.
        Falls back to empty dict if not found (graceful degradation).
        """
        if os.path.exists(REFERENCE_STATS_PATH):
            with open(REFERENCE_STATS_PATH, "r") as f:
                stats = json.load(f)
            logger.info(
                "Loaded reference stats from %s (%d numeric, %d categorical features)",
                REFERENCE_STATS_PATH,
                len(stats.get("numeric", {})),
                len(stats.get("categorical", {})),
            )
            return stats
        else:
            logger.warning(
                "Reference stats not found at %s. "
                "Run train.py to generate them. Drift monitoring disabled.",
                REFERENCE_STATS_PATH,
            )
            return {}

    # ── SQLite persistence ───────────────────────────────────────────────────

    def _init_db(self) -> None:
        """Create the drift_log table if it doesn't exist."""
        try:
            os.makedirs(os.path.dirname(DB_PATH) or ".", exist_ok=True)
            with sqlite3.connect(DB_PATH) as conn:
                conn.execute("""
                    CREATE TABLE IF NOT EXISTS drift_log (
                        id        INTEGER PRIMARY KEY AUTOINCREMENT,
                        ts        DATETIME DEFAULT CURRENT_TIMESTAMP,
                        payload   TEXT NOT NULL
                    )
                """)
                conn.commit()
        except Exception as exc:
            logger.warning("Drift DB init failed (non-fatal): %s", exc)

    def log_request(self, data: dict) -> None:
        """Persist one prediction request to the drift_log table."""
        try:
            with sqlite3.connect(DB_PATH) as conn:
                conn.execute(
                    "INSERT INTO drift_log (payload) VALUES (?)",
                    (json.dumps(data),)
                )
                conn.commit()
        except Exception as exc:
            logger.debug("Drift log write failed (non-fatal): %s", exc)

    def _load_recent(self, n: int = 500) -> list:
        """Load the most recent n records from the drift_log table."""
        try:
            with sqlite3.connect(DB_PATH) as conn:
                rows = conn.execute(
                    "SELECT payload FROM drift_log ORDER BY id DESC LIMIT ?", (n,)
                ).fetchall()
            return [json.loads(r[0]) for r in rows]
        except Exception as exc:
            logger.debug("Drift log read failed: %s", exc)
            return []

    # ── Drift check ──────────────────────────────────────────────────────────

    def check_drift(self, n: int = 500) -> dict:
        """
        Compare recent requests against training reference statistics.
        Returns a structured report with per-feature drift scores.
        """
        if not self.reference_stats:
            return {
                "status": "unavailable",
                "reason": "Reference stats not loaded. Run train.py first.",
            }

        rows = self._load_recent(n)
        if len(rows) < MIN_SAMPLES:
            return {
                "status": "insufficient_data",
                "samples": len(rows),
                "required": MIN_SAMPLES,
            }

        df = pd.DataFrame(rows)
        report = {
            "samples":  len(rows),
            "features": {},
            "alerts":   [],
            "warnings": [],
        }

        # Numeric — PSI
        for feature, ref in self.reference_stats.get("numeric", {}).items():
            if feature not in df.columns:
                continue
            values = pd.to_numeric(df[feature], errors="coerce").dropna().values
            if len(values) == 0:
                continue
            counts, _ = np.histogram(values, bins=ref["bins"])
            actual_pcts = (counts / max(counts.sum(), 1)).tolist()
            psi = _psi(ref["bin_pcts"], actual_pcts)
            report["features"][feature] = {"psi": psi, "type": "numeric"}

            if psi > PSI_ALERT:
                msg = f"ALERT: Significant drift in '{feature}' (PSI={psi})"
                report["alerts"].append(msg)
                logger.warning(msg)
            elif psi > PSI_WARNING:
                msg = f"WARNING: Moderate drift in '{feature}' (PSI={psi})"
                report["warnings"].append(msg)
                logger.warning(msg)

        # Categorical — Chi-Squared
        for feature, ref_dist in self.reference_stats.get("categorical", {}).items():
            if feature not in df.columns:
                continue
            actual_dist = df[feature].astype(str).value_counts(normalize=True).to_dict()
            chi2 = _chi2(ref_dist, actual_dist)
            report["features"][feature] = {"chi2": chi2, "type": "categorical"}

            if chi2 > CHI2_ALERT:
                msg = f"ALERT: Significant drift in '{feature}' (chi2={chi2})"
                report["alerts"].append(msg)
                logger.warning(msg)
            elif chi2 > CHI2_WARNING:
                msg = f"WARNING: Moderate drift in '{feature}' (chi2={chi2})"
                report["warnings"].append(msg)
                logger.warning(msg)

        report["status"] = (
            "drift_detected" if report["alerts"]
            else "drift_warning" if report["warnings"]
            else "no_drift"
        )
        return report


# Singleton used by main.py
drift_monitor = DriftMonitor()