Fix notebook: HF Space URL, /step envelope, health check retry on cold start

- ENV_BASE_URL now points to ajsaxena-deceit.hf.space by default (USE_LOCAL_DOCKER=False)
- /step calls now use {"action": {...}} envelope as required by OpenEnv
- /reset response unpacked from {"observation": {...}} wrapper
- Health check retries 12x with 10s sleep to handle HF cold start

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

training/sanity_run.ipynb

@@ -30,46 +30,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# ============================================================\n",
-    "# SANITY RUN CONFIG (Phase 3)\n",
-    "# ============================================================\n",
-    "TRAINING_STEPS       = 50\n",
-    "ROLLOUTS_PER_PROMPT  = 4\n",
-    "BATCH_SIZE           = 2\n",
-    "LEARNING_RATE        = 5e-6\n",
-    "LORA_RANK            = 16\n",
-    "SAVE_STEPS           = 25\n",
-    "\n",
-    "# ============================================================\n",
-    "# FULL RUN CONFIG (Phase 5) — uncomment to activate\n",
-    "# ============================================================\n",
-    "# TRAINING_STEPS       = 500\n",
-    "# ROLLOUTS_PER_PROMPT  = 8\n",
-    "# BATCH_SIZE           = 4\n",
-    "# LEARNING_RATE        = 2e-6\n",
-    "# LORA_RANK            = 32\n",
-    "# SAVE_STEPS           = 100\n",
-    "\n",
-    "# ============================================================\n",
-    "# Environment connection — toggle here\n",
-    "# ============================================================\n",
-    "USE_LOCAL_DOCKER = True   # True = local Docker on port 8000 (default, faster)\n",
-    "                           # False = deployed HF Space (for Phase 5+)\n",
-    "\n",
-    "HF_SPACE_URL = \"https://<your-hf-username>-deceit-env.hf.space\"  # only used if above is False\n",
-    "\n",
-    "ENV_BASE_URL = \"http://localhost:8000\" if USE_LOCAL_DOCKER else HF_SPACE_URL\n",
-    "\n",
-    "# ============================================================\n",
-    "# Model & logging\n",
-    "# ============================================================\n",
-    "MODEL_NAME    = \"unsloth/Qwen2.5-0.5B-Instruct\"\n",
-    "HF_REPO_ID    = \"<your-hf-username>/deceit-qwen-0.5b-sanity\"  # checkpoint destination\n",
-    "WANDB_PROJECT = \"deceit-sanity\"\n",
-    "\n",
-    "print(f\"Config loaded. Steps={TRAINING_STEPS}, ENV={ENV_BASE_URL}\")"
-   ]
+   "source": "# ============================================================\n# SANITY RUN CONFIG (Phase 3)\n# ============================================================\nTRAINING_STEPS       = 50\nROLLOUTS_PER_PROMPT  = 4\nBATCH_SIZE           = 2\nLEARNING_RATE        = 5e-6\nLORA_RANK            = 16\nSAVE_STEPS           = 25\n\n# ============================================================\n# FULL RUN CONFIG (Phase 5) — uncomment to activate\n# ============================================================\n# TRAINING_STEPS       = 500\n# ROLLOUTS_PER_PROMPT  = 8\n# BATCH_SIZE           = 4\n# LEARNING_RATE        = 2e-6\n# LORA_RANK            = 32\n# SAVE_STEPS           = 100\n\n# ============================================================\n# Environment connection — toggle here\n# ============================================================\nUSE_LOCAL_DOCKER = False  # True = local Docker on port 8000\n                           # False = deployed HF Space (default for Colab)\n\nHF_SPACE_URL = \"https://ajsaxena-deceit.hf.space\"  # Ajsaxena/DECEIT on HF Spaces\n\nENV_BASE_URL = \"http://localhost:8000\" if USE_LOCAL_DOCKER else HF_SPACE_URL\n\n# ============================================================\n# Model & logging\n# ============================================================\nMODEL_NAME    = \"unsloth/Qwen2.5-0.5B-Instruct\"\nHF_REPO_ID    = \"Ajsaxena/deceit-qwen-0.5b-sanity\"  # checkpoint destination\nWANDB_PROJECT = \"deceit-sanity\"\n\nprint(f\"Config loaded. Steps={TRAINING_STEPS}, ENV={ENV_BASE_URL}\")"
   },
   {
    "cell_type": "markdown",
@@ -167,14 +128,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import requests\n",
-    "\n",
-    "# Verify env is reachable\n",
-    "resp = requests.get(f\"{ENV_BASE_URL}/health\", timeout=10)\n",
-    "print(f\"Health check: {resp.status_code} — {resp.json()}\")\n",
-    "assert resp.status_code == 200, f\"Env not reachable at {ENV_BASE_URL}\""
-   ]
+   "source": "import requests\nimport time\n\n# Verify env is reachable — retries for HF Space cold start (up to 2 min)\nprint(f\"Connecting to {ENV_BASE_URL} ...\")\nfor attempt in range(12):\n    try:\n        resp = requests.get(f\"{ENV_BASE_URL}/health\", timeout=15)\n        if resp.status_code == 200:\n            print(f\"✓ Health check passed: {resp.json()}\")\n            break\n        else:\n            print(f\"  Attempt {attempt+1}: status {resp.status_code}, retrying...\")\n    except Exception as e:\n        print(f\"  Attempt {attempt+1}: {e}, retrying in 10s...\")\n    time.sleep(10)\nelse:\n    raise RuntimeError(f\"Env not reachable at {ENV_BASE_URL} after 12 attempts\")"
   },
   {
    "cell_type": "markdown",
@@ -270,96 +224,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "def run_rollout(model, tokenizer, base_url: str, verbose: bool = False) -> dict:\n",
-    "    \"\"\"Run one full episode and return trajectory + total reward.\"\"\"\n",
-    "    # Reset environment\n",
-    "    resp = requests.post(f\"{base_url}/reset\", json={}, timeout=15)\n",
-    "    resp.raise_for_status()\n",
-    "    obs = resp.json()\n",
-    "\n",
-    "    question   = obs.get(\"question\", \"\")\n",
-    "    context    = obs.get(\"context\", [])\n",
-    "    max_turns  = obs.get(\"max_turns\", 3)\n",
-    "\n",
-    "    total_reward = 0.0\n",
-    "    steps        = 0\n",
-    "    parse_fails  = 0\n",
-    "    trajectory   = []\n",
-    "\n",
-    "    for turn in range(max_turns):\n",
-    "        # Build prompt for this turn\n",
-    "        context_str = \"\\n\".join(context) if context else \"\"\n",
-    "        user_content = f\"Question: {question}\"\n",
-    "        if context_str:\n",
-    "            user_content += f\"\\n\\n{context_str}\"\n",
-    "        user_content += f\"\\n\\nTurn {turn + 1} of {max_turns}. Respond in JSON.\"\n",
-    "\n",
-    "        messages = [\n",
-    "            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "            {\"role\": \"user\",   \"content\": user_content},\n",
-    "        ]\n",
-    "        prompt = tokenizer.apply_chat_template(\n",
-    "            messages, tokenize=False, add_generation_prompt=True\n",
-    "        )\n",
-    "        inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n",
-    "\n",
-    "        with torch.no_grad():\n",
-    "            output_ids = model.generate(\n",
-    "                **inputs,\n",
-    "                max_new_tokens=256,\n",
-    "                do_sample=True,\n",
-    "                temperature=0.7,\n",
-    "                pad_token_id=tokenizer.eos_token_id,\n",
-    "            )\n",
-    "        generated = tokenizer.decode(\n",
-    "            output_ids[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True\n",
-    "        )\n",
-    "\n",
-    "        # Parse action\n",
-    "        try:\n",
-    "            action = parse_action(generated)\n",
-    "        except Exception:\n",
-    "            action = PARSE_FAIL_ACTION.copy()\n",
-    "            parse_fails += 1\n",
-    "\n",
-    "        # Force final on last turn\n",
-    "        if turn == max_turns - 1:\n",
-    "            action[\"is_final\"] = True\n",
-    "\n",
-    "        if verbose:\n",
-    "            print(f\"  Turn {turn+1}: is_final={action['is_final']} answer='{action['answer']}' confidence={action['confidence']:.2f}\")\n",
-    "\n",
-    "        # Step environment\n",
-    "        step_resp = requests.post(f\"{base_url}/step\", json=action, timeout=30)\n",
-    "        step_resp.raise_for_status()\n",
-    "        step_obs = step_resp.json()\n",
-    "\n",
-    "        reward   = step_obs.get(\"reward\", 0.0)\n",
-    "        done     = step_obs.get(\"done\", False)\n",
-    "        context  = step_obs.get(\"context\", [])\n",
-    "\n",
-    "        total_reward += reward\n",
-    "        steps += 1\n",
-    "        trajectory.append({\n",
-    "            \"turn\": turn + 1, \"action\": action, \"reward\": reward,\n",
-    "            \"done\": done, \"metadata\": step_obs.get(\"metadata\", {})\n",
-    "        })\n",
-    "\n",
-    "        if done:\n",
-    "            break\n",
-    "\n",
-    "    return {\n",
-    "        \"question\":     question,\n",
-    "        \"total_reward\": total_reward,\n",
-    "        \"steps\":        steps,\n",
-    "        \"parse_fails\":  parse_fails,\n",
-    "        \"trajectory\":   trajectory,\n",
-    "    }\n",
-    "\n",
-    "\n",
-    "print(\"Rollout function ready.\")"
-   ]
+   "source": "def run_rollout(model, tokenizer, base_url: str, verbose: bool = False) -> dict:\n    \"\"\"Run one full episode and return trajectory + total reward.\"\"\"\n    resp = requests.post(f\"{base_url}/reset\", json={}, timeout=30)\n    resp.raise_for_status()\n    obs = resp.json()\n    # OpenEnv wraps observation in {\"observation\": {...}}\n    obs_data   = obs.get(\"observation\", obs)\n    question   = obs_data.get(\"question\", \"\")\n    context    = obs_data.get(\"context\", [])\n    max_turns  = obs_data.get(\"max_turns\", 3)\n\n    total_reward = 0.0\n    steps        = 0\n    parse_fails  = 0\n    trajectory   = []\n\n    for turn in range(max_turns):\n        context_str = \"\\n\".join(context) if context else \"\"\n        user_content = f\"Question: {question}\"\n        if context_str:\n            user_content += f\"\\n\\n{context_str}\"\n        user_content += f\"\\n\\nTurn {turn + 1} of {max_turns}. Respond in JSON.\"\n\n        messages = [\n            {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n            {\"role\": \"user\",   \"content\": user_content},\n        ]\n        prompt = tokenizer.apply_chat_template(\n            messages, tokenize=False, add_generation_prompt=True\n        )\n        inputs = tokenizer(prompt, return_tensors=\"pt\").to(model.device)\n\n        with torch.no_grad():\n            output_ids = model.generate(\n                **inputs,\n                max_new_tokens=256,\n                do_sample=True,\n                temperature=0.7,\n                pad_token_id=tokenizer.eos_token_id,\n            )\n        generated = tokenizer.decode(\n            output_ids[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True\n        )\n\n        try:\n            action = parse_action(generated)\n        except Exception:\n            action = PARSE_FAIL_ACTION.copy()\n            parse_fails += 1\n\n        if turn == max_turns - 1:\n            action[\"is_final\"] = True\n\n        if verbose:\n            print(f\"  Turn {turn+1}: is_final={action['is_final']} answer='{action['answer']}' confidence={action['confidence']:.2f}\")\n\n        # OpenEnv /step expects {\"action\": {...}}\n        step_resp = requests.post(f\"{base_url}/step\", json={\"action\": action}, timeout=30)\n        step_resp.raise_for_status()\n        step_obs = step_resp.json()\n\n        # Response is {\"observation\": {...}, \"reward\": ..., \"done\": ...}\n        step_obs_data = step_obs.get(\"observation\", step_obs)\n        reward  = step_obs.get(\"reward\", 0.0) or 0.0\n        done    = step_obs.get(\"done\", False)\n        context = step_obs_data.get(\"context\", [])\n\n        total_reward += reward\n        steps += 1\n        trajectory.append({\n            \"turn\": turn + 1, \"action\": action, \"reward\": reward,\n            \"done\": done, \"metadata\": step_obs_data.get(\"metadata\", {})\n        })\n\n        if done:\n            break\n\n    return {\n        \"question\":     question,\n        \"total_reward\": total_reward,\n        \"steps\":        steps,\n        \"parse_fails\":  parse_fails,\n        \"trajectory\":   trajectory,\n    }\n\n\nprint(\"Rollout function ready.\")"
   },
   {
    "cell_type": "markdown",
@@ -467,97 +332,7 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "import threading\n",
-    "\n",
-    "_env_lock = threading.Lock()\n",
-    "\n",
-    "def grpo_reward_fn(completions, prompts=None, **kwargs):\n",
-    "    \"\"\"GRPO reward function: run one rollout per completion, return list of rewards.\n",
-    "    \n",
-    "    GRPO passes a list of completions (generated texts) for the same prompt.\n",
-    "    Each gets an independent rollout in the environment.\n",
-    "    \"\"\"\n",
-    "    rewards = []\n",
-    "    parse_fail_count = 0\n",
-    "\n",
-    "    for completion_text in completions:\n",
-    "        # Parse the initial action from the model's first completion\n",
-    "        try:\n",
-    "            action = parse_action(completion_text)\n",
-    "        except Exception:\n",
-    "            action = PARSE_FAIL_ACTION.copy()\n",
-    "            parse_fail_count += 1\n",
-    "\n",
-    "        try:\n",
-    "            with _env_lock:\n",
-    "                # Reset for fresh episode\n",
-    "                reset_resp = requests.post(f\"{ENV_BASE_URL}/reset\", json={}, timeout=15)\n",
-    "                reset_resp.raise_for_status()\n",
-    "                obs = reset_resp.json()\n",
-    "                max_turns = obs.get(\"max_turns\", 3)\n",
-    "\n",
-    "                # If model committed on turn 1, just step once\n",
-    "                # If not final, continue rolling out with greedy decoding\n",
-    "                total_reward = 0.0\n",
-    "                current_action = action\n",
-    "                context = obs.get(\"context\", [])\n",
-    "                question = obs.get(\"question\", \"\")\n",
-    "\n",
-    "                for turn in range(max_turns):\n",
-    "                    if turn == max_turns - 1:\n",
-    "                        current_action[\"is_final\"] = True\n",
-    "\n",
-    "                    step_resp = requests.post(f\"{ENV_BASE_URL}/step\", json=current_action, timeout=30)\n",
-    "                    step_resp.raise_for_status()\n",
-    "                    step_obs = step_resp.json()\n",
-    "\n",
-    "                    total_reward += step_obs.get(\"reward\", 0.0)\n",
-    "                    done = step_obs.get(\"done\", False)\n",
-    "                    context = step_obs.get(\"context\", [])\n",
-    "\n",
-    "                    if done:\n",
-    "                        break\n",
-    "\n",
-    "                    # Continue rollout with model for subsequent turns\n",
-    "                    context_str = \"\\n\".join(context)\n",
-    "                    user_content = f\"Question: {question}\\n\\n{context_str}\\n\\nTurn {turn+2} of {max_turns}. Respond in JSON.\"\n",
-    "                    messages = [\n",
-    "                        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "                        {\"role\": \"user\",   \"content\": user_content},\n",
-    "                    ]\n",
-    "                    next_prompt = tokenizer.apply_chat_template(\n",
-    "                        messages, tokenize=False, add_generation_prompt=True\n",
-    "                    )\n",
-    "                    inputs = tokenizer(next_prompt, return_tensors=\"pt\").to(model.device)\n",
-    "                    with torch.no_grad():\n",
-    "                        out_ids = model.generate(\n",
-    "                            **inputs, max_new_tokens=256,\n",
-    "                            do_sample=False,  # greedy for subsequent turns\n",
-    "                            pad_token_id=tokenizer.eos_token_id,\n",
-    "                        )\n",
-    "                    next_text = tokenizer.decode(\n",
-    "                        out_ids[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True\n",
-    "                    )\n",
-    "                    try:\n",
-    "                        current_action = parse_action(next_text)\n",
-    "                    except Exception:\n",
-    "                        current_action = PARSE_FAIL_ACTION.copy()\n",
-    "\n",
-    "        except Exception as e:\n",
-    "            print(f\"  [reward_fn] Episode error: {e}\")\n",
-    "            total_reward = -1.3  # worst possible reward on crash\n",
-    "\n",
-    "        rewards.append(total_reward)\n",
-    "\n",
-    "    if parse_fail_count > 0:\n",
-    "        print(f\"  [reward_fn] Parse failures: {parse_fail_count}/{len(completions)}\")\n",
-    "\n",
-    "    return rewards\n",
-    "\n",
-    "\n",
-    "print(\"GRPO reward function ready.\")"
-   ]
+   "source": "import threading\n\n_env_lock = threading.Lock()\n\ndef grpo_reward_fn(completions, prompts=None, **kwargs):\n    \"\"\"GRPO reward function: run one rollout per completion, return list of rewards.\"\"\"\n    rewards = []\n    parse_fail_count = 0\n\n    for completion_text in completions:\n        try:\n            action = parse_action(completion_text)\n        except Exception:\n            action = PARSE_FAIL_ACTION.copy()\n            parse_fail_count += 1\n\n        try:\n            with _env_lock:\n                reset_resp = requests.post(f\"{ENV_BASE_URL}/reset\", json={}, timeout=30)\n                reset_resp.raise_for_status()\n                obs = reset_resp.json()\n                obs_data  = obs.get(\"observation\", obs)\n                max_turns = obs_data.get(\"max_turns\", 3)\n                question  = obs_data.get(\"question\", \"\")\n                context   = obs_data.get(\"context\", [])\n\n                total_reward   = 0.0\n                current_action = action\n\n                for turn in range(max_turns):\n                    if turn == max_turns - 1:\n                        current_action[\"is_final\"] = True\n\n                    # OpenEnv /step expects {\"action\": {...}}\n                    step_resp = requests.post(\n                        f\"{ENV_BASE_URL}/step\",\n                        json={\"action\": current_action},\n                        timeout=30,\n                    )\n                    step_resp.raise_for_status()\n                    step_obs      = step_resp.json()\n                    step_obs_data = step_obs.get(\"observation\", step_obs)\n\n                    reward  = step_obs.get(\"reward\", 0.0) or 0.0\n                    done    = step_obs.get(\"done\", False)\n                    context = step_obs_data.get(\"context\", [])\n                    total_reward += reward\n\n                    if done:\n                        break\n\n                    # Subsequent turns: greedy decoding\n                    context_str  = \"\\n\".join(context)\n                    user_content = f\"Question: {question}\\n\\n{context_str}\\n\\nTurn {turn+2} of {max_turns}. Respond in JSON.\"\n                    messages = [\n                        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n                        {\"role\": \"user\",   \"content\": user_content},\n                    ]\n                    next_prompt = tokenizer.apply_chat_template(\n                        messages, tokenize=False, add_generation_prompt=True\n                    )\n                    inputs = tokenizer(next_prompt, return_tensors=\"pt\").to(model.device)\n                    with torch.no_grad():\n                        out_ids = model.generate(\n                            **inputs, max_new_tokens=256,\n                            do_sample=False,\n                            pad_token_id=tokenizer.eos_token_id,\n                        )\n                    next_text = tokenizer.decode(\n                        out_ids[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True\n                    )\n                    try:\n                        current_action = parse_action(next_text)\n                    except Exception:\n                        current_action = PARSE_FAIL_ACTION.copy()\n\n        except Exception as e:\n            print(f\"  [reward_fn] Episode error: {e}\")\n            total_reward = -1.3\n\n        rewards.append(total_reward)\n\n    if parse_fail_count > 0:\n        print(f\"  [reward_fn] Parse failures: {parse_fail_count}/{len(completions)}\")\n\n    return rewards\n\n\nprint(\"GRPO reward function ready.\")"
   },
   {
    "cell_type": "markdown",
@@ -793,4 +568,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}