diff --git a/kernel_ai/ml/config.py b/kernel_ai/ml/config.py index 7192130..67a63ae 100644 --- a/kernel_ai/ml/config.py +++ b/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) diff --git a/kernel_ai/ml/model.py b/kernel_ai/ml/model.py new file mode 100644 index 0000000..ddf9509 --- /dev/null +++ b/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", {}), + ) diff --git a/kernel_ai/ml/train.py b/kernel_ai/ml/train.py new file mode 100644 index 0000000..ffada2e --- /dev/null +++ b/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() diff --git a/kernel_ai/ml/worker.py b/kernel_ai/ml/worker.py index 8cb6506..fd132fa 100644 --- a/kernel_ai/ml/worker.py +++ b/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. diff --git a/requirements.txt b/requirements.txt index cf20b81..1c0ff5a 100755 --- a/requirements.txt +++ b/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