"""
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] + "..."
# ---------------------------------------------------------------------------
# Bootstrap package (notebook may only download this file)
# ---------------------------------------------------------------------------
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"
# Prefer the freshly downloaded tree over any earlier import.
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 # noqa: F401
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"),
)
# ---------------------------------------------------------------------------
# Secrets
# ---------------------------------------------------------------------------
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
# Private dataset secrets.json (API-friendly path)
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"),
)
# ---------------------------------------------------------------------------
# Model resolve
# ---------------------------------------------------------------------------
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}")
# ---------------------------------------------------------------------------
# Lightweight HF agent rollout (checkpoint eval)
# ---------------------------------------------------------------------------
_TOOL_CALL_RE = re.compile(
r"\s*(\{.*?\})\s*",
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]]:
# Bound methods so held-out eval can call fn(**kwargs); TRL uses class functions itself.
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: , ```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))
# Whole-string JSON object / list fallback
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
# unwrap {"type":"function","function":{...}}
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:
# conversational: [[{role, content, tool_calls?}]]
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:
# Last resort so the GRPO group is not all "DONE"/no-tools.
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): # type: ignore[misc,valid-type]
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-first-tool: if the model refuses tools on early turns, inject ls(".")
# so the episode can continue and shaped rewards can differentiate.
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: # noqa: BLE001
result = f"TOOL_ERROR: {e}"
except Exception as e: # noqa: BLE001
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: # noqa: BLE001
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
# ---------------------------------------------------------------------------
# Langfuse checkpoint eval
# ---------------------------------------------------------------------------
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"))
# Clean root (no bare trace_context — that creates phantom orphan parents).
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):
# conversational messages
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: # noqa: BLE001
_vlog("train.completion.wandb_failed", error=str(e))
return [0.0] * len(completions)
# ---------------------------------------------------------------------------
# Callback: eval + langfuse THEN save
# ---------------------------------------------------------------------------
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()
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
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
# 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):
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]]
# Keep only rows with docker image + eval script
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)
# 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),
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: # 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
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import GRPOConfig, GRPOTrainer
from deepagents_rl.train.swe_env import DeepAgentSWEEnv
# Probe grader
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))
# T4 (sm75) cannot run TRL's bf16 chunked logprob kernels (need sm80+).
# Also: fp16 AMP GradScaler crashes if any grads stay bf16
# (`_amp_foreach_non_finite_check_and_unscale_cuda` not implemented for BFloat16).
# On T4: no AMP (fp16=False, bf16=False); keep 4-bit weights + fp16 compute dtype.
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 # avoid GradScaler + QLoRA dtype mismatch on T4
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)
# Keep trainable LoRA weights in fp32 so optimizer/grad-norm never sees bf16 under AMP-off.
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,
)
# Nudge Qwen chat template toward required tool calls when supported.
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,
# Do not set reward_weights here: env get_reward is appended after init,
# so a fixed [0,1] length mismatches and crashes GRPOTrainer.
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], # logs only; env get_reward still scores SWE
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()