| """
|
| Kaggle agent GRPO: Qwen + deepagents SWE harness + remote Docker reward.
|
|
|
| Hard requirements:
|
| - W&B metrics every step
|
| - BEFORE every checkpoint: held-out agent eval -> Langfuse -> THEN save/push HF
|
| - HF adapter revisions tagged grpo-step-{n}
|
| - SWE_GRADER_URL must reach local scripts/swe_grader_server.py (tunnel)
|
|
|
| Secrets (dataset / User Secrets / env):
|
| HF_TOKEN, WANDB_API_KEY, LANGFUSE_PUBLIC_KEY, LANGFUSE_SECRET_KEY
|
| SWE_GRADER_URL, optional SWE_GRADER_TOKEN
|
| """
|
| from __future__ import annotations
|
|
|
| import inspect
|
| import json
|
| import os
|
| import re
|
| import sys
|
| 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] + "..."
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def _ensure_package() -> None:
|
| """Always refresh deepagents_rl from HF so Kaggle never runs a stale copy."""
|
| _vlog("bootstrap.ensure_package.begin")
|
| token = os.environ.get("HF_TOKEN")
|
| if not token:
|
| raise RuntimeError("HF_TOKEN required to download deepagents_rl sources")
|
|
|
| from huggingface_hub import snapshot_download
|
|
|
| root = Path("/kaggle/working/deepagents_rl_src")
|
| _vlog("bootstrap.ensure_package.download_from_hf", local_dir=str(root))
|
| snapshot_download(
|
| repo_id=os.getenv("HF_REPO_ID", "vanishingradient/deepagents-rl-coding"),
|
| repo_type="model",
|
| token=token,
|
| allow_patterns=["src/deepagents_rl/**"],
|
| local_dir=str(root),
|
| force_download=True,
|
| )
|
| pkg_parent = root / "src"
|
|
|
| while str(pkg_parent) in sys.path:
|
| sys.path.remove(str(pkg_parent))
|
| sys.path.insert(0, str(pkg_parent))
|
| for mod in list(sys.modules):
|
| if mod == "deepagents_rl" or mod.startswith("deepagents_rl."):
|
| del sys.modules[mod]
|
| import deepagents_rl
|
| from deepagents_rl.train import swe_env as _swe_env
|
|
|
| _vlog(
|
| "bootstrap.ensure_package.done",
|
| path=str(pkg_parent),
|
| swe_env=str(Path(_swe_env.__file__).resolve()),
|
| has_phantom_alias=hasattr(_swe_env, "_PHANTOM_DIRS"),
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def load_secrets() -> None:
|
| _vlog("secrets.load.begin")
|
| try:
|
| from kaggle_secrets import UserSecretsClient
|
|
|
| us = UserSecretsClient()
|
| for key in (
|
| "HF_TOKEN",
|
| "WANDB_API_KEY",
|
| "LANGFUSE_PUBLIC_KEY",
|
| "LANGFUSE_SECRET_KEY",
|
| "LANGFUSE_HOST",
|
| "LANGFUSE_BASE_URL",
|
| "SWE_GRADER_URL",
|
| "SWE_GRADER_TOKEN",
|
| "WANDB_PROJECT",
|
| ):
|
| try:
|
| val = us.get_secret(key)
|
| if val:
|
| os.environ[key] = val
|
| except Exception:
|
| pass
|
| except Exception:
|
| pass
|
|
|
|
|
| secret_candidates = [Path("/kaggle/input/deepagents-rl-secrets/secrets.json")]
|
| input_root = Path("/kaggle/input")
|
| if input_root.exists():
|
| secret_candidates.extend(input_root.rglob("secrets.json"))
|
| for p in secret_candidates:
|
| if p.exists():
|
| try:
|
| data = json.loads(p.read_text(encoding="utf-8"))
|
| for k, v in data.items():
|
| if v and not os.getenv(k):
|
| os.environ[k] = str(v)
|
| except Exception:
|
| pass
|
| break
|
|
|
| missing = [
|
| k
|
| for k in (
|
| "HF_TOKEN",
|
| "WANDB_API_KEY",
|
| "LANGFUSE_PUBLIC_KEY",
|
| "LANGFUSE_SECRET_KEY",
|
| "SWE_GRADER_URL",
|
| )
|
| if not os.getenv(k)
|
| ]
|
| if missing:
|
| raise RuntimeError(
|
| f"Missing required secrets/env: {missing}. "
|
| "Need SWE_GRADER_URL pointing at the local Docker grader tunnel."
|
| )
|
| os.environ.setdefault("LANGFUSE_HOST", "https://cloud.langfuse.com")
|
| os.environ.setdefault("LANGFUSE_BASE_URL", os.environ["LANGFUSE_HOST"])
|
| 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"),
|
| )
|
|
|
|
|
|
|
|
|
|
|
|
|
| CANDIDATES = [
|
| "Qwen/Qwen2.5-Coder-1.5B-Instruct",
|
| "Qwen/Qwen2.5-Coder-0.5B-Instruct",
|
| "Qwen/Qwen2.5-1.5B-Instruct",
|
| "HuggingFaceTB/SmolLM2-1.7B-Instruct",
|
| "google/gemma-2-2b-it",
|
| ]
|
|
|
|
|
| def resolve_model() -> str:
|
| from huggingface_hub import HfApi
|
|
|
| override = os.getenv("TRAIN_MODEL_ID")
|
| api = HfApi(token=os.environ["HF_TOKEN"])
|
| tried = []
|
| for repo in ([override] if override else []) + CANDIDATES:
|
| if not repo or repo in tried:
|
| continue
|
| tried.append(repo)
|
| try:
|
| api.model_info(repo)
|
| print(f"[resolve] using {repo}")
|
| out = Path("/kaggle/working/resolved_model.json")
|
| out.parent.mkdir(parents=True, exist_ok=True)
|
| out.write_text(
|
| json.dumps({"train_model_id": repo, "tried": tried}, indent=2),
|
| encoding="utf-8",
|
| )
|
| return repo
|
| except Exception as e:
|
| print(f"[resolve] skip {repo}: {e}")
|
| raise RuntimeError(f"No model resolved; tried {tried}")
|
|
|
|
|
|
|
|
|
|
|
|
|
| _TOOL_CALL_RE = re.compile(
|
| r"<tool_call>\s*(\{.*?\})\s*</tool_call>",
|
| re.DOTALL,
|
| )
|
| _JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL)
|
| _BARE_JSON_RE = re.compile(
|
| r"\{\s*\"name\"\s*:\s*\"[^\"]+\"\s*,\s*\"arguments\"\s*:\s*\{.*?\}\s*\}",
|
| re.DOTALL,
|
| )
|
|
|
|
|
| def _env_tool_callables(env: Any) -> list[Callable[..., str]]:
|
|
|
| names = [
|
| name
|
| for name, _ in inspect.getmembers(type(env), predicate=inspect.isfunction)
|
| if name not in ("reset", "get_reward") and not name.startswith("_")
|
| ]
|
| return [getattr(env, name) for name in names]
|
|
|
|
|
| def _parse_one_tool_obj(obj: dict[str, Any]) -> dict[str, Any] | None:
|
| name = obj.get("name") or obj.get("function", {}).get("name")
|
| args = (
|
| obj.get("arguments")
|
| or obj.get("parameters")
|
| or obj.get("function", {}).get("arguments")
|
| or {}
|
| )
|
| if isinstance(args, str):
|
| try:
|
| args = json.loads(args)
|
| except json.JSONDecodeError:
|
| args = {"input": args}
|
| if not name:
|
| return None
|
| return {"name": str(name), "arguments": args if isinstance(args, dict) else {}}
|
|
|
|
|
| def _parse_tool_calls(text: str) -> list[dict[str, Any]]:
|
| """Parse Qwen-style tool calls: <tool_call>, ```json fences, or bare {name,arguments}."""
|
| calls: list[dict[str, Any]] = []
|
| seen: set[str] = set()
|
| raw = text or ""
|
|
|
| candidates: list[str] = []
|
| candidates.extend(m.group(1) for m in _TOOL_CALL_RE.finditer(raw))
|
| candidates.extend(m.group(1) for m in _JSON_FENCE_RE.finditer(raw))
|
| candidates.extend(m.group(0) for m in _BARE_JSON_RE.finditer(raw))
|
|
|
|
|
| stripped = raw.strip()
|
| if stripped.startswith("{") or stripped.startswith("["):
|
| candidates.append(stripped)
|
|
|
| for blob in candidates:
|
| try:
|
| obj = json.loads(blob)
|
| except json.JSONDecodeError:
|
| continue
|
| items = obj if isinstance(obj, list) else [obj]
|
| for item in items:
|
| if not isinstance(item, dict):
|
| continue
|
|
|
| if "function" in item and isinstance(item["function"], dict):
|
| parsed = _parse_one_tool_obj(item["function"])
|
| else:
|
| parsed = _parse_one_tool_obj(item)
|
| if parsed is None:
|
| continue
|
| key = json.dumps(parsed, sort_keys=True)
|
| if key in seen:
|
| continue
|
| seen.add(key)
|
| calls.append(parsed)
|
| return calls
|
|
|
|
|
| def _parsed_to_openai_tool_calls(parsed: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
| out = []
|
| for i, p in enumerate(parsed):
|
| args = p.get("arguments") if isinstance(p.get("arguments"), dict) else {}
|
| out.append(
|
| {
|
| "id": f"call_text_{i}",
|
| "type": "function",
|
| "function": {
|
| "name": str(p["name"]),
|
| "arguments": args,
|
| },
|
| }
|
| )
|
| return out
|
|
|
|
|
| def _inject_text_tool_calls_into_completions(
|
| completions: list[Any],
|
| *,
|
| force_ls_if_empty: bool = True,
|
| ) -> tuple[list[Any], int]:
|
| """TRL only runs tools from native ``tool_calls``. Qwen often emits markdown JSON — promote it."""
|
| injected = 0
|
| fixed: list[Any] = []
|
| for comp in completions:
|
|
|
| if (
|
| isinstance(comp, list)
|
| and comp
|
| and isinstance(comp[0], dict)
|
| and comp[0].get("role") == "assistant"
|
| ):
|
| msg = dict(comp[0])
|
| existing = msg.get("tool_calls") or []
|
| if existing:
|
| fixed.append([msg] + list(comp[1:]))
|
| continue
|
| content = str(msg.get("content") or "")
|
| parsed = _parse_tool_calls(content)
|
| if not parsed and force_ls_if_empty:
|
|
|
| parsed = [{"name": "ls", "arguments": {"path": "."}}]
|
| msg["content"] = (content + "\n").strip() + "\n(forced tool: ls path='.')"
|
| if parsed:
|
| msg["tool_calls"] = _parsed_to_openai_tool_calls(parsed)
|
| injected += 1
|
| _vlog(
|
| "train.inject_text_tool_calls",
|
| n=len(parsed),
|
| names=[p["name"] for p in parsed],
|
| forced=not bool(_parse_tool_calls(content)),
|
| )
|
| fixed.append([msg] + list(comp[1:]))
|
| else:
|
| fixed.append(comp)
|
| return fixed, injected
|
|
|
|
|
| def make_tool_repair_grpo_trainer(base_cls: type):
|
| """Subclass GRPOTrainer to execute markdown/JSON tool blobs + forced ls."""
|
|
|
| class ToolRepairGRPOTrainer(base_cls):
|
| def _tool_call_loop(self, prompts, prompt_ids, completion_ids, completions, logprobs, images, multimodal_fields):
|
| force = os.getenv("DARL_FORCE_FIRST_TOOL", "1").strip() not in ("0", "false", "no")
|
| completions, n_inj = _inject_text_tool_calls_into_completions(
|
| completions, force_ls_if_empty=force
|
| )
|
| if n_inj:
|
| _vlog("train.tool_repair.injected", n=n_inj, force=force)
|
| return super()._tool_call_loop(
|
| prompts,
|
| prompt_ids,
|
| completion_ids,
|
| completions,
|
| logprobs,
|
| images,
|
| multimodal_fields,
|
| )
|
|
|
| ToolRepairGRPOTrainer.__name__ = "ToolRepairGRPOTrainer"
|
| return ToolRepairGRPOTrainer
|
|
|
|
|
| def run_heldout_agent_eval(
|
| model,
|
| tokenizer,
|
| problems: list[dict],
|
| *,
|
| max_turns: int = 8,
|
| max_new_tokens: int = 512,
|
| ) -> list[dict]:
|
| """Greedy multi-turn tool loop + remote Docker reward (checkpoint metric)."""
|
| import torch
|
|
|
| from deepagents_rl.train.swe_env import DeepAgentSWEEnv
|
|
|
| rows = []
|
| model.eval()
|
| _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)
|
| else:
|
| messages = [{"role": "user", "content": str(prompt or "")}]
|
|
|
| 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,
|
| tools=tools,
|
| tokenize=False,
|
| add_generation_prompt=True,
|
| )
|
| except TypeError:
|
| prompt_text = tokenizer.apply_chat_template(
|
| messages, tokenize=False, add_generation_prompt=True
|
| )
|
| except Exception:
|
| prompt_text = messages[-1]["content"] if messages else ""
|
|
|
| def _generate_once() -> str:
|
| inputs_local = tokenizer(prompt_text, return_tensors="pt")
|
| inputs_local = {k: v.to(model.device) for k, v in inputs_local.items()}
|
| with torch.no_grad():
|
| out_ids = model.generate(
|
| **inputs_local,
|
| max_new_tokens=max_new_tokens,
|
| do_sample=False,
|
| pad_token_id=tokenizer.eos_token_id,
|
| )
|
| gen_ids = out_ids[0][inputs_local["input_ids"].shape[-1] :]
|
| return tokenizer.decode(gen_ids, skip_special_tokens=True)
|
|
|
| child_ctx = None
|
| if hasattr(env, "_child_trace_context"):
|
| child_ctx = env._child_trace_context()
|
| if lf is not None and child_ctx is not None:
|
| with lf.start_as_current_observation(
|
| name=f"llm-turn-{turn}",
|
| as_type="generation",
|
| trace_context=child_ctx,
|
| input={"messages_tail": messages[-2:], "turn": turn},
|
| metadata={
|
| "phase": "ckpt_eval",
|
| "turn": turn,
|
| "trace_id": env._trace_id,
|
| },
|
| ) as gen_span:
|
| completion = _generate_once()
|
| calls = _parse_tool_calls(completion)
|
| gen_span.update(
|
| output={
|
| "completion": completion[:4000],
|
| "n_tool_calls": len(calls),
|
| "tool_names": [c["name"] for c in calls],
|
| }
|
| )
|
| else:
|
| 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:
|
|
|
|
|
| force = os.getenv("DARL_FORCE_FIRST_TOOL", "1").strip() not in (
|
| "0",
|
| "false",
|
| "no",
|
| )
|
| if force and turn < 2:
|
| _vlog(
|
| "heldout_eval.force_first_tool",
|
| task_id=task_id,
|
| turn=turn,
|
| tool="ls",
|
| )
|
| calls = [{"name": "ls", "arguments": {"path": "."}}]
|
| messages.append(
|
| {
|
| "role": "assistant",
|
| "content": (
|
| "(system) No tool call detected — forcing ls(path='.') "
|
| "so you can continue with real paths."
|
| ),
|
| }
|
| )
|
| else:
|
| _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']}"
|
| else:
|
| try:
|
| result = fn(**call["arguments"])
|
| except TypeError:
|
| try:
|
| result = fn(*call["arguments"].values())
|
| except Exception as e:
|
| result = f"TOOL_ERROR: {e}"
|
| except Exception as e:
|
| result = f"TOOL_ERROR: {e}"
|
| messages.append(
|
| {
|
| "role": "tool",
|
| "name": call["name"],
|
| "content": str(result)[:4000],
|
| }
|
| )
|
|
|
| reward = float(env.get_reward())
|
| report = getattr(env, "_last_report", {}) 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:
|
| traceback.print_exc()
|
| _vlog("heldout_eval.problem.error", task_id=task_id, error=str(e))
|
| rows.append(
|
| {
|
| "task_id": task_id,
|
| "prompt": str(p.get("prompt"))[:2000],
|
| "completion": f"EVAL_ERROR: {e}",
|
| "pass_frac": 0.0,
|
| "reward_total": 0.0,
|
| "resolved": False,
|
| "resolution": "ERROR",
|
| }
|
| )
|
| finally:
|
| try:
|
| env._end_trace()
|
| except Exception:
|
| pass
|
| try:
|
| env._cleanup()
|
| except Exception:
|
| pass
|
| _vlog("heldout_eval.done", n=len(rows))
|
| return rows
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def langfuse_log_ckpt_eval(
|
| *,
|
| step: int,
|
| hf_revision: str,
|
| train_model_id: str,
|
| results: list[dict[str, Any]],
|
| wandb_run_id: str | None,
|
| ) -> dict[str, Any]:
|
| from langfuse import Langfuse, get_client
|
|
|
| Langfuse(
|
| public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
|
| secret_key=os.environ["LANGFUSE_SECRET_KEY"],
|
| host=os.getenv("LANGFUSE_HOST") or os.getenv("LANGFUSE_BASE_URL"),
|
| )
|
| lf = get_client()
|
| if not lf.auth_check():
|
| raise RuntimeError("Langfuse auth failed during ckpt eval")
|
|
|
| n = len(results)
|
| mean_pass = sum(float(r.get("pass_frac") or 0) for r in results) / n if n else 0.0
|
| mean_reward = sum(float(r.get("reward_total") or 0) for r in results) / n if n else 0.0
|
| n_perfect = sum(1 for r in results if float(r.get("pass_frac") or 0) >= 1.0)
|
| n_resolved = sum(1 for r in results if r.get("resolved"))
|
|
|
|
|
| with lf.start_as_current_observation(
|
| name=f"kaggle-ckpt-eval:step-{step}",
|
| as_type="evaluator",
|
| metadata={
|
| "step": step,
|
| "hf_revision": hf_revision,
|
| "train_model_id": train_model_id,
|
| "wandb_run_id": wandb_run_id,
|
| "phase": "ckpt_eval_aggregate",
|
| "harness": "deepagents-swe",
|
| },
|
| input={"step": step, "hf_revision": hf_revision, "n": n},
|
| end_on_exit=True,
|
| ) as root:
|
| agg_trace_id = str(getattr(root, "trace_id", "") or "")
|
| root_span_id = str(getattr(root, "id", "") or "")
|
| for r in results:
|
| child_ctx = (
|
| {"trace_id": agg_trace_id, "parent_span_id": root_span_id}
|
| if agg_trace_id and root_span_id
|
| else None
|
| )
|
| with lf.start_as_current_observation(
|
| name=f"ckpt-summary:{r['task_id']}",
|
| as_type="span",
|
| trace_context=child_ctx,
|
| input={"prompt": str(r.get("prompt", ""))[:2000]},
|
| ) as span:
|
| span.update(
|
| output={
|
| "completion": str(r.get("completion", ""))[:2000],
|
| "pass_frac": r.get("pass_frac"),
|
| "reward_total": r.get("reward_total"),
|
| "resolved": r.get("resolved"),
|
| "n_tool_calls": r.get("n_tool_calls"),
|
| "tools_used": r.get("tools_used"),
|
| "rollout_trace_id": r.get("langfuse_trace_id"),
|
| }
|
| )
|
| summary = {
|
| "mean_pass_frac": mean_pass,
|
| "mean_reward_total": mean_reward,
|
| "n_perfect": n_perfect,
|
| "n_resolved": n_resolved,
|
| "n_problems": n,
|
| "step": step,
|
| "hf_revision": hf_revision,
|
| }
|
| root.update(output=summary)
|
| lf.score_current_trace(name="mean_pass_frac", value=mean_pass)
|
| lf.score_current_trace(name="mean_reward_total", value=mean_reward)
|
| lf.score_current_trace(name="n_perfect", value=float(n_perfect))
|
| lf.score_current_trace(name="n_resolved", value=float(n_resolved))
|
| lf.score_current_trace(name="checkpoint_step", value=float(step))
|
| lf.flush()
|
| _vlog(
|
| "langfuse.ckpt_eval",
|
| step=step,
|
| rev=hf_revision,
|
| mean_reward=mean_reward,
|
| resolved=n_resolved,
|
| n=n,
|
| trace_id=agg_trace_id,
|
| )
|
| return {"langfuse_trace_id": agg_trace_id, **summary}
|
|
|
|
|
| def log_completions_reward(completions: list[Any], **kwargs: Any) -> list[float]:
|
| """Side-channel: log raw model completions to W&B/console (reward always 0)."""
|
| previews = []
|
| for i, c in enumerate(completions):
|
| if isinstance(c, list):
|
|
|
| bits = []
|
| for m in c[-4:]:
|
| if isinstance(m, dict):
|
| role = m.get("role")
|
| content = str(m.get("content") or "")[:400]
|
| tcalls = m.get("tool_calls")
|
| bits.append(f"{role}: {content}" + (f" tool_calls={tcalls}" if tcalls else ""))
|
| else:
|
| bits.append(str(m)[:400])
|
| preview = " | ".join(bits)
|
| else:
|
| preview = str(c)[:800]
|
| previews.append(preview)
|
| _vlog("train.completion", i=i, preview=preview[:500], has_tool_calls=("tool_calls" in preview))
|
|
|
| try:
|
| import wandb
|
|
|
| if wandb.run is not None:
|
| n_with_tools = sum(1 for p in previews if "tool_calls" in p)
|
| table = wandb.Table(columns=["i", "preview", "has_tool_calls"])
|
| for i, p in enumerate(previews[:16]):
|
| table.add_data(i, p[:1500], "tool_calls" in p)
|
| wandb.log(
|
| {
|
| "train/completion_previews": table,
|
| "train/n_completions": len(previews),
|
| "train/n_completions_with_tool_calls_str": n_with_tools,
|
| }
|
| )
|
| except Exception as e:
|
| _vlog("train.completion.wandb_failed", error=str(e))
|
|
|
| return [0.0] * len(completions)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| def make_eval_before_save_callback(
|
| *,
|
| eval_problems: list[dict],
|
| train_model_id: str,
|
| hub_model_id: str,
|
| save_steps: int,
|
| n_eval: int,
|
| max_eval_turns: int,
|
| ):
|
| from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
|
|
|
| class EvalBeforeCheckpointCallback(TrainerCallback):
|
| def __init__(self) -> None:
|
| self._last_eval_step = -1
|
| self.trainer = None
|
|
|
| def on_step_end(
|
| self,
|
| args: TrainingArguments,
|
| state: TrainerState,
|
| control: TrainerControl,
|
| **kwargs,
|
| ):
|
| step = state.global_step
|
| if save_steps <= 0 or step <= 0 or step % save_steps != 0:
|
| return control
|
| if step == self._last_eval_step:
|
| return control
|
| self._last_eval_step = step
|
|
|
| model = kwargs.get("model")
|
| if model is None and self.trainer is not None:
|
| model = self.trainer.model
|
| tokenizer = (
|
| kwargs.get("processing_class")
|
| or kwargs.get("tokenizer")
|
| or globals().get("_TOKENIZER")
|
| )
|
| if tokenizer is None and self.trainer is not None:
|
| tokenizer = getattr(self.trainer, "processing_class", None) or getattr(
|
| self.trainer, "tokenizer", None
|
| )
|
| if model is None or tokenizer is None:
|
| raise RuntimeError("ckpt-eval requires model+tokenizer")
|
|
|
| rev = f"grpo-step-{step}"
|
| problems = eval_problems[:n_eval]
|
| _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
|
| if wandb.run is not None:
|
| wandb_run_id = wandb.run.id
|
| mean_r = sum(float(r["reward_total"]) for r in results) / max(len(results), 1)
|
| mean_tools = sum(float(r.get("n_tool_calls") or 0) for r in results) / max(
|
| len(results), 1
|
| )
|
| table = wandb.Table(
|
| columns=[
|
| "task_id",
|
| "reward_total",
|
| "pass_frac",
|
| "resolved",
|
| "resolution",
|
| "n_tool_calls",
|
| "tools_used",
|
| "langfuse_trace_id",
|
| "completion",
|
| ]
|
| )
|
| for r in results:
|
| table.add_data(
|
| str(r.get("task_id") or ""),
|
| float(r.get("reward_total") or 0),
|
| float(r.get("pass_frac") or 0),
|
| bool(r.get("resolved")),
|
| str(r.get("resolution") or ""),
|
| int(r.get("n_tool_calls") or 0),
|
| ",".join(r.get("tools_used") or []),
|
| str(r.get("langfuse_trace_id") or ""),
|
| str(r.get("completion") or "")[:1500],
|
| )
|
| wandb.log(
|
| {
|
| "eval/mean_reward_total": mean_r,
|
| "eval/mean_pass_frac": sum(
|
| float(r["pass_frac"]) for r in results
|
| )
|
| / max(len(results), 1),
|
| "eval/n_resolved": sum(1 for r in results if r.get("resolved")),
|
| "eval/mean_n_tool_calls": mean_tools,
|
| "eval/checkpoint_step": step,
|
| "eval/episodes": table,
|
| },
|
| step=step,
|
| )
|
|
|
| summary = langfuse_log_ckpt_eval(
|
| step=step,
|
| hf_revision=rev,
|
| train_model_id=train_model_id,
|
| results=results,
|
| wandb_run_id=wandb_run_id,
|
| )
|
| out = Path(args.output_dir) / f"ckpt_eval_step_{step}.json"
|
| out.write_text(
|
| json.dumps(summary | {"results": results}, indent=2), encoding="utf-8"
|
| )
|
|
|
| try:
|
| model.push_to_hub(
|
| hub_model_id,
|
| token=os.environ["HF_TOKEN"],
|
| revision=rev,
|
| commit_message=(
|
| f"Agent GRPO {rev}; mean_reward={summary['mean_reward_total']:.3f}"
|
| ),
|
| private=False,
|
| )
|
| print(f"[hf] pushed {hub_model_id}@{rev}")
|
| if wandb.run is not None:
|
| wandb.log(
|
| {
|
| "hf/revision": rev,
|
| "hf/mean_reward_at_push": summary["mean_reward_total"],
|
| "langfuse/ckpt_trace_id": summary.get("langfuse_trace_id"),
|
| },
|
| step=step,
|
| )
|
| except Exception as e:
|
| print(f"[hf] push failed (local ckpt still ok): {e}")
|
| traceback.print_exc()
|
| except Exception:
|
| print("[ckpt-eval] FAILED — re-raising")
|
| traceback.print_exc()
|
| raise
|
| return control
|
|
|
| return EvalBeforeCheckpointCallback()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| class Cfg:
|
| """Plain class (not dataclass) — safe under importlib.exec_module on Py3.13."""
|
|
|
| def __init__(
|
| self,
|
| dataset_id: str = "vanishingradient/deepagents-rl-swe-lite",
|
| hub_model_id: str = "vanishingradient/deepagents-rl-coding",
|
| train_split: str = "dev",
|
| eval_split: str = "test",
|
| max_steps: int = 20,
|
| save_steps: int = 5,
|
| num_generations: int = 2,
|
| learning_rate: float = 1e-6,
|
| max_completion_length: int = 1024,
|
| max_seq_length: int = 4096,
|
| max_tool_calling_iterations: int = 12,
|
| n_eval: int = 2,
|
| 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
|
| self.train_split = train_split
|
| self.eval_split = eval_split
|
| self.max_steps = max_steps
|
| self.save_steps = save_steps
|
| self.num_generations = num_generations
|
| self.learning_rate = learning_rate
|
| self.max_completion_length = max_completion_length
|
| self.max_seq_length = max_seq_length
|
| self.max_tool_calling_iterations = max_tool_calling_iterations
|
| self.n_eval = n_eval
|
| self.max_eval_turns = max_eval_turns
|
| self.max_train_instances = max_train_instances
|
| self.lora_r = lora_r
|
|
|
| self.mode = (mode or os.getenv("DARL_MODE") or "full").strip().lower()
|
|
|
|
|
| def _load_swe_splits(cfg: Cfg):
|
| from datasets import Dataset, load_dataset
|
|
|
| from deepagents_rl.train.swe_env import prepare_swe_grpo_rows
|
|
|
| ds = load_dataset(cfg.dataset_id, token=os.environ["HF_TOKEN"])
|
| train_split = cfg.train_split if cfg.train_split in ds else "train"
|
| eval_split = cfg.eval_split if cfg.eval_split in ds else (
|
| "test" if "test" in ds else train_split
|
| )
|
|
|
| train_recs = [dict(r) for r in ds[train_split]]
|
| eval_recs = [dict(r) for r in ds[eval_split]]
|
|
|
|
|
| def _ok(r: dict) -> bool:
|
| img = r.get("image") or r.get("docker_image") or ""
|
| return bool(img) and bool(r.get("eval_script"))
|
|
|
| train_recs = [r for r in train_recs if _ok(r)][: cfg.max_train_instances]
|
| eval_recs = [r for r in eval_recs if _ok(r)]
|
|
|
| train_rows = prepare_swe_grpo_rows(train_recs)
|
| eval_rows = prepare_swe_grpo_rows(eval_recs)
|
| if not train_rows:
|
| raise RuntimeError(f"No usable SWE train rows in {cfg.dataset_id}/{train_split}")
|
| if not eval_rows:
|
| eval_rows = train_rows[: max(1, cfg.n_eval)]
|
|
|
| print(
|
| f"[data] train={len(train_rows)} eval={len(eval_rows)} "
|
| f"splits={train_split}/{eval_split} dataset={cfg.dataset_id}"
|
| )
|
| 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)
|
|
|
|
|
| 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},
|
| )
|
|
|
|
|
| env = DeepAgentSWEEnv()
|
| try:
|
| reset_kwargs = {
|
| k: v
|
| for k, v in problem.items()
|
| if k
|
| not in (
|
| "prompt",
|
| "problem_statement",
|
| )
|
| }
|
|
|
| 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")
|
|
|
| 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",
|
| )
|
| }
|
|
|
| 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:
|
| 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
|
|
|
|
|
| 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),
|
| mode=getattr(cfg, "mode", "full"),
|
| )
|
| 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:
|
| _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
|
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
| from trl import GRPOConfig, GRPOTrainer
|
|
|
| from deepagents_rl.train.swe_env import DeepAgentSWEEnv
|
|
|
|
|
| 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']}"
|
| _vlog("main.grader_health_check", url=grader)
|
| health = requests.get(f"{grader}/health", headers=headers, timeout=30)
|
| health.raise_for_status()
|
| _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))
|
|
|
|
|
|
|
|
|
|
|
| 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()
|
| use_fp16 = False
|
| compute_dtype = torch.bfloat16 if use_bf16 else torch.float16
|
| _vlog(
|
| "main.dtype_choice",
|
| bf16=use_bf16,
|
| fp16=use_fp16,
|
| compute_dtype=str(compute_dtype),
|
| capability=(major, minor),
|
| device=torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
|
| )
|
| bnb = BitsAndBytesConfig(
|
| load_in_4bit=True,
|
| bnb_4bit_compute_dtype=compute_dtype,
|
| 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
|
| )
|
| if tokenizer.pad_token is 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,
|
| device_map="auto",
|
| 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,
|
| 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",
|
| ],
|
| )
|
| model = get_peft_model(model, lora)
|
|
|
| for name, param in model.named_parameters():
|
| if param.requires_grad and param.dtype != torch.float32:
|
| _vlog("main.cast_trainable_fp32", name=name, from_dtype=str(param.dtype))
|
| param.data = param.data.float()
|
|
|
| output_dir = "/kaggle/working/grpo-out"
|
| Path(output_dir).mkdir(parents=True, exist_ok=True)
|
| globals()["_TOKENIZER"] = tokenizer
|
|
|
| wandb.init(
|
| project=os.environ["WANDB_PROJECT"],
|
| name="kaggle-grpo-swe",
|
| job_type="train",
|
| config={
|
| "train_model_id": model_id,
|
| "dataset_id": cfg.dataset_id,
|
| "harness": "deepagents-swe",
|
| "max_steps": cfg.max_steps,
|
| "save_steps": cfg.save_steps,
|
| "n_eval": cfg.n_eval,
|
| "num_generations": cfg.num_generations,
|
| "max_tool_calling_iterations": cfg.max_tool_calling_iterations,
|
| "hub_model_id": cfg.hub_model_id,
|
| "swe_grader_url": grader,
|
| "eval_before_every_checkpoint": True,
|
| },
|
| )
|
|
|
| _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,
|
| eval_problems[: cfg.n_eval],
|
| max_turns=cfg.max_eval_turns,
|
| )
|
| base_summary = langfuse_log_ckpt_eval(
|
| step=0,
|
| hf_revision="grpo-step-0-baseline",
|
| train_model_id=model_id,
|
| results=baseline,
|
| wandb_run_id=wandb.run.id if wandb.run else None,
|
| )
|
| Path(output_dir, "ckpt_eval_step_0.json").write_text(
|
| json.dumps(base_summary | {"results": baseline}, indent=2), encoding="utf-8"
|
| )
|
| base_table = wandb.Table(
|
| columns=[
|
| "task_id",
|
| "reward_total",
|
| "pass_frac",
|
| "resolved",
|
| "resolution",
|
| "n_tool_calls",
|
| "tools_used",
|
| "langfuse_trace_id",
|
| "completion",
|
| ]
|
| )
|
| for r in baseline:
|
| base_table.add_data(
|
| str(r.get("task_id") or ""),
|
| float(r.get("reward_total") or 0),
|
| float(r.get("pass_frac") or 0),
|
| bool(r.get("resolved")),
|
| str(r.get("resolution") or ""),
|
| int(r.get("n_tool_calls") or 0),
|
| ",".join(r.get("tools_used") or []),
|
| str(r.get("langfuse_trace_id") or ""),
|
| str(r.get("completion") or "")[:1500],
|
| )
|
| wandb.log(
|
| {
|
| "eval/mean_pass_frac": base_summary["mean_pass_frac"],
|
| "eval/mean_reward_total": base_summary["mean_reward_total"],
|
| "eval/n_resolved": base_summary.get("n_resolved", 0),
|
| "eval/mean_n_tool_calls": sum(float(r.get("n_tool_calls") or 0) for r in baseline)
|
| / max(len(baseline), 1),
|
| "eval/checkpoint_step": 0,
|
| "eval/episodes": base_table,
|
| "langfuse/ckpt_trace_id": base_summary.get("langfuse_trace_id"),
|
| },
|
| step=0,
|
| )
|
|
|
|
|
| chat_kwargs: dict[str, Any] = {"tool_choice": "required"}
|
| gen_kwargs: dict[str, Any] = {
|
| "temperature": float(os.getenv("DARL_GEN_TEMPERATURE", "0.8")),
|
| "top_p": 0.95,
|
| }
|
| args = GRPOConfig(
|
| output_dir=output_dir,
|
| run_name="kaggle-grpo-swe",
|
| learning_rate=cfg.learning_rate,
|
| per_device_train_batch_size=1,
|
| gradient_accumulation_steps=4,
|
| num_generations=cfg.num_generations,
|
| max_completion_length=cfg.max_completion_length,
|
| max_tool_calling_iterations=cfg.max_tool_calling_iterations,
|
| max_steps=cfg.max_steps,
|
| logging_steps=1,
|
| save_steps=cfg.save_steps,
|
| save_total_limit=3,
|
| report_to=["wandb"],
|
| bf16=use_bf16,
|
| fp16=use_fp16,
|
| remove_unused_columns=False,
|
|
|
|
|
| hub_model_id=cfg.hub_model_id,
|
| hub_token=os.environ["HF_TOKEN"],
|
| push_to_hub=False,
|
| chat_template_kwargs=chat_kwargs,
|
| generation_kwargs=gen_kwargs,
|
| )
|
| _vlog(
|
| "main.grpo_config_tool_nudge",
|
| chat_template_kwargs=chat_kwargs,
|
| generation_kwargs=gen_kwargs,
|
| )
|
|
|
| cb = make_eval_before_save_callback(
|
| eval_problems=eval_problems,
|
| train_model_id=model_id,
|
| hub_model_id=cfg.hub_model_id,
|
| save_steps=cfg.save_steps,
|
| n_eval=cfg.n_eval,
|
| max_eval_turns=cfg.max_eval_turns,
|
| )
|
|
|
| TrainerCls = make_tool_repair_grpo_trainer(GRPOTrainer)
|
| _vlog(
|
| "main.trainer_init",
|
| environment="DeepAgentSWEEnv",
|
| trainer=TrainerCls.__name__,
|
| max_steps=cfg.max_steps,
|
| num_generations=cfg.num_generations,
|
| max_tool_calling_iterations=cfg.max_tool_calling_iterations,
|
| )
|
| trainer = TrainerCls(
|
| model=model,
|
| reward_funcs=[log_completions_reward],
|
| args=args,
|
| train_dataset=train,
|
| processing_class=tokenizer,
|
| callbacks=[cb],
|
| environment_factory=DeepAgentSWEEnv,
|
| )
|
| 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"
|
| print(f"=== FINAL AGENT EVAL step={final_step} ===")
|
| final_results = run_heldout_agent_eval(
|
| model,
|
| tokenizer,
|
| eval_problems[: cfg.n_eval],
|
| max_turns=cfg.max_eval_turns,
|
| )
|
| final_summary = langfuse_log_ckpt_eval(
|
| step=final_step,
|
| hf_revision=rev,
|
| train_model_id=model_id,
|
| results=final_results,
|
| wandb_run_id=wandb.run.id if wandb.run else None,
|
| )
|
| model.push_to_hub(
|
| cfg.hub_model_id,
|
| token=os.environ["HF_TOKEN"],
|
| revision=rev,
|
| commit_message=(
|
| f"Final agent GRPO {rev}; mean_reward={final_summary['mean_reward_total']:.3f}"
|
| ),
|
| )
|
| model.push_to_hub(
|
| cfg.hub_model_id,
|
| token=os.environ["HF_TOKEN"],
|
| commit_message=(
|
| f"Final agent GRPO main; mean_reward={final_summary['mean_reward_total']:.3f}"
|
| ),
|
| )
|
| wandb.log(
|
| {
|
| "eval/final_mean_reward_total": final_summary["mean_reward_total"],
|
| "hf/final_revision": rev,
|
| }
|
| )
|
| wandb.finish()
|
| print("TRAIN DONE", final_summary)
|
|
|
|
|
| _TOKENIZER = None
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
