fix: robust notebook setup (no magic shell) + local CWD auto-detect

- Use subprocess+shutil for Colab git clone; surface clone errors
- Local mode finds repo root without /content
- Cell 3 auto-chdirs from training/ if needed
- Clear errors when project files are missing

Made-with: Cursor

training/train_grpo.ipynb

@@ -23,20 +23,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [
-        {
-          "ename": "",
-          "evalue": "",
-          "output_type": "error",
-          "traceback": [
-            "\u001b[1;31mRunning cells with '.venv (Python 3.13.1)' requires the ipykernel package.\n",
-            "\u001b[1;31mInstall 'ipykernel' into the Python environment. \n",
-            "\u001b[1;31mCommand: '/Users/vaibhavkhandare/Projects/mernstack/openenv-course/viraltest/.venv/bin/python -m pip install ipykernel -U --force-reinstall'"
-          ]
-        }
-      ],
       "source": [
         "# Cell 1: Install dependencies\n",
         "!pip install -q torch torchvision torchaudio\n",
@@ -44,41 +31,155 @@
         "!pip install -q matplotlib pandas\n",
         "!pip install -q pydantic httpx\n",
         "!pip install -q \"openenv-core[core]>=0.2.2\""
+      ],
+      "execution_count": 5,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.3\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.0.1\u001b[0m\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+            "zsh:1: 4.45.0 not found\n",
+            "\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.3\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.0.1\u001b[0m\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+            "\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.3\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.0.1\u001b[0m\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
+            "\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.3\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m26.0.1\u001b[0m\n",
+            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
+          ]
+        }
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
-        "# Cell 2: Clone the repo and set up paths\n",
-        "import os, sys\n",
-        "REPO_DIR = \"/content/viral-posts-env\"\n",
+        "# Cell 2: Resolve repo path (Colab: fresh clone. Local: auto-detect project root)\n",
+        "import os\n",
+        "import sys\n",
+        "import shutil\n",
+        "import subprocess\n",
+        "from pathlib import Path\n",
+        "\n",
         "REPO_BRANCH = \"hack1\"\n",
-        "if not os.path.exists(REPO_DIR):\n",
-        "    !git clone --branch {REPO_BRANCH} --depth 1 https://github.com/VaibhavKhandare/viral-posts-env.git {REPO_DIR}\n",
-        "os.chdir(REPO_DIR)\n",
-        "sys.path.insert(0, REPO_DIR)\n",
+        "REPO_URL = \"https://github.com/VaibhavKhandare/viral-posts-env.git\"\n",
+        "COLAB_REPO = Path(\"/content/viral-posts-env\")\n",
+        "\n",
+        "\n",
+        "def _is_repo_root(p: Path) -> bool:\n",
+        "    return (p / \"server\" / \"viraltest_environment.py\").is_file() and (p / \"models.py\").is_file()\n",
+        "\n",
+        "\n",
+        "def _find_local_root() -> Path:\n",
+        "    here = Path.cwd().resolve()\n",
+        "    for cand in (here, here.parent, here.parent.parent):\n",
+        "        if _is_repo_root(cand):\n",
+        "            return cand\n",
+        "    raise FileNotFoundError(\n",
+        "        \"Could not find project root. cd into viral-posts-env or run this notebook in Google Colab.\"\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "# --- Colab: always clone a clean copy (avoids stale 7-day code) ---\n",
+        "if Path(\"/content\").is_dir():\n",
+        "    if COLAB_REPO.exists():\n",
+        "        shutil.rmtree(COLAB_REPO, ignore_errors=True)\n",
+        "    p = subprocess.run(\n",
+        "        [\n",
+        "            \"git\", \"clone\", \"--branch\", REPO_BRANCH, \"--depth\", \"1\",\n",
+        "            REPO_URL, str(COLAB_REPO),\n",
+        "        ],\n",
+        "        capture_output=True,\n",
+        "        text=True,\n",
+        "    )\n",
+        "    if p.returncode != 0:\n",
+        "        raise RuntimeError(\n",
+        "            \"git clone failed. Check network and branch name.\\n\"\n",
+        "            f\"stdout:\\n{p.stdout}\\nstderr:\\n{p.stderr}\"\n",
+        "        )\n",
+        "    if not COLAB_REPO.is_dir():\n",
+        "        raise FileNotFoundError(f\"Clone did not create {COLAB_REPO}\")\n",
+        "    os.chdir(COLAB_REPO)\n",
+        "    print(\"Mode: Colab (fresh clone)\")\n",
+        "else:\n",
+        "    # --- Local machine: do not use /content ---\n",
+        "    root = _find_local_root()\n",
+        "    os.chdir(root)\n",
+        "    print(\"Mode: local\")\n",
+        "    print(f\"Repo root: {root}\")\n",
+        "\n",
+        "REPO_DIR = str(Path.cwd().resolve())\n",
+        "if REPO_DIR not in sys.path:\n",
+        "    sys.path.insert(0, REPO_DIR)\n",
         "\n",
         "PLOTS_DIR = os.path.join(REPO_DIR, \"plots\")\n",
         "os.makedirs(PLOTS_DIR, exist_ok=True)\n",
+        "\n",
+        "try:\n",
+        "    commit = subprocess.check_output(\n",
+        "        [\"git\", \"rev-parse\", \"--short\", \"HEAD\"],\n",
+        "        stderr=subprocess.DEVNULL,\n",
+        "        text=True,\n",
+        "    ).strip()\n",
+        "except Exception:\n",
+        "    commit = \"n/a\"\n",
+        "\n",
         "print(f\"Working dir: {os.getcwd()}\")\n",
+        "print(f\"Branch: {REPO_BRANCH}\")\n",
+        "print(f\"Commit: {commit}\")\n",
         "print(f\"Plots dir: {PLOTS_DIR}\")"
+      ],
+      "execution_count": 6,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "text": [
+            "fatal: could not create leading directories of '/content/viral-posts-env': Read-only file system\n"
+          ]
+        },
+        {
+          "output_type": "error",
+          "ename": "FileNotFoundError",
+          "evalue": "[Errno 2] No such file or directory: '/content/viral-posts-env'",
+          "traceback": [
+            "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
+            "\u001b[31mFileNotFoundError\u001b[39m                         Traceback (most recent call last)",
+            "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[6]\u001b[39m\u001b[32m, line 13\u001b[39m\n\u001b[32m      9\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m os.path.exists(REPO_DIR):\n\u001b[32m     10\u001b[39m     get_ipython().system(\u001b[33m'rm -rf {REPO_DIR}'\u001b[39m)\n\u001b[32m     11\u001b[39m \n\u001b[32m     12\u001b[39m get_ipython().system(\u001b[33m'git clone --branch {REPO_BRANCH} --depth 1 {REPO_URL} {REPO_DIR}'\u001b[39m)\n\u001b[32m---> \u001b[39m\u001b[32m13\u001b[39m os.chdir(REPO_DIR)\n\u001b[32m     14\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m REPO_DIR \u001b[38;5;28;01mnot\u001b[39;00m \u001b[38;5;28;01min\u001b[39;00m sys.path:\n\u001b[32m     15\u001b[39m     sys.path.insert(\u001b[32m0\u001b[39m, REPO_DIR)\n\u001b[32m     16\u001b[39m \n",
+            "\u001b[31mFileNotFoundError\u001b[39m: [Errno 2] No such file or directory: '/content/viral-posts-env'"
+          ]
+        }
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
-        "# Cell 3: Imports\n",
-        "import json, random, time, textwrap, copy\n",
+        "# Cell 3: Imports (with runtime validation)\n",
+        "import json, random, time, textwrap, copy, os, sys\n",
         "from pathlib import Path\n",
         "from typing import Any, Dict, List, Optional, Tuple\n",
         "from collections import defaultdict\n",
         "\n",
+        "# Find repo root if notebook was opened from training/ and Cell 2 was skipped\n",
+        "if not Path(\"server/viraltest_environment.py\").is_file():\n",
+        "    for cand in (Path.cwd(), Path.cwd().parent, Path.cwd().parent.parent):\n",
+        "        if (cand / \"server\" / \"viraltest_environment.py\").is_file():\n",
+        "            os.chdir(cand)\n",
+        "            s = str(cand.resolve())\n",
+        "            if s not in sys.path:\n",
+        "                sys.path.insert(0, s)\n",
+        "            print(\"Auto chdir to repo root:\", s)\n",
+        "            break\n",
+        "    else:\n",
+        "        raise RuntimeError(\n",
+        "            \"Project files not found. Run **Cell 2** first (Colab), or run from repo root.\\n\"\n",
+        "            f\"  cwd = {os.getcwd()!r}\\n\"\n",
+        "        )\n",
+        "\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
         "import pandas as pd\n",
@@ -97,8 +198,16 @@
         "TASKS = [\"monthly_engage\", \"monthly_strategic\", \"monthly_competitive\"]\n",
         "\n",
         "print(f\"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}\")\n",
-        "print(f\"Tags: {len(TAG_POOL)}, Topics: {len(ALL_TOPICS)}, Horizon: {TASK_HORIZON} days\")"
-      ]
+        "print(f\"Tags: {len(TAG_POOL)}, Topics: {len(ALL_TOPICS)}, Horizon: {TASK_HORIZON} days\")\n",
+        "\n",
+        "# Hard stop if stale repo/code is loaded\n",
+        "assert TASK_HORIZON == 30, (\n",
+        "    f\"Expected TASK_HORIZON=30, got {TASK_HORIZON}. \"\n",
+        "    \"Restart runtime and run from Cell 1 again (clean clone on hack1).\"\n",
+        ")"
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -111,9 +220,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 4: Define heuristic agents + episode runner\n",
         "_rng = random.Random(42)\n",
@@ -190,24 +297,36 @@
         "            \"rewards\": rewards, \"energies\": energies}\n",
         "\n",
         "print(\"Agents and episode runner defined.\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
-        "# Cell 5: Run baselines\n",
+        "# Cell 5: Run baselines (safe)\n",
         "print(\"Running heuristic baselines (5 agents × 3 tasks)...\")\n",
         "print(\"=\" * 70)\n",
         "\n",
+        "required = [\"BASELINE_AGENTS\", \"run_episode\", \"TASKS\", \"random\"]\n",
+        "missing = [k for k in required if k not in globals()]\n",
+        "if missing:\n",
+        "    raise RuntimeError(\n",
+        "        f\"Missing prerequisites: {missing}. Run notebook from top (Cell 1 -> Cell 5).\"\n",
+        "    )\n",
+        "\n",
         "baseline_results = {}\n",
         "for name, fn in BASELINE_AGENTS.items():\n",
         "    baseline_results[name] = {}\n",
         "    for task in TASKS:\n",
         "        _rng = random.Random(42)\n",
-        "        result = run_episode(task, fn, seed=42)\n",
+        "        try:\n",
+        "            result = run_episode(task, fn, seed=42)\n",
+        "        except Exception as e:\n",
+        "            raise RuntimeError(\n",
+        "                f\"Baseline failed for agent={name}, task={task}: {type(e).__name__}: {e}\"\n",
+        "            ) from e\n",
         "        baseline_results[name][task] = result\n",
         "        print(f\"  {name:>12s} | {task:>22s} | score={result['grader_score']:.4f} \"\n",
         "              f\"| energy={result['final_energy']:.2f}\")\n",
@@ -219,13 +338,13 @@
         "for name in BASELINE_AGENTS:\n",
         "    scores = [baseline_results[name][t][\"grader_score\"] for t in TASKS]\n",
         "    print(f\"{name:<14s} {scores[0]:>10.4f} {scores[1]:>12.4f} {scores[2]:>14.4f} {sum(scores)/3:>8.4f}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 6: Baseline plots\n",
         "fig, axes = plt.subplots(1, 3, figsize=(16, 5), sharey=True)\n",
@@ -243,7 +362,9 @@
         "fig.tight_layout()\n",
         "fig.savefig(f\"{PLOTS_DIR}/baseline_leaderboard.png\", dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -256,9 +377,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 7: Load model\n",
         "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
@@ -282,13 +401,13 @@
         "model.eval()\n",
         "print(f\"Model loaded. Device: {model.device}\")\n",
         "print(f\"Memory: {torch.cuda.memory_allocated()/1e9:.1f} GB\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 8: LLM agent functions\n",
         "SYSTEM_PROMPT = textwrap.dedent(\"\"\"\\\n",
@@ -404,7 +523,9 @@
         "            \"burned_out\": obs.creator_energy <= 0}\n",
         "\n",
         "print(\"LLM agent functions defined.\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -417,9 +538,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 9: Run untrained model\n",
         "print(\"Running UNTRAINED base model on all tasks...\")\n",
@@ -436,7 +555,9 @@
         "print(\"BEFORE TRAINING:\")\n",
         "for t in TASKS:\n",
         "    print(f\"  {t}: grader={before_results[t]['grader_score']:.4f}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -455,9 +576,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 10: Attach LoRA adapter\n",
         "from peft import LoraConfig, get_peft_model, TaskType\n",
@@ -472,13 +591,13 @@
         "model.enable_input_require_grads()\n",
         "peft_model = get_peft_model(model, lora_config)\n",
         "peft_model.print_trainable_parameters()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 11: Training loop\n",
         "from trl import SFTTrainer, SFTConfig\n",
@@ -569,7 +688,9 @@
         "elapsed = time.time() - t_start\n",
         "print(f\"\\nTraining complete in {elapsed/60:.1f} min\")\n",
         "print(pd.DataFrame(training_log).to_string(index=False))"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -582,9 +703,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 12: Run trained model\n",
         "print(\"Running TRAINED model on all tasks...\")\n",
@@ -602,7 +721,9 @@
         "print(\"AFTER TRAINING:\")\n",
         "for t in TASKS:\n",
         "    print(f\"  {t}: grader={after_results[t]['grader_score']:.4f}\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -613,9 +734,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 13: Training curves\n",
         "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
@@ -637,13 +756,13 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/reward_curve.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 14: Before vs After\n",
         "task_labels = [t.replace('monthly_', '').title() for t in TASKS]\n",
@@ -673,13 +792,13 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/before_after.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 15: Trajectory comparison\n",
         "fig, axes = plt.subplots(2, 3, figsize=(16, 8))\n",
@@ -703,7 +822,9 @@
         "fig.tight_layout()\n",
         "fig.savefig(f'{PLOTS_DIR}/training_trajectories.png', dpi=150, bbox_inches='tight')\n",
         "plt.show()"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "markdown",
@@ -714,9 +835,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 16: Final summary\n",
         "print(\"=\" * 67)\n",
@@ -753,13 +872,13 @@
         "\n",
         "print(f\"\\nSaved to {PLOTS_DIR}/\")\n",
         "print(\"All results are from real LoRA weight updates on real environment runs.\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     },
     {
       "cell_type": "code",
-      "execution_count": null,
       "metadata": {},
-      "outputs": [],
       "source": [
         "# Cell 17: Save adapter\n",
         "save_path = \"./viraltest_trained_adapter\"\n",
@@ -767,7 +886,9 @@
         "tokenizer.save_pretrained(save_path)\n",
         "print(f\"LoRA adapter saved to {save_path}\")\n",
         "print(\"Load with: PeftModel.from_pretrained(base_model, save_path)\")"
-      ]
+      ],
+      "execution_count": null,
+      "outputs": []
     }
   ],
   "metadata": {
@@ -779,10 +900,18 @@
       "name": "python3"
     },
     "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
       "name": "python",
-      "version": "3.13.1"
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.14.2"
     }
   },
   "nbformat": 4,
   "nbformat_minor": 4
-}
+}