From 29b2945c37f2f00eb42f7d5baa87016491b58806 Mon Sep 17 00:00:00 2001 From: Aleksei Fedorov Date: Mon, 22 Jun 2026 17:43:04 +0000 Subject: [PATCH] ML module init --- kernel_ai/api/rest.py | 1 + kernel_ai/http/api.py | 2 + kernel_ai/http/api_handlers/kernel.py | 35 ++++ kernel_ai/ml/__init__.py | 19 ++ kernel_ai/ml/__main__.py | 6 + kernel_ai/ml/baseline.py | 91 +++++++++ kernel_ai/ml/config.py | 64 +++++++ kernel_ai/ml/features.py | 258 ++++++++++++++++++++++++++ kernel_ai/ml/store.py | 182 ++++++++++++++++++ kernel_ai/ml/worker.py | 135 ++++++++++++++ requirements.txt | 1 + 11 files changed, 794 insertions(+) create mode 100644 kernel_ai/ml/__init__.py create mode 100644 kernel_ai/ml/__main__.py create mode 100644 kernel_ai/ml/baseline.py create mode 100644 kernel_ai/ml/config.py create mode 100644 kernel_ai/ml/features.py create mode 100644 kernel_ai/ml/store.py create mode 100644 kernel_ai/ml/worker.py diff --git a/kernel_ai/api/rest.py b/kernel_ai/api/rest.py index 42c2971..027e267 100644 --- a/kernel_ai/api/rest.py +++ b/kernel_ai/api/rest.py @@ -28,6 +28,7 @@ ("/proc-timeline-branches", "get_proc_timeline_branches", h.get_proc_timeline_branches, None), ("/execution-context", "get_execution_context", h.get_execution_context, None), ("/kernel-dna", "kernel_dna", h.kernel_dna, None), + ("/ml-anomalies", "ml_anomalies", h.ml_anomalies, None), ("/crypto-realtime", "crypto_realtime", h.crypto_realtime, None), ("/security-realtime", "security_realtime", h.security_realtime, None), ("/processes-realtime", "processes_realtime", h.processes_realtime, None), diff --git a/kernel_ai/http/api.py b/kernel_ai/http/api.py index afec3aa..6642670 100644 --- a/kernel_ai/http/api.py +++ b/kernel_ai/http/api.py @@ -6,6 +6,7 @@ io_pulse, kernel_data, kernel_dna, + ml_anomalies, nginx_files, process_kernel_map, sentry_test, @@ -63,6 +64,7 @@ "isolation_context", "kernel_data", "kernel_dna", + "ml_anomalies", "network_stack_realtime", "nginx_files", "process_kernel_map", diff --git a/kernel_ai/http/api_handlers/kernel.py b/kernel_ai/http/api_handlers/kernel.py index 646c772..f17db72 100644 --- a/kernel_ai/http/api_handlers/kernel.py +++ b/kernel_ai/http/api_handlers/kernel.py @@ -82,6 +82,41 @@ def kernel_dna(): return api_json(_telemetry.get_kernel_dna_data) +def ml_anomalies(): + """Recent ML-detected anomalies (Stage 1 baselines) for Kernel DNA. + + Read-only: the Flask app never computes anomalies, it only reads what the + isolated ML worker wrote to the shared store. If the store is unreachable, + the underlying reader returns an empty list rather than failing the page. + """ + + def _payload(): + from kernel_ai.ml.config import MLConfig + from kernel_ai.ml.store import fetch_recent_anomalies + + try: + since = int(request.args.get("since_seconds", 120)) + except (TypeError, ValueError): + since = 120 + since = max(5, min(since, 3600)) + try: + limit = int(request.args.get("limit", 100)) + except (TypeError, ValueError): + limit = 100 + limit = max(1, min(limit, 500)) + + cfg = MLConfig() + anomalies = fetch_recent_anomalies(cfg.dsn, since_seconds=since, limit=limit) + return { + "timestamp": datetime.now().isoformat(), + "since_seconds": since, + "count": len(anomalies), + "anomalies": anomalies, + } + + return api_json(_payload) + + def sentry_test(): """Temporary endpoint for manual Sentry verification.""" if not current_app.config.get("SENTRY_TEST_ENDPOINT_ENABLED", False): diff --git a/kernel_ai/ml/__init__.py b/kernel_ai/ml/__init__.py new file mode 100644 index 0000000..df4b40e --- /dev/null +++ b/kernel_ai/ml/__init__.py @@ -0,0 +1,19 @@ +"""Kernel-AI anomaly detection pipeline (Stage 1: statistical baselines). + +This package is intentionally decoupled from the Flask request path: + + collect features -> update baseline -> score -> persist anomalies + +The long-running detector lives in :mod:`kernel_ai.ml.worker` and runs as a +separate process (``python -m kernel_ai.ml``). The Flask app only *reads* +already-computed anomalies from the shared store, so a fault in the ML side can +never block or crash the main service. + +Stages (see docs/KERNEL_DNA_AIML_ROADMAP.md): + Stage 1 (this code) - online EWMA baselines + robust z-score, pure-Python. + Stage 2+ - scikit-learn models, MLflow tracking (added later). +""" + +from kernel_ai.ml.config import MLConfig + +__all__ = ["MLConfig"] diff --git a/kernel_ai/ml/__main__.py b/kernel_ai/ml/__main__.py new file mode 100644 index 0000000..9fec3ba --- /dev/null +++ b/kernel_ai/ml/__main__.py @@ -0,0 +1,6 @@ +"""Entrypoint: ``python -m kernel_ai.ml`` runs the Stage 1 detector loop.""" + +from kernel_ai.ml.worker import main + +if __name__ == "__main__": + main() diff --git a/kernel_ai/ml/baseline.py b/kernel_ai/ml/baseline.py new file mode 100644 index 0000000..f853e0e --- /dev/null +++ b/kernel_ai/ml/baseline.py @@ -0,0 +1,91 @@ +"""Stage 1 detector: online EWMA baseline + robust z-score. + +For every feature we keep a slowly-adapting estimate of its *normal* value +(EWMA mean) and its *normal* spread (EWMA variance, West's incremental form). +A live value far above the baseline (high positive z-score) is what we treat as +an attack-shaped mutation: a sudden burst of processes, syscalls, retransmits, +page faults, IRQs, etc. + +Pure Python on purpose — O(1) memory per feature, no numpy/sklearn — so it runs +comfortably inside a tight memory budget. Heavier models arrive at Stage 2. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + + +@dataclass +class _Stat: + mean: float = 0.0 + var: float = 0.0 + count: int = 0 + + +@dataclass(frozen=True) +class Score: + name: str + value: float + mean: float + std: float + z: float + warm: bool + + +class EwmaBaseline: + """Per-feature exponentially-weighted mean/variance with z-score scoring.""" + + def __init__(self, alpha: float, warmup_samples: int) -> None: + self.alpha = alpha + self.warmup = warmup_samples + self._stats: dict[str, _Stat] = {} + + def _update(self, st: _Stat, x: float) -> None: + if st.count == 0: + st.mean = x + st.var = 0.0 + else: + diff = x - st.mean + incr = self.alpha * diff + st.mean += incr + # West's incremental EWMA variance. + st.var = (1.0 - self.alpha) * (st.var + diff * incr) + st.count += 1 + + def update_and_score(self, features: dict[str, float], min_std: dict[str, float]) -> dict[str, Score]: + """Score the current values against the baseline, *then* fold them in. + + Scoring before updating means a real spike is measured against the + pre-spike baseline (it hasn't yet been absorbed), which is what we want. + """ + scores: dict[str, Score] = {} + for name, value in features.items(): + st = self._stats.setdefault(name, _Stat()) + warm = st.count < self.warmup + std = math.sqrt(max(0.0, st.var)) + floor = max(min_std.get(name, 0.0), 1e-9) + std_eff = max(std, floor) + z = (value - st.mean) / std_eff if st.count > 0 else 0.0 + scores[name] = Score(name=name, value=value, mean=st.mean, std=std, z=z, warm=warm) + self._update(st, value) + return scores + + # --- warm-restart support: snapshot / restore baseline across restarts --- + + def export_state(self) -> list[dict]: + return [ + {"name": n, "mean": s.mean, "var": s.var, "count": s.count} + for n, s in self._stats.items() + ] + + def load_state(self, rows: list[dict]) -> None: + for row in rows or []: + try: + self._stats[row["name"]] = _Stat( + mean=float(row["mean"]), + var=float(row["var"]), + count=int(row["count"]), + ) + except (KeyError, TypeError, ValueError): + continue diff --git a/kernel_ai/ml/config.py b/kernel_ai/ml/config.py new file mode 100644 index 0000000..7192130 --- /dev/null +++ b/kernel_ai/ml/config.py @@ -0,0 +1,64 @@ +"""Configuration for the ML anomaly pipeline. + +Everything is driven by environment variables so the same code runs locally +(local Postgres) and on PROD (managed Postgres) by only swapping the DSN. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + + +def _env_float(name: str, default: float) -> float: + try: + return float(os.getenv(name, str(default))) + except (TypeError, ValueError): + return default + + +def _env_int(name: str, default: int) -> int: + try: + return int(os.getenv(name, str(default))) + except (TypeError, ValueError): + return default + + +@dataclass(frozen=True) +class MLConfig: + """Resolved ML pipeline settings (read once at process start).""" + + # Postgres connection string. On deploy, override with the PROD value. + # Example: postgresql://user:pass@host:5432/dbname + dsn: str = os.getenv( + "KERNEL_AI_ML_DSN", + "postgresql://kernel_ai:kernel_ai_dev_pw@127.0.0.1:5432/kernel_ai_ml", + ) + + # Sampling cadence of the detector loop (seconds between feature snapshots). + interval_sec: float = _env_float("KERNEL_AI_ML_INTERVAL_SEC", 2.0) + + # EWMA smoothing: alpha = 2 / (window + 1). Larger window = slower, calmer + # baseline. ~60 gives a baseline that adapts over a couple of minutes. + baseline_window: int = _env_int("KERNEL_AI_ML_BASELINE_WINDOW", 60) + + # Samples to observe before we start emitting anomalies (let the baseline + # settle so we don't fire on the cold-start transient). + warmup_samples: int = _env_int("KERNEL_AI_ML_WARMUP", 30) + + # Robust z-score thresholds. A feature whose current value sits this many + # standard deviations above its baseline becomes a mutation. + z_warn: float = _env_float("KERNEL_AI_ML_Z_WARN", 4.0) + z_crit: float = _env_float("KERNEL_AI_ML_Z_CRIT", 7.0) + + # Persist the raw feature snapshot each tick (useful for later training / + # drift analysis). Disable to keep the DB tiny. + store_features: bool = os.getenv("KERNEL_AI_ML_STORE_FEATURES", "true").lower() == "true" + + # How long to keep rows (hours). The worker prunes older data each cycle. + 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) + + @property + def alpha(self) -> float: + return 2.0 / (max(2, self.baseline_window) + 1.0) diff --git a/kernel_ai/ml/features.py b/kernel_ai/ml/features.py new file mode 100644 index 0000000..ae70ba8 --- /dev/null +++ b/kernel_ai/ml/features.py @@ -0,0 +1,258 @@ +"""Feature extraction from procfs. + +Turns raw kernel counters into a flat ``{name: float}`` feature vector once per +tick. Monotonic counters (context switches, page faults, retransmits, ...) are +converted into per-second *rates* using the previous snapshot, because rates — +not absolute totals — are what reveal an attack-shaped spike. + +Each feature also carries metadata (subsystem + helix position + a noise floor) +so a detected anomaly can be mapped straight onto the Kernel DNA visualization. +""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass + + +@dataclass(frozen=True) +class FeatureSpec: + name: str + subsystem: str + # Position on the DNA helix (0..1), aligned with the gene bands: + # sched 0.0-0.2, net 0.2-0.4, fs 0.4-0.6, mm 0.6-0.8, drivers 0.8-1.0. + position: float + # Absolute noise floor for the std used in the z-score. Prevents a quiet, + # near-constant feature from firing on microscopic deviations. + min_std: float + label: str + + +FEATURE_SPECS: dict[str, FeatureSpec] = { + "proc_count": FeatureSpec("proc_count", "sched", 0.06, 3.0, "processes"), + "procs_running": FeatureSpec("procs_running", "sched", 0.10, 1.0, "runnable now"), + "procs_blocked": FeatureSpec("procs_blocked", "sched", 0.14, 1.0, "blocked (D)"), + "ctxt_per_sec": FeatureSpec("ctxt_per_sec", "sched", 0.18, 50.0, "context switches/s"), + "run_queue": FeatureSpec("run_queue", "sched", 0.16, 1.0, "run-queue depth"), + "load1": FeatureSpec("load1", "sched", 0.08, 0.2, "loadavg 1m"), + "tcp_retrans_per_sec": FeatureSpec("tcp_retrans_per_sec", "net", 0.30, 0.5, "TCP retrans/s"), + "tcp_inseg_per_sec": FeatureSpec("tcp_inseg_per_sec", "net", 0.26, 20.0, "TCP in segs/s"), + "tcp_outseg_per_sec": FeatureSpec("tcp_outseg_per_sec", "net", 0.34, 20.0, "TCP out segs/s"), + "net_softirq_per_sec": FeatureSpec("net_softirq_per_sec", "net", 0.38, 20.0, "NET softirq/s"), + "block_softirq_per_sec": FeatureSpec("block_softirq_per_sec", "fs", 0.46, 5.0, "BLOCK softirq/s"), + "pgfault_per_sec": FeatureSpec("pgfault_per_sec", "mm", 0.62, 50.0, "minor faults/s"), + "pgmajfault_per_sec": FeatureSpec("pgmajfault_per_sec", "mm", 0.68, 1.0, "major faults/s"), + "pgscan_direct_per_sec": FeatureSpec("pgscan_direct_per_sec", "mm", 0.72, 5.0, "direct reclaim/s"), + "swap_io_per_sec": FeatureSpec("swap_io_per_sec", "mm", 0.76, 1.0, "swap io/s"), + "psi_mem_some10": FeatureSpec("psi_mem_some10", "mm", 0.64, 0.5, "mem PSI some10"), + "psi_mem_full10": FeatureSpec("psi_mem_full10", "mm", 0.78, 0.3, "mem PSI full10"), + "hardirq_per_sec": FeatureSpec("hardirq_per_sec", "drivers", 0.90, 50.0, "hard IRQ/s"), + "cpu_busy_pct": FeatureSpec("cpu_busy_pct", "drivers", 0.86, 3.0, "cpu busy %"), +} + + +def _read_kv(path: str) -> dict[str, int]: + out: dict[str, int] = {} + try: + with open(path, "r", encoding="utf-8", errors="ignore") as f: + for raw in f: + parts = raw.split() + if len(parts) == 2: + try: + out[parts[0]] = int(parts[1]) + except ValueError: + continue + except OSError: + return {} + return out + + +def _read_psi_mem() -> tuple[float, float]: + some = full = 0.0 + try: + with open("/proc/pressure/memory", "r", encoding="utf-8", errors="ignore") as f: + for raw in f: + parts = raw.split() + if not parts: + continue + target = parts[0] + for tok in parts[1:]: + if tok.startswith("avg10="): + try: + val = float(tok.split("=", 1)[1]) + except ValueError: + val = 0.0 + if target == "some": + some = val + elif target == "full": + full = val + except OSError: + pass + return some, full + + +def _read_loadavg() -> tuple[float, int]: + try: + with open("/proc/loadavg", "r", encoding="utf-8", errors="ignore") as f: + parts = f.read().split() + load1 = float(parts[0]) + runnable = int(parts[3].split("/", 1)[0]) if len(parts) >= 4 and "/" in parts[3] else 0 + return load1, runnable + except (OSError, ValueError, IndexError): + return 0.0, 0 + + +def _read_stat() -> dict: + out = {"ctxt": 0, "procs_running": 0, "procs_blocked": 0, "cpu_busy": 0, "cpu_total": 0} + try: + with open("/proc/stat", "r", encoding="utf-8", errors="ignore") as f: + for line in f: + if line.startswith("cpu "): + nums = [int(x) for x in line.split()[1:] if x.isdigit()] + if len(nums) >= 4: + idle = nums[3] + (nums[4] if len(nums) > 4 else 0) + total = sum(nums) + out["cpu_busy"] = total - idle + out["cpu_total"] = total + elif line.startswith("ctxt "): + out["ctxt"] = int(line.split()[1]) + elif line.startswith("procs_running "): + out["procs_running"] = int(line.split()[1]) + elif line.startswith("procs_blocked "): + out["procs_blocked"] = int(line.split()[1]) + except (OSError, ValueError): + pass + return out + + +def _read_tcp_snmp() -> dict[str, int]: + out = {"RetransSegs": 0, "InSegs": 0, "OutSegs": 0} + try: + with open("/proc/net/snmp", "r", encoding="utf-8", errors="ignore") as f: + lines = [ln.strip() for ln in f if ln.startswith("Tcp:")] + if len(lines) >= 2: + headers = lines[-2].split()[1:] + values = lines[-1].split()[1:] + for key, value in zip(headers, values): + if key in out: + try: + out[key] = int(value) + except ValueError: + pass + except OSError: + pass + return out + + +def _read_softirq_totals() -> dict[str, int]: + totals: dict[str, int] = {} + try: + with open("/proc/softirqs", "r", encoding="utf-8", errors="ignore") as f: + for raw in f.readlines()[1:]: + if ":" not in raw: + continue + left, right = raw.split(":", 1) + totals[left.strip()] = sum(int(t) for t in right.split() if t.isdigit()) + except OSError: + return {} + return totals + + +def _read_hardirq_total() -> int: + total = 0 + try: + with open("/proc/interrupts", "r", encoding="utf-8", errors="ignore") as f: + for raw in f.readlines()[1:]: + if ":" not in raw: + continue + _, right = raw.split(":", 1) + for tok in right.split(): + if tok.isdigit(): + total += int(tok) + else: + break + except OSError: + return 0 + return total + + +def _count_procs() -> int: + try: + return sum(1 for d in os.listdir("/proc") if d.isdigit()) + except OSError: + return 0 + + +class FeatureExtractor: + """Stateful procfs reader: holds the previous counter snapshot to derive + per-second rates. Call :meth:`collect` once per tick.""" + + def __init__(self) -> None: + self._prev: dict[str, float] | None = None + self._prev_ts: float | None = None + + def collect(self) -> dict[str, float] | None: + now = time.time() + stat = _read_stat() + vmstat = _read_kv("/proc/vmstat") + tcp = _read_tcp_snmp() + softirq = _read_softirq_totals() + hardirq = _read_hardirq_total() + psi_some, psi_full = _read_psi_mem() + load1, run_queue = _read_loadavg() + + raw = { + "ctxt": float(stat["ctxt"]), + "pgfault": float(vmstat.get("pgfault", 0)), + "pgmajfault": float(vmstat.get("pgmajfault", 0)), + "pgscan_direct": float(vmstat.get("pgscan_direct", 0)), + "swap_io": float(vmstat.get("pswpin", 0) + vmstat.get("pswpout", 0)), + "tcp_retrans": float(tcp["RetransSegs"]), + "tcp_inseg": float(tcp["InSegs"]), + "tcp_outseg": float(tcp["OutSegs"]), + "net_softirq": float(softirq.get("NET_RX", 0) + softirq.get("NET_TX", 0)), + "block_softirq": float(softirq.get("BLOCK", 0)), + "hardirq": float(hardirq), + "cpu_busy": float(stat["cpu_busy"]), + "cpu_total": float(stat["cpu_total"]), + } + + prev, prev_ts = self._prev, self._prev_ts + self._prev, self._prev_ts = raw, now + + # First sample: no previous counters, so rates are undefined. Skip it. + if prev is None or prev_ts is None: + return None + dt = now - prev_ts + if dt <= 0: + return None + + def rate(key: str) -> float: + return max(0.0, (raw[key] - prev.get(key, raw[key])) / dt) + + cpu_total_delta = raw["cpu_total"] - prev.get("cpu_total", raw["cpu_total"]) + cpu_busy_delta = raw["cpu_busy"] - prev.get("cpu_busy", raw["cpu_busy"]) + cpu_busy_pct = (cpu_busy_delta / cpu_total_delta * 100.0) if cpu_total_delta > 0 else 0.0 + + return { + "proc_count": float(_count_procs()), + "procs_running": float(stat["procs_running"]), + "procs_blocked": float(stat["procs_blocked"]), + "ctxt_per_sec": rate("ctxt"), + "run_queue": float(run_queue), + "load1": load1, + "tcp_retrans_per_sec": rate("tcp_retrans"), + "tcp_inseg_per_sec": rate("tcp_inseg"), + "tcp_outseg_per_sec": rate("tcp_outseg"), + "net_softirq_per_sec": rate("net_softirq"), + "block_softirq_per_sec": rate("block_softirq"), + "pgfault_per_sec": rate("pgfault"), + "pgmajfault_per_sec": rate("pgmajfault"), + "pgscan_direct_per_sec": rate("pgscan_direct"), + "swap_io_per_sec": rate("swap_io"), + "psi_mem_some10": psi_some, + "psi_mem_full10": psi_full, + "hardirq_per_sec": rate("hardirq"), + "cpu_busy_pct": max(0.0, min(100.0, cpu_busy_pct)), + } diff --git a/kernel_ai/ml/store.py b/kernel_ai/ml/store.py new file mode 100644 index 0000000..571c36c --- /dev/null +++ b/kernel_ai/ml/store.py @@ -0,0 +1,182 @@ +"""Postgres-backed store for the ML pipeline. + +Schema (all created on demand): + ml_feature_snapshots - one JSONB row per tick (raw features, for training) + ml_anomalies - detected mutations served to the Kernel DNA UI + ml_baseline_state - EWMA state, persisted for warm restarts + +The worker owns one long-lived connection. The Flask read path opens a fresh +short-lived connection per call (thread-safe, low volume), and tolerates the DB +being unreachable by returning empty results instead of raising. +""" + +from __future__ import annotations + +import logging + +import psycopg +from psycopg.types.json import Json + +logger = logging.getLogger("kernel_ai.ml.store") + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS ml_feature_snapshots ( + ts timestamptz NOT NULL DEFAULT now(), + features jsonb NOT NULL +); +CREATE INDEX IF NOT EXISTS ml_feature_snapshots_ts_idx ON ml_feature_snapshots (ts); + +CREATE TABLE IF NOT EXISTS ml_anomalies ( + id bigserial PRIMARY KEY, + ts timestamptz NOT NULL DEFAULT now(), + source text NOT NULL DEFAULT 'stage1_baseline', + feature text NOT NULL, + subsystem text, + type text NOT NULL, + severity text NOT NULL, + score double precision NOT NULL, + value double precision, + baseline_mean double precision, + baseline_std double precision, + position double precision, + message text, + meta jsonb +); +CREATE INDEX IF NOT EXISTS ml_anomalies_ts_idx ON ml_anomalies (ts); + +CREATE TABLE IF NOT EXISTS ml_baseline_state ( + name text PRIMARY KEY, + mean double precision NOT NULL, + var double precision NOT NULL, + count integer NOT NULL, + updated_at timestamptz NOT NULL DEFAULT now() +); +""" + + +def connect(dsn: str, *, autocommit: bool = True) -> psycopg.Connection: + return psycopg.connect(dsn, autocommit=autocommit) + + +def ensure_schema(conn: psycopg.Connection) -> None: + with conn.cursor() as cur: + cur.execute(_SCHEMA) + + +class PostgresStore: + """Write side, used by the worker over a single persistent connection.""" + + def __init__(self, dsn: str) -> None: + self.dsn = dsn + self.conn = connect(dsn) + ensure_schema(self.conn) + + def insert_feature_snapshot(self, features: dict[str, float]) -> None: + with self.conn.cursor() as cur: + cur.execute( + "INSERT INTO ml_feature_snapshots (features) VALUES (%s)", + (Json(features),), + ) + + def insert_anomalies(self, anomalies: list[dict]) -> None: + if not anomalies: + return + with self.conn.cursor() as cur: + cur.executemany( + """ + INSERT INTO ml_anomalies + (source, feature, subsystem, type, severity, score, value, + baseline_mean, baseline_std, position, message, meta) + VALUES + (%(source)s, %(feature)s, %(subsystem)s, %(type)s, %(severity)s, + %(score)s, %(value)s, %(baseline_mean)s, %(baseline_std)s, + %(position)s, %(message)s, %(meta)s) + """, + [ + { + "source": a.get("source", "stage1_baseline"), + "feature": a["feature"], + "subsystem": a.get("subsystem"), + "type": a["type"], + "severity": a["severity"], + "score": a["score"], + "value": a.get("value"), + "baseline_mean": a.get("baseline_mean"), + "baseline_std": a.get("baseline_std"), + "position": a.get("position"), + "message": a.get("message"), + "meta": Json(a.get("meta") or {}), + } + for a in anomalies + ], + ) + + def save_baseline(self, rows: list[dict]) -> None: + if not rows: + return + with self.conn.cursor() as cur: + cur.executemany( + """ + INSERT INTO ml_baseline_state (name, mean, var, count, updated_at) + VALUES (%(name)s, %(mean)s, %(var)s, %(count)s, now()) + ON CONFLICT (name) DO UPDATE + SET mean = EXCLUDED.mean, + var = EXCLUDED.var, + count = EXCLUDED.count, + updated_at = now() + """, + rows, + ) + + def load_baseline(self) -> list[dict]: + with self.conn.cursor() as cur: + cur.execute("SELECT name, mean, var, count FROM ml_baseline_state") + return [ + {"name": r[0], "mean": r[1], "var": r[2], "count": r[3]} + for r in cur.fetchall() + ] + + def prune(self, features_hours: int, anomalies_hours: int) -> None: + with self.conn.cursor() as cur: + cur.execute( + "DELETE FROM ml_feature_snapshots WHERE ts < now() - make_interval(hours => %s)", + (features_hours,), + ) + cur.execute( + "DELETE FROM ml_anomalies WHERE ts < now() - make_interval(hours => %s)", + (anomalies_hours,), + ) + + def close(self) -> None: + try: + self.conn.close() + except Exception: # noqa: BLE001 - best-effort cleanup + pass + + +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: + with connect(dsn) as conn, conn.cursor() as cur: + cur.execute( + """ + SELECT ts, source, feature, subsystem, type, severity, score, value, + baseline_mean, baseline_std, position, message + FROM ml_anomalies + WHERE ts > now() - make_interval(secs => %s) + ORDER BY ts DESC + LIMIT %s + """, + (since_seconds, limit), + ) + cols = [d.name for d in cur.description] + out = [] + for row in cur.fetchall(): + rec = dict(zip(cols, row)) + if rec.get("ts") is not None: + rec["ts"] = rec["ts"].isoformat() + out.append(rec) + return out + except Exception as exc: # noqa: BLE001 - API must degrade gracefully + logger.warning("fetch_recent_anomalies failed: %s", exc) + return [] diff --git a/kernel_ai/ml/worker.py b/kernel_ai/ml/worker.py new file mode 100644 index 0000000..8cb6506 --- /dev/null +++ b/kernel_ai/ml/worker.py @@ -0,0 +1,135 @@ +"""Stage 1 detector loop. + + collect features -> score vs baseline -> emit anomalies -> persist + +Runs as its own process so it is fully isolated from the Flask request path. +Anomalies are written to Postgres in a shape that maps directly onto Kernel DNA +mutations (type / severity / message / position), enriched with the statistical +evidence (z-score, baseline) that triggered them. +""" + +from __future__ import annotations + +import logging +import signal +import time + +from kernel_ai.ml.baseline import EwmaBaseline, Score +from kernel_ai.ml.config import MLConfig +from kernel_ai.ml.features import FEATURE_SPECS, FeatureExtractor +from kernel_ai.ml.store import PostgresStore + +logger = logging.getLogger("kernel_ai.ml.worker") + +# Persist baseline state / prune old rows every N ticks (not every tick). +_HOUSEKEEPING_EVERY = 30 + + +def _build_anomalies(scores: dict[str, Score], cfg: MLConfig) -> list[dict]: + """Turn high positive z-scores into Kernel DNA mutation records.""" + out: list[dict] = [] + for name, sc in scores.items(): + if sc.warm: + continue + # Attacks present as bursts: we flag upward deviations only. + if sc.z < cfg.z_warn or sc.value <= sc.mean: + continue + spec = FEATURE_SPECS.get(name) + severity = "high" if sc.z >= cfg.z_crit else "medium" + out.append( + { + "source": "stage1_baseline", + "feature": name, + "subsystem": spec.subsystem if spec else None, + "type": f"baseline_spike:{name}", + "severity": severity, + "score": round(sc.z, 3), + "value": round(sc.value, 3), + "baseline_mean": round(sc.mean, 3), + "baseline_std": round(sc.std, 3), + "position": spec.position if spec else 0.5, + "message": ( + f"{(spec.label if spec else name)} spike: " + f"{sc.value:.1f} vs baseline {sc.mean:.1f}±{sc.std:.1f} (z={sc.z:.1f})" + ), + "meta": {"stage": 1, "z": round(sc.z, 3), "alpha": round(cfg.alpha, 4)}, + } + ) + return out + + +class MLWorker: + def __init__(self, cfg: MLConfig | None = None) -> None: + self.cfg = cfg or MLConfig() + self.extractor = FeatureExtractor() + self.baseline = EwmaBaseline(alpha=self.cfg.alpha, warmup_samples=self.cfg.warmup_samples) + self.store = PostgresStore(self.cfg.dsn) + self._running = True + self._min_std = {n: s.min_std for n, s in FEATURE_SPECS.items()} + + def stop(self, *_args) -> None: + self._running = False + + def _tick(self) -> int: + features = self.extractor.collect() + if not features: + return 0 + scores = self.baseline.update_and_score(features, self._min_std) + anomalies = _build_anomalies(scores, self.cfg) + if self.cfg.store_features: + self.store.insert_feature_snapshot(features) + if anomalies: + self.store.insert_anomalies(anomalies) + return len(anomalies) + + def run(self) -> None: + signal.signal(signal.SIGTERM, self.stop) + signal.signal(signal.SIGINT, self.stop) + + restored = self.store.load_baseline() + if restored: + self.baseline.load_state(restored) + logger.info("restored baseline state for %d features", len(restored)) + + logger.info( + "ML worker started: interval=%.1fs window=%d warmup=%d z_warn=%.1f z_crit=%.1f", + self.cfg.interval_sec, self.cfg.baseline_window, self.cfg.warmup_samples, + self.cfg.z_warn, self.cfg.z_crit, + ) + + ticks = 0 + while self._running: + start = time.time() + try: + n = self._tick() + ticks += 1 + if n: + logger.info("tick %d: emitted %d anomalies", ticks, n) + 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) + except Exception as exc: # noqa: BLE001 - keep the loop alive + logger.exception("tick failed: %s", exc) + # Reconnect on DB hiccups rather than dying. + try: + self.store = PostgresStore(self.cfg.dsn) + except Exception: # noqa: BLE001 + time.sleep(2.0) + + elapsed = time.time() - start + time.sleep(max(0.0, self.cfg.interval_sec - elapsed)) + + # Graceful shutdown: persist what we learned. + try: + self.store.save_baseline(self.baseline.export_state()) + finally: + self.store.close() + logger.info("ML worker stopped after %d ticks", ticks) + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s %(message)s", + ) + MLWorker().run() diff --git a/requirements.txt b/requirements.txt index 81c01db..cf20b81 100755 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ openai==0.28.1 python-dotenv==1.0.0 gunicorn==21.2.0 sentry-sdk +psycopg[binary]==3.2.3