src/pipeline/runner.py

"""Async pipeline: generator -> scorer -> broadcaster, with structured logging."""

import logging
import time
from collections.abc import AsyncIterator, Awaitable, Callable
from typing import Any

from src.detector.anomaly import AnomalyDetector
from src.detector.drift import DriftDetector
from src.detector.metrics import RunningMetrics
from src.stream.generator import (
    BASELINE_PRESSURE,
    BASELINE_TEMP,
    BASELINE_VIBRATION,
    NOISE_SIGMA,
)
from src.stream.models import Observation

logger = logging.getLogger(__name__)


def _drift_signal(obs: Observation) -> float:
    """Average normalized deviation of all three sensors from Phase A baselines."""
    return (
        (obs.temperature - BASELINE_TEMP) / NOISE_SIGMA
        + (obs.pressure - BASELINE_PRESSURE) / NOISE_SIGMA
        + (obs.vibration - BASELINE_VIBRATION) / NOISE_SIGMA
    ) / 3.0


def _broadcast_message(
    obs: Observation,
    anomaly_score: float,
    alert: bool,
    drift_event: bool,
    running_precision: float,
    running_recall: float,
    running_f1: float,
    latency_ms: float,
    observation_index: int,
    total_anomalies_detected: int,
    total_drift_events: int,
) -> dict[str, Any]:
    return {
        "timestamp": obs.timestamp,
        "temperature": obs.temperature,
        "pressure": obs.pressure,
        "vibration": obs.vibration,
        "label": obs.label,
        "phase": obs.phase,
        "anomaly_score": anomaly_score,
        "alert": alert,
        "drift_event": drift_event,
        "running_precision": running_precision,
        "running_recall": running_recall,
        "running_f1": running_f1,
        "latency_ms": latency_ms,
        "observation_index": observation_index,
        "total_anomalies_detected": total_anomalies_detected,
        "total_drift_events": total_drift_events,
    }


async def run_pipeline(
    stream: AsyncIterator[Observation],
    broadcaster: Callable[[dict[str, Any]], Awaitable[None]],
    *,
    anomaly_threshold: float = 0.5,
    n_trees: int = 10,
    height: int = 6,
    window_size: int = 100,
    drift_delta: float = 0.002,
    drift_grace_period: int = 30,
    state: dict[str, Any] | None = None,
) -> None:
    """Run the detection pipeline: score each observation, detect drift, broadcast messages."""
    detector = AnomalyDetector(
        threshold=anomaly_threshold,
        n_trees=n_trees,
        height=height,
        window_size=window_size,
    )
    drift_detector = DriftDetector(delta=drift_delta, grace_period=drift_grace_period)
    metrics = RunningMetrics()
    observation_index = 0
    total_anomalies = 0
    total_drifts = 0

    logger.info("pipeline_start", extra={"threshold": anomaly_threshold, "n_trees": n_trees})

    _baseline_frozen = False

    async for obs in stream:
        t0 = time.perf_counter()

        if obs.phase == "A":
            detector.learn_scaler(obs)
        elif not _baseline_frozen:
            detector.freeze_baseline()
            _baseline_frozen = True

        score = detector.score(obs)
        drift_event_flag = drift_detector.update(_drift_signal(obs))

        if drift_event_flag:
            total_drifts += 1
            logger.info(
                "drift_detected",
                extra={
                    "observation_index": observation_index,
                    "phase": obs.phase,
                    "drift_signal": round(_drift_signal(obs), 4),
                },
            )

        alert = score > anomaly_threshold
        if alert:
            total_anomalies += 1

        metrics.update(ground_truth=(obs.label != "normal"), predicted=alert)

        latency_ms = (time.perf_counter() - t0) * 1000

        if state is not None:
            state["total_observations"] = observation_index + 1
            state["total_anomalies_detected"] = total_anomalies
            state["total_drift_events"] = total_drifts
            state["current_precision"] = metrics.precision
            state["current_recall"] = metrics.recall
            state["current_f1"] = metrics.f1

        msg = _broadcast_message(
            obs=obs,
            anomaly_score=score,
            alert=alert,
            drift_event=drift_event_flag,
            running_precision=metrics.precision,
            running_recall=metrics.recall,
            running_f1=metrics.f1,
            latency_ms=round(latency_ms, 3),
            observation_index=observation_index,
            total_anomalies_detected=total_anomalies,
            total_drift_events=total_drifts,
        )
        await broadcaster(msg)
        observation_index += 1

    logger.info(
        "pipeline_end",
        extra={
            "total_observations": observation_index,
            "total_anomalies": total_anomalies,
            "total_drifts": total_drifts,
            "final_precision": round(metrics.precision, 4),
            "final_recall": round(metrics.recall, 4),
            "final_f1": round(metrics.f1, 4),
        },
    )