diff --git a/index.html b/index.html
index 944d411..4a2ffc0 100755
--- a/index.html
+++ b/index.html
@@ -92,7 +92,7 @@ Linux Kernel Ring 0 Visualization
-
+
diff --git a/kernel_ai/ml/config.py b/kernel_ai/ml/config.py
index 67a63ae..0995664 100644
--- a/kernel_ai/ml/config.py
+++ b/kernel_ai/ml/config.py
@@ -87,6 +87,22 @@ class MLConfig:
# sustained anomaly).
if_cooldown_sec: float = _env_float("KERNEL_AI_ML_IF_COOLDOWN_SEC", 15.0)
+ # --- Stage 3 (drift + auto-retrain) ---
+ # Window of recent snapshots used to measure drift (minutes).
+ drift_window_min: int = _env_int("KERNEL_AI_ML_DRIFT_WINDOW_MIN", 30)
+ # Drift trips when the live flag rate exceeds expected (contamination) by
+ # this multiple, or when the mean per-feature distribution shift (in train
+ # std units) exceeds the z threshold.
+ drift_rate_mult: float = _env_float("KERNEL_AI_ML_DRIFT_RATE_MULT", 5.0)
+ drift_feature_z: float = _env_float("KERNEL_AI_ML_DRIFT_FEATURE_Z", 3.0)
+ # Poison guard: exclude snapshots within +/- this many seconds of a
+ # high-severity anomaly from the retraining set (don't learn the attack).
+ poison_guard_sec: int = _env_int("KERNEL_AI_ML_POISON_GUARD_SEC", 120)
+ # A retrained model is rejected (kept previous) if its training flag rate is
+ # outside this sane band (degenerate model: flags nothing or everything).
+ retrain_min_flag_rate: float = _env_float("KERNEL_AI_ML_RETRAIN_MIN_FLAG", 0.001)
+ retrain_max_flag_rate: float = _env_float("KERNEL_AI_ML_RETRAIN_MAX_FLAG", 0.30)
+
@property
def alpha(self) -> float:
return 2.0 / (max(2, self.baseline_window) + 1.0)
diff --git a/kernel_ai/ml/drift.py b/kernel_ai/ml/drift.py
new file mode 100644
index 0000000..1c144a3
--- /dev/null
+++ b/kernel_ai/ml/drift.py
@@ -0,0 +1,131 @@
+"""Stage 3 drift monitor.
+
+Measures, without any labels, whether the live system has drifted away from the
+"normal" the current model was trained on. Two complementary signals:
+
+ 1. flag_rate - fraction of recent snapshots the model flags as anomalous.
+ If this sits far above the model's `contamination` for a
+ sustained period (without a real incident), normal has moved.
+ 2. feature_drift - mean per-feature shift of recent data vs the training
+ distribution, measured in training-std units (a PSI-like z).
+
+Either signal crossing its threshold marks `drifted = True`, which the retrain
+orchestrator uses to decide whether to refit.
+"""
+
+from __future__ import annotations
+
+import logging
+import statistics
+
+from kernel_ai.ml.config import MLConfig
+from kernel_ai.ml.features import FEATURE_SPECS
+from kernel_ai.ml.model import IsolationForestModel
+from kernel_ai.ml.store import fetch_recent_feature_dicts, insert_drift
+
+logger = logging.getLogger("kernel_ai.ml.drift")
+
+# Cap each feature's drift z so a single quiet, near-constant metric can't
+# dominate the aggregate (a tiny absolute move on a ~0-variance feature would
+# otherwise blow up to astronomical z).
+_MAX_FEATURE_Z = 25.0
+
+
+def _feature_drift(recent: list[dict], feature_stats: dict[str, dict]) -> tuple[float, dict]:
+ """Mean absolute shift of recent feature means vs training means, in train
+ std units. Returns (aggregate_score, per_feature_detail).
+
+ The denominator floor uses the same per-feature noise floor as the Stage 1
+ z-score (FEATURE_SPECS[...].min_std), so quiet features need a *meaningful*
+ move — not a microscopic one — to register as drift.
+ """
+ if not recent or not feature_stats:
+ return 0.0, {}
+ per_feature = {}
+ z_values = []
+ for name, st in feature_stats.items():
+ vals = [float(r.get(name, 0.0)) for r in recent]
+ if not vals:
+ continue
+ recent_mean = statistics.fmean(vals)
+ train_mean = float(st.get("mean", 0.0))
+ train_std = float(st.get("std", 0.0))
+ spec = FEATURE_SPECS.get(name)
+ noise_floor = spec.min_std if spec else 0.0
+ floor = max(train_std, noise_floor, abs(train_mean) * 0.05, 1e-6)
+ z = min(_MAX_FEATURE_Z, abs(recent_mean - train_mean) / floor)
+ per_feature[name] = round(z, 3)
+ z_values.append(z)
+ aggregate = statistics.fmean(z_values) if z_values else 0.0
+ return aggregate, per_feature
+
+
+def compute_drift(cfg: MLConfig | None = None, *, persist: bool = True) -> dict:
+ cfg = cfg or MLConfig()
+ try:
+ model = IsolationForestModel.load(cfg.model_path)
+ except Exception as exc: # noqa: BLE001
+ logger.warning("drift: no model to compare against (%s)", exc)
+ return {"available": False, "reason": "no_model"}
+
+ recent = fetch_recent_feature_dicts(cfg.dsn, minutes=cfg.drift_window_min)
+ n = len(recent)
+ contamination = float(model.meta.get("contamination", cfg.if_contamination))
+
+ if n == 0:
+ result = {
+ "available": True, "n_recent": 0, "flag_rate": 0.0,
+ "expected_rate": contamination, "feature_drift": 0.0, "drifted": False,
+ "detail": {"reason": "no_recent_data"},
+ }
+ if persist:
+ insert_drift(cfg.dsn, result)
+ return result
+
+ matrix = [[float(r.get(name, 0.0)) for name in model.feature_names] for r in recent]
+ preds = model.model.predict(matrix)
+ flagged = sum(1 for p in preds if p == -1)
+ flag_rate = flagged / n
+
+ feature_drift, per_feature = _feature_drift(recent, model.meta.get("feature_stats", {}))
+
+ rate_drift = flag_rate > contamination * cfg.drift_rate_mult
+ dist_drift = feature_drift > cfg.drift_feature_z
+ drifted = bool(rate_drift or dist_drift)
+
+ # Surface the biggest-shifting features for explainability.
+ top = sorted(per_feature.items(), key=lambda kv: kv[1], reverse=True)[:5]
+
+ result = {
+ "available": True,
+ "n_recent": n,
+ "flag_rate": round(flag_rate, 4),
+ "expected_rate": round(contamination, 4),
+ "feature_drift": round(feature_drift, 4),
+ "drifted": drifted,
+ "detail": {
+ "rate_drift": rate_drift,
+ "dist_drift": dist_drift,
+ "rate_mult_threshold": cfg.drift_rate_mult,
+ "feature_z_threshold": cfg.drift_feature_z,
+ "top_features": dict(top),
+ },
+ }
+ if persist:
+ insert_drift(cfg.dsn, result)
+ logger.info(
+ "drift: n=%d flag_rate=%.3f (exp %.3f) feature_drift=%.2f drifted=%s",
+ n, flag_rate, contamination, feature_drift, drifted,
+ )
+ return result
+
+
+def main() -> None:
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
+ import json
+
+ print(json.dumps(compute_drift(), indent=2))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/kernel_ai/ml/retrain.py b/kernel_ai/ml/retrain.py
new file mode 100644
index 0000000..2b64f09
--- /dev/null
+++ b/kernel_ai/ml/retrain.py
@@ -0,0 +1,65 @@
+"""Stage 3 auto-retrain orchestrator (run by a systemd timer).
+
+ measure drift -> decide -> retrain on clean recent data -> register
+
+Decision:
+ * default: always retrain (the timer cadence IS the schedule).
+ * --only-if-drift: retrain only when the drift monitor trips (use with a
+ more frequent timer to react faster without needless refits).
+
+Safety:
+ * training excludes high-severity anomaly windows (poison guard), and
+ * a degenerate model (flags ~nothing/~everything) is rejected, keeping the
+ previous one. Both live in train.py; this just orchestrates and logs.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+
+from kernel_ai.ml.config import MLConfig
+from kernel_ai.ml.drift import compute_drift
+from kernel_ai.ml import train as train_mod
+
+logger = logging.getLogger("kernel_ai.ml.retrain")
+
+
+def run(*, only_if_drift: bool, min_samples: int) -> int:
+ cfg = MLConfig()
+ drift = compute_drift(cfg, persist=True)
+
+ if only_if_drift and drift.get("available") and not drift.get("drifted"):
+ logger.info("no drift detected (flag_rate=%s) - skipping retrain", drift.get("flag_rate"))
+ return 0
+
+ try:
+ metrics = train_mod.train(
+ cfg,
+ min_samples=min_samples,
+ contamination=cfg.if_contamination,
+ trees=cfg.if_n_estimators,
+ exclude_anomalous=True,
+ enforce_guardrails=True,
+ )
+ except SystemExit as exc:
+ # Soft-skip (not enough clean data, or guardrail tripped): keep previous
+ # model and don't mark the timer run as failed.
+ logger.warning("retrain skipped: %s", exc)
+ return 0
+
+ logger.info("retrain complete: %s", metrics)
+ return 0
+
+
+def main() -> None:
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
+ parser = argparse.ArgumentParser(description="Drift-aware auto-retrain")
+ parser.add_argument("--only-if-drift", action="store_true")
+ parser.add_argument("--min-samples", type=int, default=100)
+ args = parser.parse_args()
+ raise SystemExit(run(only_if_drift=args.only_if_drift, min_samples=args.min_samples))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/kernel_ai/ml/store.py b/kernel_ai/ml/store.py
index 571c36c..c8443af 100644
--- a/kernel_ai/ml/store.py
+++ b/kernel_ai/ml/store.py
@@ -51,6 +51,17 @@
count integer NOT NULL,
updated_at timestamptz NOT NULL DEFAULT now()
);
+
+CREATE TABLE IF NOT EXISTS ml_drift (
+ ts timestamptz NOT NULL DEFAULT now(),
+ flag_rate double precision,
+ expected_rate double precision,
+ feature_drift double precision,
+ n_recent integer,
+ drifted boolean,
+ detail jsonb
+);
+CREATE INDEX IF NOT EXISTS ml_drift_ts_idx ON ml_drift (ts);
"""
@@ -154,6 +165,78 @@ def close(self) -> None:
pass
+def fetch_recent_feature_dicts(dsn: str, *, minutes: int = 30, limit: int = 5000) -> list[dict]:
+ """Recent feature snapshots (as dicts) for drift measurement."""
+ with connect(dsn) as conn, conn.cursor() as cur:
+ cur.execute(
+ """
+ SELECT features FROM ml_feature_snapshots
+ WHERE ts > now() - make_interval(mins => %s)
+ ORDER BY ts DESC LIMIT %s
+ """,
+ (minutes, limit),
+ )
+ return [r[0] for r in cur.fetchall() if isinstance(r[0], dict)]
+
+
+def fetch_training_snapshots(
+ dsn: str,
+ *,
+ limit: int = 50000,
+ exclude_anomalous: bool = True,
+ guard_sec: int = 120,
+) -> list[dict]:
+ """Snapshots for (re)training. With ``exclude_anomalous`` we drop any
+ snapshot within +/- ``guard_sec`` of a HIGH-severity anomaly, so a real
+ incident can't poison the model's idea of "normal"."""
+ with connect(dsn) as conn, conn.cursor() as cur:
+ if exclude_anomalous:
+ cur.execute(
+ """
+ SELECT s.features
+ FROM ml_feature_snapshots s
+ WHERE NOT EXISTS (
+ SELECT 1 FROM ml_anomalies a
+ WHERE a.severity = 'high'
+ AND a.ts BETWEEN s.ts - make_interval(secs => %s)
+ AND s.ts + make_interval(secs => %s)
+ )
+ ORDER BY s.ts DESC LIMIT %s
+ """,
+ (guard_sec, guard_sec, limit),
+ )
+ else:
+ cur.execute(
+ "SELECT features FROM ml_feature_snapshots ORDER BY ts DESC LIMIT %s",
+ (limit,),
+ )
+ return [r[0] for r in cur.fetchall() if isinstance(r[0], dict)]
+
+
+def insert_drift(dsn: str, record: dict) -> None:
+ """Persist one drift measurement (best-effort)."""
+ try:
+ with connect(dsn) as conn, conn.cursor() as cur:
+ cur.execute(
+ """
+ INSERT INTO ml_drift
+ (flag_rate, expected_rate, feature_drift, n_recent, drifted, detail)
+ VALUES (%(flag_rate)s, %(expected_rate)s, %(feature_drift)s,
+ %(n_recent)s, %(drifted)s, %(detail)s)
+ """,
+ {
+ "flag_rate": record.get("flag_rate"),
+ "expected_rate": record.get("expected_rate"),
+ "feature_drift": record.get("feature_drift"),
+ "n_recent": record.get("n_recent"),
+ "drifted": record.get("drifted"),
+ "detail": Json(record.get("detail") or {}),
+ },
+ )
+ except Exception as exc: # noqa: BLE001
+ logger.warning("insert_drift failed: %s", exc)
+
+
def fetch_recent_anomalies(dsn: str, *, since_seconds: int = 120, limit: int = 100) -> list[dict]:
"""Read recent anomalies for the API. Never raises on DB trouble."""
try:
diff --git a/kernel_ai/ml/train.py b/kernel_ai/ml/train.py
index ffada2e..213ea1d 100644
--- a/kernel_ai/ml/train.py
+++ b/kernel_ai/ml/train.py
@@ -23,7 +23,7 @@
from kernel_ai.ml.config import MLConfig
from kernel_ai.ml.features import FEATURE_SPECS
from kernel_ai.ml.model import IsolationForestModel
-from kernel_ai.ml.store import connect
+from kernel_ai.ml.store import fetch_training_snapshots
logger = logging.getLogger("kernel_ai.ml.train")
@@ -31,20 +31,20 @@
FEATURE_ORDER = list(FEATURE_SPECS.keys())
-def load_matrix(dsn: str, limit: int = 50000) -> list[list[float]]:
- """Load feature snapshots into a dense matrix in canonical column order."""
- with connect(dsn) as conn, conn.cursor() as cur:
- cur.execute(
- "SELECT features FROM ml_feature_snapshots ORDER BY ts DESC LIMIT %s",
- (limit,),
- )
- rows = cur.fetchall()
- matrix = []
- for (features,) in rows:
- if not isinstance(features, dict):
- continue
- matrix.append([float(features.get(name, 0.0)) for name in FEATURE_ORDER])
- return matrix
+def _dicts_to_matrix(rows: list[dict]) -> list[list[float]]:
+ return [[float(r.get(name, 0.0)) for name in FEATURE_ORDER] for r in rows]
+
+
+def _feature_stats(matrix: list[list[float]]) -> dict[str, dict]:
+ """Per-feature mean/std over the training set, stored in the model meta so
+ the drift monitor can later compare live data against the trained 'normal'."""
+ stats: dict[str, dict] = {}
+ for col, name in enumerate(FEATURE_ORDER):
+ values = [row[col] for row in matrix]
+ mean = statistics.fmean(values) if values else 0.0
+ std = statistics.pstdev(values) if len(values) > 1 else 0.0
+ stats[name] = {"mean": mean, "std": std}
+ return stats
def _percentile(values: list[float], pct: float) -> float:
@@ -55,36 +55,61 @@ def _percentile(values: list[float], pct: float) -> float:
return ordered[k]
-def train(cfg: MLConfig, *, min_samples: int, contamination: float, trees: int) -> dict:
- matrix = load_matrix(cfg.dsn)
+def train(
+ cfg: MLConfig,
+ *,
+ min_samples: int,
+ contamination: float,
+ trees: int,
+ exclude_anomalous: bool = True,
+ enforce_guardrails: bool = True,
+) -> dict:
+ rows = fetch_training_snapshots(
+ cfg.dsn,
+ exclude_anomalous=exclude_anomalous,
+ guard_sec=cfg.poison_guard_sec,
+ )
+ matrix = _dicts_to_matrix(rows)
n = len(matrix)
if n < min_samples:
raise SystemExit(
- f"Not enough samples to train: have {n}, need >= {min_samples}. "
+ f"Not enough clean samples to train: have {n}, need >= {min_samples}. "
f"Let the worker collect more (it runs every {cfg.interval_sec:.0f}s)."
)
model = IsolationForestModel(feature_names=FEATURE_ORDER)
model.fit(matrix, contamination=contamination, n_estimators=trees)
+ # Stash training distribution so drift can be measured against it later.
+ model.meta["feature_stats"] = _feature_stats(matrix)
# Evaluate on the training set: how the score is distributed and what
# fraction would be flagged (should sit near `contamination`).
scores = model.score_matrix(matrix)
preds = model.model.predict(matrix)
flagged = sum(1 for p in preds if p == -1)
+ flag_rate = flagged / n if n else 0.0
metrics = {
"n_samples": float(n),
"n_features": float(len(FEATURE_ORDER)),
- "flag_rate": flagged / n if n else 0.0,
+ "flag_rate": flag_rate,
"score_mean": statistics.fmean(scores) if scores else 0.0,
"score_p95": _percentile(scores, 95),
"score_max": max(scores) if scores else 0.0,
+ "excluded_anomalous": 1.0 if exclude_anomalous else 0.0,
}
+ # Guardrail: refuse to promote a degenerate model (flags ~nothing or
+ # ~everything). The previously saved artifact is left untouched.
+ if enforce_guardrails and not (cfg.retrain_min_flag_rate <= flag_rate <= cfg.retrain_max_flag_rate):
+ raise SystemExit(
+ f"Refusing to save model: flag_rate={flag_rate:.3f} outside sane band "
+ f"[{cfg.retrain_min_flag_rate}, {cfg.retrain_max_flag_rate}]. Keeping previous model."
+ )
+
# Persist the artifact the worker loads (independent of MLflow availability).
os.makedirs(os.path.dirname(cfg.model_path), exist_ok=True)
model.save(cfg.model_path)
- logger.info("saved model artifact -> %s", cfg.model_path)
+ logger.info("saved model artifact -> %s (flag_rate=%.3f, n=%d)", cfg.model_path, flag_rate, n)
# MLflow tracking + registry (best-effort: training must not hard-fail if
# MLflow has a hiccup, the artifact is already saved above).