Smoke mode + detailed train entry

kaggle/train_grpo.py

@@ -775,6 +775,7 @@ class Cfg:
         max_eval_turns: int = 8,
         max_train_instances: int = 16,
         lora_r: int = 16,
+        mode: str | None = None,
     ) -> None:
         self.dataset_id = dataset_id
         self.hub_model_id = hub_model_id
@@ -791,6 +792,8 @@ class Cfg:
         self.max_eval_turns = max_eval_turns
         self.max_train_instances = max_train_instances
         self.lora_r = lora_r
+        # mode: "smoke" = 1 instance + reward checks then exit; "full" = GRPO train
+        self.mode = (mode or os.getenv("DARL_MODE") or "full").strip().lower()
 
 
 def _load_swe_splits(cfg: Cfg):
@@ -829,10 +832,275 @@ def _load_swe_splits(cfg: Cfg):
     return Dataset.from_list(train_rows), eval_rows
 
 
+def run_smoke_one_case(cfg: Cfg) -> dict[str, Any]:
+    """One SWE instance end-to-end: env tools + empty-patch reward + 1 agent rollout.
+
+    Hard-fails if grader/env/reward plumbing is broken. Soft-warns if the small
+    model makes zero tool calls or scores 0 (common before training).
+    """
+    import torch
+    import wandb
+    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+    from deepagents_rl.train.swe_env import DeepAgentSWEEnv
+
+    out_dir = Path("/kaggle/working/grpo-out")
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # Force single-case data
+    cfg.max_train_instances = 1
+    cfg.n_eval = 1
+    cfg.max_eval_turns = max(cfg.max_eval_turns, 10)
+    train, eval_problems = _load_swe_splits(cfg)
+    if eval_problems:
+        problem = dict(eval_problems[0])
+    else:
+        problem = dict(train[0])
+    instance_id = problem.get("instance_id") or problem.get("task_id")
+    _vlog("smoke.instance", instance_id=instance_id, keys=list(problem.keys())[:20])
+
+    checks: dict[str, Any] = {
+        "instance_id": instance_id,
+        "grader_health": False,
+        "env_reset": False,
+        "ls_root_ok": False,
+        "ls_src_aliased": False,
+        "empty_patch_reward_ok": False,
+        "empty_patch_reward": None,
+        "empty_patch_report": {},
+        "agent_reward": None,
+        "agent_n_tool_calls": 0,
+        "agent_tools_used": [],
+        "agent_resolved": False,
+        "langfuse_trace_id": "",
+        "hard_ok": False,
+        "soft_ok": False,
+        "errors": [],
+    }
+
+    import requests
+
+    grader = os.environ["SWE_GRADER_URL"].rstrip("/")
+    headers = {"ngrok-skip-browser-warning": "1"}
+    if os.environ.get("SWE_GRADER_TOKEN"):
+        headers["Authorization"] = f"Bearer {os.environ['SWE_GRADER_TOKEN']}"
+    health = requests.get(f"{grader}/health", headers=headers, timeout=30)
+    health.raise_for_status()
+    checks["grader_health"] = True
+    _vlog("smoke.grader_health", body=health.json())
+
+    os.environ["WANDB_PROJECT"] = os.getenv("WANDB_PROJECT", "deepagents-rl-coding")
+    wandb.login(key=os.environ["WANDB_API_KEY"])
+    wandb.init(
+        project=os.environ["WANDB_PROJECT"],
+        name="kaggle-grpo-swe-smoke",
+        job_type="smoke",
+        config={"mode": "smoke", "instance_id": instance_id},
+    )
+
+    # --- Phase 1: deterministic env + empty-patch grader ---
+    env = DeepAgentSWEEnv()
+    try:
+        reset_kwargs = {
+            k: v
+            for k, v in problem.items()
+            if k
+            not in (
+                "prompt",
+                "problem_statement",
+            )
+        }
+        # Ensure docker_image key for reset
+        if "docker_image" not in reset_kwargs and problem.get("docker_image"):
+            reset_kwargs["docker_image"] = problem["docker_image"]
+        tip = env.reset(**reset_kwargs)
+        checks["env_reset"] = True
+        _vlog("smoke.env.reset", tip_preview=str(tip)[:200])
+
+        ls_root = env.ls(".")
+        checks["ls_root_ok"] = bool(ls_root) and "path_not_found" not in ls_root.lower()
+        _vlog("smoke.ls_root", preview=str(ls_root)[:500], ok=checks["ls_root_ok"])
+
+        ls_src = env.ls("src")
+        # Must NOT be bare ERROR path_not_found — alias/fallback should list root
+        checks["ls_src_aliased"] = (
+            "path_not_found" not in ls_src.lower()
+            or "repo root" in ls_src.lower()
+            or "showing repo root" in ls_src.lower()
+            or "listing repo root" in ls_src.lower()
+            or "not used in this repo" in ls_src.lower()
+        )
+        _vlog("smoke.ls_src", preview=str(ls_src)[:500], ok=checks["ls_src_aliased"])
+
+        empty_reward = float(env.get_reward())
+        empty_report = getattr(env, "_last_report", {}) or {}
+        checks["empty_patch_reward"] = empty_reward
+        checks["empty_patch_report"] = {
+            k: empty_report.get(k)
+            for k in (
+                "pass_frac",
+                "resolved",
+                "resolution",
+                "error",
+                "reward_total",
+            )
+        }
+        # Hard OK if grader returned a numeric reward without transport failure
+        err = str(empty_report.get("error") or "")
+        transport_fail = any(
+            x in err.lower()
+            for x in ("connection", "timeout", "refused", "ngrok", "502", "503", "404")
+        )
+        checks["empty_patch_reward_ok"] = (not transport_fail) and isinstance(
+            empty_reward, float
+        )
+        checks["langfuse_trace_id"] = getattr(env, "_trace_id", "") or ""
+        _vlog(
+            "smoke.empty_patch_reward",
+            reward=empty_reward,
+            report=checks["empty_patch_report"],
+            ok=checks["empty_patch_reward_ok"],
+        )
+    except Exception as e:  # noqa: BLE001
+        checks["errors"].append(f"env_phase: {type(e).__name__}: {e}")
+        _vlog("smoke.env_phase.error", error=str(e))
+        traceback.print_exc()
+    finally:
+        try:
+            env._end_trace()
+        except Exception:
+            pass
+        try:
+            env._cleanup()
+        except Exception:
+            pass
+
+    # --- Phase 2: one agent rollout (model may score 0; still validates loop) ---
+    model_id = resolve_model()
+    major, minor = (0, 0)
+    if torch.cuda.is_available():
+        major, minor = torch.cuda.get_device_capability(0)
+    use_bf16 = major >= 8 and torch.cuda.is_bf16_supported()
+    compute_dtype = torch.bfloat16 if use_bf16 else torch.float16
+    bnb = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=compute_dtype,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id, token=os.environ["HF_TOKEN"], trust_remote_code=True
+    )
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        quantization_config=bnb,
+        device_map="auto",
+        token=os.environ["HF_TOKEN"],
+        trust_remote_code=True,
+    )
+    model = prepare_model_for_kbit_training(model)
+    model = get_peft_model(
+        model,
+        LoraConfig(
+            r=cfg.lora_r,
+            lora_alpha=cfg.lora_r * 2,
+            lora_dropout=0.0,
+            bias="none",
+            task_type="CAUSAL_LM",
+            target_modules=[
+                "q_proj",
+                "k_proj",
+                "v_proj",
+                "o_proj",
+                "gate_proj",
+                "up_proj",
+                "down_proj",
+            ],
+        ),
+    )
+
+    agent_rows = run_heldout_agent_eval(
+        model,
+        tokenizer,
+        [problem],
+        max_turns=cfg.max_eval_turns,
+        max_new_tokens=512,
+    )
+    row = agent_rows[0] if agent_rows else {}
+    checks["agent_reward"] = float(row.get("reward_total") or 0.0)
+    checks["agent_n_tool_calls"] = int(row.get("n_tool_calls") or 0)
+    checks["agent_tools_used"] = list(row.get("tools_used") or [])
+    checks["agent_resolved"] = bool(row.get("resolved"))
+    if row.get("langfuse_trace_id"):
+        checks["langfuse_trace_id"] = row["langfuse_trace_id"]
+    checks["agent_completion_preview"] = str(row.get("completion") or "")[:800]
+    _vlog(
+        "smoke.agent_done",
+        reward=checks["agent_reward"],
+        n_tool_calls=checks["agent_n_tool_calls"],
+        tools=checks["agent_tools_used"],
+        resolved=checks["agent_resolved"],
+    )
+
+    checks["hard_ok"] = all(
+        [
+            checks["grader_health"],
+            checks["env_reset"],
+            checks["ls_root_ok"],
+            checks["ls_src_aliased"],
+            checks["empty_patch_reward_ok"],
+            not checks["errors"],
+        ]
+    )
+    checks["soft_ok"] = checks["agent_n_tool_calls"] > 0
+
+    report_path = out_dir / "smoke_report.json"
+    report_path.write_text(json.dumps(checks, indent=2, default=str), encoding="utf-8")
+    _vlog("smoke.report_written", path=str(report_path), hard_ok=checks["hard_ok"])
+
+    wandb.log(
+        {
+            "smoke/hard_ok": 1.0 if checks["hard_ok"] else 0.0,
+            "smoke/soft_ok": 1.0 if checks["soft_ok"] else 0.0,
+            "smoke/empty_patch_reward": float(checks["empty_patch_reward"] or 0.0),
+            "smoke/agent_reward": float(checks["agent_reward"] or 0.0),
+            "smoke/agent_n_tool_calls": float(checks["agent_n_tool_calls"]),
+            "smoke/instance_id": str(instance_id or ""),
+            "smoke/ls_root_ok": 1.0 if checks["ls_root_ok"] else 0.0,
+            "smoke/ls_src_aliased": 1.0 if checks["ls_src_aliased"] else 0.0,
+        }
+    )
+    wandb.finish()
+
+    print("=== SMOKE REPORT ===", flush=True)
+    print(json.dumps(checks, indent=2, default=str), flush=True)
+    if not checks["hard_ok"]:
+        raise RuntimeError(
+            "SMOKE HARD FAIL — grader/env/reward plumbing broken. "
+            f"errors={checks['errors']} report={report_path}"
+        )
+    if not checks["soft_ok"]:
+        print(
+            "SMOKE SOFT WARN: agent made 0 tool calls (prompt/model issue). "
+            "Plumbing OK — full GRPO may still train tool use.",
+            flush=True,
+        )
+    else:
+        print("SMOKE PASS (hard+soft). Safe to run full GRPO.", flush=True)
+    return checks
+
+
 def main(cfg: Cfg | None = None) -> None:
     global _TOKENIZER
     cfg = cfg or Cfg()
-    _vlog("main.begin", cfg_dataset=getattr(cfg, "dataset_id", None))
+    _vlog(
+        "main.begin",
+        cfg_dataset=getattr(cfg, "dataset_id", None),
+        mode=getattr(cfg, "mode", "full"),
+    )
     load_secrets()
     _ensure_package()
 
@@ -851,6 +1119,12 @@ def main(cfg: Cfg | None = None) -> None:
     except Exception as e:  # noqa: BLE001
         _vlog("main.logging_configure_failed", error=str(e))
 
+    if str(getattr(cfg, "mode", "full")).lower() == "smoke":
+        _vlog("main.branch_smoke")
+        run_smoke_one_case(cfg)
+        _vlog("main.smoke_done_exit")
+        return
+
     import torch
     import wandb
     from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training