Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -22,18 +22,39 @@ import traceback
 from pathlib import Path
 from typing import Any, Callable
 
+
+def _vlog(msg: str, **fields: Any) -> None:
+    extra = ""
+    if fields:
+        try:
+            extra = " | " + json.dumps(fields, default=str)[:1500]
+        except Exception:
+            extra = f" | {fields!r}"[:1500]
+    print(f"[darl] {msg}{extra}", flush=True)
+
+
+def _truncate_preview(obj: Any, n: int = 400) -> str:
+    try:
+        s = obj if isinstance(obj, str) else json.dumps(obj, default=str)
+    except Exception:
+        s = repr(obj)
+    return s if len(s) <= n else s[:n] + "..."
+
+
 # ---------------------------------------------------------------------------
 # Bootstrap package (notebook may only download this file)
 # ---------------------------------------------------------------------------
 
 
 def _ensure_package() -> None:
+    _vlog("bootstrap.ensure_package.begin")
     try:
         import deepagents_rl  # noqa: F401
 
+        _vlog("bootstrap.ensure_package.already_importable")
         return
     except ImportError:
-        pass
+        _vlog("bootstrap.ensure_package.download_from_hf")
 
     token = os.environ.get("HF_TOKEN")
     if not token:
@@ -55,6 +76,8 @@ def _ensure_package() -> None:
         sys.path.insert(0, str(pkg_parent))
     import deepagents_rl  # noqa: F401
 
+    _vlog("bootstrap.ensure_package.done", path=str(pkg_parent))
+
 
 # ---------------------------------------------------------------------------
 # Secrets
@@ -62,6 +85,7 @@ def _ensure_package() -> None:
 
 
 def load_secrets() -> None:
+    _vlog("secrets.load.begin")
     try:
         from kaggle_secrets import UserSecretsClient
 
@@ -123,6 +147,13 @@ def load_secrets() -> None:
     os.environ.setdefault("WANDB_PROJECT", "deepagents-rl-coding")
     os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "0")
     os.environ.setdefault("TRL_EXPERIMENTAL_SILENCE", "1")
+    os.environ.setdefault("DARL_LOG_LEVEL", "DEBUG")
+    _vlog(
+        "secrets.load.done",
+        has_grader=bool(os.getenv("SWE_GRADER_URL")),
+        grader_host=(os.getenv("SWE_GRADER_URL") or "")[:60],
+        wandb_project=os.getenv("WANDB_PROJECT"),
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -217,12 +248,16 @@ def run_heldout_agent_eval(
 
     rows = []
     model.eval()
-    for p in problems:
+    _vlog("heldout_eval.begin", n_problems=len(problems), max_turns=max_turns)
+    for pi, p in enumerate(problems):
         env = DeepAgentSWEEnv()
+        task_id = p.get("task_id") or p.get("instance_id")
+        _vlog("heldout_eval.problem.begin", i=pi, task_id=task_id)
         try:
             env.reset(**{k: v for k, v in p.items() if k != "prompt"})
             tools = _env_tool_callables(env)
             tool_map = {t.__name__: t for t in tools}
+            _vlog("heldout_eval.tools", tools=list(tool_map.keys()))
             prompt = p.get("prompt")
             if isinstance(prompt, list):
                 messages = list(prompt)
@@ -232,6 +267,7 @@ def run_heldout_agent_eval(
             final_text = ""
             lf = getattr(env, "_lf", None)
             for turn in range(max_turns):
+                _vlog("heldout_eval.llm_turn.begin", task_id=task_id, turn=turn)
                 try:
                     prompt_text = tokenizer.apply_chat_template(
                         messages,
@@ -279,11 +315,26 @@ def run_heldout_agent_eval(
                     completion = _generate_once()
                     calls = _parse_tool_calls(completion)
 
+                _vlog(
+                    "heldout_eval.llm_turn.end",
+                    task_id=task_id,
+                    turn=turn,
+                    completion_preview=completion[:400],
+                    n_tool_calls=len(calls),
+                    tool_names=[c["name"] for c in calls],
+                )
                 final_text = completion
                 messages.append({"role": "assistant", "content": completion})
                 if not calls:
+                    _vlog("heldout_eval.stop_no_tools", task_id=task_id, turn=turn)
                     break
                 for call in calls:
+                    _vlog(
+                        "heldout_eval.tool_dispatch",
+                        task_id=task_id,
+                        tool=call["name"],
+                        arguments=_truncate_preview(call.get("arguments")),
+                    )
                     fn = tool_map.get(call["name"])
                     if fn is None:
                         result = f"Unknown tool: {call['name']}"
@@ -307,31 +358,34 @@ def run_heldout_agent_eval(
 
             reward = float(env.get_reward())
             report = getattr(env, "_last_report", {}) or {}
-            rows.append(
-                {
-                    "task_id": p.get("task_id") or p.get("instance_id"),
-                    "prompt": (
-                        prompt
-                        if isinstance(prompt, str)
-                        else json.dumps(prompt)[:2000]
-                    ),
-                    "completion": final_text[:4000],
-                    "pass_frac": float(report.get("pass_frac") or reward),
-                    "reward_total": reward,
-                    "resolved": bool(report.get("resolved")),
-                    "resolution": report.get("resolution"),
-                    "n_tool_calls": len(getattr(env, "_tool_calls", []) or []),
-                    "tools_used": [
-                        c.get("name") for c in (getattr(env, "_tool_calls", []) or [])
-                    ],
-                    "langfuse_trace_id": getattr(env, "_trace_id", "") or "",
-                }
-            )
+            row = {
+                "task_id": task_id,
+                "prompt": (
+                    prompt
+                    if isinstance(prompt, str)
+                    else json.dumps(prompt)[:2000]
+                ),
+                "completion": final_text[:4000],
+                "pass_frac": float(report.get("pass_frac") or reward),
+                "reward_total": reward,
+                "resolved": bool(report.get("resolved")),
+                "resolution": report.get("resolution"),
+                "n_tool_calls": len(getattr(env, "_tool_calls", []) or []),
+                "tools_used": [
+                    c.get("name") for c in (getattr(env, "_tool_calls", []) or [])
+                ],
+                "langfuse_trace_id": getattr(env, "_trace_id", "") or "",
+            }
+            rows.append(row)
+            _vlog("heldout_eval.problem.done", **{k: row[k] for k in (
+                "task_id", "reward_total", "pass_frac", "resolved", "n_tool_calls", "tools_used"
+            )})
         except Exception as e:  # noqa: BLE001
             traceback.print_exc()
+            _vlog("heldout_eval.problem.error", task_id=task_id, error=str(e))
             rows.append(
                 {
-                    "task_id": p.get("task_id") or p.get("instance_id"),
+                    "task_id": task_id,
                     "prompt": str(p.get("prompt"))[:2000],
                     "completion": f"EVAL_ERROR: {e}",
                     "pass_frac": 0.0,
@@ -340,6 +394,7 @@ def run_heldout_agent_eval(
                     "resolution": "ERROR",
                 }
             )
+    _vlog("heldout_eval.done", n=len(rows))
         finally:
             try:
                 env._end_trace()
@@ -488,11 +543,21 @@ def make_eval_before_save_callback(
 
             rev = f"grpo-step-{step}"
             problems = eval_problems[:n_eval]
-            print(f"\n=== AGENT EVAL BEFORE CHECKPOINT step={step} rev={rev} n={len(problems)} ===")
+            _vlog("ckpt_eval.begin", step=step, rev=rev, n=len(problems))
+            print(
+                f"\n=== AGENT EVAL BEFORE CHECKPOINT step={step} rev={rev} n={len(problems)} ===",
+                flush=True,
+            )
             try:
                 results = run_heldout_agent_eval(
                     model, tokenizer, problems, max_turns=max_eval_turns
                 )
+                _vlog(
+                    "ckpt_eval.rollouts_done",
+                    step=step,
+                    mean_reward=sum(float(r["reward_total"]) for r in results)
+                    / max(len(results), 1),
+                )
                 import wandb
 
                 wandb_run_id = None
@@ -668,9 +733,25 @@ def _load_swe_splits(cfg: Cfg):
 def main(cfg: Cfg | None = None) -> None:
     global _TOKENIZER
     cfg = cfg or Cfg()
+    _vlog("main.begin", cfg_dataset=getattr(cfg, "dataset_id", None))
     load_secrets()
     _ensure_package()
 
+    try:
+        from deepagents_rl.observability.logging import configure_logging
+
+        out_log = Path("/kaggle/working/grpo-out")
+        out_log.mkdir(parents=True, exist_ok=True)
+        configure_logging(
+            level=os.getenv("DARL_LOG_LEVEL", "DEBUG"),
+            log_file=out_log / "run.log",
+            events_jsonl=out_log / "events.jsonl",
+            force=True,
+        )
+        _vlog("main.logging_configured", log_file=str(out_log / "run.log"))
+    except Exception as e:  # noqa: BLE001
+        _vlog("main.logging_configure_failed", error=str(e))
+
     import torch
     import wandb
     from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
@@ -686,15 +767,20 @@ def main(cfg: Cfg | None = None) -> None:
     headers = {"ngrok-skip-browser-warning": "1"}
     if os.environ.get("SWE_GRADER_TOKEN"):
         headers["Authorization"] = f"Bearer {os.environ['SWE_GRADER_TOKEN']}"
+    _vlog("main.grader_health_check", url=grader)
     health = requests.get(f"{grader}/health", headers=headers, timeout=30)
     health.raise_for_status()
-    print("[grader] health", health.json())
+    _vlog("main.grader_health_ok", body=health.json())
 
     os.environ["WANDB_PROJECT"] = os.getenv("WANDB_PROJECT", "deepagents-rl-coding")
+    _vlog("main.wandb_login")
     wandb.login(key=os.environ["WANDB_API_KEY"])
 
+    _vlog("main.resolve_model")
     model_id = resolve_model()
+    _vlog("main.load_dataset")
     train, eval_problems = _load_swe_splits(cfg)
+    _vlog("main.dataset_ready", train_n=len(train), eval_n=len(eval_problems))
 
     bnb = BitsAndBytesConfig(
         load_in_4bit=True,
@@ -702,6 +788,7 @@ def main(cfg: Cfg | None = None) -> None:
         bnb_4bit_quant_type="nf4",
         bnb_4bit_use_double_quant=True,
     )
+    _vlog("main.load_tokenizer", model_id=model_id)
     tokenizer = AutoTokenizer.from_pretrained(
         model_id, token=os.environ["HF_TOKEN"], trust_remote_code=True
     )
@@ -709,6 +796,7 @@ def main(cfg: Cfg | None = None) -> None:
         tokenizer.pad_token = tokenizer.eos_token
     _TOKENIZER = tokenizer
 
+    _vlog("main.load_model_4bit", model_id=model_id, cuda=torch.cuda.is_available())
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         quantization_config=bnb,
@@ -716,6 +804,7 @@ def main(cfg: Cfg | None = None) -> None:
         token=os.environ["HF_TOKEN"],
         trust_remote_code=True,
     )
+    _vlog("main.prepare_lora")
     model = prepare_model_for_kbit_training(model)
     lora = LoraConfig(
         r=cfg.lora_r,
@@ -758,7 +847,8 @@ def main(cfg: Cfg | None = None) -> None:
         },
     )
 
-    print("=== BASELINE AGENT EVAL (step 0) ===")
+    _vlog("main.baseline_eval.begin", n_eval=cfg.n_eval)
+    print("=== BASELINE AGENT EVAL (step 0) ===", flush=True)
     baseline = run_heldout_agent_eval(
         model,
         tokenizer,
@@ -845,6 +935,13 @@ def main(cfg: Cfg | None = None) -> None:
         max_eval_turns=cfg.max_eval_turns,
     )
 
+    _vlog(
+        "main.trainer_init",
+        environment="DeepAgentSWEEnv",
+        max_steps=cfg.max_steps,
+        num_generations=cfg.num_generations,
+        max_tool_calling_iterations=cfg.max_tool_calling_iterations,
+    )
     trainer = GRPOTrainer(
         model=model,
         args=args,
@@ -855,7 +952,9 @@ def main(cfg: Cfg | None = None) -> None:
     )
     cb.trainer = trainer
 
+    _vlog("main.trainer.train.begin")
     trainer.train()
+    _vlog("main.trainer.train.done", global_step=int(trainer.state.global_step))
 
     final_step = int(trainer.state.global_step)
     rev = f"grpo-step-{final_step}-final"