feat(training): A10G-optimised pipeline — auto train.py, Dockerfile.train, GH Action sync

- Add fastmcp to pyproject.toml (was imported but not declared — env server Dockerfile failed)
- Add train.py: fully automated SFT→GRPO→push script for HF Space auto-run
- Add Dockerfile.train: training Space image (JupyterLab on :7860 or auto train.py)
- Add train-entrypoint.sh: START_MODE=jupyter|auto switch
- Add .github/workflows/sync_hf_space.yml: push main → HF env-server Space on every commit
- Rewrite train_grpo.ipynb: remove google.colab, fix for HF Spaces env vars, remove
unnecessary env server subprocess, tune hyperparams for 46 GB VRAM
- grpo.py: pass max_turns through hyperparams (default 4 for fast GRPO episodes)

Timing targets on A10G Large:
SFT (2 epochs, batch 4, grad_accum 4): ~45 min
GRPO (100 steps, 4 rollouts, max_turns=4, 256 completion tokens): ~90 min
Total: ~2.5 hours

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.github/workflows/sync_hf_space.yml

@@ -0,0 +1,43 @@
+name: Sync → HF Space (env server)
+
+# Pushes the repo to the env-server HF Space on every commit to main.
+# The Space rebuilds its Docker image automatically, which downloads
+# scenarios from HF Hub and restarts the env server.
+#
+# Required GitHub secrets:
+#   HF_TOKEN      - HuggingFace write token (Settings → Secrets → Actions)
+#   HF_USERNAME   - your HuggingFace username
+#   ENV_SPACE_NAME - name of your env-server Space (e.g. "ci-triage-env")
+#
+# The training Space is NOT auto-synced here (rebuilding mid-training would
+# kill a running job). Manually push to it when you want to update.
+
+on:
+  push:
+    branches: [main]
+
+jobs:
+  sync-env-space:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repo (full history)
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          lfs: true
+
+      - name: Push to HF Space
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+          HF_USERNAME: ${{ secrets.HF_USERNAME }}
+          ENV_SPACE_NAME: ${{ secrets.ENV_SPACE_NAME }}
+        run: |
+          git config --global user.email "github-action@ci-triage"
+          git config --global user.name "CI Triage Sync"
+
+          REMOTE="https://${HF_USERNAME}:${HF_TOKEN}@huggingface.co/spaces/${HF_USERNAME}/${ENV_SPACE_NAME}"
+          git remote add hf-env "$REMOTE" 2>/dev/null || git remote set-url hf-env "$REMOTE"
+
+          # Force-push main → Space repo (Space will auto-rebuild Docker image)
+          git push hf-env HEAD:main --force
+          echo "✓ Pushed to https://huggingface.co/spaces/${HF_USERNAME}/${ENV_SPACE_NAME}"

Dockerfile.train

@@ -0,0 +1,48 @@
+# Training Space Dockerfile — JupyterLab + auto-run on port 7860.
+#
+# Two modes (controlled by START_MODE env var in Space settings):
+#   START_MODE=jupyter  → opens JupyterLab so you can run train_grpo.ipynb manually
+#   START_MODE=auto     → runs train.py immediately, no interaction needed
+#
+# HF Space secrets to set:
+#   HF_TOKEN, HF_USERNAME, WANDB_API_KEY
+#   HF_SCENARIOS_REPO, HF_SFT_DATASET_REPO, HF_MODEL_REPO (optional)
+#   GRPO_STEPS (optional, default 100)
+
+FROM pytorch/pytorch:2.4.0-cuda12.1-cudnn9-devel
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        git curl build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /workspace
+
+# 1. Install unsloth (must come after torch, hence not in pyproject extras)
+RUN pip install --no-cache-dir \
+    "unsloth[cu121-torch240] @ git+https://github.com/unslothai/unsloth.git"
+
+# 2. Install project + all training deps
+COPY pyproject.toml ./
+COPY src/ src/
+RUN pip install --no-cache-dir -e ".[data,training]"
+
+# 3. JupyterLab for interactive mode
+RUN pip install --no-cache-dir jupyterlab ipywidgets
+
+# 4. Copy notebooks and training scripts
+COPY notebooks/ notebooks/
+COPY train.py ./
+
+# Persistent storage expected at /data (attach 20 GB disk in Space settings)
+RUN mkdir -p /data/checkpoints /data/scenarios /data/sft_dataset
+
+EXPOSE 7860
+
+ENV START_MODE=jupyter
+
+COPY train-entrypoint.sh /train-entrypoint.sh
+RUN chmod +x /train-entrypoint.sh
+ENTRYPOINT ["/train-entrypoint.sh"]

notebooks/train_grpo.ipynb

@@ -2,176 +2,250 @@
  "cells": [
   {
    "cell_type": "markdown",
+   "id": "intro",
    "metadata": {},
    "source": [
-    "# CI-Triage-Env — GRPO Training Notebook\n",
+    "# CI-Triage-Env — GRPO Training\n",
     "\n",
-    "Colab-runnable end-to-end training pipeline:\n",
+    "**Hardware target**: A10G Large (46 GB VRAM, 12 vCPU) — HuggingFace Space\n",
+    "\n",
+    "Pipeline:\n",
     "1. Install dependencies\n",
-    "2. Pull scenario corpus from HF Hub\n",
-    "3. Start env server\n",
-    "4. SFT warmstart on C3 trajectory dataset\n",
-    "5. GRPO smoke test (100 steps)\n",
-    "6. Full GRPO (3000 steps)\n",
-    "7. Push adapter to HF Hub\n",
-    "\n",
-    "**Prerequisites**: `HF_TOKEN`, `OPENAI_API_KEY`, `WANDB_API_KEY` set as Colab secrets."
+    "2. Authenticate (HF + W&B)\n",
+    "3. Pull scenario corpus from HF Hub\n",
+    "4. Pull SFT dataset from HF Hub\n",
+    "5. SFT warmstart (~45 min)\n",
+    "6. GRPO fine-tuning (~90 min for 100 steps)\n",
+    "7. Push final model to HF Hub\n",
+    "\n",
+    "**Set these as Space secrets** (Settings → Variables and secrets):\n",
+    "- `HF_TOKEN` — HuggingFace write token\n",
+    "- `HF_USERNAME` — your HF username\n",
+    "- `WANDB_API_KEY` — Weights & Biases key (get free at wandb.ai)\n",
+    "\n",
+    "**Time budget**: SFT≈45 min + GRPO≈90 min = ~2.5 hours on A10G Large.\n",
+    "Monitor training live at https://wandb.ai (project: `ci-triage-env`)."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-install",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 1: Install dependencies\n",
-    "!pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n",
-    "!pip install -q unsloth trl transformers accelerate peft\n",
-    "!pip install -q wandb datasets huggingface_hub openai httpx fastapi uvicorn pydantic jsonschema\n",
-    "!pip install -q -e .  # install ci_triage_env package in editable mode"
+    "# Cell 1 — Install deps (run once; ~10 min including unsloth compile)\n",
+    "import subprocess, sys\n",
+    "\n",
+    "def run(cmd):\n",
+    "    result = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n",
+    "    if result.returncode != 0:\n",
+    "        print(result.stderr[-2000:])\n",
+    "        raise RuntimeError(f'Command failed: {cmd}')\n",
+    "    return result.stdout\n",
+    "\n",
+    "# PyTorch is pre-installed in the Space Docker image; install the rest\n",
+    "run('pip install -q \"unsloth[cu121-torch240] @ git+https://github.com/unslothai/unsloth.git\"')\n",
+    "run('pip install -q trl>=0.11 transformers>=4.45 accelerate>=0.30 peft')\n",
+    "run('pip install -q wandb datasets huggingface_hub')\n",
+    "run('pip install -q -e /workspace')  # install ci_triage_env package\n",
+    "print('All dependencies installed.')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-auth",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 2: Environment setup\n",
+    "# Cell 2 — Authenticate\n",
     "import os\n",
-    "from google.colab import userdata\n",
+    "from huggingface_hub import login\n",
+    "import wandb\n",
     "\n",
-    "os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')\n",
-    "os.environ['WANDB_API_KEY'] = userdata.get('WANDB_API_KEY')\n",
-    "os.environ['OPENAI_API_KEY'] = userdata.get('OPENAI_API_KEY')\n",
-    "os.environ['WANDB_PROJECT'] = 'ci-triage-env'\n",
+    "HF_TOKEN    = os.environ['HF_TOKEN']\n",
+    "HF_USERNAME = os.environ['HF_USERNAME']\n",
+    "WANDB_KEY   = os.environ.get('WANDB_API_KEY', '')\n",
     "\n",
-    "import wandb\n",
-    "wandb.login()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Cell 3: Pull scenario corpus from HF dataset hub\n",
-    "# Replace YOUR_ORG with your HuggingFace org/username\n",
-    "HF_DATASET_REPO = 'YOUR_ORG/ci-triage-scenarios'\n",
-    "HF_MODEL_REPO  = 'YOUR_ORG/ci-triage-trained-qwen3.5-4b'\n",
+    "login(token=HF_TOKEN)\n",
+    "if WANDB_KEY:\n",
+    "    wandb.login(key=WANDB_KEY)\n",
+    "    os.environ['WANDB_PROJECT'] = 'ci-triage-env'\n",
+    "else:\n",
+    "    os.environ['WANDB_DISABLED'] = 'true'\n",
+    "    print('W&B disabled — set WANDB_API_KEY secret to enable')\n",
     "\n",
-    "from huggingface_hub import snapshot_download\n",
-    "scen_dir = snapshot_download(HF_DATASET_REPO, repo_type='dataset',\n",
-    "                             local_dir='data_artifacts/scenarios')\n",
-    "print(f'Scenarios downloaded to {scen_dir}')"
+    "# Repo names (edit if you used different names)\n",
+    "SCENARIOS_REPO   = f'{HF_USERNAME}/ci-triage-scenarios'\n",
+    "SFT_DATASET_REPO = f'{HF_USERNAME}/ci-triage-sft'\n",
+    "MODEL_REPO       = f'{HF_USERNAME}/ci-triage-agent'\n",
+    "print(f'Authenticated as {HF_USERNAME}')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-scenarios",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 4: Start env server in background\n",
-    "import subprocess, time\n",
-    "server_proc = subprocess.Popen(\n",
-    "    ['python', '-m', 'ci_triage_env.env.server'],\n",
-    "    stdout=subprocess.PIPE, stderr=subprocess.PIPE\n",
-    ")\n",
-    "time.sleep(4)  # give server time to start\n",
-    "print('Env server started, PID:', server_proc.pid)"
+    "# Cell 3 — Download scenario corpus from HF Hub\n",
+    "from pathlib import Path\n",
+    "from huggingface_hub import snapshot_download\n",
+    "\n",
+    "SCEN_DIR = Path('/data/scenarios')\n",
+    "SCEN_DIR.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "existing = list(SCEN_DIR.rglob('*.json'))\n",
+    "if existing:\n",
+    "    print(f'Scenarios already present: {len(existing)} files — skipping download')\n",
+    "else:\n",
+    "    snapshot_download(\n",
+    "        repo_id=SCENARIOS_REPO,\n",
+    "        repo_type='dataset',\n",
+    "        local_dir=str(SCEN_DIR),\n",
+    "        token=HF_TOKEN,\n",
+    "    )\n",
+    "    n = len(list(SCEN_DIR.rglob('*.json')))\n",
+    "    print(f'Downloaded {n} scenario files')\n",
+    "\n",
+    "train_dir = SCEN_DIR / 'train'\n",
+    "print(f'Train scenarios: {len(list(train_dir.rglob(\"*.json\")))}')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-sft-ds",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 5: Generate SFT trajectories (skip if already done)\n",
-    "import os\n",
-    "if not os.path.exists('data_artifacts/sft_dataset'):\n",
-    "    from ci_triage_env.training.trajectory_gen import main as traj_main\n",
-    "    traj_main([\n",
-    "        '--count', '600',\n",
-    "        '--model', 'gpt-4o-mini',\n",
-    "        '--budget', '25.0',\n",
-    "        '--output', 'data_artifacts/sft_dataset/',\n",
-    "    ])\n",
+    "# Cell 4 — Download SFT dataset from HF Hub\n",
+    "from datasets import load_dataset, load_from_disk\n",
+    "\n",
+    "SFT_DS_DIR = Path('/data/sft_dataset')\n",
+    "\n",
+    "if SFT_DS_DIR.exists():\n",
+    "    ds = load_from_disk(str(SFT_DS_DIR))\n",
+    "    print(f'SFT dataset already present: {len(ds)} examples')\n",
     "else:\n",
-    "    print('SFT dataset already exists, skipping generation.')"
+    "    ds = load_dataset(SFT_DATASET_REPO, split='train', token=HF_TOKEN)\n",
+    "    SFT_DS_DIR.mkdir(parents=True, exist_ok=True)\n",
+    "    ds.save_to_disk(str(SFT_DS_DIR))\n",
+    "    print(f'Downloaded {len(ds)} SFT examples')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-sft",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 6: SFT warmstart\n",
+    "# Cell 5 — SFT warmstart\n",
+    "# Estimated time: ~45 min on A10G Large\n",
+    "# Optimised: batch_size=4, grad_accum=4 → effective batch 16, bf16 via unsloth\n",
     "from ci_triage_env.training.sft import run_sft\n",
     "\n",
-    "run_sft(\n",
-    "    dataset_path='data_artifacts/sft_dataset/',\n",
-    "    output_dir='checkpoints/sft/',\n",
-    "    num_epochs=3,\n",
-    "    per_device_batch_size=1,\n",
-    "    gradient_accumulation_steps=4,\n",
-    ")\n",
-    "print('SFT complete. Checkpoint at checkpoints/sft/')"
+    "SFT_CKPT = Path('/data/checkpoints/sft')\n",
+    "\n",
+    "if SFT_CKPT.exists():\n",
+    "    print(f'SFT checkpoint found at {SFT_CKPT} — skipping (delete to retrain)')\n",
+    "else:\n",
+    "    run_sft(\n",
+    "        dataset_path=str(SFT_DS_DIR),\n",
+    "        output_dir=str(SFT_CKPT),\n",
+    "        num_epochs=2,\n",
+    "        per_device_batch_size=4,       # A10G Large has 46 GB — fits 4 sequences\n",
+    "        gradient_accumulation_steps=4, # effective batch = 16\n",
+    "    )\n",
+    "    print(f'SFT done → {SFT_CKPT}')\n",
+    "\n",
+    "    # Push immediately so checkpoint is safe even if GRPO fails\n",
+    "    from huggingface_hub import upload_folder\n",
+    "    upload_folder(\n",
+    "        folder_path=str(SFT_CKPT),\n",
+    "        repo_id=MODEL_REPO + '-sft',\n",
+    "        repo_type='model',\n",
+    "        token=HF_TOKEN,\n",
+    "        commit_message='SFT warmstart checkpoint (Qwen3.5-4B + LoRA)',\n",
+    "    )\n",
+    "    print(f'SFT checkpoint pushed to {MODEL_REPO}-sft')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-grpo",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 7: GRPO smoke test (100 steps, ~30 min)\n",
+    "# Cell 6 — GRPO fine-tuning\n",
+    "# Estimated time: ~90 min for 100 steps on A10G Large\n",
+    "#\n",
+    "# Why 100 steps? Each step = 4 multi-turn rollouts (max 4 tool calls each).\n",
+    "# Sequential rollout with model.generate() is the bottleneck: ~50 sec/step.\n",
+    "# Increase GRPO_STEPS if you have more time budget.\n",
+    "#\n",
+    "# MockEnvClient is used in-process — no server needed, full speed.\n",
+    "\n",
+    "from ci_triage_env.training.mock_env_client import MockEnvClient\n",
     "from ci_triage_env.training.grpo import run_grpo\n",
     "\n",
+    "GRPO_CKPT  = Path('/data/checkpoints/grpo')\n",
+    "GRPO_STEPS = 100  # increase to 200 if you have ~3 hours total\n",
+    "\n",
+    "env_client = MockEnvClient(scenarios_dir=str(SCEN_DIR / 'train'))\n",
+    "print(f'MockEnvClient loaded {len(env_client.scenario_ids)} train scenarios')\n",
+    "print(f'Starting GRPO — {GRPO_STEPS} steps, ~{GRPO_STEPS * 50 // 60} min estimated')\n",
+    "print('Monitor: https://wandb.ai (project: ci-triage-env)')\n",
+    "\n",
     "run_grpo(\n",
-    "    sft_checkpoint_dir='checkpoints/sft/',\n",
-    "    output_dir='checkpoints/grpo_smoke/',\n",
-    "    total_steps=100,\n",
-    ")\n",
-    "print('Smoke test complete. Check W&B for reward curve.')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Cell 8: Full GRPO run (3000 steps, ~30h wall-clock)\n",
-    "# Monitor: https://wandb.ai/<entity>/ci-triage-env\n",
-    "# Hard-stop rules: see plan/branch-c-reward-training/phase-c4.md\n",
-    "run_grpo(\n",
-    "    sft_checkpoint_dir='checkpoints/sft/',\n",
-    "    output_dir='checkpoints/grpo_full/',\n",
-    "    total_steps=3000,\n",
+    "    sft_checkpoint_dir=str(SFT_CKPT),\n",
+    "    output_dir=str(GRPO_CKPT),\n",
+    "    total_steps=GRPO_STEPS,\n",
+    "    env_client=env_client,\n",
+    "    scenarios_train_path=str(SCEN_DIR / 'train'),\n",
+    "    hyperparams={\n",
+    "        # ── training update (fast) ──────────────────────\n",
+    "        'per_device_train_batch_size': 1,\n",
+    "        'gradient_accumulation_steps': 4,   # effective batch = 4\n",
+    "        'learning_rate': 5e-6,\n",
+    "        'kl_coef': 0.04,\n",
+    "        # ── rollout generation (bottleneck) ────────────\n",
+    "        'num_generations': 4,               # 4 rollouts per training sample\n",
+    "        'max_prompt_length': 2048,\n",
+    "        'max_completion_length': 256,        # short = fast; CI responses are concise\n",
+    "        'temperature': 0.8,\n",
+    "        'top_p': 0.95,\n",
+    "        # ── logging ────────────────────────────────────\n",
+    "        'logging_steps': 5,\n",
+    "        'save_steps': 50,\n",
+    "        'report_to': 'wandb' if WANDB_KEY else 'none',\n",
+    "    },\n",
     ")\n",
-    "print('Full GRPO complete.')"
+    "print(f'GRPO done → {GRPO_CKPT}')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-push",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Cell 9: Push trained adapter to HF Hub\n",
+    "# Cell 7 — Push final model to HF Hub\n",
     "from huggingface_hub import upload_folder\n",
     "\n",
     "upload_folder(\n",
-    "    folder_path='checkpoints/grpo_full/',\n",
-    "    repo_id=HF_MODEL_REPO,\n",
+    "    folder_path=str(GRPO_CKPT),\n",
+    "    repo_id=MODEL_REPO,\n",
     "    repo_type='model',\n",
-    "    commit_message='CI-Triage-Env GRPO-trained Qwen3.5-4B adapter',\n",
+    "    token=HF_TOKEN,\n",
+    "    commit_message=f'GRPO-trained adapter — {GRPO_STEPS} steps on A10G Large',\n",
     ")\n",
-    "print(f'Adapter pushed to https://huggingface.co/{HF_MODEL_REPO}')"
+    "print(f'Final model: https://huggingface.co/{MODEL_REPO}')"
    ]
   }
  ],

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "huggingface_hub>=0.23",
     "jsonschema>=4.21",
     "openenv-core>=0.2.3",
+    "fastmcp>=0.4",
 ]
 
 [project.optional-dependencies]

src/ci_triage_env/training/grpo.py

@@ -64,10 +64,12 @@ def run_grpo(
     train_dir = Path(scenarios_train_path)
     scenario_ids = [p.stem for p in train_dir.rglob("*.json")] if train_dir.exists() else []
 
+    max_turns = hp.pop("max_turns", 4)   # short episodes for faster GRPO
     rollout = TrainingRollout(
         env_client=env_client,
         scenarios_train=scenario_ids,
         weights=weights_override,
+        max_turns=max_turns,
     )
 
     model, tokenizer = load_model_for_sft(model_name=sft_checkpoint_dir)

train-entrypoint.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [[ "${START_MODE:-jupyter}" == "auto" ]]; then
+    echo "[train-entrypoint] START_MODE=auto — running train.py"
+    exec python /workspace/train.py
+else
+    echo "[train-entrypoint] START_MODE=jupyter — launching JupyterLab on :7860"
+    exec jupyter lab \
+        --ip=0.0.0.0 \
+        --port=7860 \
+        --no-browser \
+        --allow-root \
+        --NotebookApp.token="" \
+        --NotebookApp.password="" \
+        --notebook-dir=/workspace
+fi

train.py

@@ -0,0 +1,158 @@
+"""Automated end-to-end training script for HF Spaces.
+
+Runs: scenario download → SFT warmstart → GRPO fine-tuning → push to HF Hub.
+All config comes from environment variables (set as Space secrets).
+
+Optimised for A10G Large (46 GB VRAM, 12 vCPU).
+
+Required env vars:
+    HF_TOKEN         - HuggingFace write token
+    HF_USERNAME      - your HF username
+    WANDB_API_KEY    - Weights & Biases API key
+
+Optional:
+    HF_SCENARIOS_REPO   - default: {HF_USERNAME}/ci-triage-scenarios
+    HF_SFT_DATASET_REPO - default: {HF_USERNAME}/ci-triage-sft
+    HF_MODEL_REPO       - default: {HF_USERNAME}/ci-triage-agent
+    GRPO_STEPS          - default: 100  (set lower to finish faster, higher for more training)
+    SKIP_SFT            - set to "1" to skip SFT and jump straight to GRPO (if checkpoint exists)
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+# ── resolve config ────────────────────────────────────────────────────────────
+HF_TOKEN    = os.environ["HF_TOKEN"]
+HF_USERNAME = os.environ["HF_USERNAME"]
+WANDB_KEY   = os.environ.get("WANDB_API_KEY", "")
+
+SCENARIOS_REPO   = os.environ.get("HF_SCENARIOS_REPO",   f"{HF_USERNAME}/ci-triage-scenarios")
+SFT_DATASET_REPO = os.environ.get("HF_SFT_DATASET_REPO", f"{HF_USERNAME}/ci-triage-sft")
+MODEL_REPO       = os.environ.get("HF_MODEL_REPO",       f"{HF_USERNAME}/ci-triage-agent")
+GRPO_STEPS       = int(os.environ.get("GRPO_STEPS", "100"))
+SKIP_SFT         = os.environ.get("SKIP_SFT", "0") == "1"
+
+DATA_ROOT  = Path("/data")
+SCEN_DIR   = DATA_ROOT / "scenarios"
+SFT_DS_DIR = DATA_ROOT / "sft_dataset"
+SFT_CKPT   = DATA_ROOT / "checkpoints" / "sft"
+GRPO_CKPT  = DATA_ROOT / "checkpoints" / "grpo"
+
+# ── auth ──────────────────────────────────────────────────────────────────────
+from huggingface_hub import login
+login(token=HF_TOKEN)
+
+if WANDB_KEY:
+    import wandb
+    wandb.login(key=WANDB_KEY)
+    os.environ["WANDB_PROJECT"] = "ci-triage-env"
+else:
+    os.environ["WANDB_DISABLED"] = "true"
+
+# ── Step 1: download scenario corpus ─────────────────────────────────────────
+if not SCEN_DIR.exists() or not any(SCEN_DIR.rglob("*.json")):
+    print(f"\n[1/4] Downloading scenarios from {SCENARIOS_REPO} …")
+    from huggingface_hub import snapshot_download
+    snapshot_download(
+        repo_id=SCENARIOS_REPO,
+        repo_type="dataset",
+        local_dir=str(SCEN_DIR),
+        token=HF_TOKEN,
+    )
+else:
+    n = sum(1 for _ in SCEN_DIR.rglob("*.json"))
+    print(f"\n[1/4] Scenarios already present ({n} files) — skipping download.")
+
+train_scen = list(SCEN_DIR.rglob("train/**/*.json")) or list(SCEN_DIR.rglob("*.json"))
+print(f"      Train scenarios available: {len(train_scen)}")
+
+# ── Step 2: download SFT dataset ─────────────────────────────────────────────
+if not SFT_DS_DIR.exists():
+    print(f"\n[2/4] Downloading SFT dataset from {SFT_DATASET_REPO} …")
+    from datasets import load_dataset
+    ds = load_dataset(SFT_DATASET_REPO, split="train", token=HF_TOKEN)
+    SFT_DS_DIR.mkdir(parents=True, exist_ok=True)
+    ds.save_to_disk(str(SFT_DS_DIR))
+    print(f"      {len(ds)} SFT examples saved.")
+else:
+    from datasets import load_from_disk
+    ds = load_from_disk(str(SFT_DS_DIR))
+    print(f"\n[2/4] SFT dataset already present ({len(ds)} examples) — skipping download.")
+
+# ── Step 3: SFT warmstart ─────────────────────────────────────────────────────
+if SKIP_SFT and SFT_CKPT.exists():
+    print(f"\n[3/4] SKIP_SFT=1 and checkpoint found at {SFT_CKPT} — skipping SFT.")
+else:
+    print(f"\n[3/4] SFT warmstart — {len(ds)} examples, A10G-optimised settings …")
+    from ci_triage_env.training.sft import run_sft
+    run_sft(
+        dataset_path=str(SFT_DS_DIR),
+        output_dir=str(SFT_CKPT),
+        num_epochs=2,
+        per_device_batch_size=4,      # 46 GB → fit 4 sequences comfortably
+        gradient_accumulation_steps=4, # effective batch = 16
+    )
+    print(f"      SFT done → {SFT_CKPT}")
+
+    # Push SFT checkpoint immediately so it's saved even if GRPO fails
+    print("      Pushing SFT checkpoint to HF Hub …")
+    from huggingface_hub import upload_folder
+    upload_folder(
+        folder_path=str(SFT_CKPT),
+        repo_id=MODEL_REPO + "-sft",
+        repo_type="model",
+        token=HF_TOKEN,
+        commit_message="SFT warmstart checkpoint",
+    )
+
+# ── Step 4: GRPO fine-tuning ──────────────────────────────────────────────────
+print(f"\n[4/4] GRPO training — {GRPO_STEPS} steps, MockEnvClient in-process …")
+print("      Monitoring: https://wandb.ai (search project ci-triage-env)")
+
+from ci_triage_env.training.mock_env_client import MockEnvClient
+from ci_triage_env.training.grpo import run_grpo
+
+env_client = MockEnvClient(scenarios_dir=str(SCEN_DIR / "train"))
+print(f"      Loaded {len(env_client.scenario_ids)} train scenarios into MockEnvClient")
+
+# A10G Large optimised hyperparams.
+# max_turns=4 + max_completion_length=256 keeps each rollout to ~15 sec so
+# 100 steps × 4 rollouts ≈ 100 min total — fits the 2-3 hour budget.
+run_grpo(
+    sft_checkpoint_dir=str(SFT_CKPT),
+    output_dir=str(GRPO_CKPT),
+    total_steps=GRPO_STEPS,
+    env_client=env_client,
+    scenarios_train_path=str(SCEN_DIR / "train"),
+    hyperparams={
+        "per_device_train_batch_size": 1,
+        "gradient_accumulation_steps": 4,   # effective batch = 4
+        "num_generations": 4,
+        "max_prompt_length": 2048,
+        "max_completion_length": 256,
+        "learning_rate": 5e-6,
+        "kl_coef": 0.04,
+        "temperature": 0.8,
+        "top_p": 0.95,
+        "logging_steps": 5,
+        "save_steps": 50,
+        "report_to": "wandb" if WANDB_KEY else "none",
+    },
+)
+print(f"      GRPO done → {GRPO_CKPT}")
+
+# ── Push final model ──────────────────────────────────────────────────────────
+print(f"\n[done] Pushing final model to {MODEL_REPO} …")
+from huggingface_hub import upload_folder
+upload_folder(
+    folder_path=str(GRPO_CKPT),
+    repo_id=MODEL_REPO,
+    repo_type="model",
+    token=HF_TOKEN,
+    commit_message=f"GRPO-trained adapter — {GRPO_STEPS} steps",
+)
+print(f"       Model at: https://huggingface.co/{MODEL_REPO}")
+print("\nTraining complete.")