{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d3e36380",
   "metadata": {},
   "source": [
    "# IncidentCommander — RL Training (Colab T4 / A100)\n",
    "\n",
    "**Stack:** Unsloth · Qwen2.5-1.5B-Instruct (4-bit QLoRA actor) · Qwen2.5-72B-Instruct critic via Hugging Face Inference Providers · PPO + GAE\n",
    "\n",
    "This notebook trains the IncidentCommander SRE agent against the in-repo simulator (**381 deterministic scenarios** including 166 saboteur/Slack scenarios, cascading topology, K8s adversary, runbook traps).\n",
    "\n",
    "## Why this critic\n",
    "Claude Haiku 4.5 is **not** hosted on Hugging Face — Anthropic and HF are different vendors. Instead we use **Qwen2.5-72B-Instruct via the HF Inference Providers router**, which (a) is free with any HF account, (b) is ~48× the size of the actor, giving the value head genuine compute headroom, and (c) is routinely served through Together / Nebius for free credits. The same code path also accepts `meta-llama/Meta-Llama-3.1-70B-Instruct` or `mistralai/Mistral-Large-2407`.\n",
    "\n",
    "## How to run\n",
    "1. **Runtime → Change runtime type → T4 GPU (or A100)**\n",
    "2. Run the cells top-to-bottom.\n",
    "3. Set your `HF_TOKEN` when prompted (free tier works)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a85e3201",
   "metadata": {},
   "source": [
    "## 1 · Verify GPU + install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c77e6201",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi || echo 'No GPU — switch the runtime to T4 or A100.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7b2bc58",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q --upgrade pip\n",
    "%pip install -q 'unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git'\n",
    "%pip install -q --no-deps 'xformers<0.0.27' trl peft accelerate bitsandbytes\n",
    "%pip install -q 'huggingface_hub>=0.25' pydantic==2.* httpx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1743d04",
   "metadata": {},
   "source": [
    "## 2 · Get the IncidentCommander repo into Colab\n",
    "\n",
    "**Option A (recommended):** push the repo to GitHub (`scripts/push_to_remotes.ps1` does that for you) and set `IC_REPO_URL` below.\n",
    "\n",
    "**Option B:** zip the local repo (`Compress-Archive -Path .\\rl-agent, .\\colab -DestinationPath ic.zip`) and drag it into Colab's `/content/` folder as `incident-commander.zip`.\n",
    "\n",
    "**Option C:** clone from a Hugging Face Space repo via `IC_HF_SPACE`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c065ac22",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, subprocess, sys, zipfile\n",
    "REPO_DIR = '/content/incident-commander'\n",
    "\n",
    "# Default to the user's HF Space so the notebook works zero-config.\n",
    "# Override either of these env vars to clone from somewhere else.\n",
    "os.environ.setdefault('IC_HF_SPACE', 'sagnik-mukherjee/incodent-commander')\n",
    "# Optional: GitHub mirror — also defaulted to the user's repo.\n",
    "os.environ.setdefault('IC_REPO_URL', 'https://github.com/r1cksync/meta-rl-hack.git')\n",
    "\n",
    "IC_REPO_URL = os.environ.get('IC_REPO_URL', '')\n",
    "if IC_REPO_URL and not os.path.isdir(REPO_DIR):\n",
    "    try:\n",
    "        subprocess.run(['git','clone','--depth','1', IC_REPO_URL, REPO_DIR], check=True)\n",
    "    except subprocess.CalledProcessError:\n",
    "        print(f'GitHub clone failed for {IC_REPO_URL}; will try the HF Space next.')\n",
    "\n",
    "if not os.path.isdir(REPO_DIR) and os.path.exists('/content/incident-commander.zip'):\n",
    "    os.makedirs(REPO_DIR, exist_ok=True)\n",
    "    with zipfile.ZipFile('/content/incident-commander.zip') as z:\n",
    "        z.extractall(REPO_DIR)\n",
    "\n",
    "IC_HF_SPACE = os.environ.get('IC_HF_SPACE', '')\n",
    "if not os.path.isdir(REPO_DIR) and IC_HF_SPACE:\n",
    "    subprocess.run(['git','clone',\n",
    "                    f'https://huggingface.co/spaces/{IC_HF_SPACE}', REPO_DIR], check=True)\n",
    "\n",
    "assert os.path.isdir(REPO_DIR), 'No repo present — pick Option A/B/C above.'\n",
    "%cd /content/incident-commander\n",
    "sys.path.insert(0, '/content/incident-commander')\n",
    "sys.path.insert(0, '/content/incident-commander/rl-agent')\n",
    "print('Repo ready at', REPO_DIR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ebb2e2c",
   "metadata": {},
   "source": [
    "## 3 · Hugging Face token\n",
    "\n",
    "`HF_TOKEN` powers (a) actor weight downloads and (b) the Qwen2.5-72B critic over the Inference Providers router.\n",
    "\n",
    "**Two ways to provide it** — the cell below tries them in order:\n",
    "\n",
    "1. **Colab Secret (recommended)** — open the 🔑 icon in Colab's left sidebar → **Add new secret** → name `HF_TOKEN`, paste your `hf_…` token, toggle **Notebook access** on. The cell picks it up automatically with no prompt.\n",
    "2. **Inline `getpass`** — if no secret is set, the cell falls back to a hidden prompt where you paste the token.\n",
    "\n",
    "> The token you wanted to use first (`hf_IBf…Mhl`) cannot be checked into this notebook because Hugging Face's pre-receive hook blocks any file that contains a token string. Paste it via secret or `getpass` instead. **Rotate it after this run** — it has been shared in plaintext."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09f781f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, getpass\n",
    "\n",
    "# 1) Try Colab's built-in secrets manager (left sidebar → Secrets → add HF_TOKEN).\n",
    "try:\n",
    "    from google.colab import userdata           # type: ignore\n",
    "    _t = userdata.get('HF_TOKEN')\n",
    "    if _t:\n",
    "        os.environ['HF_TOKEN'] = _t\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "# 2) Otherwise prompt for it (paste your hf_… token).\n",
    "if not os.environ.get('HF_TOKEN'):\n",
    "    os.environ['HF_TOKEN'] = getpass.getpass('Paste your HF token (hf_…): ')\n",
    "\n",
    "os.environ['HUGGING_FACE_HUB_TOKEN'] = os.environ['HF_TOKEN']\n",
    "os.environ['INCIDENT_COMMANDER_MOCK'] = 'true'\n",
    "\n",
    "from huggingface_hub import login, whoami\n",
    "login(os.environ['HF_TOKEN'], add_to_git_credential=False)\n",
    "me = whoami(token=os.environ['HF_TOKEN'])\n",
    "print(f\"HF login OK as: {me.get('name', me)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fa944ce",
   "metadata": {},
   "source": [
    "## 4 · Smoke-test the simulator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9882eba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, glob, collections\n",
    "sys.path.insert(0, '/content/incident-commander/rl-agent')\n",
    "from simulator import SimState, dispatch, load_scenario\n",
    "files = sorted(glob.glob('rl-agent/scenarios/sim/*/*.json'))\n",
    "by_diff = collections.Counter(p.split('/')[-2] for p in files)\n",
    "ok = 0\n",
    "for p in files:\n",
    "    s = SimState(); scn = load_scenario(p, s)\n",
    "    last = None\n",
    "    for step in scn['correct_action_chain']:\n",
    "        last = dispatch(step, s)\n",
    "    if last and last.ok: ok += 1\n",
    "print(f'{ok}/{len(files)} scenario chains pass.  By difficulty: {dict(by_diff)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b71958d",
   "metadata": {},
   "source": [
    "## 5 · Sanity-check the HF Inference critic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8d7aea7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from colab.train_lib import LLMCritic\n",
    "critic = LLMCritic(provider='hf', model='Qwen/Qwen2.5-72B-Instruct')\n",
    "v_good = critic.value(\n",
    "    observation='{\"task\":\"users_db is dead, agent has not failed over yet\"}',\n",
    "    action={'id': 'platform.failover_replica', 'params': {'target': 'users_db'}})\n",
    "v_bad  = critic.value(\n",
    "    observation='{\"task\":\"users_db is unhealthy, runbook NOT read\"}',\n",
    "    action={'id': 'platform.restart_cluster', 'params': {'target': 'users_db'}})\n",
    "print(f'failover V={v_good:.2f}    restart-brute V={v_bad:.2f}')\n",
    "assert v_good >= v_bad, 'Critic should rate failover above brute-force restart.'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2df84fc0",
   "metadata": {},
   "source": [
    "## 6 · Configure + run training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b10d4f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings, logging\n",
    "warnings.filterwarnings('ignore', category=FutureWarning, module='transformers')\n",
    "warnings.filterwarnings('ignore', message='.*max_new_tokens.*max_length.*')\n",
    "warnings.filterwarnings('ignore', message='.*attention mask API.*')\n",
    "logging.getLogger('transformers').setLevel(logging.ERROR)\n",
    "\n",
    "from colab.train_lib import CFG, train_loop\n",
    "\n",
    "# REAL training run — 120 PPO updates × 6 rollouts/update ≈ 11.5 k transitions\n",
    "# (~70–90 min on a T4, ~25 min on an A100). Bump `rollouts_per_update` and\n",
    "# `total_updates` together if you want even longer.\n",
    "CFG.update({\n",
    "    'total_updates':        120,\n",
    "    'rollouts_per_update':  6,\n",
    "    'max_steps_per_ep':     16,\n",
    "    'critic_provider':      'hf',\n",
    "    'critic_model':         'Qwen/Qwen2.5-72B-Instruct',\n",
    "    'lr':                   1e-5,\n",
    "    'kl_coef':              0.02,\n",
    "    'clip_eps':             0.20,\n",
    "    'gae_lambda':           0.92,\n",
    "    'checkpoint_every':     20,\n",
    "    'run_name':             'real01',\n",
    "    'tasks': [\n",
    "        # Mix difficulty + saboteur/Slack templates so the policy sees the\n",
    "        # full task distribution every update. All IDs verified against the\n",
    "        # rl-agent/scenarios/sim/ tree.\n",
    "        'sim_easy_lambda_throttle_001',\n",
    "        'sim_easy_lambda_throttle_010',\n",
    "        'sim_med_eb_lambda_016',\n",
    "        'sim_med_eb_lambda_021',\n",
    "        'sim_hard_apigw_chain_001',\n",
    "        'sim_hard_ddb_chain_021',\n",
    "        'sim_hard_iam_chain_011',\n",
    "        'sim_advanced_cascade_users_db_001',\n",
    "        'sim_advanced_runbook_trap_postgres_001',\n",
    "        'sim_advanced_trolley_orders_db_001',\n",
    "        'sim_advanced_saboteur_duel_001',\n",
    "        'sim_advanced_slack_redherring_001',\n",
    "        'sim_gen_app_leak_checkout_007',\n",
    "        'sim_gen_app_leak_payments_019',\n",
    "        'sim_gen_db_duel_users_db_003',\n",
    "        'sim_gen_db_duel_orders_db_015',\n",
    "        'sim_gen_redherring_payments_013',\n",
    "        'sim_gen_redherring_auth_001',\n",
    "        'sim_gen_cascade_payments_db_004',\n",
    "        'sim_gen_cascade_users_db_023',\n",
    "        'sim_gen_cache_warm_session_cache_004',\n",
    "        'sim_gen_peak_frontend_001',\n",
    "        'sim_gen_restore_payments_db_001',\n",
    "    ],\n",
    "})\n",
    "\n",
    "print(f\"Starting REAL run: {CFG['total_updates']} updates × \"\n",
    "      f\"{CFG['rollouts_per_update']} rollouts × \"\n",
    "      f\"{CFG['max_steps_per_ep']} steps  ≈  \"\n",
    "      f\"{CFG['total_updates'] * CFG['rollouts_per_update'] * CFG['max_steps_per_ep']:,} transitions max\")\n",
    "\n",
    "log_path = train_loop()\n",
    "print('Training log:', log_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "436e2f1f",
   "metadata": {},
   "source": [
    "## 7 · Quick visualization of training curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0de13c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, matplotlib.pyplot as plt\n",
    "data = json.load(open(log_path))\n",
    "u = [e['update']      for e in data['updates']]\n",
    "r = [e['mean_reward'] for e in data['updates']]\n",
    "v = [e['mean_value']  for e in data['updates']]\n",
    "k = [e['ppo']['kl']   for e in data['updates']]\n",
    "fig, ax = plt.subplots(1, 3, figsize=(14, 3.5))\n",
    "ax[0].plot(u, r, color='#3fb950'); ax[0].set_title('mean_reward'); ax[0].grid(alpha=.3)\n",
    "ax[1].plot(u, v, color='#58a6ff'); ax[1].set_title('mean_value (Qwen-72B critic)'); ax[1].grid(alpha=.3)\n",
    "ax[2].plot(u, k, color='#f85149'); ax[2].set_title('PPO KL'); ax[2].grid(alpha=.3)\n",
    "for a in ax: a.set_xlabel('update')\n",
    "plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aed4e7bb",
   "metadata": {},
   "source": [
    "## 8 · Replay artifact with the trained agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c09d3e2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "from colab.train_lib import IncidentRolloutCollector, QwenActor, LLMCritic, CFG\n",
    "actor = QwenActor(model_name=CFG['actor_model'], max_seq_len=CFG['max_seq_len'],\n",
    "                  lora_r=CFG['lora_r'], lora_alpha=CFG['lora_alpha'],\n",
    "                  lora_dropout=CFG['lora_dropout'])\n",
    "ckpts = sorted(glob.glob('colab/logs/adapter_*_final')) or sorted(glob.glob('colab/logs/adapter_*'))\n",
    "if ckpts:\n",
    "    actor.model.load_adapter(ckpts[-1], adapter_name='trained')\n",
    "    actor.model.set_adapter('trained')\n",
    "    print('Loaded', ckpts[-1])\n",
    "critic = LLMCritic(provider=CFG['critic_provider'], model=CFG['critic_model'])\n",
    "collector = IncidentRolloutCollector(actor, critic,\n",
    "                                      tasks=['sim_advanced_saboteur_duel_001'],\n",
    "                                      max_steps_per_ep=14)\n",
    "_ = collector.collect(1)\n",
    "replays = sorted(glob.glob('rl-agent/replays/*.html'))\n",
    "print('Latest replay:', replays[-1] if replays else 'none')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "264bdb62",
   "metadata": {},
   "source": [
    "## 9 · Push trained adapter + logs back to Hugging Face\n",
    "\n",
    "Defaults to `IC_HF_USER=sagnik-mukherjee`, so this cell will create / update the public model repo **`sagnik-mukherjee/incident-commander-actor`** with the adapter, replays, and training logs. Override `os.environ['IC_HF_USER']` if you want a different account.\n",
    "\n",
    "> Requires a **Write**-scope HF token. The pre-filled token (`hf_IBf…Mhl`) needs Write to push successfully — if it is Read-only, regenerate one with Write at https://huggingface.co/settings/tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ada14fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, glob\n",
    "from huggingface_hub import HfApi, create_repo\n",
    "\n",
    "os.environ.setdefault('IC_HF_USER', 'sagnik-mukherjee')\n",
    "IC_HF_USER = os.environ['IC_HF_USER']\n",
    "\n",
    "api  = HfApi(token=os.environ['HF_TOKEN'])\n",
    "repo = f'{IC_HF_USER}/incident-commander-actor'\n",
    "create_repo(repo, exist_ok=True, repo_type='model', token=os.environ['HF_TOKEN'])\n",
    "\n",
    "finals = sorted(glob.glob('colab/logs/adapter_*_final'))\n",
    "if finals:\n",
    "    api.upload_folder(folder_path=finals[-1], repo_id=repo, repo_type='model',\n",
    "                      path_in_repo='adapter')\n",
    "api.upload_folder(folder_path='colab/logs', repo_id=repo, repo_type='model',\n",
    "                  path_in_repo='logs', allow_patterns=['*.json'])\n",
    "if os.path.isdir('rl-agent/replays'):\n",
    "    api.upload_folder(folder_path='rl-agent/replays', repo_id=repo,\n",
    "                      repo_type='model', path_in_repo='replays',\n",
    "                      allow_patterns=['*.html'])\n",
    "print(f'Pushed → https://huggingface.co/{repo}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d83a70e2",
   "metadata": {},
   "source": [
    "## 10 · Download artifacts to your laptop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1e43326",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "from google.colab import files\n",
    "shutil.make_archive('/content/ic_artifacts', 'zip',\n",
    "                    root_dir='/content/incident-commander', base_dir='colab/logs')\n",
    "shutil.make_archive('/content/ic_replays',  'zip',\n",
    "                    root_dir='/content/incident-commander', base_dir='rl-agent/replays')\n",
    "files.download('/content/ic_artifacts.zip')\n",
    "files.download('/content/ic_replays.zip')"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}