Training pipeline scripts

training/__init__.py

@@ -0,0 +1,13 @@
+"""
+KernelX Intelligence Layer — Model Training Pipeline
+
+Stages:
+  0. Environment setup (requirements.txt)
+  1. Data preprocessing (data/preprocess.py)
+  2. World Model SFT training (models/train_world_model.py)
+  3. Strategist RL training via GRPO (models/train_strategist.py)
+  4. Monitoring (logged during training)
+  5. Export & quantization to GGUF (models/export_gguf.py)
+  6. Inference engine (inference/strategy_engine.py)
+  7. Gradio demo (demo/app.py)
+"""

training/data/__init__.py

@@ -0,0 +1,23 @@
+"""KernelX training data preprocessing utilities."""
+
+from .preprocess import (
+    CONFIG,
+    ACTIVE_FEATURES,
+    FEATURE_NAMES,
+    SYMLOG_FEATURES,
+    IDX_CPU,
+    IDX_PRIO,
+    IDX_STATIC_PRIO,
+    IDX_NORMAL_PRIO,
+    IDX_EXEC_NS,
+    IDX_VRUNTIME,
+    IDX_MIGRATIONS,
+    IDX_CPUS_ALLOWED,
+    IDX_CTX_SWITCHES,
+    IDX_WAIT_US,
+    symmetric_log,
+    preprocess_features,
+    extract_active,
+    format_state,
+    preprocess_record,
+)

training/data/preprocess.py

@@ -0,0 +1,223 @@
+"""
+KernelX Intelligence Layer — Data Ingestion and Preprocessing (Stage 1)
+
+Reads raw state_transitions.jsonl from the bridge's TrajectoryManager,
+applies feature scaling (symlog for huge counters), drops sparse-zero
+features, and produces train/val/test splits for World Model and
+Strategist training.
+
+Usage:
+    python -m training.data.preprocess --input data/state_transitions.jsonl
+"""
+
+import json
+import argparse
+import sys
+from pathlib import Path
+from typing import List, Dict, Any
+
+import numpy as np
+
+# ---------------------------------------------------------------------------
+# Configuration (loaded from preprocessing_config.json)
+# ---------------------------------------------------------------------------
+
+CONFIG_PATH = Path(__file__).parent / "preprocessing_config.json"
+
+def load_config() -> dict:
+    with open(CONFIG_PATH) as f:
+        return json.load(f)
+
+CONFIG = load_config()
+
+SYMLOG_FEATURES = CONFIG["symlog_features"]          # [4, 5, 6]
+ACTIVE_FEATURES = CONFIG["active_features"]           # [0,1,2,3,4,5,6,7,12,23]
+FEATURE_NAMES   = CONFIG["feature_names"]             # 10 short names
+SPARSE_ZERO     = CONFIG["sparse_zero_features"]      # indices to drop
+
+# ---------------------------------------------------------------------------
+# Feature index constants (positions within the ACTIVE 10D vector)
+# ---------------------------------------------------------------------------
+
+IDX_CPU           = 0   # raw index 0
+IDX_PRIO          = 1   # raw index 1
+IDX_STATIC_PRIO   = 2   # raw index 2
+IDX_NORMAL_PRIO   = 3   # raw index 3
+IDX_EXEC_NS       = 4   # raw index 4 (symlog)
+IDX_VRUNTIME      = 5   # raw index 5 (symlog)
+IDX_MIGRATIONS    = 6   # raw index 6 (symlog)
+IDX_CPUS_ALLOWED  = 7   # raw index 7
+IDX_CTX_SWITCHES  = 8   # raw index 12
+IDX_WAIT_US       = 9   # raw index 23
+
+# ---------------------------------------------------------------------------
+# Scaling functions
+# ---------------------------------------------------------------------------
+
+def symmetric_log(x: float) -> float:
+    """sgn(x) * ln(1 + |x|) — compresses huge values while preserving sign."""
+    return float(np.sign(x) * np.log1p(np.abs(x)))
+
+
+def preprocess_features(raw_features: List[float]) -> List[float]:
+    """Transform a raw 24D feature vector: symlog the big counters."""
+    f = list(raw_features)  # copy
+    for idx in SYMLOG_FEATURES:
+        f[idx] = symmetric_log(f[idx])
+    return f
+
+
+def extract_active(scaled_features: List[float]) -> List[float]:
+    """Keep only the active (non-zero, information-carrying) features."""
+    return [scaled_features[i] for i in ACTIVE_FEATURES]
+
+# ---------------------------------------------------------------------------
+# Prompt formatting for LLM consumption
+# ---------------------------------------------------------------------------
+
+def format_state(active_vector: List[float]) -> str:
+    """Convert a 10D active feature vector into a compact text string."""
+    parts = []
+    for name, val in zip(FEATURE_NAMES, active_vector):
+        if val == int(val):
+            parts.append(f"{name}:{int(val)}")
+        else:
+            parts.append(f"{name}:{val:.2f}")
+    return " | ".join(parts)
+
+# ---------------------------------------------------------------------------
+# Record processing
+# ---------------------------------------------------------------------------
+
+def preprocess_record(record: Dict[str, Any]) -> Dict[str, Any]:
+    """Transform a single raw JSONL record into training-ready format."""
+    s_t_scaled = preprocess_features(record["state_t"]["features"])
+    s_t1_scaled = preprocess_features(record["state_t_next"]["features"])
+
+    s_t_active = extract_active(s_t_scaled)
+    s_t1_active = extract_active(s_t1_scaled)
+
+    return {
+        "state": s_t_active,
+        "action": record["action"],
+        "reward": record["reward"],
+        "next_state": s_t1_active,
+        "pid": record["state_t"]["pid"],
+        "cpu": record["state_t"]["cpu"],
+        "timestamp": record["state_t"]["timestamp"],
+    }
+
+# ---------------------------------------------------------------------------
+# Dataset audit
+# ---------------------------------------------------------------------------
+
+def audit_dataset(records: List[Dict]) -> None:
+    """Print per-feature statistics for the raw dataset."""
+    all_features = []
+    for r in records:
+        all_features.append(r["state_t"]["features"])
+        all_features.append(r["state_t_next"]["features"])
+
+    arr = np.array(all_features, dtype=np.float64)
+
+    print(f"\nTotal transitions: {len(records)}")
+    print(f"Total feature vectors: {len(all_features)}")
+    print(f"\n{'Idx':<5} {'Min':<22} {'Max':<22} {'Mean':<22} {'Std':<22} {'Zeros%':<10}")
+    print("-" * 103)
+    for i in range(24):
+        col = arr[:, i]
+        zero_pct = (col == 0).sum() / len(col) * 100
+        print(f"{i:<5} {col.min():<22.2f} {col.max():<22.2f} {col.mean():<22.2f} {col.std():<22.2f} {zero_pct:<10.1f}")
+
+    print(f"\nNaN count: {np.isnan(arr).sum()}")
+    print(f"Inf count: {np.isinf(arr).sum()}")
+
+    actions = [r["action"] for r in records]
+    rewards = [r["reward"] for r in records]
+    print(f"\nAction — unique values: {sorted(set(actions))}")
+    print(f"Reward — min: {min(rewards)}, max: {max(rewards)}, mean: {np.mean(rewards):.2f}, std: {np.std(rewards):.2f}")
+
+
+# ---------------------------------------------------------------------------
+# Train / Val / Test split (chronological)
+# ---------------------------------------------------------------------------
+
+def split_chronological(processed: List[Dict], train_ratio=0.8, val_ratio=0.1):
+    """Split processed records chronologically (NOT randomly)."""
+    processed.sort(key=lambda x: x["timestamp"])
+    n = len(processed)
+    train_end = int(n * train_ratio)
+    val_end = int(n * (train_ratio + val_ratio))
+    return processed[:train_end], processed[train_end:val_end], processed[val_end:]
+
+# ---------------------------------------------------------------------------
+# Main pipeline
+# ---------------------------------------------------------------------------
+
+def run_pipeline(input_path: str, output_dir: str, audit: bool = True):
+    """Full preprocessing pipeline: audit -> scale -> split -> save."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load raw data
+    print(f"Loading raw data from {input_path} ...")
+    records = []
+    with open(input_path) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                records.append(json.loads(line))
+
+    if not records:
+        print("ERROR: No records found in input file.")
+        sys.exit(1)
+
+    print(f"Loaded {len(records)} raw transitions.")
+
+    # Audit
+    if audit:
+        audit_dataset(records)
+
+    # Preprocess
+    print("\nPreprocessing (symlog scaling + active feature extraction) ...")
+    processed = [preprocess_record(r) for r in records]
+
+    # Verify transform on first record
+    sample = processed[0]
+    print(f"\nSample preprocessed state (10D):")
+    print(f"  {format_state(sample['state'])}")
+    print(f"Sample preprocessed next_state:")
+    print(f"  {format_state(sample['next_state'])}")
+
+    # Save full processed dataset
+    processed_path = output_dir / "processed_transitions.jsonl"
+    with open(processed_path, "w") as f:
+        for p in processed:
+            f.write(json.dumps(p) + "\n")
+    print(f"\nSaved {len(processed)} processed records to {processed_path}")
+
+    # Split
+    train, val, test = split_chronological(processed)
+    for split_name, split_data in [("train", train), ("val", val), ("test", test)]:
+        split_path = output_dir / f"{split_name}.jsonl"
+        with open(split_path, "w") as f:
+            for item in split_data:
+                f.write(json.dumps(item) + "\n")
+        print(f"{split_name}: {len(split_data)} records -> {split_path}")
+
+    print("\nPreprocessing complete.")
+    return train, val, test
+
+
+def main():
+    parser = argparse.ArgumentParser(description="KernelX data preprocessing pipeline")
+    parser.add_argument("--input", required=True, help="Path to raw state_transitions.jsonl")
+    parser.add_argument("--output-dir", default=str(Path(__file__).parent), help="Output directory for processed data")
+    parser.add_argument("--no-audit", action="store_true", help="Skip the dataset audit step")
+    args = parser.parse_args()
+
+    run_pipeline(args.input, args.output_dir, audit=not args.no_audit)
+
+
+if __name__ == "__main__":
+    main()

training/data/preprocessing_config.json

@@ -0,0 +1,26 @@
+{
+  "symlog_features": [4, 5, 6],
+  "active_features": [0, 1, 2, 3, 4, 5, 6, 7, 12, 23],
+  "feature_names": ["cpu", "prio", "sprio", "nprio", "exec_ns", "vrt", "migr", "cpus", "csw", "wt_us"],
+  "passthrough_features": [0, 7, 12, 23],
+  "priority_features": [1, 2, 3],
+  "sparse_zero_features": [8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22],
+  "feature_descriptions": {
+    "0": "bpf_get_smp_processor_id() - CPU core ID",
+    "1": "task->prio - Dynamic priority (0-139)",
+    "2": "task->static_prio - Static priority (nice-based)",
+    "3": "task->normal_prio - Normal priority",
+    "4": "task->se.sum_exec_runtime - Total CPU time in ns (symlog-scaled)",
+    "5": "task->se.vruntime - CFS virtual runtime (symlog-scaled)",
+    "6": "task->se.nr_migrations - CPU migration count (symlog-scaled)",
+    "7": "task->nr_cpus_allowed - CPU affinity mask size",
+    "12": "context switch count from cpu_stats",
+    "23": "wait time in microseconds ((now - start_ts) / 1000)"
+  },
+  "model": {
+    "name": "HuggingFaceTB/SmolLM2-360M-Instruct",
+    "backup": "Qwen/Qwen2.5-0.5B-Instruct",
+    "max_seq_length": 512,
+    "target_inference_ms": 50
+  }
+}

training/demo/app.py

@@ -0,0 +1,322 @@
+"""
+KernelX Intelligence Layer — Gradio Demo (Stage 7)
+
+Judge-facing demo that shows:
+  1. Baseline SmolLM2-360M output (untrained) on kernel states
+  2. World Model predictions vs actual next states
+  3. Strategist scheduling actions (action-only, no rationale for speed)
+  4. Before/after metrics comparison
+
+Usage:
+    python -m training.demo.app \
+        --strategist-model training/models/strategist_merged/strategist-q4km.gguf \
+        --test-data training/data/test.jsonl
+
+    # Without trained models (shows heuristic baseline):
+    python -m training.demo.app --test-data training/data/test.jsonl --no-model
+"""
+
+import argparse
+import json
+import re
+import sys
+import time
+
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+from training.data.preprocess import (
+    FEATURE_NAMES, format_state, load_config,
+    IDX_WAIT_US, IDX_CTX_SWITCHES, IDX_EXEC_NS,
+)
+from training.environment.rewards import RewardComputer
+
+CONFIG = load_config()
+
+# ---------------------------------------------------------------------------
+# Heuristic baseline policy (for comparison)
+# ---------------------------------------------------------------------------
+
+def heuristic_policy(state: list) -> float:
+    """Simple rule-based policy for baseline comparison."""
+    wait_us = state[IDX_WAIT_US]
+    csw = state[IDX_CTX_SWITCHES]
+
+    if wait_us > 15:
+        return -0.6
+    elif csw > 10:
+        return -0.3
+    else:
+        return 0.05
+
+# ---------------------------------------------------------------------------
+# Model wrappers
+# ---------------------------------------------------------------------------
+
+class StrategistWrapper:
+    """Wraps a GGUF model for use in the demo."""
+
+    def __init__(self, model_path: str):
+        from llama_cpp import Llama
+        self.llm = Llama(model_path=model_path, n_ctx=512, n_threads=4, verbose=False)
+
+    def predict(self, state: list, pid: int, cpu: int) -> tuple:
+        state_str = format_state(state)
+        prompt = (
+            "<|system|>You are a Linux kernel scheduling strategist. "
+            "Given the current system state, output a scheduling action.<|end|>\n"
+            f"<|user|>[STATE] {state_str}\n"
+            f"[PID] {pid} [CPU] {cpu}\n"
+            "[ACTION]<|end|>\n"
+            "<|assistant|>"
+        )
+
+        start = time.perf_counter()
+        output = self.llm(prompt, max_tokens=8, temperature=0.2)
+        latency = (time.perf_counter() - start) * 1000
+
+        text = output["choices"][0]["text"]
+        action_match = re.search(r"([-+]?\d*\.?\d+)", text)
+
+        action = float(action_match.group(1)) if action_match else 0.0
+        action = max(-1.0, min(1.0, action))
+
+        return action, latency
+
+
+class WorldModelWrapper:
+    """Wraps a GGUF or HF model for world model predictions."""
+
+    def __init__(self, model_path: str):
+        from llama_cpp import Llama
+        self.llm = Llama(model_path=model_path, n_ctx=512, n_threads=4, verbose=False)
+
+    def predict_next_state(self, state: list, action: float, pid: int) -> list:
+        state_str = format_state(state)
+        prompt = (
+            "<|system|>You are a Linux kernel simulator. "
+            "Predict the next system state.<|end|>\n"
+            f"<|user|>[STATE] {state_str}\n"
+            f"[ACTION] {action:.4f}\n"
+            f"[PID] {pid}\n"
+            "Predict [NEXT_STATE]<|end|>\n"
+            "<|assistant|>"
+        )
+
+        output = self.llm(prompt, max_tokens=128, temperature=0.1)
+        text = output["choices"][0]["text"]
+
+        # Parse predicted state
+        values = []
+        for part in text.split("|"):
+            part = part.strip()
+            if ":" in part:
+                try:
+                    values.append(float(part.split(":")[1]))
+                except ValueError:
+                    pass
+
+        if len(values) == len(FEATURE_NAMES):
+            return values
+        return state  # fallback: return same state
+
+# ---------------------------------------------------------------------------
+# Demo functions
+# ---------------------------------------------------------------------------
+
+def run_comparison(
+    record: dict,
+    strategist: Optional[StrategistWrapper],
+    reward_computer: RewardComputer,
+) -> dict:
+    """Run baseline vs trained comparison on a single transition."""
+    state = record["state"]
+    next_state = record["next_state"]
+    pid = record["pid"]
+    cpu = record["cpu"]
+
+    # Heuristic baseline
+    h_action = heuristic_policy(state)
+    h_reward = reward_computer.compute_total(
+        state=state, action=h_action, prev_action=0.0,
+        next_state=next_state,
+    )
+
+    result = {
+        "state": format_state(state),
+        "actual_next_state": format_state(next_state),
+        "pid": pid,
+        "cpu": cpu,
+        "heuristic": {
+            "action": h_action,
+            "reward": h_reward,
+        },
+    }
+
+    # Trained strategist
+    if strategist:
+        s_action, s_latency = strategist.predict(state, pid, cpu)
+        s_reward = reward_computer.compute_total(
+            state=state, action=s_action, prev_action=0.0,
+            next_state=next_state,
+        )
+        result["strategist"] = {
+            "action": s_action,
+            "latency_ms": s_latency,
+            "reward": s_reward,
+        }
+
+    return result
+
+# ---------------------------------------------------------------------------
+# Gradio app
+# ---------------------------------------------------------------------------
+
+def build_gradio_app(
+    test_records: list,
+    strategist: Optional[StrategistWrapper],
+    world_model: Optional[WorldModelWrapper],
+):
+    """Build and return the Gradio interface."""
+    import gradio as gr
+
+    reward_computer = RewardComputer(alpha=1.0, beta=2.0, gamma=0.5)
+
+    def analyze_state(record_idx: int):
+        """Run analysis on a selected test record."""
+        idx = int(record_idx) % len(test_records)
+        record = test_records[idx]
+        result = run_comparison(record, strategist, reward_computer)
+
+        # Format output
+        lines = []
+        lines.append(f"## State #{idx}")
+        lines.append(f"**PID:** {result['pid']}  |  **CPU:** {result['cpu']}")
+        lines.append(f"**Current State:** `{result['state']}`")
+        lines.append(f"**Actual Next State:** `{result['actual_next_state']}`")
+        lines.append("")
+
+        # Heuristic
+        h = result["heuristic"]
+        lines.append("### Heuristic Baseline")
+        lines.append(f"- **Action:** {h['action']:.4f}")
+        lines.append(f"- **Total Reward:** {h['reward']['total']:.4f}")
+        lines.append(f"  - Throughput: {h['reward']['throughput']:.4f}")
+        lines.append(f"  - Latency: {h['reward']['latency']:.4f}")
+        lines.append(f"  - Stability: {h['reward']['stability']:.4f}")
+        lines.append("")
+
+        # Strategist
+        if "strategist" in result:
+            s = result["strategist"]
+            lines.append("### Trained Strategist")
+            lines.append(f"- **Action:** {s['action']:.4f}")
+            lines.append(f"- **Inference Latency:** {s['latency_ms']:.1f}ms")
+            lines.append(f"- **Total Reward:** {s['reward']['total']:.4f}")
+            lines.append(f"  - Throughput: {s['reward']['throughput']:.4f}")
+            lines.append(f"  - Latency: {s['reward']['latency']:.4f}")
+            lines.append(f"  - Stability: {s['reward']['stability']:.4f}")
+            lines.append("")
+
+            # Comparison
+            delta = s["reward"]["total"] - h["reward"]["total"]
+            direction = "better" if delta > 0 else "worse"
+            lines.append(f"### Comparison")
+            lines.append(f"Strategist is **{delta:+.4f}** reward ({direction} than heuristic)")
+        else:
+            lines.append("*No trained model loaded. Run with --strategist-model to compare.*")
+
+        return "\n".join(lines)
+
+    def batch_comparison():
+        """Run comparison across all test records and report aggregate stats."""
+        h_rewards = []
+        s_rewards = []
+        s_latencies = []
+
+        n = min(50, len(test_records))  # cap for speed
+        for i in range(n):
+            result = run_comparison(test_records[i], strategist, reward_computer)
+            h_rewards.append(result["heuristic"]["reward"]["total"])
+            if "strategist" in result:
+                s_rewards.append(result["strategist"]["reward"]["total"])
+                s_latencies.append(result["strategist"]["latency_ms"])
+
+        s_mean = f"{np.mean(s_rewards):.4f}" if s_rewards else "N/A"
+        s_std = f"{np.std(s_rewards):.4f}" if s_rewards else "N/A"
+
+        lines = [f"## Batch Comparison ({n} samples)\n"]
+        lines.append("| Metric | Heuristic | Strategist |")
+        lines.append("|--------|-----------|------------|")
+        lines.append(f"| Mean Reward | {np.mean(h_rewards):.4f} | {s_mean} |")
+        lines.append(f"| Std Reward | {np.std(h_rewards):.4f} | {s_std} |")
+        if s_latencies:
+            lines.append(f"| Mean Latency | N/A | {np.mean(s_latencies):.1f}ms |")
+            lines.append(f"| P95 Latency | N/A | {np.percentile(s_latencies, 95):.1f}ms |")
+        if s_rewards:
+            win_rate = sum(1 for s, h in zip(s_rewards, h_rewards) if s > h) / len(s_rewards)
+            lines.append(f"| Win Rate | — | {win_rate*100:.1f}% |")
+
+        return "\n".join(lines)
+
+    # Build Gradio UI
+    with gr.Blocks(title="KernelX Intelligence Layer") as app:
+        gr.Markdown("# KernelX Intelligence Layer Demo")
+        gr.Markdown("Compare heuristic baseline vs trained Strategist on real kernel states.")
+
+        with gr.Row():
+            record_slider = gr.Slider(
+                minimum=0, maximum=len(test_records) - 1,
+                step=1, value=0, label="Test Record Index"
+            )
+            analyze_btn = gr.Button("Analyze", variant="primary")
+
+        output_md = gr.Markdown()
+        analyze_btn.click(fn=analyze_state, inputs=[record_slider], outputs=[output_md])
+
+        gr.Markdown("---")
+        batch_btn = gr.Button("Run Batch Comparison (50 samples)")
+        batch_output = gr.Markdown()
+        batch_btn.click(fn=batch_comparison, outputs=[batch_output])
+
+    return app
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="KernelX Gradio Demo")
+    parser.add_argument("--test-data", required=True, help="Path to test.jsonl")
+    parser.add_argument("--strategist-model", default=None, help="GGUF Strategist model")
+    parser.add_argument("--world-model", default=None, help="GGUF World Model")
+    parser.add_argument("--no-model", action="store_true", help="Run without trained models")
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--share", action="store_true", help="Create public Gradio link")
+    args = parser.parse_args()
+
+    # Load test data
+    records = [json.loads(l) for l in open(args.test_data) if l.strip()]
+    print(f"Loaded {len(records)} test records")
+
+    # Load models
+    strategist = None
+    world_model = None
+
+    if not args.no_model:
+        if args.strategist_model:
+            print(f"Loading Strategist: {args.strategist_model}")
+            strategist = StrategistWrapper(args.strategist_model)
+        if args.world_model:
+            print(f"Loading World Model: {args.world_model}")
+            world_model = WorldModelWrapper(args.world_model)
+
+    app = build_gradio_app(records, strategist, world_model)
+    app.launch(server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+    main()

training/environment/__init__.py

@@ -0,0 +1,4 @@
+"""KernelX RL environment and reward functions."""
+
+from .rewards import RewardComputer
+from .environment import KernelSchedulerEnv, KernelState, KernelAction

training/environment/environment.py

@@ -0,0 +1,131 @@
+"""
+KernelX Intelligence Layer — RL Environment (OpenEnv structure)
+
+Provides reset/step interface for training the Strategist policy via GRPO.
+Replays recorded transitions from the preprocessed JSONL dataset and
+computes multi-objective rewards.
+"""
+
+import json
+import random
+from dataclasses import dataclass, field
+from typing import List, Tuple
+
+from .rewards import RewardComputer
+
+
+@dataclass
+class KernelState:
+    """Observation wrapper for the RL environment."""
+    features: List[float]   # active features (10D after preprocessing)
+    pid: int
+    cpu: int
+    timestep: int
+    prev_action: float
+
+
+@dataclass
+class KernelAction:
+    """Action output from the Strategist."""
+    value: float            # scheduling weight in [-1.0, 1.0]
+
+
+class KernelSchedulerEnv:
+    """Offline RL environment that replays recorded kernel transitions.
+
+    Each episode starts at a random position in the dataset and runs for
+    max_steps transitions.  The reward is computed from the multi-objective
+    RewardComputer.
+    """
+
+    def __init__(
+        self,
+        data_path: str = "training/data/train.jsonl",
+        max_steps: int = 10,
+        alpha: float = 1.0,
+        beta: float = 2.0,
+        gamma: float = 0.5,
+    ):
+        self.records = [json.loads(l) for l in open(data_path) if l.strip()]
+        self.max_steps = max_steps
+        self.reward_computer = RewardComputer(alpha=alpha, beta=beta, gamma=gamma)
+
+        # Episode state
+        self.timestep = 0
+        self.current_idx = 0
+        self.prev_action = 0.0
+
+        if len(self.records) < max_steps + 1:
+            raise ValueError(
+                f"Dataset has {len(self.records)} records but max_steps={max_steps} "
+                f"requires at least {max_steps + 1}"
+            )
+
+    def reset(self) -> KernelState:
+        """Start a fresh episode from a random point in the dataset."""
+        self.timestep = 0
+        self.current_idx = random.randint(0, len(self.records) - self.max_steps - 1)
+        self.prev_action = 0.0
+        return self._get_state()
+
+    def step(self, action: KernelAction) -> Tuple[KernelState, dict, bool]:
+        """Apply action, compute reward, advance to next state.
+
+        Returns:
+            next_state: The new KernelState after the transition
+            reward_breakdown: Dict with 'total' and per-component rewards
+            done: Whether the episode has ended
+        """
+        current = self.records[self.current_idx + self.timestep]
+        next_idx = self.current_idx + self.timestep + 1
+        next_rec = self.records[next_idx] if next_idx < len(self.records) else current
+
+        reward_breakdown = self.reward_computer.compute_total(
+            state=current["state"],
+            action=action,
+            prev_action=self.prev_action,
+            next_state=next_rec["state"],
+        )
+
+        self.timestep += 1
+        self.prev_action = action.value
+        done = self.timestep >= self.max_steps
+
+        return self._get_state(), reward_breakdown, done
+
+    def _get_state(self) -> KernelState:
+        """Read the current state from the dataset."""
+        rec = self.records[self.current_idx + self.timestep]
+        return KernelState(
+            features=rec["state"],
+            pid=rec["pid"],
+            cpu=rec["cpu"],
+            timestep=self.timestep,
+            prev_action=self.prev_action,
+        )
+
+    def simulate(self, state_features: list, action_value: float) -> list:
+        """Lightweight next-state lookup for reward_fn during GRPO.
+
+        Finds the nearest recorded state in the dataset and returns
+        its recorded next_state. This is a simple approximation;
+        the World Model provides higher-fidelity simulation.
+        """
+        import numpy as np
+
+        state_arr = np.array(state_features)
+        best_dist = float("inf")
+        best_next = state_features  # fallback
+
+        # Sample a subset to keep this fast
+        sample_size = min(500, len(self.records))
+        indices = random.sample(range(len(self.records)), sample_size)
+
+        for idx in indices:
+            rec = self.records[idx]
+            dist = float(np.linalg.norm(state_arr - np.array(rec["state"])))
+            if dist < best_dist:
+                best_dist = dist
+                best_next = rec["next_state"]
+
+        return best_next

training/environment/rewards.py

@@ -0,0 +1,142 @@
+"""
+KernelX Intelligence Layer — Multi-Objective Reward Function
+
+Decomposes R_t = alpha * log(throughput + 1) - beta * delta_wait - gamma * |a_t - a_{t-1}|
+into independent, inspectable reward components.
+
+Throughput proxy: delta(sum_exec_runtime) since IPC is not yet collected
+from PMU sidecars (indices 13-22 are zero in current data).
+"""
+
+import numpy as np
+
+# Active feature indices (positions within the 10D active vector)
+IDX_CPU           = 0
+IDX_PRIO          = 1
+IDX_STATIC_PRIO   = 2
+IDX_NORMAL_PRIO   = 3
+IDX_EXEC_NS       = 4   # symlog-scaled sum_exec_runtime
+IDX_VRUNTIME      = 5   # symlog-scaled vruntime
+IDX_MIGRATIONS    = 6   # symlog-scaled nr_migrations
+IDX_CPUS_ALLOWED  = 7
+IDX_CTX_SWITCHES  = 8
+IDX_WAIT_US       = 9   # wait time in microseconds
+
+
+class RewardComputer:
+    """Multi-objective reward for kernel scheduling decisions.
+
+    Components:
+        throughput  — reward for CPU progress (delta exec_runtime)
+        latency     — penalty for increased wait time
+        stability   — penalty for jittery action changes
+        format      — reward for action value in valid range
+    """
+
+    def __init__(self, alpha: float = 1.0, beta: float = 2.0, gamma: float = 0.5):
+        self.alpha = alpha
+        self.beta = beta
+        self.gamma = gamma
+
+    def throughput_reward(self, state: list, next_state: list) -> float:
+        """Reward for throughput improvement.
+
+        Uses delta(sum_exec_runtime) as proxy for IPC (not yet available).
+        Index 4 = sum_exec_runtime (symlog-scaled).
+        Positive delta means the process accumulated CPU time = progress.
+        """
+        exec_delta = next_state[IDX_EXEC_NS] - state[IDX_EXEC_NS]
+        return self.alpha * float(np.log(max(0.0, exec_delta) + 1))
+
+    def latency_reward(self, state: list, next_state: list) -> float:
+        """Penalty for increased wait time.
+
+        Index 9 = wait_time in microseconds (raw, not symlog-scaled).
+        Positive delta means wait time increased = bad.
+        """
+        wait_delta = next_state[IDX_WAIT_US] - state[IDX_WAIT_US]
+        return -self.beta * max(0.0, wait_delta)
+
+    def stability_reward(self, action: float, prev_action: float) -> float:
+        """Penalty for jittery scheduling changes."""
+        return -self.gamma * abs(action - prev_action)
+
+    def format_reward(self, action_value: float) -> float:
+        """Reward for action value in valid range."""
+        return 1.0 if -1.0 <= action_value <= 1.0 else 0.0
+
+    def compute_total(
+        self,
+        state: list,
+        action,  # KernelAction or has .value attribute
+        prev_action: float,
+        next_state: list,
+    ) -> dict:
+        """Compute all reward components and return breakdown."""
+        action_val = action.value if hasattr(action, "value") else float(action)
+
+        r_throughput = self.throughput_reward(state, next_state)
+        r_latency = self.latency_reward(state, next_state)
+        r_stability = self.stability_reward(action_val, prev_action)
+        r_format = self.format_reward(action_val)
+
+        total = r_throughput + r_latency + r_stability + r_format
+
+        return {
+            "total": total,
+            "throughput": r_throughput,
+            "latency": r_latency,
+            "stability": r_stability,
+            "format": r_format,
+        }
+
+    def calibrate(self, records: list, n_samples: int = 200) -> dict:
+        """Run reward function on random actions to verify healthy distribution.
+
+        Returns stats about reward distribution for tuning alpha/beta/gamma.
+        """
+        import random
+
+        totals = []
+        components = {"throughput": [], "latency": [], "stability": [], "format": []}
+
+        samples = random.sample(records, min(n_samples, len(records)))
+        for rec in samples:
+            fake_action_val = random.uniform(-1.0, 1.0)
+            fake_prev = random.uniform(-1.0, 1.0)
+
+            r_t = self.throughput_reward(rec["state"], rec["next_state"])
+            r_l = self.latency_reward(rec["state"], rec["next_state"])
+            r_s = self.stability_reward(fake_action_val, fake_prev)
+            r_f = self.format_reward(fake_action_val)
+
+            total = r_t + r_l + r_s + r_f
+            totals.append(total)
+            components["throughput"].append(r_t)
+            components["latency"].append(r_l)
+            components["stability"].append(r_s)
+            components["format"].append(r_f)
+
+        stats = {
+            "total": {
+                "mean": float(np.mean(totals)),
+                "std": float(np.std(totals)),
+                "min": float(np.min(totals)),
+                "max": float(np.max(totals)),
+                "zero_rate": float(np.mean([1 if t == 0 else 0 for t in totals])),
+            }
+        }
+        for name, vals in components.items():
+            stats[name] = {
+                "mean": float(np.mean(vals)),
+                "std": float(np.std(vals)),
+                "min": float(np.min(vals)),
+                "max": float(np.max(vals)),
+            }
+
+        print("\n=== Reward Calibration ===")
+        for name, s in stats.items():
+            print(f"  {name:>12}: mean={s['mean']:+.3f}  std={s['std']:.3f}  "
+                  f"range=[{s['min']:+.3f}, {s['max']:+.3f}]")
+
+        return stats

training/inference/__init__.py

@@ -0,0 +1 @@
+"""KernelX inference engine for real-time scheduling."""

training/inference/benchmark_latency.py

@@ -0,0 +1,147 @@
+"""
+KernelX Intelligence Layer — Latency Benchmark
+
+Measures end-to-end inference latency of the quantized Strategist model
+on the target CPU hardware. Reports mean, P50, P95, P99, and max latency.
+
+Usage:
+    python -m training.inference.benchmark_latency \
+        --model training/models/strategist_merged/strategist-q4km.gguf \
+        --samples 200
+"""
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+from training.data.preprocess import FEATURE_NAMES, format_state, load_config
+from training.inference.strategy_engine import build_inference_prompt, parse_output
+
+CONFIG = load_config()
+
+
+def benchmark(
+    model_path: str,
+    test_data_path: str = None,
+    n_samples: int = 200,
+    n_threads: int = 4,
+    temperature: float = 0.2,
+    max_tokens: int = 64,
+    warmup: int = 5,
+):
+    """Run latency benchmark on the quantized GGUF model."""
+    from llama_cpp import Llama
+
+    print(f"Loading model: {model_path}")
+    llm = Llama(model_path=model_path, n_ctx=512, n_threads=n_threads, verbose=False)
+
+    # Build test prompts
+    if test_data_path:
+        records = [json.loads(l) for l in open(test_data_path) if l.strip()]
+    else:
+        # Synthetic test data
+        records = []
+        for i in range(n_samples):
+            state = [float(i % 16)]  # cpu
+            state += [120.0, 120.0, 120.0]  # priorities
+            state += [20.0 + i * 0.1, 28.0 + i * 0.01, 8.0 + i * 0.001]  # symlog'd
+            state += [16.0, float(i % 50), float(5 + i % 30)]  # cpus, csw, wt_us
+            records.append({"state": state, "pid": 1000 + i, "cpu": i % 16})
+
+    records = records[:n_samples]
+    prompts = []
+    for rec in records:
+        prompts.append(build_inference_prompt(rec["state"], rec["pid"], rec["cpu"]))
+
+    # Warmup
+    print(f"Warming up ({warmup} iterations) ...")
+    for i in range(warmup):
+        llm(prompts[i % len(prompts)], max_tokens=max_tokens, temperature=temperature)
+
+    # Benchmark
+    print(f"\nBenchmarking {len(prompts)} samples ...")
+    latencies = []
+    format_ok = 0
+    token_counts = []
+
+    for prompt in prompts:
+        start = time.perf_counter()
+        output = llm(prompt, max_tokens=max_tokens, temperature=temperature)
+        elapsed = time.perf_counter() - start
+        latencies.append(elapsed)
+
+        text = output["choices"][0]["text"]
+        tokens = output["usage"]["completion_tokens"]
+        token_counts.append(tokens)
+
+        action_val = parse_output(text)
+        if -1.0 <= action_val <= 1.0:
+            format_ok += 1
+
+    latencies_ms = np.array(latencies) * 1000
+
+    # Report
+    target = CONFIG["model"]["target_inference_ms"]
+    print(f"\n{'='*50}")
+    print(f" KernelX Latency Benchmark")
+    print(f"{'='*50}")
+    print(f" Model:    {Path(model_path).name}")
+    print(f" Threads:  {n_threads}")
+    print(f" Samples:  {len(prompts)}")
+    print(f" Target:   <{target}ms")
+    print(f"{'='*50}")
+    print(f" Mean:     {np.mean(latencies_ms):>8.1f} ms")
+    print(f" Median:   {np.median(latencies_ms):>8.1f} ms")
+    print(f" P95:      {np.percentile(latencies_ms, 95):>8.1f} ms")
+    print(f" P99:      {np.percentile(latencies_ms, 99):>8.1f} ms")
+    print(f" Max:      {np.max(latencies_ms):>8.1f} ms")
+    print(f" Min:      {np.min(latencies_ms):>8.1f} ms")
+    print(f" Std:      {np.std(latencies_ms):>8.1f} ms")
+    print(f"{'='*50}")
+    print(f" Tokens/s: {np.sum(token_counts) / np.sum(latencies):>8.1f}")
+    print(f" Format OK:{format_ok}/{len(prompts)} ({format_ok/len(prompts)*100:.1f}%)")
+    print(f"{'='*50}")
+
+    p95 = np.percentile(latencies_ms, 95)
+    if p95 <= target:
+        print(f" VERDICT:  PASS (P95 {p95:.1f}ms <= {target}ms)")
+    else:
+        print(f" VERDICT:  FAIL (P95 {p95:.1f}ms > {target}ms)")
+
+    return {
+        "mean_ms": float(np.mean(latencies_ms)),
+        "p95_ms": float(p95),
+        "p99_ms": float(np.percentile(latencies_ms, 99)),
+        "format_ok_pct": format_ok / len(prompts) * 100,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark KernelX inference latency")
+    parser.add_argument("--model", required=True, help="Path to GGUF model")
+    parser.add_argument("--test-data", default=None, help="Test JSONL (optional)")
+    parser.add_argument("--samples", type=int, default=200)
+    parser.add_argument("--threads", type=int, default=4)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-tokens", type=int, default=8)
+    parser.add_argument("--warmup", type=int, default=5)
+    args = parser.parse_args()
+
+    benchmark(
+        model_path=args.model,
+        test_data_path=args.test_data,
+        n_samples=args.samples,
+        n_threads=args.threads,
+        temperature=args.temperature,
+        max_tokens=args.max_tokens,
+        warmup=args.warmup,
+    )
+
+
+if __name__ == "__main__":
+    main()

training/inference/strategy_engine.py

@@ -0,0 +1,307 @@
+"""
+KernelX Intelligence Layer — Inference Engine (Stage 6)
+
+Three-thread architecture for real-time kernel scheduling:
+  Thread 1 (Telemetry):  reads latest state from shared memory
+  Thread 2 (Strategist): runs LLM inference every cycle
+  Thread 3 (Updater):    writes action to shared memory command slot
+
+Uses llama.cpp via llama-cpp-python for sub-50ms CPU inference
+with the quantized GGUF Strategist model.
+
+Usage:
+    python -m training.inference.strategy_engine \
+        --model training/models/strategist_merged/strategist-q4km.gguf \
+        --shm-path /dev/shm/kernelx_state
+"""
+
+import argparse
+import json
+import mmap
+import os
+import re
+import struct
+import sys
+import threading
+import time
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+from training.data.preprocess import (
+    ACTIVE_FEATURES, FEATURE_NAMES, SYMLOG_FEATURES,
+    symmetric_log, format_state, load_config,
+)
+
+CONFIG = load_config()
+
+# ---------------------------------------------------------------------------
+# Shared memory layout (must match bridge/src/main.rs HUDState)
+# ---------------------------------------------------------------------------
+# features:           [u64; 24]  = 192 bytes (offset 0)
+# current_action:     f32        = 4   bytes (offset 192)
+# active_pid:         u32        = 4   bytes (offset 196)
+# is_clamped:         u32        = 4   bytes (offset 200)
+# reasoning:          [u8; 128]  = 128 bytes (offset 204)
+# p99_wait_us:        u64        = 8   bytes (offset 332)
+#                                  Total: 340 bytes
+
+SHM_SIZE = 340
+FEATURES_OFFSET = 0
+FEATURES_SIZE = 192     # 24 * 8 bytes (u64)
+ACTION_OFFSET = 192
+PID_OFFSET = 196
+CLAMPED_OFFSET = 200
+REASONING_OFFSET = 204
+REASONING_SIZE = 128
+P99_OFFSET = 332
+
+# ---------------------------------------------------------------------------
+# Prompt building (mirrors train_strategist.py)
+# ---------------------------------------------------------------------------
+
+def build_inference_prompt(active_features: list, pid: int, cpu: int) -> str:
+    state_str = format_state(active_features)
+    return (
+        "<|system|>You are a Linux kernel scheduling strategist. "
+        "Given the current system state, output a scheduling action.<|end|>\n"
+        f"<|user|>[STATE] {state_str}\n"
+        f"[PID] {pid} [CPU] {cpu}\n"
+        "[ACTION]<|end|>\n"
+        "<|assistant|>"
+    )
+
+
+def parse_output(text: str) -> float:
+    """Parse action float from model output."""
+    action_match = re.search(r"\[ACTION\]\s*([-+]?\d*\.?\d+)", text)
+    if not action_match:
+        action_match = re.search(r"([-+]?\d*\.?\d+)", text)
+    if not action_match:
+        return 0.0
+
+    action_val = float(action_match.group(1))
+    return max(-1.0, min(1.0, action_val))
+
+# ---------------------------------------------------------------------------
+# Shared memory reader/writer
+# ---------------------------------------------------------------------------
+
+def read_features_from_shm(shm: mmap.mmap) -> tuple:
+    """Read 24D raw features + PID from shared memory.
+
+    Returns:
+        (raw_features_24d: list[float], pid: int, cpu: int)
+    """
+    shm.seek(FEATURES_OFFSET)
+    raw_bytes = shm.read(FEATURES_SIZE)
+    raw_features = list(np.frombuffer(raw_bytes, dtype=np.uint64).astype(np.float64))
+
+    shm.seek(PID_OFFSET)
+    pid = struct.unpack("<I", shm.read(4))[0]
+
+    # CPU is in features[0] (from bpf_get_smp_processor_id)
+    cpu = int(raw_features[0]) if raw_features else 0
+
+    return raw_features, pid, cpu
+
+
+def preprocess_for_inference(raw_features: list) -> list:
+    """Apply symlog scaling and extract active features for the LLM."""
+    f = list(raw_features)
+    for idx in SYMLOG_FEATURES:
+        f[idx] = symmetric_log(f[idx])
+    return [f[i] for i in ACTIVE_FEATURES]
+
+
+def write_action_to_shm(shm: mmap.mmap, action: float):
+    """Write action value to shared memory."""
+    shm.seek(ACTION_OFFSET)
+    shm.write(struct.pack("<f", action))
+
+# ---------------------------------------------------------------------------
+# Strategy Engine (three-thread architecture)
+# ---------------------------------------------------------------------------
+
+class StrategyEngine:
+    """Real-time scheduling inference engine.
+
+    Reads kernel state from shared memory, runs the quantized Strategist
+    model, and writes the scheduling action back.
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        shm_path: str = "/dev/shm/kernelx_state",
+        n_threads: int = 2,
+        poll_interval_ms: float = 10.0,
+        update_interval_ms: float = 50.0,
+        temperature: float = 0.2,
+        max_tokens: int = 8,
+    ):
+        from llama_cpp import Llama
+
+        self.model_path = model_path
+        self.shm_path = shm_path
+        self.poll_interval = poll_interval_ms / 1000.0
+        self.update_interval = update_interval_ms / 1000.0
+        self.temperature = temperature
+        self.max_tokens = max_tokens
+
+        # Load model
+        print(f"[StrategyEngine] Loading model: {model_path}")
+        self.llm = Llama(
+            model_path=model_path,
+            n_ctx=512,
+            n_threads=n_threads,
+            verbose=False,
+        )
+
+        # Shared state (protected by lock)
+        self.lock = threading.Lock()
+        self.current_features = [0.0] * len(ACTIVE_FEATURES)
+        self.current_pid = 0
+        self.current_cpu = 0
+        self.latest_action = 0.0
+
+        # Control
+        self.running = False
+        self.shm = None
+
+        # Metrics
+        self.inference_count = 0
+        self.inference_latencies = []
+
+    def _open_shm(self) -> mmap.mmap:
+        """Open shared memory file for read/write."""
+        fd = os.open(self.shm_path, os.O_RDWR)
+        return mmap.mmap(fd, SHM_SIZE, mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE)
+
+    def start(self):
+        """Start all three threads."""
+        if not os.path.exists(self.shm_path):
+            print(f"[StrategyEngine] WARNING: SHM {self.shm_path} not found.")
+            print("[StrategyEngine] Running in dry-run mode (no SHM I/O).")
+            self.shm = None
+        else:
+            self.shm = self._open_shm()
+            print(f"[StrategyEngine] Connected to SHM: {self.shm_path}")
+
+        self.running = True
+
+        threads = [
+            threading.Thread(target=self._telemetry_loop, name="telemetry", daemon=True),
+            threading.Thread(target=self._strategist_loop, name="strategist", daemon=True),
+            threading.Thread(target=self._update_loop, name="updater", daemon=True),
+        ]
+
+        for t in threads:
+            t.start()
+            print(f"[StrategyEngine] Started {t.name} thread")
+
+        print("[StrategyEngine] All threads running. Press Ctrl+C to stop.")
+
+        try:
+            while self.running:
+                time.sleep(1.0)
+                # Periodic stats
+                if self.inference_count > 0 and self.inference_count % 20 == 0:
+                    recent = self.inference_latencies[-20:]
+                    avg_ms = np.mean(recent) * 1000
+                    print(f"[StrategyEngine] Inferences: {self.inference_count}, "
+                          f"Avg latency: {avg_ms:.1f}ms, "
+                          f"Action: {self.latest_action:.4f}")
+        except KeyboardInterrupt:
+            self.stop()
+
+    def stop(self):
+        """Stop all threads."""
+        print("\n[StrategyEngine] Shutting down ...")
+        self.running = False
+        if self.shm:
+            self.shm.close()
+
+    def _telemetry_loop(self):
+        """Thread 1: Read latest state from shared memory."""
+        while self.running:
+            if self.shm:
+                try:
+                    raw_features, pid, cpu = read_features_from_shm(self.shm)
+                    active = preprocess_for_inference(raw_features)
+                    with self.lock:
+                        self.current_features = active
+                        self.current_pid = pid
+                        self.current_cpu = cpu
+                except Exception:
+                    pass
+            time.sleep(self.poll_interval)
+
+    def _strategist_loop(self):
+        """Thread 2: Run LLM inference every cycle."""
+        while self.running:
+            with self.lock:
+                features = list(self.current_features)
+                pid = self.current_pid
+                cpu = self.current_cpu
+
+            prompt = build_inference_prompt(features, pid, cpu)
+
+            start = time.perf_counter()
+            output = self.llm(prompt, max_tokens=self.max_tokens,
+                              temperature=self.temperature)
+            elapsed = time.perf_counter() - start
+
+            text = output["choices"][0]["text"]
+            action = parse_output(text)
+
+            with self.lock:
+                self.latest_action = action
+
+            self.inference_count += 1
+            self.inference_latencies.append(elapsed)
+
+    def _update_loop(self):
+        """Thread 3: Write action to shared memory command slot."""
+        while self.running:
+            with self.lock:
+                action = self.latest_action
+
+            if self.shm:
+                try:
+                    write_action_to_shm(self.shm, action)
+                except Exception:
+                    pass
+
+            time.sleep(self.update_interval)
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="KernelX Strategy Engine")
+    parser.add_argument("--model", required=True, help="Path to GGUF model")
+    parser.add_argument("--shm-path", default="/dev/shm/kernelx_state")
+    parser.add_argument("--threads", type=int, default=2, help="llama.cpp threads")
+    parser.add_argument("--poll-ms", type=float, default=10.0)
+    parser.add_argument("--update-ms", type=float, default=50.0)
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--max-tokens", type=int, default=64)
+    args = parser.parse_args()
+
+    engine = StrategyEngine(
+        model_path=args.model,
+        shm_path=args.shm_path,
+        n_threads=args.threads,
+        poll_interval_ms=args.poll_ms,
+        update_interval_ms=args.update_ms,
+        temperature=args.temperature,
+        max_tokens=args.max_tokens,
+    )
+    engine.start()
+
+
+if __name__ == "__main__":
+    main()

training/requirements.txt

@@ -0,0 +1,12 @@
+torch>=2.0
+transformers>=4.40
+trl>=0.12
+unsloth
+peft>=0.13
+datasets>=2.18
+accelerate>=0.30
+llama-cpp-python>=0.3
+wandb
+pandas
+numpy
+gradio>=4.0

training/run_pipeline.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""
+KernelX Intelligence Layer — Full Training Pipeline Runner
+
+Runs all stages end-to-end: preprocess -> World Model SFT -> Strategist GRPO -> export.
+
+Usage:
+    # Full pipeline
+    python training/run_pipeline.py \
+        --raw-data data/state_transitions.jsonl \
+        --output-root training
+
+    # Resume from a specific stage
+    python training/run_pipeline.py \
+        --raw-data data/state_transitions.jsonl \
+        --output-root training \
+        --start-stage 3
+
+Stages:
+    1 = Preprocess data
+    2 = Train World Model (SFT)
+    3 = Train Strategist (warm-start SFT + GRPO)
+    4 = Export & quantize to GGUF
+    5 = Validate quantized model
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+def main():
+    parser = argparse.ArgumentParser(description="KernelX full training pipeline")
+    parser.add_argument("--raw-data", required=True, help="Path to raw state_transitions.jsonl")
+    parser.add_argument("--output-root", default="training", help="Root output directory")
+    parser.add_argument("--start-stage", type=int, default=1, help="Stage to start from (1-5)")
+    parser.add_argument("--end-stage", type=int, default=5, help="Stage to end at (1-5)")
+    parser.add_argument("--epochs-world", type=int, default=3, help="World Model training epochs")
+    parser.add_argument("--epochs-strategist", type=int, default=3, help="Strategist GRPO epochs")
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--quantize", default="Q4_K_M")
+    parser.add_argument("--wandb", action="store_true")
+    parser.add_argument("--curriculum", action="store_true")
+    args = parser.parse_args()
+
+    root = Path(args.output_root)
+    data_dir = root / "data"
+    models_dir = root / "models"
+
+    train_path = data_dir / "train.jsonl"
+    val_path = data_dir / "val.jsonl"
+    test_path = data_dir / "test.jsonl"
+
+    # ------------------------------------------------------------------
+    # Stage 1: Preprocess
+    # ------------------------------------------------------------------
+    if args.start_stage <= 1 <= args.end_stage:
+        print("\n" + "=" * 60)
+        print(" STAGE 1: Data Preprocessing")
+        print("=" * 60)
+
+        from training.data.preprocess import run_pipeline
+        run_pipeline(args.raw_data, str(data_dir), audit=True)
+
+    # ------------------------------------------------------------------
+    # Stage 2: World Model SFT
+    # ------------------------------------------------------------------
+    if args.start_stage <= 2 <= args.end_stage:
+        print("\n" + "=" * 60)
+        print(" STAGE 2: World Model Training (SFT)")
+        print("=" * 60)
+
+        from training.models.train_world_model import train, evaluate_world_model
+
+        wm_dir = models_dir / "world_model_final"
+        model, tokenizer = train(
+            train_path=str(train_path),
+            val_path=str(val_path),
+            output_dir=str(wm_dir),
+            num_epochs=args.epochs_world,
+            batch_size=args.batch_size,
+            use_wandb=args.wandb,
+        )
+
+        if test_path.exists():
+            evaluate_world_model(model, tokenizer, str(test_path))
+
+    # ------------------------------------------------------------------
+    # Stage 3: Strategist Training (warm-start + GRPO)
+    # ------------------------------------------------------------------
+    if args.start_stage <= 3 <= args.end_stage:
+        print("\n" + "=" * 60)
+        print(" STAGE 3: Strategist Training (GRPO)")
+        print("=" * 60)
+
+        from training.models.train_strategist import (
+            run_warmstart, run_grpo, inspect_generations,
+        )
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        records = [json.loads(l) for l in open(train_path) if l.strip()]
+
+        # Phase 1: Warm-start
+        ws_dir = models_dir / "strategist_warmstart"
+        model, tokenizer = run_warmstart(
+            records=records,
+            output_dir=str(ws_dir),
+            use_wandb=args.wandb,
+        )
+
+        # Phase 2: GRPO
+        strat_dir = models_dir / "strategist_final"
+        model, tokenizer = run_grpo(
+            model=model,
+            tokenizer=tokenizer,
+            train_records=records,
+            output_dir=str(strat_dir),
+            num_epochs=args.epochs_strategist,
+            use_curriculum=args.curriculum,
+            use_wandb=args.wandb,
+        )
+
+        # Inspect
+        inspect_generations(model, tokenizer, records, n=10)
+
+    # ------------------------------------------------------------------
+    # Stage 4: Export & Quantize
+    # ------------------------------------------------------------------
+    if args.start_stage <= 4 <= args.end_stage:
+        print("\n" + "=" * 60)
+        print(" STAGE 4: Export & Quantize to GGUF")
+        print("=" * 60)
+
+        from training.models.export_gguf import merge_lora, convert_to_gguf
+
+        strat_dir = models_dir / "strategist_final"
+        merged_dir = models_dir / "strategist_merged"
+        merged_path = merge_lora(str(strat_dir), str(merged_dir))
+        convert_to_gguf(merged_path, str(merged_dir), args.quantize)
+
+    # ------------------------------------------------------------------
+    # Stage 5: Validate
+    # ------------------------------------------------------------------
+    if args.start_stage <= 5 <= args.end_stage:
+        print("\n" + "=" * 60)
+        print(" STAGE 5: Validate Quantized Model")
+        print("=" * 60)
+
+        from training.models.export_gguf import validate_gguf
+
+        quant_name = f"strategist-{args.quantize.lower().replace('_', '')}.gguf"
+        gguf_path = models_dir / "strategist_merged" / quant_name
+
+        if gguf_path.exists():
+            validate_gguf(str(gguf_path), str(test_path) if test_path.exists() else None)
+        else:
+            print(f"GGUF not found at {gguf_path}. Run stage 4 first.")
+
+    # ------------------------------------------------------------------
+    print("\n" + "=" * 60)
+    print(" Pipeline complete!")
+    print("=" * 60)
+    print(f"\nArtifacts in: {root}/")
+    print(f"  Data:   {data_dir}/")
+    print(f"  Models: {models_dir}/")
+    print(f"\nNext steps:")
+    print(f"  1. Run demo:     python -m training.demo.app --test-data {test_path}")
+    print(f"  2. Run engine:   python -m training.inference.strategy_engine --model <gguf>")
+    print(f"  3. Benchmark:    python -m training.inference.benchmark_latency --model <gguf>")
+
+
+if __name__ == "__main__":
+    main()