kernel_dna_isoforest

kernel_ai/ml/config.py

@@ -8,6 +8,10 @@
 
 import os
 from dataclasses import dataclass
+from pathlib import Path
+
+# Project root (parent of the ``kernel_ai`` package), used for default paths.
+_PROJECT_ROOT = Path(__file__).resolve().parents[2]
 
 
 def _env_float(name: str, default: float) -> float:
@@ -59,6 +63,30 @@ class MLConfig:
     retain_features_hours: int = _env_int("KERNEL_AI_ML_RETAIN_FEATURES_H", 48)
     retain_anomalies_hours: int = _env_int("KERNEL_AI_ML_RETAIN_ANOMALIES_H", 168)
 
+    # --- Stage 2 (IsolationForest) ---
+    # Enable the second-opinion model in the worker. If the model file is
+    # missing, Stage 2 stays dormant regardless of this flag.
+    enable_stage2: bool = os.getenv("KERNEL_AI_ML_STAGE2", "true").lower() == "true"
+    # Saved model artifact the worker loads (decoupled from MLflow).
+    model_path: str = os.getenv(
+        "KERNEL_AI_ML_MODEL_PATH",
+        str(_PROJECT_ROOT / "models" / "isoforest_latest.joblib"),
+    )
+    # MLflow tracking store. sqlite gives both tracking + a model registry while
+    # staying file-light (no always-on server). Override on PROD if desired.
+    mlflow_uri: str = os.getenv(
+        "KERNEL_AI_MLFLOW_URI",
+        f"sqlite:///{_PROJECT_ROOT / 'mlflow.db'}",
+    )
+    mlflow_experiment: str = os.getenv("KERNEL_AI_MLFLOW_EXPERIMENT", "kernel_dna_anomaly")
+    mlflow_model_name: str = os.getenv("KERNEL_AI_MLFLOW_MODEL", "kernel_dna_isoforest")
+    # IsolationForest training defaults.
+    if_contamination: float = _env_float("KERNEL_AI_ML_IF_CONTAMINATION", 0.02)
+    if_n_estimators: int = _env_int("KERNEL_AI_ML_IF_TREES", 200)
+    # Min seconds between IsolationForest mutations (avoid per-tick spam during a
+    # sustained anomaly).
+    if_cooldown_sec: float = _env_float("KERNEL_AI_ML_IF_COOLDOWN_SEC", 15.0)
+
     @property
     def alpha(self) -> float:
         return 2.0 / (max(2, self.baseline_window) + 1.0)

kernel_ai/ml/model.py

@@ -0,0 +1,90 @@
+"""Stage 2 model: scikit-learn IsolationForest wrapper.
+
+IsolationForest learns the *shape* of normal multivariate behaviour and flags
+points that are easy to "isolate" (few random splits) as anomalies. Unlike the
+Stage 1 per-feature z-score, it can catch anomalies that only show up as an
+unusual *combination* of features (e.g. high syscalls + low CPU + high retrans).
+
+This module is deliberately MLflow-free: it only knows how to fit / save / load /
+score. Experiment tracking lives in :mod:`kernel_ai.ml.train`; the inference
+worker just loads the saved artifact via joblib. Tree-based -> no feature
+scaling needed, which keeps the artifact simple.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import joblib
+
+
+@dataclass
+class IsolationForestModel:
+    feature_names: list[str]
+    model: Any = None
+    meta: dict = field(default_factory=dict)
+
+    def fit(
+        self,
+        matrix: list[list[float]],
+        *,
+        contamination: float = 0.02,
+        n_estimators: int = 200,
+        random_state: int = 42,
+    ) -> "IsolationForestModel":
+        from sklearn.ensemble import IsolationForest
+
+        clf = IsolationForest(
+            n_estimators=n_estimators,
+            contamination=contamination,
+            random_state=random_state,
+            n_jobs=1,
+        )
+        clf.fit(matrix)
+        self.model = clf
+        self.meta.update(
+            {
+                "contamination": contamination,
+                "n_estimators": n_estimators,
+                "n_samples": len(matrix),
+                "n_features": len(self.feature_names),
+            }
+        )
+        return self
+
+    def _vectorize(self, features: dict[str, float]) -> list[float]:
+        return [float(features.get(name, 0.0)) for name in self.feature_names]
+
+    def score_one(self, features: dict[str, float]) -> tuple[bool, float]:
+        """Return (is_anomaly, anomaly_score).
+
+        anomaly_score is ``-decision_function``: higher = more anomalous, and
+        crosses 0 at the model's learned normal/anomalous boundary.
+        """
+        if self.model is None:
+            return False, 0.0
+        vec = [self._vectorize(features)]
+        is_anomaly = bool(self.model.predict(vec)[0] == -1)
+        anomaly_score = float(-self.model.decision_function(vec)[0])
+        return is_anomaly, anomaly_score
+
+    def score_matrix(self, matrix: list[list[float]]) -> list[float]:
+        if self.model is None or not matrix:
+            return []
+        return [float(-s) for s in self.model.decision_function(matrix)]
+
+    def save(self, path: str) -> None:
+        joblib.dump(
+            {"feature_names": self.feature_names, "model": self.model, "meta": self.meta},
+            path,
+        )
+
+    @classmethod
+    def load(cls, path: str) -> "IsolationForestModel":
+        blob = joblib.load(path)
+        return cls(
+            feature_names=blob["feature_names"],
+            model=blob["model"],
+            meta=blob.get("meta", {}),
+        )

kernel_ai/ml/train.py

@@ -0,0 +1,140 @@
+"""Stage 2 training: fit IsolationForest on collected feature snapshots.
+
+Flow:
+    Postgres ml_feature_snapshots  ->  matrix  ->  IsolationForest
+        -> evaluate (flag rate, score distribution)
+        -> log params/metrics/model to MLflow  (sqlite backend = tracking + registry)
+        -> save artifact to models/isoforest_latest.joblib  (what the worker loads)
+
+MLflow runs in a *file-light* sqlite store, so there's no always-on server.
+Inspect runs on demand with:  mlflow ui --backend-store-uri sqlite:///mlflow.db
+
+Usage:
+    python -m kernel_ai.ml.train [--min-samples N] [--contamination C] [--trees T]
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+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 connect
+
+logger = logging.getLogger("kernel_ai.ml.train")
+
+# Canonical feature ordering shared by training and inference.
+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 _percentile(values: list[float], pct: float) -> float:
+    if not values:
+        return 0.0
+    ordered = sorted(values)
+    k = max(0, min(len(ordered) - 1, int(round((pct / 100.0) * (len(ordered) - 1)))))
+    return ordered[k]
+
+
+def train(cfg: MLConfig, *, min_samples: int, contamination: float, trees: int) -> dict:
+    matrix = load_matrix(cfg.dsn)
+    n = len(matrix)
+    if n < min_samples:
+        raise SystemExit(
+            f"Not enough 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)
+
+    # 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)
+    metrics = {
+        "n_samples": float(n),
+        "n_features": float(len(FEATURE_ORDER)),
+        "flag_rate": flagged / n if n else 0.0,
+        "score_mean": statistics.fmean(scores) if scores else 0.0,
+        "score_p95": _percentile(scores, 95),
+        "score_max": max(scores) if scores else 0.0,
+    }
+
+    # 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)
+
+    # MLflow tracking + registry (best-effort: training must not hard-fail if
+    # MLflow has a hiccup, the artifact is already saved above).
+    try:
+        import mlflow
+        import mlflow.sklearn
+
+        mlflow.set_tracking_uri(cfg.mlflow_uri)
+        mlflow.set_experiment(cfg.mlflow_experiment)
+        with mlflow.start_run() as run:
+            mlflow.log_params(
+                {
+                    "contamination": contamination,
+                    "n_estimators": trees,
+                    "n_samples": n,
+                    "n_features": len(FEATURE_ORDER),
+                    "baseline_window": cfg.baseline_window,
+                    "feature_order": ",".join(FEATURE_ORDER),
+                }
+            )
+            mlflow.log_metrics(metrics)
+            mlflow.sklearn.log_model(
+                model.model,
+                artifact_path="isoforest",
+                registered_model_name=cfg.mlflow_model_name,
+            )
+            logger.info("logged MLflow run %s (experiment=%s)", run.info.run_id, cfg.mlflow_experiment)
+    except Exception as exc:  # noqa: BLE001 - tracking is optional
+        logger.warning("MLflow logging skipped: %s", exc)
+
+    return metrics
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
+    cfg = MLConfig()
+    parser = argparse.ArgumentParser(description="Train Stage 2 IsolationForest")
+    parser.add_argument("--min-samples", type=int, default=100)
+    parser.add_argument("--contamination", type=float, default=cfg.if_contamination)
+    parser.add_argument("--trees", type=int, default=cfg.if_n_estimators)
+    args = parser.parse_args()
+
+    metrics = train(
+        cfg,
+        min_samples=args.min_samples,
+        contamination=args.contamination,
+        trees=args.trees,
+    )
+    logger.info("training done: %s", metrics)
+
+
+if __name__ == "__main__":
+    main()

kernel_ai/ml/worker.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import logging
+import os
 import signal
 import time
 
@@ -58,6 +59,39 @@ def _build_anomalies(scores: dict[str, Score], cfg: MLConfig) -> list[dict]:
     return out
 
 
+def _build_isoforest_anomaly(score: float, scores: dict[str, Score], cfg: MLConfig) -> dict:
+    """Build a mutation from an IsolationForest verdict.
+
+    The forest judges the whole feature vector, so we borrow the Stage 1
+    z-scores to point at the single most-deviating feature for the helix
+    position / subsystem and a human-readable cause.
+    """
+    top = max(scores.values(), key=lambda s: abs(s.z), default=None)
+    feature = top.name if top else "vector"
+    spec = FEATURE_SPECS.get(feature)
+    severity = "high" if score > 0.1 else "medium"
+    label = spec.label if spec else feature
+    return {
+        "source": "stage2_isoforest",
+        "feature": feature,
+        "subsystem": spec.subsystem if spec else None,
+        "type": f"isoforest:{feature}",
+        "severity": severity,
+        "score": round(score, 4),
+        "value": round(top.value, 3) if top else None,
+        "baseline_mean": round(top.mean, 3) if top else None,
+        "baseline_std": round(top.std, 3) if top else None,
+        "position": spec.position if spec else 0.5,
+        "message": (
+            f"IsolationForest flagged an unusual system state "
+            f"(top deviation: {label}, z={top.z:.1f}, if_score={score:.3f})"
+            if top else
+            f"IsolationForest flagged an unusual system state (if_score={score:.3f})"
+        ),
+        "meta": {"stage": 2, "if_score": round(score, 4)},
+    }
+
+
 class MLWorker:
     def __init__(self, cfg: MLConfig | None = None) -> None:
         self.cfg = cfg or MLConfig()
@@ -66,6 +100,35 @@ def __init__(self, cfg: MLConfig | None = None) -> None:
         self.store = PostgresStore(self.cfg.dsn)
         self._running = True
         self._min_std = {n: s.min_std for n, s in FEATURE_SPECS.items()}
+        # Stage 2 model (loaded lazily; absent until train.py has produced it).
+        self.model = None
+        self._model_mtime: float | None = None
+        self._last_if_emit = 0.0
+        self._maybe_load_model()
+
+    def _maybe_load_model(self) -> None:
+        """Load / hot-reload the IsolationForest artifact if present and changed.
+
+        Importing sklearn/joblib only happens once a model file exists, so a
+        Stage-1-only deployment never pays the memory cost.
+        """
+        if not self.cfg.enable_stage2:
+            return
+        path = self.cfg.model_path
+        try:
+            mtime = os.path.getmtime(path)
+        except OSError:
+            return  # no model yet -> Stage 2 stays dormant
+        if self._model_mtime is not None and mtime <= self._model_mtime:
+            return
+        try:
+            from kernel_ai.ml.model import IsolationForestModel
+
+            self.model = IsolationForestModel.load(path)
+            self._model_mtime = mtime
+            logger.info("loaded Stage 2 model: %s (%s)", path, self.model.meta)
+        except Exception as exc:  # noqa: BLE001 - keep running on Stage 1 only
+            logger.warning("failed to load Stage 2 model: %s", exc)
 
     def stop(self, *_args) -> None:
         self._running = False
@@ -76,6 +139,21 @@ def _tick(self) -> int:
             return 0
         scores = self.baseline.update_and_score(features, self._min_std)
         anomalies = _build_anomalies(scores, self.cfg)
+
+        # Stage 2 second opinion: the forest can catch unusual *combinations*
+        # the per-feature z-score misses. Rate-limited so a sustained anomaly
+        # doesn't spam one mutation per tick.
+        if self.model is not None:
+            try:
+                is_anom, if_score = self.model.score_one(features)
+            except Exception as exc:  # noqa: BLE001
+                is_anom, if_score = False, 0.0
+                logger.warning("isoforest scoring failed: %s", exc)
+            now = time.time()
+            if is_anom and (now - self._last_if_emit) >= self.cfg.if_cooldown_sec:
+                anomalies.append(_build_isoforest_anomaly(if_score, scores, self.cfg))
+                self._last_if_emit = now
+
         if self.cfg.store_features:
             self.store.insert_feature_snapshot(features)
         if anomalies:
@@ -108,6 +186,8 @@ def run(self) -> None:
                 if ticks % _HOUSEKEEPING_EVERY == 0:
                     self.store.save_baseline(self.baseline.export_state())
                     self.store.prune(self.cfg.retain_features_hours, self.cfg.retain_anomalies_hours)
+                    # Pick up a freshly retrained model without a restart.
+                    self._maybe_load_model()
             except Exception as exc:  # noqa: BLE001 - keep the loop alive
                 logger.exception("tick failed: %s", exc)
                 # Reconnect on DB hiccups rather than dying.

requirements.txt

@@ -6,3 +6,6 @@ python-dotenv==1.0.0
 gunicorn==21.2.0
 sentry-sdk
 psycopg[binary]==3.2.3
+scikit-learn==1.5.2
+joblib==1.4.2
+mlflow==2.17.2