test(calibration): add trigger unit tests + reset EMA after retrain

Addresses review on llm-d#32:
- Extract the drift-trigger decision logic into CalibrationTrigger (pure,
  unit-testable with fixed inputs); wire continuous_coverage_loop to it.
- Reset the EMA and counter once a triggered retrain lands, and suppress
  re-firing while it is in flight, so sustained drift no longer queues
  back-to-back retrains during the retrain window.
- Unit tests: steady drift fires; one-off/noisy do not; steady calibrated
  never fires; 0-100 coverage scale (not 0-1); worse-of-ttft/tpot; and no
  re-fire until the retrain lands, then a clean reset.
- Commit the manual drift harness under tests/drift_harness/ for the e2e
  work in llm-d#15.

Signed-off-by: Greg Neighbors <26003+gkneighb@users.noreply.github.com>

Author: gkneighb

src/llm_d_latency_predictor/calibration.py

@@ -0,0 +1,80 @@
+"""Calibration-drift detector for continuous coverage evaluation.
+
+The decision logic is pure (no threading, sleeps, or servers) so it can be
+unit-tested with fixed inputs. `continuous_coverage_loop` feeds one coverage
+evaluation per interval into `CalibrationTrigger.update`, which returns True on
+the step that should request a retrain.
+
+See llm-d/llm-d-latency-predictor#19 / #32.
+"""
+
+from __future__ import annotations
+
+# Sentinel distinct from None, since last_retrain_time can legitimately be None
+# (model never retrained). _UNSET means "not waiting on a triggered retrain".
+_UNSET = object()
+
+
+class CalibrationTrigger:
+    """Detects sustained calibration drift and decides when to request a retrain.
+
+    Coverage is expected on a 0-100 scale (matching `quantile_coverage`), compared
+    against `target_pct` (e.g. 90 for p90) -- NOT a 0-1 fraction. An EMA of the
+    absolute deviation smooths transient noise; a retrain is requested only after
+    the EMA stays above `threshold` for `k` consecutive evaluations.
+
+    After firing, the detector suppresses further firing until the requested
+    retrain lands (detected by a change in `last_retrain_time`), then resets its
+    EMA and counter so post-retrain history starts clean. This prevents a slow
+    retrain from queuing back-to-back retrains while the model is still drifted.
+    """
+
+    def __init__(self, target_pct: float, threshold: float, k: int, ema_alpha: float):
+        self.target_pct = target_pct
+        self.threshold = threshold
+        self.k = k
+        self.ema_alpha = ema_alpha
+        self.ttft_ema = 0.0
+        self.tpot_ema = 0.0
+        self.consecutive_bad = 0
+        self._pending_retrain = _UNSET
+
+    @property
+    def max_dev(self) -> float:
+        return max(self.ttft_ema, self.tpot_ema)
+
+    @property
+    def awaiting_retrain(self) -> bool:
+        return self._pending_retrain is not _UNSET
+
+    def update(self, ttft_cov: float | None, tpot_cov: float | None, last_retrain_time) -> bool:
+        """Fold in one coverage evaluation. Returns True if a retrain should fire now."""
+        a = self.ema_alpha
+        if ttft_cov is not None:
+            self.ttft_ema = (1 - a) * self.ttft_ema + a * abs(ttft_cov - self.target_pct)
+        if tpot_cov is not None:
+            self.tpot_ema = (1 - a) * self.tpot_ema + a * abs(tpot_cov - self.target_pct)
+
+        # A triggered retrain is in flight: keep tracking the EMA (for /metrics) but
+        # do not re-fire. Once the retrain lands, reset so history starts fresh.
+        if self._pending_retrain is not _UNSET:
+            if last_retrain_time != self._pending_retrain:
+                self.reset()
+            return False
+
+        if self.max_dev > self.threshold:
+            self.consecutive_bad += 1
+            if self.consecutive_bad >= self.k:
+                self.consecutive_bad = 0
+                self._pending_retrain = last_retrain_time
+                return True
+        else:
+            self.consecutive_bad = 0
+        return False
+
+    def reset(self) -> None:
+        """Clear EMA, counter, and pending-retrain state."""
+        self.ttft_ema = 0.0
+        self.tpot_ema = 0.0
+        self.consecutive_bad = 0
+        self._pending_retrain = _UNSET

src/llm_d_latency_predictor/training_server.py

@@ -32,6 +32,8 @@
 from sklearn.linear_model import BayesianRidge
 from sklearn.preprocessing import StandardScaler
 
+from llm_d_latency_predictor.calibration import CalibrationTrigger
+
 try:
     import xgboost as xgb
 
@@ -1918,42 +1920,35 @@ def continuous_coverage_loop():
     # Wait long enough for the first train() to populate models + initial coverage.
     time.sleep(15)
     target_pct = settings.QUANTILE_ALPHA * 100
-    alpha = settings.CALIBRATION_EMA_ALPHA
-    ttft_dev_ema = 0.0
-    tpot_dev_ema = 0.0
-    consecutive_bad = 0
+    trigger = CalibrationTrigger(
+        target_pct=target_pct,
+        threshold=settings.CALIBRATION_TRIGGER_THRESHOLD,
+        k=settings.CALIBRATION_TRIGGER_K,
+        ema_alpha=settings.CALIBRATION_EMA_ALPHA,
+    )
 
     logging.info(
         f"Continuous coverage loop started "
         f"(eval_interval={settings.COVERAGE_EVAL_INTERVAL_SEC}s, "
         f"target={target_pct:.0f}%, threshold={settings.CALIBRATION_TRIGGER_THRESHOLD:.2f}pp, "
-        f"k={settings.CALIBRATION_TRIGGER_K}, ema_alpha={alpha:.2f})"
+        f"k={settings.CALIBRATION_TRIGGER_K}, ema_alpha={settings.CALIBRATION_EMA_ALPHA:.2f})"
     )
 
     while not predictor._shutdown_event.is_set():
         try:
             ttft_cov, tpot_cov = predictor.evaluate_current_coverage()
-            if ttft_cov is not None:
-                ttft_dev_ema = (1 - alpha) * ttft_dev_ema + alpha * abs(ttft_cov - target_pct)
-            if tpot_cov is not None:
-                tpot_dev_ema = (1 - alpha) * tpot_dev_ema + alpha * abs(tpot_cov - target_pct)
-            max_dev = max(ttft_dev_ema, tpot_dev_ema)
-            if max_dev > settings.CALIBRATION_TRIGGER_THRESHOLD:
-                consecutive_bad += 1
-                logging.info(
-                    f"Coverage drift: ttft_cov={ttft_cov} tpot_cov={tpot_cov} "
-                    f"ttft_dev_ema={ttft_dev_ema:.2f} tpot_dev_ema={tpot_dev_ema:.2f} "
-                    f"consecutive_bad={consecutive_bad}/{settings.CALIBRATION_TRIGGER_K}"
+            fired = trigger.update(ttft_cov, tpot_cov, predictor.last_retrain_time)
+            logging.info(
+                f"Coverage drift: ttft_cov={ttft_cov} tpot_cov={tpot_cov} "
+                f"max_dev={trigger.max_dev:.2f} consecutive_bad={trigger.consecutive_bad}/{settings.CALIBRATION_TRIGGER_K}"
+                f"{' (awaiting retrain)' if trigger.awaiting_retrain else ''}"
+            )
+            if fired:
+                logging.warning(
+                    f"Calibration trigger fired (max_dev={trigger.max_dev:.2f}pp > "
+                    f"threshold={settings.CALIBRATION_TRIGGER_THRESHOLD:.2f}pp). Requesting immediate retrain."
                 )
-                if consecutive_bad >= settings.CALIBRATION_TRIGGER_K:
-                    logging.warning(
-                        f"Calibration trigger fired (max_dev={max_dev:.2f}pp > "
-                        f"threshold={settings.CALIBRATION_TRIGGER_THRESHOLD:.2f}pp). Requesting immediate retrain."
-                    )
-                    predictor._calibration_trigger.set()
-                    consecutive_bad = 0  # one trigger per drift event
-            else:
-                consecutive_bad = 0
+                predictor._calibration_trigger.set()
         except Exception:
             logging.error("Error in continuous coverage loop", exc_info=True)
         predictor._shutdown_event.wait(timeout=settings.COVERAGE_EVAL_INTERVAL_SEC)

tests/drift_harness/README.md

@@ -0,0 +1,27 @@
+# Calibration drift harness
+
+Manual harness for exercising continuous coverage evaluation and the
+calibration-triggered retrain end-to-end against a live `training_server`.
+Committed here so the automated CI/e2e work (#15) can reuse it rather than
+re-derive it. The unit tests for the trigger *logic* live in
+`tests/test_calibration_trigger.py` and need no server.
+
+## Scripts
+
+- `synth_workload.py` — drives `/add_training_data_bulk` with a synthetic
+  linear-noise latency model, switching regimes at `--drift-at` to inject drift.
+  Drift modes: `slope` (multiply slope coefficients — dispersion-dominant on the
+  heavy-tailed features), `location` (shift intercepts), `noise` (scale sigma).
+- `scrape_metrics.py` — polls `/metrics` and records the coverage (and
+  quantile-loss) series to CSV over time.
+
+## Example
+
+```bash
+# terminal 1: run the server with the coverage loop enabled
+LATENCY_COVERAGE_EVAL_INTERVAL_SEC=5 uvicorn llm_d_latency_predictor.training_server:app --port 8000
+
+# terminal 2: baseline for 90s, then inject drift; watch the trigger fire
+python tests/drift_harness/synth_workload.py --duration 300 --drift-at 90 --rate 150
+python tests/drift_harness/scrape_metrics.py --duration 300 --output coverage.csv
+```

tests/drift_harness/scrape_metrics.py

@@ -0,0 +1,113 @@
+"""
+Polls the training_server's /metrics endpoint, extracts the
+ttft_coverage_percent and tpot_coverage_percent values (each a deque of up
+to 5 most-recent test-set coverage measurements), and appends one row per
+poll to a CSV.
+
+Output columns:
+  elapsed_sec, ttft_cov_latest, ttft_cov_mean5, tpot_cov_latest, tpot_cov_mean5
+
+Each coverage value is a fraction in [0, 1] (the actual fraction of test
+samples for which y_true <= y_pred). For a p90 model the well-calibrated
+value is 0.9.
+
+Usage:
+  python scrape_metrics.py --duration 1800 --output coverage.csv
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import re
+import sys
+import time
+import urllib.request
+
+COV_RE = re.compile(
+    r'(?P<name>(?:ttft|tpot)_coverage_percent)\{idx="(?P<idx>\d+)"\}\s+(?P<value>[\d.eE+-]+)'
+)
+QL_RE = re.compile(
+    r'(?P<name>(?:ttft|tpot)_quantile_loss)\{idx="(?P<idx>\d+)"\}\s+(?P<value>[\d.eE+-]+)'
+)
+
+
+def fetch_metrics(url: str) -> str:
+    with urllib.request.urlopen(url, timeout=5) as r:
+        return r.read().decode()
+
+
+def _parse_series(body: str, pattern: re.Pattern, ttft_name: str) -> tuple[list[float], list[float]]:
+    """Return (ttft_values, tpot_values) for a metric family, ordered by idx (oldest→newest)."""
+    ttft: dict[int, float] = {}
+    tpot: dict[int, float] = {}
+    for m in pattern.finditer(body):
+        idx = int(m["idx"])
+        value = float(m["value"])
+        if m["name"] == ttft_name:
+            ttft[idx] = value
+        else:
+            tpot[idx] = value
+    return (
+        [ttft[i] for i in sorted(ttft)],
+        [tpot[i] for i in sorted(tpot)],
+    )
+
+
+def parse_coverage(body: str) -> tuple[list[float], list[float]]:
+    return _parse_series(body, COV_RE, "ttft_coverage_percent")
+
+
+def parse_quantile_loss(body: str) -> tuple[list[float], list[float]]:
+    return _parse_series(body, QL_RE, "ttft_quantile_loss")
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--server", default="http://localhost:8000")
+    ap.add_argument("--duration", type=float, default=1800.0)
+    ap.add_argument("--interval", type=float, default=5.0)
+    ap.add_argument("--output", default="coverage.csv")
+    args = ap.parse_args()
+
+    url = args.server.rstrip("/") + "/metrics"
+    start = time.time()
+
+    def latest(xs):
+        return xs[-1] if xs else ""
+
+    def mean5(xs):
+        return sum(xs) / len(xs) if xs else ""
+
+    with open(args.output, "w", newline="") as f:
+        w = csv.writer(f)
+        w.writerow([
+            "elapsed_sec",
+            "ttft_cov_latest", "ttft_cov_mean5", "tpot_cov_latest", "tpot_cov_mean5",
+            "ttft_ql_latest", "ttft_ql_mean5", "tpot_ql_latest", "tpot_ql_mean5",
+        ])
+        while True:
+            elapsed = time.time() - start
+            if elapsed > args.duration:
+                break
+            try:
+                body = fetch_metrics(url)
+                ttft_c, tpot_c = parse_coverage(body)
+                ttft_q, tpot_q = parse_quantile_loss(body)
+            except Exception as e:
+                print(f"[{elapsed:7.1f}s] scrape error: {e}", file=sys.stderr, flush=True)
+                ttft_c, tpot_c, ttft_q, tpot_q = [], [], [], []
+            w.writerow([
+                f"{elapsed:.1f}",
+                latest(ttft_c), mean5(ttft_c), latest(tpot_c), mean5(tpot_c),
+                latest(ttft_q), mean5(ttft_q), latest(tpot_q), mean5(tpot_q),
+            ])
+            f.flush()
+            time.sleep(args.interval)
+
+    print(f"wrote {args.output}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())