fix: space_runner installs unsloth at click-time on live GPU

training/space_runner.py

@@ -1,8 +1,10 @@
 """
-Gradio UI for the training Space.
+Gradio UI for the HF training Space.
 Training does NOT start automatically — user must click "Start Training".
+Unsloth is installed at click-time so it picks up the A100's CUDA correctly.
 """
 import os
+import subprocess
 import sys
 import threading
 import time
@@ -14,7 +16,7 @@ from dotenv import load_dotenv
 load_dotenv()
 
 _log_lines: list[str] = []
-_training_status = "idle"  # idle | running | complete | failed
+_training_status = "idle"  # idle | installing | running | complete | failed
 
 
 def _append_log(msg: str):
@@ -22,78 +24,91 @@ def _append_log(msg: str):
     _log_lines.append(f"[{ts}] {msg}")
 
 
-def _run_training():
-    global _training_status
-    _training_status = "running"
-    _append_log("Training started.")
-    try:
-        # Redirect stdout so log lines appear in the UI
-        import io
-        import contextlib
+class _Tee:
+    def __init__(self, orig):
+        self._orig = orig
 
-        sys.path.insert(0, str(Path(__file__).parent.parent))
-        from training.train_grpo import main
-
-        # Capture print output
-        old_stdout = sys.stdout
-        old_stderr = sys.stderr
+    def write(self, msg):
+        if msg.strip():
+            _append_log(msg.rstrip())
+        self._orig.write(msg)
 
-        class Tee:
-            def __init__(self, orig):
-                self._orig = orig
+    def flush(self):
+        self._orig.flush()
 
-            def write(self, msg):
-                if msg.strip():
-                    _append_log(msg.rstrip())
-                self._orig.write(msg)
 
-            def flush(self):
-                self._orig.flush()
-
-        sys.stdout = Tee(old_stdout)
-        sys.stderr = Tee(old_stderr)
+def _run_training():
+    global _training_status
 
-        try:
-            main()
-        finally:
-            sys.stdout = old_stdout
-            sys.stderr = old_stderr
+    # Step 1 — install unsloth on the live GPU
+    _training_status = "installing"
+    _append_log("Installing unsloth on GPU hardware...")
+    try:
+        subprocess.check_call(
+            [sys.executable, "-m", "pip", "install", "--quiet",
+             "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"],
+            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+        )
+        _append_log("unsloth installed successfully.")
+    except subprocess.CalledProcessError as e:
+        _training_status = "failed"
+        _append_log(f"❌ unsloth install failed: {e}")
+        return
 
+    # Step 2 — run training
+    _training_status = "running"
+    _append_log("Starting GRPO training...")
+    old_out, old_err = sys.stdout, sys.stderr
+    sys.stdout = _Tee(old_out)
+    sys.stderr = _Tee(old_err)
+    try:
+        sys.path.insert(0, str(Path(__file__).parent.parent))
+        from training.train_grpo import main
+        main()
         _training_status = "complete"
-        _append_log("✅ Training complete. Check WandB for curves.")
+        _append_log("✅ Training complete! Model pushed to HF Hub. Check WandB for curves.")
     except Exception as e:
         _training_status = "failed"
         _append_log(f"❌ Training failed: {e}")
+    finally:
+        sys.stdout = old_out
+        sys.stderr = old_err
 
 
 def start_training():
     global _training_status
-    if _training_status == "running":
-        return "⚠️ Training is already running.", _get_logs()
+    if _training_status in ("installing", "running"):
+        return "⚠️ Already in progress.", _get_logs()
     if _training_status == "complete":
         return "✅ Training already complete.", _get_logs()
 
-    missing = []
-    if not os.getenv("WANDB_API_KEY"):
-        missing.append("WANDB_API_KEY")
-    if not os.getenv("HF_TOKEN"):
-        missing.append("HF_TOKEN")
-    if not os.getenv("HF_HUB_MODEL_ID"):
-        missing.append("HF_HUB_MODEL_ID")
+    missing = [k for k in ("WANDB_API_KEY", "HF_TOKEN", "HF_HUB_MODEL_ID")
+               if not os.getenv(k)]
     if missing:
-        return f"❌ Missing secrets: {', '.join(missing)}. Set them in Space Settings → Variables and secrets.", _get_logs()
+        return f"❌ Missing secrets: {', '.join(missing)}", _get_logs()
 
+    import torch
+    if not torch.cuda.is_available():
+        return "❌ No GPU detected. Upgrade Space hardware to A100 first.", _get_logs()
+
+    gpu = torch.cuda.get_device_name(0)
+    _append_log(f"GPU detected: {gpu}")
     threading.Thread(target=_run_training, daemon=True).start()
-    return "🚀 Training started! Logs updating below...", _get_logs()
+    return f"🚀 Started on {gpu}. Installing unsloth...", _get_logs()
 
 
 def _get_logs() -> str:
-    return "\n".join(_log_lines[-80:]) if _log_lines else "No logs yet."
+    return "\n".join(_log_lines[-100:]) if _log_lines else "No logs yet."
 
 
 def get_status() -> str:
-    icons = {"idle": "⏸️ Idle", "running": "🔄 Training in progress...",
-             "complete": "✅ Complete", "failed": "❌ Failed"}
+    icons = {
+        "idle": "⏸️  Idle — ready to start",
+        "installing": "⚙️  Installing unsloth on GPU...",
+        "running": "🔄  Training in progress...",
+        "complete": "✅  Training complete",
+        "failed": "❌  Failed — check logs",
+    }
     return icons.get(_training_status, _training_status)
 
 
@@ -106,32 +121,29 @@ def refresh():
 with gr.Blocks(title="ContextCorruption Training") as demo:
     gr.Markdown("""
 # ContextCorruption-Env — GRPO Training
-**Qwen2-1.5B-Instruct** fine-tuned to identify corrupted documents and resist misleading context.
+Fine-tuning **Qwen2-1.5B-Instruct** to identify corrupted documents.
 
-Before starting, ensure these secrets are set in **Space Settings → Variables and secrets**:
-- `WANDB_API_KEY`
-- `HF_TOKEN`
-- `HF_HUB_MODEL_ID` (e.g. `Siddh12334/qwen-1.5b-context-corruption`)
+**Before clicking Start, confirm:**
+- Space hardware is set to **A100 Large** (Settings → Space hardware)
+- Secrets are set: `WANDB_API_KEY` · `HF_TOKEN` · `HF_HUB_MODEL_ID`
 """)
 
-    status_box = gr.Textbox(label="Status", value="⏸️ Idle", interactive=False)
-    log_box = gr.Textbox(label="Training Logs", lines=20, interactive=False,
-                         value="Waiting to start...")
-    msg_box = gr.Textbox(label="Message", interactive=False)
+    status_box = gr.Textbox(label="Status", value=get_status(), interactive=False)
+    log_box    = gr.Textbox(label="Live Logs", lines=25, interactive=False, value="Waiting...")
+    msg_box    = gr.Textbox(label="", interactive=False)
 
     with gr.Row():
-        start_btn = gr.Button("🚀 Start Training", variant="primary", scale=2)
-        refresh_btn = gr.Button("🔄 Refresh Logs", scale=1)
+        start_btn   = gr.Button("🚀 Start Training", variant="primary", scale=2)
+        refresh_btn = gr.Button("🔄 Refresh", scale=1)
 
     gr.Markdown("""
 ---
-**Config:** 500 episodes · 3 epochs · Qwen2-1.5B · LoRA r=16 · A10G ~1.5 hrs · ~$2
+**Estimated time on A100:** ~30–45 min · **Cost:** ~$3–5 from your HF credits
+After training, model is pushed to HF Hub and WandB has the reward/loss curves.
 """)
 
     start_btn.click(fn=start_training, outputs=[msg_box, log_box])
     refresh_btn.click(fn=refresh, outputs=[status_box, log_box])
-
-    # Auto-refresh every 10s while running
     demo.load(fn=refresh, outputs=[status_box, log_box], every=10)