diff --git a/kernel_ai/ml/retrain.py b/kernel_ai/ml/retrain.py index 2b64f09..299d2ec 100644 --- a/kernel_ai/ml/retrain.py +++ b/kernel_ai/ml/retrain.py @@ -1,16 +1,19 @@ -"""Stage 3 auto-retrain orchestrator (run by a systemd timer). +"""Stage 3 + Stage 4 auto-retrain orchestrator (run by a systemd timer). - measure drift -> decide -> retrain on clean recent data -> register + measure drift -> decide -> IsolationForest on clean data -> STIDE profile 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). + * default: always retrain IsolationForest (the timer cadence IS the schedule). + * --only-if-drift: skip IsolationForest when the drift monitor is quiet. + * STIDE (``stide_latest.joblib``) is always refreshed from ``ml_syscall_ngrams`` + when Stage 4 is on — including ``--stide-only`` and IF soft-skips — so the + audit/L2 vocabulary does not wait on feature drift. Safety: - * training excludes high-severity anomaly windows (poison guard), and + * IF 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. + * STIDE uses a frequency poison guard inside ``build_profile``. """ from __future__ import annotations @@ -25,40 +28,81 @@ logger = logging.getLogger("kernel_ai.ml.retrain") -def run(*, only_if_drift: bool, min_samples: int) -> int: +def _refresh_stide(cfg: MLConfig, metrics: dict) -> None: + """Best-effort STIDE rebuild from accumulated audit/socket n-grams.""" + if not cfg.enable_stage4: + return + # train() already rebuilt STIDE on a successful IF pass — skip duplicate work. + if metrics.get("seq_vocab_kept") is not None: + return + try: + from kernel_ai.ml.sequence import build_profile + + seq_meta = build_profile(cfg) + metrics["seq_vocab_kept"] = float(seq_meta.get("vocab_kept", 0)) + logger.info("STIDE profile refreshed: %s", seq_meta) + except SystemExit as exc: + logger.info("STIDE profile not rebuilt: %s", exc) + except Exception as exc: # noqa: BLE001 - sequence profile is optional + logger.warning("STIDE profile build failed: %s", exc) + + +def run(*, only_if_drift: bool, min_samples: int, stide_only: bool = False) -> int: cfg = MLConfig() - drift = compute_drift(cfg, persist=True) + metrics: dict = {} - 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")) + if stide_only: + _refresh_stide(cfg, metrics) + logger.info("stide-only retrain complete: %s", metrics) 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 + drift = compute_drift(cfg, persist=True) + skip_if = bool(only_if_drift and drift.get("available") and not drift.get("drifted")) + if skip_if: + logger.info( + "no drift detected (flag_rate=%s) - skipping IsolationForest retrain", + drift.get("flag_rate"), + ) + else: + 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 + # IsolationForest and still try STIDE below. + logger.warning("IsolationForest retrain skipped: %s", exc) + metrics = {} + + _refresh_stide(cfg, metrics) 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 = argparse.ArgumentParser(description="Drift-aware auto-retrain (IF + STIDE)") parser.add_argument("--only-if-drift", action="store_true") parser.add_argument("--min-samples", type=int, default=100) + parser.add_argument( + "--stide-only", + action="store_true", + help="only rebuild stide_latest.joblib from ml_syscall_ngrams (no IsolationForest)", + ) args = parser.parse_args() - raise SystemExit(run(only_if_drift=args.only_if_drift, min_samples=args.min_samples)) + raise SystemExit( + run( + only_if_drift=args.only_if_drift, + min_samples=args.min_samples, + stide_only=args.stide_only, + ) + ) if __name__ == "__main__": diff --git a/kernel_ai/ml/sequence.py b/kernel_ai/ml/sequence.py index 6698d2f..767e88f 100644 --- a/kernel_ai/ml/sequence.py +++ b/kernel_ai/ml/sequence.py @@ -140,6 +140,14 @@ def update_stream(self, events) -> int: def recent(self) -> list[str]: return list(self._recent) + def recent_tokens(self) -> list[str]: + """Reconstruct overlapping syscall tokens from the rolling n-gram window. + + Stage 8 Markov trains on syscall names (from expanded n-grams); scoring + must use the same alphabet, not the ``a|b|c`` keys STIDE stores. + """ + return ngrams_to_tokens(self.recent(), n=self.n) + def recent_by_pid(self, *, min_len: int) -> list[tuple[int, list[str]]]: """Return ``(pid, ngram_window)`` for pids with enough recent n-grams.""" need = max(1, int(min_len)) @@ -156,6 +164,27 @@ def drain_pending(self) -> dict[str, int]: return pending +def ngrams_to_tokens(ngrams: list[str], *, n: int = 3) -> list[str]: + """Stitch sliding n-gram keys back into a syscall token stream.""" + if not ngrams: + return [] + first = [t for t in str(ngrams[0]).split(_SEP) if t] + if not first: + return [] + out = list(first) + width = max(2, int(n)) + for key in ngrams[1:]: + parts = [t for t in str(key).split(_SEP) if t] + if not parts: + continue + # Overlapping slide: keep only the new trailing token when width matches. + if len(parts) == width and len(out) >= width - 1 and out[-(width - 1) :] == parts[:-1]: + out.append(parts[-1]) + else: + out.extend(parts) + return out + + @dataclass class StideModel: """A "normal" n-gram vocabulary with window-mismatch scoring.""" diff --git a/kernel_ai/ml/worker.py b/kernel_ai/ml/worker.py index 145a298..2b1092f 100644 --- a/kernel_ai/ml/worker.py +++ b/kernel_ai/ml/worker.py @@ -322,16 +322,16 @@ def _tick_sequence(self) -> dict | None: ) def _tick_stage8(self) -> dict | None: - """Stage 8 stub: score the Stage 4 rolling window if a model is ready.""" + """Stage 8: score reconstructed syscall tokens if a Markov model is ready.""" if self.deep_scorer is None or self.seq_tracker is None: return None self.deep_scorer.maybe_reload() if not self.deep_scorer.ready: return None - window = self.seq_tracker.recent() - if len(window) < max(8, self.cfg.stage8_window // 4): + tokens = self.seq_tracker.recent_tokens() + if len(tokens) < max(8, self.cfg.stage8_window // 4): return None - tokens = window[-self.cfg.stage8_window :] + tokens = tokens[-self.cfg.stage8_window :] score = self.deep_scorer.score_tokens(tokens) if not score: return None diff --git a/tests/test_ml_collectors.py b/tests/test_ml_collectors.py index b1de8c7..46f5038 100644 --- a/tests/test_ml_collectors.py +++ b/tests/test_ml_collectors.py @@ -19,6 +19,19 @@ def test_resolve_syscall_nr_arch_collision(): assert resolve_syscall_name(117, AUDIT_ARCH_AARCH64) == "ptrace" +def test_ngrams_to_tokens_stitches_overlap(): + from kernel_ai.ml.sequence import ngrams_to_tokens + + keys = ["clone|openat|execve", "openat|execve|connect", "execve|connect|setuid"] + assert ngrams_to_tokens(keys, n=3) == [ + "clone", + "openat", + "execve", + "connect", + "setuid", + ] + + def test_encode_decode_roundtrip(): events = [ SyscallEvent(ts=1.0, pid=10, uid=0, comm="bash", syscall="clone"),