{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ARIA \u2014 DevOps Incident Response: GRPO Training (RunPod A40/A100)\n",
    "\n",
    "**Model:** `unsloth/Meta-Llama-3.1-8B-Instruct` (4-bit quantized)  \n",
    "**Tasks:** `easy` \u2192 `medium` \u2192 `hard` \u2192 `bonus` (full curriculum)  \n",
    "**Episodes:** 40 per task (160 total)  \n",
    "**Expected runtime:** ~4\u20135 hours on A40 48GB | ~3\u20134 hours on A100 40GB  \n",
    "**Cost:** ~$2\u20134 total on RunPod  \n",
    "\n",
    "This is the full-scale training run. Use this for the flagship result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RunPod Setup\n",
    "1. Go to **runpod.io** \u2192 Secure Cloud \u2192 Deploy\n",
    "2. Select: **RTX A40 (48GB)** or **A100 PCIe (40GB)**\n",
    "3. Template: **RunPod PyTorch 2.1** (has CUDA pre-installed)\n",
    "4. Storage: 50GB is enough\n",
    "5. Once pod starts \u2192 click **Jupyter Lab** \u2192 upload this notebook\n",
    "6. Set your HF token in Cell 2\n",
    "7. Run all cells \u2014 it will finish in 4-5 hours unattended"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 1: Install dependencies \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "import subprocess, sys, os\n",
    "\n",
    "# Set logits flag BEFORE any import\n",
    "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n",
    "\n",
    "subprocess.run(['pip', 'install', '-q',\n",
    "    'unsloth',\n",
    "    'mergekit',                    # required by trl>=0.9\n",
    "    'trl>=0.9.0',\n",
    "    'transformers>=4.48.0',        # needs 4.48+ for CompileConfig\n",
    "    'accelerate>=0.26.0',\n",
    "    'peft>=0.10.0',\n",
    "    'bitsandbytes',\n",
    "    'requests',\n",
    "    'matplotlib',\n",
    "    'huggingface_hub',\n",
    "    'scipy'\n",
    "], capture_output=True, text=True)\n",
    "\n",
    "# Clear stale module cache\n",
    "for mod in list(sys.modules.keys()):\n",
    "    if any(x in mod for x in ['trl','unsloth','transformers','peft']):\n",
    "        del sys.modules[mod]\n",
    "\n",
    "# Verify \u2014 unsloth must be imported first\n",
    "import unsloth\n",
    "from unsloth import FastLanguageModel\n",
    "import transformers, peft, torch\n",
    "\n",
    "print(f'\u2705 unsloth {unsloth.__version__}')\n",
    "print(f'\u2705 transformers {transformers.__version__}')\n",
    "print(f'\u2705 torch {torch.__version__} | CUDA: {torch.cuda.is_available()}')\n",
    "print(f'\u2705 UNSLOTH_RETURN_LOGITS = {os.environ[\"UNSLOTH_RETURN_LOGITS\"]}')\n",
    "print('\u2705 All dependencies installed')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 2: Config \u2014 SET YOUR HF TOKEN HERE \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "import os\n",
    "os.environ['UNSLOTH_RETURN_LOGITS'] = '1'\n",
    "\n",
    "HF_TOKEN = os.environ.get('HF_TOKEN', 'YOUR_HF_WRITE_TOKEN_HERE')  # set as env var or paste here\n",
    "os.environ['HF_TOKEN'] = HF_TOKEN\n",
    "\n",
    "from huggingface_hub import login\n",
    "login(token=HF_TOKEN, add_to_git_credential=False)\n",
    "print('\u2705 Logged in to HuggingFace')\n",
    "\n",
    "CONFIG = {\n",
    "    # Model \u2014 8B for A100\n",
    "    'model_name': 'unsloth/Meta-Llama-3.1-8B-Instruct',\n",
    "    'max_seq_length': 3072,\n",
    "    'load_in_4bit': True,\n",
    "\n",
    "    # Environment\n",
    "    'env_url': 'https://arijit-07-devops-incident-response.hf.space',\n",
    "    'tasks': ['easy', 'medium', 'hard', 'bonus'],\n",
    "    'episodes_per_task': 40,\n",
    "    'max_steps_per_episode': 12,        # reduced from 20 \u2014 tighter episodes\n",
    "\n",
    "    # Training \u2014 conservative to prevent catastrophic forgetting\n",
    "    'learning_rate': 5e-6,              # FIXED: was 1e-5, caused degradation\n",
    "    'grpo_group_size': 6,\n",
    "    'lora_rank': 32,\n",
    "    'lora_alpha': 64,\n",
    "    'max_grad_norm': 0.5,\n",
    "    'kl_coeff': 0.05,                   # NEW: prevents catastrophic forgetting\n",
    "\n",
    "    # Output\n",
    "    'hf_repo': 'Arijit-07/aria-devops-llama8b',\n",
    "    'output_dir': '/workspace/aria-llama8b',\n",
    "    'save_every_n_episodes': 20,\n",
    "}\n",
    "\n",
    "import torch\n",
    "print(f'GPU: {torch.cuda.get_device_name(0)}')\n",
    "print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')\n",
    "print(f'Model: {CONFIG[\"model_name\"]} | Tasks: {CONFIG[\"tasks\"]}')\n",
    "print(f'LR: {CONFIG[\"learning_rate\"]} | KL: {CONFIG[\"kl_coeff\"]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 3: Environment Client \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "import requests, json, time, random\n",
    "\n",
    "BASE_URL = CONFIG['env_url']\n",
    "\n",
    "def env_reset(task_id, seed=None):\n",
    "    payload = {'task_id': task_id}\n",
    "    if seed is not None: payload['seed'] = seed\n",
    "    for attempt in range(3):\n",
    "        try:\n",
    "            r = requests.post(f'{BASE_URL}/reset', json=payload, timeout=30)\n",
    "            r.raise_for_status()\n",
    "            return r.json()\n",
    "        except:\n",
    "            if attempt == 2: raise\n",
    "            time.sleep(5)\n",
    "\n",
    "def env_step(action):\n",
    "    for attempt in range(3):\n",
    "        try:\n",
    "            r = requests.post(f'{BASE_URL}/step', json=action, timeout=30)\n",
    "            r.raise_for_status()\n",
    "            return r.json()\n",
    "        except:\n",
    "            if attempt == 2: raise\n",
    "            time.sleep(5)\n",
    "\n",
    "def env_state():\n",
    "    r = requests.get(f'{BASE_URL}/state', timeout=30)\n",
    "    r.raise_for_status()\n",
    "    return r.json()\n",
    "\n",
    "health = requests.get(f'{BASE_URL}/health', timeout=15).json()\n",
    "print(f'\u2705 Environment: {health}')\n",
    "test_obs = env_reset('easy', seed=0)\n",
    "print(f'\u2705 Reset OK. Services: {len(test_obs.get(\"services\", []))}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 4: System Prompt + Observation Formatter \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "CONFIG = {\n",
    "    # Model\n",
    "    'model_name': 'unsloth/Llama-3.2-3B-Instruct',\n",
    "    'max_seq_length': 2048,\n",
    "    'load_in_4bit': True,\n",
    "\n",
    "    # Environment\n",
    "    'env_url': 'https://arijit-07-devops-incident-response.hf.space',\n",
    "    'tasks': ['easy', 'medium'],       # curriculum order\n",
    "    'episodes_per_task': 60,           # 60 per task = 120 total\n",
    "    'max_steps_per_episode': 15,\n",
    "\n",
    "    # Training\n",
    "    'learning_rate': 2e-5,\n",
    "    'batch_size': 4,\n",
    "    'grad_accum': 4,\n",
    "    'grpo_group_size': 4,              # number of completions per prompt\n",
    "    'lora_rank': 16,\n",
    "    'lora_alpha': 32,\n",
    "\n",
    "    # Output\n",
    "    'hf_repo': 'Arijit-07/aria-devops-llama3b',   # change if needed\n",
    "    'output_dir': '/kaggle/working/aria-llama3b',\n",
    "    'save_every_n_episodes': 30,\n",
    "}\n",
    "\n",
    "print('Config loaded:')\n",
    "for k, v in CONFIG.items():\n",
    "    print(f'  {k}: {v}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 5: Load Llama-3.1-8B with Unsloth \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "from unsloth import FastLanguageModel\n",
    "import torch\n",
    "\n",
    "print(f'Loading {CONFIG[\"model_name\"]}...')\n",
    "\n",
    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
    "    model_name=CONFIG['model_name'],\n",
    "    max_seq_length=CONFIG['max_seq_length'],\n",
    "    dtype=None,\n",
    "    load_in_4bit=CONFIG['load_in_4bit'],\n",
    "    token=HF_TOKEN,\n",
    ")\n",
    "\n",
    "model = FastLanguageModel.get_peft_model(\n",
    "    model,\n",
    "    r=CONFIG['lora_rank'],\n",
    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
    "                    'gate_proj', 'up_proj', 'down_proj'],\n",
    "    lora_alpha=CONFIG['lora_alpha'],\n",
    "    lora_dropout=0.05,\n",
    "    bias='none',\n",
    "    use_gradient_checkpointing='unsloth',\n",
    "    random_state=42,\n",
    ")\n",
    "\n",
    "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "total = sum(p.numel() for p in model.parameters())\n",
    "print(f'\u2705 Model loaded')\n",
    "print(f'   Trainable: {trainable:,} ({100*trainable/total:.2f}%)')\n",
    "print(f'   VRAM: {torch.cuda.memory_allocated()/1e9:.2f} GB used')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 6: Action Parser + Episode Runner \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "import re\n",
    "\n",
    "def parse_action(text):\n",
    "    text = text.strip()\n",
    "    for pattern in [\n",
    "        r'```json\\s*({.*?})\\s*```',\n",
    "        r'```\\s*({.*?})\\s*```',\n",
    "        r'({\\s*\"action_type\"[^}]+})',\n",
    "    ]:\n",
    "        match = re.search(pattern, text, re.DOTALL)\n",
    "        if match:\n",
    "            try: return json.loads(match.group(1))\n",
    "            except: continue\n",
    "    try: return json.loads(text)\n",
    "    except: return {'action_type': 'noop'}\n",
    "\n",
    "def generate_action(obs, task_id, temperature=0.7):\n",
    "    messages = [\n",
    "        {'role': 'system', 'content': SYSTEM_PROMPT},\n",
    "        {'role': 'user', 'content': observation_to_prompt(obs, task_id)}\n",
    "    ]\n",
    "    input_ids = tokenizer.apply_chat_template(\n",
    "        messages, tokenize=True, add_generation_prompt=True,\n",
    "        return_tensors='pt'\n",
    "    ).to('cuda')\n",
    "    FastLanguageModel.for_inference(model)\n",
    "    with torch.no_grad():\n",
    "        out = model.generate(\n",
    "            input_ids, max_new_tokens=150,\n",
    "            temperature=temperature, do_sample=True,\n",
    "            pad_token_id=tokenizer.eos_token_id,\n",
    "        )\n",
    "    generated = tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True)\n",
    "    return parse_action(generated), generated\n",
    "\n",
    "def run_episode(task_id, seed=None, verbose=False):\n",
    "    obs = env_reset(task_id, seed=seed)\n",
    "    total_reward = 0.0\n",
    "    done = False\n",
    "    for step in range(CONFIG['max_steps_per_episode']):\n",
    "        if done: break\n",
    "        action, _ = generate_action(obs, task_id)\n",
    "        if verbose: print(f'  Step {step+1}: {action}')\n",
    "        result = env_step(action)\n",
    "        total_reward += result.get('reward', 0.0)\n",
    "        obs = result.get('observation', obs)\n",
    "        done = result.get('done', False)\n",
    "    state = env_state()\n",
    "    return state.get('current_score', total_reward)\n",
    "\n",
    "print('\u2705 Episode runner ready')\n",
    "print('Testing one episode...')\n",
    "test_score = run_episode('easy', seed=99, verbose=True)\n",
    "print(f'Test score: {test_score:.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 7: Pre-Training Baseline \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "print('Running pre-training baseline (8 episodes per task)...')\n",
    "baseline_scores = {}\n",
    "\n",
    "for task_id in CONFIG['tasks']:\n",
    "    scores = [run_episode(task_id, seed=i*7+3) for i in range(8)]\n",
    "    avg = sum(scores) / len(scores)\n",
    "    baseline_scores[task_id] = {'scores': scores, 'avg': avg}\n",
    "    print(f'  [{task_id}] baseline: {avg:.3f} (min={min(scores):.3f} max={max(scores):.3f})')\n",
    "\n",
    "print('\\n\u2705 Baseline done. Starting training...')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 8: GRPO Training Loop (FIXED \u2014 Episode-level updates + KL) \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "from torch.optim import AdamW\n",
    "from transformers import get_cosine_schedule_with_warmup\n",
    "import os, time, random, copy, json\n",
    "\n",
    "os.makedirs(CONFIG['output_dir'], exist_ok=True)\n",
    "\n",
    "# Frozen reference model for KL penalty\n",
    "ref_model = copy.deepcopy(model)\n",
    "for p in ref_model.parameters():\n",
    "    p.requires_grad = False\n",
    "ref_model.eval()\n",
    "print('\u2705 Reference model frozen for KL penalty')\n",
    "\n",
    "optimizer = AdamW(\n",
    "    [p for p in model.parameters() if p.requires_grad],\n",
    "    lr=CONFIG['learning_rate'], weight_decay=0.01\n",
    ")\n",
    "total_eps = CONFIG['episodes_per_task'] * len(CONFIG['tasks'])\n",
    "scheduler = get_cosine_schedule_with_warmup(\n",
    "    optimizer,\n",
    "    num_warmup_steps=max(1, total_eps // 10),\n",
    "    num_training_steps=total_eps\n",
    ")\n",
    "\n",
    "training_log = []\n",
    "episode_scores = {t: [] for t in CONFIG['tasks']}\n",
    "global_ep = 0\n",
    "start_time = time.time()\n",
    "\n",
    "print('=' * 65)\n",
    "print('ARIA GRPO TRAINING \u2014 Llama-3.1-8B')\n",
    "print(f'LR={CONFIG[\"learning_rate\"]} | KL={CONFIG[\"kl_coeff\"]} | Groups={CONFIG[\"grpo_group_size\"]}')\n",
    "print(f'Strategy: collect full episode \u2192 score on fresh env \u2192 update once')\n",
    "print('=' * 65)\n",
    "\n",
    "def run_episode_collect(task_id, seed):\n",
    "    \"\"\"\n",
    "    FIXED: group completions scored on FRESH env snapshots.\n",
    "    Only best action advances main episode.\n",
    "    \"\"\"\n",
    "    obs = env_reset(task_id, seed=seed)\n",
    "    trajectory = []\n",
    "    done = False\n",
    "\n",
    "    FastLanguageModel.for_inference(model)\n",
    "\n",
    "    for step in range(CONFIG['max_steps_per_episode']):\n",
    "        if done:\n",
    "            break\n",
    "\n",
    "        messages = [\n",
    "            {'role': 'system', 'content': SYSTEM_PROMPT},\n",
    "            {'role': 'user', 'content': observation_to_prompt(obs, task_id)}\n",
    "        ]\n",
    "        input_ids = tokenizer.apply_chat_template(\n",
    "            messages, tokenize=True, add_generation_prompt=True,\n",
    "            return_tensors='pt'\n",
    "        ).to('cuda')\n",
    "\n",
    "        # Generate all completions first \u2014 no env calls yet\n",
    "        group_completions, group_texts = [], []\n",
    "        for _ in range(CONFIG['grpo_group_size']):\n",
    "            with torch.no_grad():\n",
    "                out = model.generate(\n",
    "                    input_ids, max_new_tokens=128, temperature=0.9,\n",
    "                    do_sample=True, pad_token_id=tokenizer.eos_token_id,\n",
    "                )\n",
    "            gen_ids = out[0][input_ids.shape[1]:]\n",
    "            group_completions.append(gen_ids)\n",
    "            group_texts.append(tokenizer.decode(gen_ids, skip_special_tokens=True))\n",
    "\n",
    "        # Score each completion on a FRESH env snapshot\n",
    "        group_rewards = []\n",
    "        for gen_text in group_texts:\n",
    "            action = parse_action(gen_text)\n",
    "            try:\n",
    "                env_reset(task_id, seed=seed)  # fresh snapshot\n",
    "                res = env_step(action)\n",
    "                r = res.get('reward', 0.0)\n",
    "            except:\n",
    "                r = 0.0\n",
    "            if action.get('action_type', 'noop') != 'noop':\n",
    "                r += 0.02  # exploration bonus\n",
    "            group_rewards.append(r)\n",
    "\n",
    "        # Advance main episode with best action\n",
    "        best_idx = group_rewards.index(max(group_rewards))\n",
    "        best_action = parse_action(group_texts[best_idx])\n",
    "        try:\n",
    "            adv_res = env_step(best_action)\n",
    "            obs = adv_res.get('observation', obs)\n",
    "            done = adv_res.get('done', False)\n",
    "        except:\n",
    "            done = True\n",
    "\n",
    "        trajectory.append({\n",
    "            'input_ids': input_ids,\n",
    "            'completions': group_completions,\n",
    "            'rewards': group_rewards,\n",
    "        })\n",
    "\n",
    "    # Get final score from accumulated rewards\n",
    "    total_reward = sum(max(s['rewards']) for s in trajectory) if trajectory else 0.0\n",
    "    return trajectory, total_reward\n",
    "\n",
    "\n",
    "def update_from_trajectory(trajectory):\n",
    "    \"\"\"Single model update from full episode with KL penalty.\"\"\"\n",
    "    if not trajectory:\n",
    "        return 0.0\n",
    "\n",
    "    FastLanguageModel.for_training(model)\n",
    "    model.train()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    total_loss = torch.tensor(0.0).to('cuda')\n",
    "\n",
    "    for step_data in trajectory:\n",
    "        input_ids = step_data['input_ids']\n",
    "        completions = step_data['completions']\n",
    "        rewards = step_data['rewards']\n",
    "\n",
    "        rewards_t = torch.tensor(rewards, dtype=torch.float32)\n",
    "        if rewards_t.std() > 1e-8:\n",
    "            advantages = (rewards_t - rewards_t.mean()) / (rewards_t.std() + 1e-8)\n",
    "        else:\n",
    "            advantages = rewards_t - rewards_t.mean()\n",
    "\n",
    "        best_idx = rewards.index(max(rewards))\n",
    "        best_ids = completions[best_idx]\n",
    "        best_adv = advantages[best_idx]\n",
    "\n",
    "        full_ids = torch.cat([input_ids[0], best_ids]).unsqueeze(0)\n",
    "        labels = full_ids.clone()\n",
    "        labels[0, :input_ids.shape[1]] = -100\n",
    "\n",
    "        outputs = model(full_ids, labels=labels)\n",
    "        policy_loss = outputs.loss * (-best_adv)\n",
    "\n",
    "        # KL penalty vs reference model\n",
    "        with torch.no_grad():\n",
    "            ref_out = ref_model(full_ids)\n",
    "        ref_logits = ref_out.logits[:, input_ids.shape[1]-1:-1, :]\n",
    "        pol_logits = outputs.logits[:, input_ids.shape[1]-1:-1, :]\n",
    "        kl = torch.nn.functional.kl_div(\n",
    "            torch.log_softmax(pol_logits, dim=-1),\n",
    "            torch.softmax(ref_logits, dim=-1),\n",
    "            reduction='batchmean'\n",
    "        )\n",
    "        total_loss = total_loss + policy_loss + CONFIG['kl_coeff'] * kl\n",
    "\n",
    "    total_loss = total_loss / len(trajectory)\n",
    "    total_loss.backward()\n",
    "    torch.nn.utils.clip_grad_norm_(\n",
    "        [p for p in model.parameters() if p.requires_grad],\n",
    "        CONFIG['max_grad_norm']\n",
    "    )\n",
    "    optimizer.step()\n",
    "    scheduler.step()\n",
    "    return total_loss.item()\n",
    "\n",
    "\n",
    "# \u2500\u2500 Main training loop \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "for task_id in CONFIG['tasks']:\n",
    "    print(f'\\n\ud83d\udccb Task: {task_id.upper()} | Baseline: {baseline_scores[task_id][\"avg\"]:.3f}')\n",
    "    print('-' * 40)\n",
    "\n",
    "    for ep in range(CONFIG['episodes_per_task']):\n",
    "        seed = random.randint(0, 9999)\n",
    "\n",
    "        trajectory, final_score = run_episode_collect(task_id, seed)\n",
    "        loss = update_from_trajectory(trajectory)\n",
    "\n",
    "        episode_scores[task_id].append(final_score)\n",
    "        global_ep += 1\n",
    "        elapsed = (time.time() - start_time) / 60\n",
    "        recent = episode_scores[task_id][-10:]\n",
    "        rolling = sum(recent) / len(recent)\n",
    "\n",
    "        training_log.append({\n",
    "            'episode': global_ep, 'task_id': task_id,\n",
    "            'score': final_score, 'rolling_avg': rolling,\n",
    "            'loss': loss, 'elapsed_min': round(elapsed, 1)\n",
    "        })\n",
    "\n",
    "        if (ep + 1) % 5 == 0:\n",
    "            delta = rolling - baseline_scores[task_id]['avg']\n",
    "            trend = '\ud83d\udcc8' if delta > 0.02 else '\ud83d\udcc9' if delta < -0.02 else '\u27a1\ufe0f'\n",
    "            print(\n",
    "                f'  {trend} Ep {ep+1:3d}/{CONFIG[\"episodes_per_task\"]} | '\n",
    "                f'Score: {final_score:.3f} | Roll-10: {rolling:.3f} | '\n",
    "                f'vs baseline: {delta:+.3f} | Loss: {loss:.4f} | {elapsed:.0f}m'\n",
    "            )\n",
    "\n",
    "        if global_ep % CONFIG['save_every_n_episodes'] == 0:\n",
    "            ckpt = f'{CONFIG[\"output_dir\"]}/ep{global_ep}'\n",
    "            model.save_pretrained(ckpt)\n",
    "            tokenizer.save_pretrained(ckpt)\n",
    "            print(f'  \ud83d\udcbe Checkpoint ep{global_ep}')\n",
    "\n",
    "    task_avg = sum(episode_scores[task_id]) / len(episode_scores[task_id])\n",
    "    base_avg = baseline_scores[task_id]['avg']\n",
    "    delta = task_avg - base_avg\n",
    "    result = '\u2705 IMPROVED' if delta > 0.02 else '\u26a0\ufe0f FLAT' if delta > -0.02 else '\u274c DEGRADED'\n",
    "    print(f'\\n{result} {task_id}: {base_avg:.3f} \u2192 {task_avg:.3f} ({delta:+.3f})')\n",
    "\n",
    "    # Save training log so far (in case of crash)\n",
    "    with open(f'{CONFIG[\"output_dir\"]}/training_log.json', 'w') as f:\n",
    "        json.dump(training_log, f, indent=2)\n",
    "    print('  \ud83d\udcdd Training log saved')\n",
    "\n",
    "print(f'\\n\ud83c\udf89 Training complete! {(time.time()-start_time)/60:.0f} minutes')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 9: Post-Training Eval + Generalization \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "FastLanguageModel.for_inference(model)\n",
    "print('Post-training evaluation (8 episodes per task, unseen seeds)...')\n",
    "\n",
    "post_scores = {}\n",
    "for task_id in CONFIG['tasks']:\n",
    "    scores = [run_episode(task_id, seed=i*13+7) for i in range(8)]\n",
    "    avg = sum(scores) / len(scores)\n",
    "    post_scores[task_id] = {'scores': scores, 'avg': avg}\n",
    "    delta = avg - baseline_scores[task_id]['avg']\n",
    "    print(f'  [{task_id}] {baseline_scores[task_id][\"avg\"]:.3f} \u2192 {avg:.3f} '\n",
    "          f'({(\"+\" if delta>=0 else \"\")}{delta:.3f})')\n",
    "\n",
    "print('\\nZero-shot generalization (ARIA tasks \u2014 never seen in training):')\n",
    "gen_scores = {}\n",
    "for task_id in ['security', 'database', 'failover']:\n",
    "    scores = []\n",
    "    for i in range(5):\n",
    "        try: scores.append(run_episode(task_id, seed=i*17+5))\n",
    "        except: scores.append(0.0)\n",
    "    avg = sum(scores) / len(scores)\n",
    "    gen_scores[task_id] = avg\n",
    "    print(f'  [{task_id}] zero-shot: {avg:.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 10: Learning Curve Visualization \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "import numpy as np\n",
    "\n",
    "fig = plt.figure(figsize=(20, 12))\n",
    "fig.patch.set_facecolor('#0d1117')\n",
    "gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.4, wspace=0.35)\n",
    "COLORS = {'easy':'#4caf50','medium':'#ff9800','hard':'#f44336','bonus':'#9c27b0'}\n",
    "\n",
    "def style_ax(ax, title):\n",
    "    ax.set_facecolor('#161b22')\n",
    "    ax.set_title(title, color='white', fontsize=12, fontweight='bold', pad=10)\n",
    "    ax.tick_params(colors='#8b949e', labelsize=9)\n",
    "    for spine in ax.spines.values(): spine.set_color('#30363d')\n",
    "    ax.spines['top'].set_visible(False)\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.grid(True, alpha=0.1, color='#30363d')\n",
    "\n",
    "for idx, task_id in enumerate(CONFIG['tasks']):\n",
    "    row, col = divmod(idx, 3)\n",
    "    ax = fig.add_subplot(gs[row, col])\n",
    "    style_ax(ax, f'Task: {task_id.upper()}')\n",
    "    task_log = [e for e in training_log if e['task_id'] == task_id]\n",
    "    eps = [e['episode'] for e in task_log]\n",
    "    scores = [e['score'] for e in task_log]\n",
    "    rolling = [e['rolling_avg'] for e in task_log]\n",
    "    color = COLORS.get(task_id, '#58a6ff')\n",
    "    ax.plot(eps, scores, alpha=0.15, color=color, linewidth=1)\n",
    "    ax.plot(eps, rolling, color=color, linewidth=2.5, label='Rolling avg (10)')\n",
    "    ax.axhline(y=baseline_scores[task_id]['avg'], color='#f85149',\n",
    "               linestyle='--', linewidth=1.5, label='Baseline')\n",
    "    ax.axhline(y=post_scores[task_id]['avg'], color='#3fb950',\n",
    "               linestyle='--', linewidth=1.5, label='Post-training')\n",
    "    ax.set_ylim(0, 1.05)\n",
    "    ax.set_xlabel('Episode', color='#8b949e', fontsize=9)\n",
    "    ax.set_ylabel('Score', color='#8b949e', fontsize=9)\n",
    "    ax.legend(facecolor='#161b22', labelcolor='white', fontsize=8)\n",
    "\n",
    "ax5 = fig.add_subplot(gs[1, 1])\n",
    "style_ax(ax5, 'Before vs After (all tasks)')\n",
    "x = np.arange(len(CONFIG['tasks']))\n",
    "w = 0.35\n",
    "before_v = [baseline_scores[t]['avg'] for t in CONFIG['tasks']]\n",
    "after_v = [post_scores[t]['avg'] for t in CONFIG['tasks']]\n",
    "b1 = ax5.bar(x-w/2, before_v, w, label='Before', color='#f85149', alpha=0.85)\n",
    "b2 = ax5.bar(x+w/2, after_v, w, label='After', color='#3fb950', alpha=0.85)\n",
    "for bar, v in zip(b1, before_v):\n",
    "    ax5.text(bar.get_x()+bar.get_width()/2., v+0.01, f'{v:.2f}',\n",
    "             ha='center', color='white', fontsize=8)\n",
    "for bar, v in zip(b2, after_v):\n",
    "    ax5.text(bar.get_x()+bar.get_width()/2., v+0.01, f'{v:.2f}',\n",
    "             ha='center', color='white', fontsize=8)\n",
    "ax5.set_xticks(x)\n",
    "ax5.set_xticklabels(CONFIG['tasks'], color='#8b949e')\n",
    "ax5.set_ylim(0, 1.15)\n",
    "ax5.legend(facecolor='#161b22', labelcolor='white', fontsize=9)\n",
    "\n",
    "ax6 = fig.add_subplot(gs[1, 2])\n",
    "ax6.set_facecolor('#161b22')\n",
    "ax6.set_title('Summary', color='white', fontsize=12, fontweight='bold')\n",
    "ax6.axis('off')\n",
    "lines = [\n",
    "    ('Model', 'Llama-3.1-8B (Unsloth 4-bit)'),\n",
    "    ('Algorithm', 'GRPO'),\n",
    "    ('LoRA rank', str(CONFIG['lora_rank'])),\n",
    "    ('Total episodes', str(global_ep)),\n",
    "    ('', ''),\n",
    "]\n",
    "for t in CONFIG['tasks']:\n",
    "    b = baseline_scores[t]['avg']; a = post_scores[t]['avg']\n",
    "    lines.append((f'  {t}', f'{b:.2f} \u2192 {a:.2f} (+{a-b:.2f})'))\n",
    "if gen_scores:\n",
    "    lines += [('', ''), ('Zero-shot', '')]\n",
    "    for t, s in gen_scores.items():\n",
    "        lines.append((f'  {t}', f'{s:.2f}'))\n",
    "y = 0.95\n",
    "for label, val in lines:\n",
    "    if not label: y -= 0.04; continue\n",
    "    ax6.text(0.02, y, label+':', color='#8b949e', fontsize=9,\n",
    "             transform=ax6.transAxes, fontweight='bold')\n",
    "    ax6.text(0.52, y, val, color='#c9d1d9', fontsize=9, transform=ax6.transAxes)\n",
    "    y -= 0.08\n",
    "\n",
    "fig.suptitle('ARIA \u2014 DevOps Incident Response\\nGRPO Training (Llama-3.1-8B Full Curriculum)',\n",
    "             color='white', fontsize=16, fontweight='bold', y=0.98)\n",
    "plt.savefig('training_curve_8b.png', dpi=150, bbox_inches='tight', facecolor='#0d1117')\n",
    "print('\u2705 Saved training_curve_8b.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# \u2500\u2500 Cell 11: Save to HuggingFace Hub \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
    "from huggingface_hub import HfApi\n",
    "import json\n",
    "\n",
    "print(f'Pushing to {CONFIG[\"hf_repo\"]}...')\n",
    "FastLanguageModel.for_inference(model)\n",
    "\n",
    "model.save_pretrained_merged(CONFIG['output_dir'], tokenizer, save_method='merged_16bit')\n",
    "model.push_to_hub_merged(CONFIG['hf_repo'], tokenizer,\n",
    "                         save_method='merged_16bit', token=HF_TOKEN)\n",
    "print(f'\u2705 Model: https://huggingface.co/{CONFIG[\"hf_repo\"]}')\n",
    "\n",
    "api = HfApi()\n",
    "for fname in ['training_curve_8b.png']:\n",
    "    api.upload_file(path_or_fileobj=fname, path_in_repo=fname,\n",
    "                    repo_id=CONFIG['hf_repo'], token=HF_TOKEN)\n",
    "    print(f'\u2705 {fname} uploaded')\n",
    "\n",
    "with open('training_log_8b.json', 'w') as f:\n",
    "    json.dump(training_log, f, indent=2)\n",
    "api.upload_file(path_or_fileobj='training_log_8b.json',\n",
    "                path_in_repo='training_log_8b.json',\n",
    "                repo_id=CONFIG['hf_repo'], token=HF_TOKEN)\n",
    "\n",
    "print('\\n\ud83c\udf89 DONE! Shut down the RunPod instance now to stop billing.')\n",
    "print(f'   Model: https://huggingface.co/{CONFIG[\"hf_repo\"]}')\n",
    "print(f'   Curve: check training_curve_8b.png in the repo')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "| Task | Before | After | Improvement |\n",
    "|---|---|---|---|\n",
    "| easy | see baseline | see post-eval | see training curve |\n",
    "| medium | see baseline | see post-eval | see training curve |\n",
    "| hard | see baseline | see post-eval | see training curve |\n",
    "| bonus | see baseline | see post-eval | see training curve |\n",
    "\n",
    "**Zero-shot generalization** on security, database, failover tasks shows  \n",
    "the model learned transferable operational reasoning, not task memorization.\n",
    "\n",
    "Weights: `https://huggingface.co/Arijit-07/aria-devops-llama8b`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}