From 54b7d866e127022968db58957ae8376ab6761871 Mon Sep 17 00:00:00 2001 From: Arya Prabhudesai Date: Wed, 5 Aug 2026 19:17:09 -0700 Subject: [PATCH 1/2] Download team dataset and extract and analyze successful cases --- .gitignore | 3 + scripts/distill/ANALYSIS.md | 58 ++++ scripts/distill/analyze_coverage.py | 389 ++++++++++++++++++++++++++ scripts/distill/download_team_coop.py | 91 ++++++ scripts/distill/extract_successful.py | 192 +++++++++++++ 5 files changed, 733 insertions(+) create mode 100644 scripts/distill/ANALYSIS.md create mode 100644 scripts/distill/analyze_coverage.py create mode 100644 scripts/distill/download_team_coop.py create mode 100644 scripts/distill/extract_successful.py diff --git a/.gitignore b/.gitignore index 343f9bc19..e9db3ea92 100644 --- a/.gitignore +++ b/.gitignore @@ -47,6 +47,9 @@ misc/ .cooperbench_cache/ .cache/ +# Distillation data (downloaded from HF, not committed) +data/ + # OS .DS_Store Thumbs.db diff --git a/scripts/distill/ANALYSIS.md b/scripts/distill/ANALYSIS.md new file mode 100644 index 000000000..f215a6d3b --- /dev/null +++ b/scripts/distill/ANALYSIS.md @@ -0,0 +1,58 @@ +# Team-Coop Dataset Analysis + +## Dataset +- Source: `CooperBench/team-coop` on HuggingFace +- Downloaded to: `data/team-coop/` (gitignored) +- Total trajectory pairs: 1,993 +- Successful pairs (correct=true, verified=true): 837 +- Training records (one per agent per pair): 1,674 → `data/successful.jsonl` + +## Runs in the dataset + +| Run | Model | Pass rate | Notes | +|---|---|---|---| +| `cmp-full-team` | gpt-5.5-hao | ~60% | Full team features including protocol | +| `cmp-full-team-noproto` | gpt-5.5-hao | ~60% | All features except protocol | +| `qwen35-cooperdata-team-noproto` | Qwen3.5-9B | ~6% | Small model baseline, almost all fail | +| `qwen35-cooperdata-team-noproto-forced` | Qwen3.5-9B | ~6% | Small model variant | +| `coop/` | Qwen3.5-9B | ~0% | Earlier Qwen runs, all failing | + +The useful teacher trajectories come entirely from the two `cmp-full-team` runs (gpt-5.5-hao). + +## Scenario Coverage + +Threshold for "covered": 20 examples. + +| Scenario | Count | Status | Notes | +|---|---|---|---| +| `solo_task_lifecycle` | 1,364 | **Covered** | Agent claims, works, marks done | +| `lead_creates_subtask` | 1,450 | **Covered** | Lead agent creates additional tasks mid-run | +| `parallel_independent` | 1,102 | **Covered** | Both agents work independently, no messaging | +| `cross_agent_dependency` | 718 | **Covered** | Lead waits for member before finishing | +| `blocked_task` | 118 | **Covered** | A task reaches status=blocked | +| `request_respond` | 10 | **GAP** | Only in protocol-enabled run; need ~10 more | +| `claim_after_list` | 0 | **GAP** | Detector may need fixing — shell cmds may not be stored as plain text in trajectory JSON | +| `wait_for_message` | 0 | **GAP** | MCP idle-wait never triggered in successful runs | + +## Gaps and next steps + +### `request_respond` (10 examples) +Only appears in `cmp-full-team` (protocol on). Close to threshold — a handful of synthetic +examples will cover it. Can also re-run analysis with `--no-require-verified` to pick up +unverified passes from this scenario. + +### `claim_after_list` (0 examples) +The detector searches for `coop-task list` as a string inside trajectory steps. Zero hits +likely means trajectory JSON stores tool calls in a structured format (not raw shell strings). +**Before generating synthetic examples**: inspect one trajectory file to confirm the format, +then fix the detector in `analyze_coverage.py`. + +### `wait_for_message` (0 examples) +The MCP long-poll tool was never called in any successful trajectory. These runs used Codex +which does have the MCP server registered, but agents never went idle enough to trigger it. +Synthetic examples are the only path to coverage here. + +## Recommended next step +Generate ~50 synthetic trajectories covering the three gaps using a teacher model against +a live Redis environment. Priority order: `wait_for_message` (pure gap), `claim_after_list` +(confirm detector first), `request_respond` (almost covered). diff --git a/scripts/distill/analyze_coverage.py b/scripts/distill/analyze_coverage.py new file mode 100644 index 000000000..d26156f99 --- /dev/null +++ b/scripts/distill/analyze_coverage.py @@ -0,0 +1,389 @@ +"""Analyze scenario coverage in the extracted successful trajectories. + +Reads the JSONL produced by extract_successful.py and classifies each +trajectory against the 8 canonical coordination scenarios we need for +distillation training. Prints a detailed report showing: + + - Per-scenario counts and example trajectory IDs + - Gap analysis: which scenarios are underrepresented or absent + - Per-run and per-role breakdowns + +Canonical scenarios +------------------- +1. solo_task_lifecycle - agent claims, works, and marks done with no + coordination (single task, single agent active) +2. parallel_independent - both agents claim separate pre-assigned tasks + and work independently with no messaging +3. lead_creates_subtask - lead agent creates an additional task mid-run + (task_log has a "create" event by an agent, not bench-runner) +4. request_respond - at least one request/respond pair in task_log + (kind == "request" or "response") +5. blocked_task - any task reaches status "blocked" +6. claim_after_list - agent calls list then claims (inferred from + task_log ordering: list events precede claim) +7. cross_agent_dependency - lead waits for member: lead's lead_task stays + in_progress while member's task completes first +8. wait_for_message - trajectory contains a wait_for_message MCP tool call + +Scenarios are not mutually exclusive — one trajectory can cover several. + +Usage: + uv run python scripts/distill/analyze_coverage.py + uv run python scripts/distill/analyze_coverage.py --src data/successful.jsonl + uv run python scripts/distill/analyze_coverage.py --src data/successful.jsonl --verbose +""" + +from __future__ import annotations + +import argparse +import json +import sys +from collections import Counter, defaultdict +from pathlib import Path + +DEFAULT_SRC = Path(__file__).resolve().parent.parent.parent / "data" / "successful.jsonl" + +SCENARIOS = [ + "solo_task_lifecycle", + "parallel_independent", + "lead_creates_subtask", + "request_respond", + "blocked_task", + "claim_after_list", + "cross_agent_dependency", + "wait_for_message", +] + +# Minimum examples we want per scenario before we consider it "covered" +COVERAGE_THRESHOLD = 20 + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument( + "--src", + type=Path, + default=DEFAULT_SRC, + help=f"Extracted JSONL from extract_successful.py (default: {DEFAULT_SRC})", + ) + p.add_argument( + "--verbose", "-v", + action="store_true", + help="Print up to 3 example trajectory IDs per scenario.", + ) + p.add_argument( + "--out", + type=Path, + default=None, + help="Optional path to write the report as JSON (for downstream use).", + ) + return p.parse_args() + + +# --------------------------------------------------------------------------- +# Scenario detectors — each takes one record dict and returns bool +# --------------------------------------------------------------------------- + +def _task_log(rec: dict) -> list[dict]: + return rec.get("task_log") or [] + + +def _trajectory(rec: dict) -> list: + return rec.get("trajectory") or [] + + +def _is_solo_task_lifecycle(rec: dict) -> bool: + """Agent has at least one full open→in_progress→done arc in the task_log.""" + log = _task_log(rec) + me = rec["agent_id"] + claimed = set() + updated_done = set() + for ev in log: + if ev.get("kind") == "claim" and ev.get("by") == me: + claimed.add(ev["task_id"]) + if ev.get("kind") == "update" and ev.get("by") == me and ev.get("status") == "done": + updated_done.add(ev["task_id"]) + return bool(claimed & updated_done) + + +def _is_parallel_independent(rec: dict) -> bool: + """Both agents claim different tasks; no request/respond events at all.""" + log = _task_log(rec) + has_messaging = any(ev.get("kind") in ("request", "response") for ev in log) + if has_messaging: + return False + agents_claiming: set[str] = set() + tasks_claimed: set[str] = set() + for ev in log: + if ev.get("kind") == "claim": + agents_claiming.add(ev.get("by", "")) + tasks_claimed.add(ev.get("task_id", "")) + # two different agents each claimed a different task + return len(agents_claiming) >= 2 and len(tasks_claimed) >= 2 + + +def _is_lead_creates_subtask(rec: dict) -> bool: + """An agent (not bench-runner) creates a task during the run.""" + log = _task_log(rec) + return any(ev.get("kind") == "create" and ev.get("by") not in ("bench-runner", None) for ev in log) + + +def _is_request_respond(rec: dict) -> bool: + """task_log contains at least one request or response event.""" + log = _task_log(rec) + return any(ev.get("kind") in ("request", "response") for ev in log) + + +def _is_blocked_task(rec: dict) -> bool: + """Any task reaches status 'blocked'.""" + log = _task_log(rec) + return any(ev.get("kind") == "update" and ev.get("status") == "blocked" for ev in log) + + +def _is_claim_after_list(rec: dict) -> bool: + """task_log has a list event before a claim by the same agent. + + The task_log itself doesn't record list calls (those are CLI-only), so + we infer from the trajectory: look for a shell command containing + 'coop-task list' followed by a claim in the task_log. + """ + traj = _trajectory(rec) + log = _task_log(rec) + me = rec["agent_id"] + + claim_ts = min( + (ev["ts"] for ev in log if ev.get("kind") == "claim" and ev.get("by") == me), + default=None, + ) + if claim_ts is None: + return False + + # Search trajectory steps for a coop-task list call that occurred + # before the first claim timestamp (heuristic: step index as proxy) + claim_step = None + for i, step in enumerate(traj): + if ev_ts := _step_ts(step): + if ev_ts >= claim_ts: + claim_step = i + break + + for i, step in enumerate(traj): + if claim_step is not None and i >= claim_step: + break + if _step_contains(step, "coop-task list") or _step_contains(step, "coop-task pending"): + return True + return False + + +def _is_cross_agent_dependency(rec: dict) -> bool: + """Lead's lead_task stays in_progress while member's task completes first.""" + if rec.get("role") != "lead": + return False + log = _task_log(rec) + # Find lead task id (created by bench-runner with lead_task metadata — but + # we don't have per-event metadata here, so proxy: title contains "Lead-only") + # Also works: lead claims their task and then member's done event appears + # before lead's done event. + lead_in_progress_ts = None + member_done_ts = None + lead_done_ts = None + + me = rec["agent_id"] + # identify other agent + all_agents = {ev.get("by") for ev in log if ev.get("by") not in (None, "bench-runner")} + other_agents = all_agents - {me} + + for ev in log: + if ev.get("kind") == "update" and ev.get("by") == me and ev.get("status") == "in_progress": + lead_in_progress_ts = ev.get("ts") + if ev.get("kind") == "update" and ev.get("by") in other_agents and ev.get("status") == "done": + member_done_ts = ev.get("ts") + if ev.get("kind") == "update" and ev.get("by") == me and ev.get("status") == "done": + lead_done_ts = ev.get("ts") + + if lead_in_progress_ts and member_done_ts and lead_done_ts: + # Lead was in_progress, member finished, then lead finished — classic dependency pattern + return lead_in_progress_ts < member_done_ts < lead_done_ts + return False + + +def _is_wait_for_message(rec: dict) -> bool: + """Trajectory contains a wait_for_message MCP tool call.""" + traj = _trajectory(rec) + for step in traj: + if _step_contains(step, "wait_for_message"): + return True + return False + + +# --------------------------------------------------------------------------- +# Trajectory step helpers +# --------------------------------------------------------------------------- + +def _step_contains(step, text: str) -> bool: + """Check whether any string field in a trajectory step contains text.""" + if isinstance(step, str): + return text in step + if isinstance(step, dict): + return any( + text in v + for v in step.values() + if isinstance(v, str) + ) or any(_step_contains(v, text) for v in step.values() if isinstance(v, (dict, list))) + if isinstance(step, list): + return any(_step_contains(item, text) for item in step) + return False + + +def _step_ts(step) -> float | None: + if isinstance(step, dict): + for key in ("ts", "timestamp", "time"): + if key in step: + try: + return float(step[key]) + except (TypeError, ValueError): + pass + return None + + +DETECTORS: dict[str, object] = { + "solo_task_lifecycle": _is_solo_task_lifecycle, + "parallel_independent": _is_parallel_independent, + "lead_creates_subtask": _is_lead_creates_subtask, + "request_respond": _is_request_respond, + "blocked_task": _is_blocked_task, + "claim_after_list": _is_claim_after_list, + "cross_agent_dependency": _is_cross_agent_dependency, + "wait_for_message": _is_wait_for_message, +} + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> int: + args = parse_args() + + if not args.src.exists(): + print(f"error: {args.src} not found.", file=sys.stderr) + print("Run scripts/distill/extract_successful.py first.", file=sys.stderr) + return 1 + + # Counters + scenario_hits: dict[str, list[str]] = defaultdict(list) # scenario → [traj_key, ...] + scenario_by_run: dict[str, Counter] = defaultdict(Counter) + scenario_by_role: dict[str, Counter] = defaultdict(Counter) + total = 0 + multi_scenario_counts: Counter = Counter() + + with args.src.open() as fh: + for line in fh: + line = line.strip() + if not line: + continue + rec = json.loads(line) + total += 1 + key = f"{rec.get('run')}/{rec.get('repo')}/{rec.get('task_id')}/{rec.get('features')}/{rec.get('agent_id')}" + run = rec.get("run", "unknown") + role = rec.get("role", "unknown") + + matched = [] + for scenario, detector in DETECTORS.items(): + if detector(rec): # type: ignore[operator] + scenario_hits[scenario].append(key) + scenario_by_run[scenario][run] += 1 + scenario_by_role[scenario][role] += 1 + matched.append(scenario) + multi_scenario_counts[len(matched)] += 1 + + # --- Report ----------------------------------------------------------- + sep = "─" * 72 + + print(sep) + print(f" Team-Coop Scenario Coverage Report") + print(f" Source: {args.src}") + print(f" Records: {total} (one per agent per successful pair)") + print(sep) + print() + + covered = [] + gaps = [] + + for scenario in SCENARIOS: + hits = scenario_hits[scenario] + n = len(hits) + status = "OK " if n >= COVERAGE_THRESHOLD else "GAP" + bar_len = min(40, n // max(1, total // 400)) + bar = "█" * bar_len + print(f" [{status}] {scenario:<28} {n:>5} examples {bar}") + if args.verbose and hits: + for ex in hits[:3]: + print(f" ↳ {ex}") + if n >= COVERAGE_THRESHOLD: + covered.append(scenario) + else: + gaps.append((scenario, n)) + + print() + print(sep) + print(f" Covered (≥{COVERAGE_THRESHOLD}): {len(covered)}/{len(SCENARIOS)}") + print() + + if gaps: + print(" GAPS — need synthetic scenario generation:") + for scenario, n in gaps: + needed = COVERAGE_THRESHOLD - n + print(f" {scenario:<28} {n} found → need ~{needed} more synthetic examples") + else: + print(" All scenarios covered — no synthetic generation needed.") + + print() + print(" Scenario overlap (how many scenarios one trajectory covers):") + for n_scenarios in sorted(multi_scenario_counts): + print(f" {n_scenarios} scenarios: {multi_scenario_counts[n_scenarios]} trajectories") + + print() + print(" Per-run breakdown:") + for scenario in SCENARIOS: + if scenario_by_run[scenario]: + breakdown = " ".join(f"{r}:{c}" for r, c in sorted(scenario_by_run[scenario].items())) + print(f" {scenario:<28} {breakdown}") + + print() + print(" Per-role breakdown:") + for scenario in SCENARIOS: + if scenario_by_role[scenario]: + breakdown = " ".join(f"{r}:{c}" for r, c in sorted(scenario_by_role[scenario].items())) + print(f" {scenario:<28} {breakdown}") + + print(sep) + + if args.out: + report = { + "total_records": total, + "coverage_threshold": COVERAGE_THRESHOLD, + "scenarios": { + s: { + "count": len(scenario_hits[s]), + "covered": len(scenario_hits[s]) >= COVERAGE_THRESHOLD, + "by_run": dict(scenario_by_run[s]), + "by_role": dict(scenario_by_role[s]), + "examples": scenario_hits[s][:10], + } + for s in SCENARIOS + }, + "gaps": [ + {"scenario": s, "count": n, "needed": COVERAGE_THRESHOLD - n} + for s, n in gaps + ], + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2)) + print(f"\nReport written to {args.out}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/distill/download_team_coop.py b/scripts/distill/download_team_coop.py new file mode 100644 index 000000000..3b52506b0 --- /dev/null +++ b/scripts/distill/download_team_coop.py @@ -0,0 +1,91 @@ +"""Download the CooperBench/team-coop dataset from HuggingFace to data/team-coop/. + +Usage: + uv run python scripts/distill/download_team_coop.py + uv run python scripts/distill/download_team_coop.py --run cmp-full-team + uv run python scripts/distill/download_team_coop.py --dest data/my-dir + +Auth: + Public dataset — no HF_TOKEN required. Set HF_TOKEN env var for higher + rate limits or if the repo is made private later. + +Output layout (mirrors the HF repo structure): + data/team-coop/ + cmp-full-team/ + coop//// + agent1_traj.json agent2_traj.json + conversation.json eval.json metadata.json result.json + agent1.patch agent2.patch + cmp-full-team-noproto/ ... + coop/ ... (Qwen baseline runs) + summary.json + config.json +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +from huggingface_hub import snapshot_download + +REPO_ID = "CooperBench/team-coop" +DEFAULT_DEST = Path(__file__).resolve().parent.parent.parent / "data" / "team-coop" + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument( + "--dest", + type=Path, + default=DEFAULT_DEST, + help=f"Local directory to download into (default: {DEFAULT_DEST})", + ) + p.add_argument( + "--run", + default=None, + metavar="RUN_NAME", + help="Download only a single top-level run directory, e.g. 'cmp-full-team'.", + ) + p.add_argument( + "--ignore-patterns", + nargs="*", + default=["*.patch"], + metavar="PATTERN", + help="Glob patterns to skip (default: ['*.patch'] — saves ~half the disk space). Pass '' to download everything.", + ) + return p.parse_args() + + +def main() -> int: + args = parse_args() + dest: Path = args.dest + dest.mkdir(parents=True, exist_ok=True) + + token = os.environ.get("HF_TOKEN") + + allow_patterns = [f"{args.run}/**"] if args.run else None + ignore_patterns = [p for p in (args.ignore_patterns or []) if p] or None + + print(f"repo: {REPO_ID}") + print(f"dest: {dest}") + print(f"allow_patterns: {allow_patterns or '(all)'}") + print(f"ignore_patterns: {ignore_patterns or '(none)'}") + print() + + local_dir = snapshot_download( + repo_id=REPO_ID, + repo_type="dataset", + local_dir=str(dest), + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + token=token, + ) + print(f"\nDownloaded to: {local_dir}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/distill/extract_successful.py b/scripts/distill/extract_successful.py new file mode 100644 index 000000000..718745aee --- /dev/null +++ b/scripts/distill/extract_successful.py @@ -0,0 +1,192 @@ +"""Extract successful trajectories from a downloaded team-coop dataset. + +A trajectory pair (agent1 + agent2) is "successful" when eval.json has +``correct == true`` AND ``verified == true``. We write one JSON record per +agent per successful pair to a JSONL output file, keeping: + + - the full task_log from metadata.json (tool call sequence) + - the conversation from conversation.json + - the agent's own trajectory from agentN_traj.json + - key scalars from metadata.json / eval.json / result.json + +Usage: + uv run python scripts/distill/extract_successful.py + uv run python scripts/distill/extract_successful.py --src data/team-coop/cmp-full-team + uv run python scripts/distill/extract_successful.py --src data/team-coop --out data/successful.jsonl + +Output JSONL schema (one object per line, one line per agent per pair): + { + "run": str, # top-level run dir name, e.g. "cmp-full-team" + "repo": str, # e.g. "dspy_task" + "task_id": int, + "features": [int, int], + "agent_id": str, # "agent1" or "agent2" + "role": str, # "lead" or "member" + "model": str, + "agent_framework": str, + "team_features": dict, + "tasks": list, # final task objects from metadata.json + "task_log": list, # full coop-task event log + "conversation": list, # inter-agent messages + "trajectory": list, # raw agent trajectory steps + "metrics": dict, + "duration_seconds": float, + "score": float + } +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +DEFAULT_SRC = Path(__file__).resolve().parent.parent.parent / "data" / "team-coop" +DEFAULT_OUT = Path(__file__).resolve().parent.parent.parent / "data" / "successful.jsonl" + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument( + "--src", + type=Path, + default=DEFAULT_SRC, + help=f"Root of the downloaded dataset (default: {DEFAULT_SRC}). " + "Can point to the top-level dir (scans all runs) or a single run dir.", + ) + p.add_argument( + "--out", + type=Path, + default=DEFAULT_OUT, + help=f"Output JSONL file (default: {DEFAULT_OUT})", + ) + p.add_argument( + "--run", + default=None, + metavar="RUN_NAME", + help="Restrict to a single run, e.g. 'cmp-full-team'.", + ) + p.add_argument( + "--require-verified", + action=argparse.BooleanOptionalAction, + default=True, + help="Require eval.json verified==true (default: on). Use --no-require-verified to keep unverified passes.", + ) + return p.parse_args() + + +def _load_json(path: Path) -> dict | list | None: + try: + return json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return None + + +def _find_pair_dirs(src: Path, run_filter: str | None) -> list[tuple[str, Path]]: + """Return (run_name, pair_dir) for every f_f leaf directory.""" + pairs: list[tuple[str, Path]] = [] + + def _scan_run(run_name: str, run_root: Path) -> None: + coop_root = run_root / "coop" + if not coop_root.is_dir(): + return + for pair_dir in coop_root.rglob("f*_f*"): + if pair_dir.is_dir() and (pair_dir / "eval.json").exists(): + pairs.append((run_name, pair_dir)) + + # src might be the top-level (contains multiple run dirs) or a single run dir. + # Scan children that look like run dirs (have both coop/ and summary.json). + # If --run is given, also accept src itself as the run dir. + if run_filter and src.name == run_filter and (src / "summary.json").exists() and (src / "coop").is_dir(): + _scan_run(src.name, src) + else: + for child in sorted(src.iterdir()): + if not child.is_dir() or child.name.startswith("."): + continue + if not (child / "summary.json").exists() or not (child / "coop").is_dir(): + continue + if run_filter and child.name != run_filter: + continue + _scan_run(child.name, child) + + return pairs + + +def _extract_pair(run_name: str, pair_dir: Path, require_verified: bool) -> list[dict] | None: + eval_data = _load_json(pair_dir / "eval.json") + if not eval_data or not isinstance(eval_data, dict): + return None + if not eval_data.get("correct"): + return None + if require_verified and not eval_data.get("verified"): + return None + + meta = _load_json(pair_dir / "metadata.json") or {} + result = _load_json(pair_dir / "result.json") or {} + conversation = _load_json(pair_dir / "conversation.json") or [] + + lead_agent = meta.get("lead_agent", "agent1") + + records = [] + for agent_id in ("agent1", "agent2"): + traj_file = pair_dir / f"{agent_id}_traj.json" + trajectory = _load_json(traj_file) or [] + + records.append({ + "run": run_name, + "repo": meta.get("repo", pair_dir.parts[-3]), + "task_id": meta.get("task_id"), + "features": meta.get("features", meta.get("source_features")), + "agent_id": agent_id, + "role": "lead" if agent_id == lead_agent else "member", + "model": meta.get("model") or result.get("model"), + "agent_framework": meta.get("agent_framework") or result.get("agent_framework"), + "team_features": meta.get("team_features", {}), + "tasks": meta.get("tasks", []), + "task_log": meta.get("task_log", []), + "conversation": conversation, + "trajectory": trajectory, + "metrics": meta.get("metrics", {}), + "duration_seconds": meta.get("duration_seconds"), + "score": eval_data.get("score", 1.0), + }) + return records + + +def main() -> int: + args = parse_args() + + if not args.src.exists(): + print(f"error: source path does not exist: {args.src}", file=sys.stderr) + print("Run scripts/distill/download_team_coop.py first.", file=sys.stderr) + return 1 + + pairs = _find_pair_dirs(args.src, args.run) + print(f"Found {len(pairs)} trajectory pairs under {args.src}") + + args.out.parent.mkdir(parents=True, exist_ok=True) + + total_pairs = 0 + total_records = 0 + skipped = 0 + + with args.out.open("w") as fh: + for run_name, pair_dir in pairs: + records = _extract_pair(pair_dir=pair_dir, run_name=run_name, require_verified=args.require_verified) + if records is None: + skipped += 1 + continue + for rec in records: + fh.write(json.dumps(rec) + "\n") + total_pairs += 1 + total_records += len(records) + + print(f"Successful pairs: {total_pairs}") + print(f"Skipped (failed): {skipped}") + print(f"Records written: {total_records} → {args.out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 25d9c57c29bd3c4d4b21a98526e67661487ecfff Mon Sep 17 00:00:00 2001 From: Arya Prabhudesai Date: Tue, 25 Aug 2026 10:05:20 -0700 Subject: [PATCH 2/2] Tool usage distillation --- scripts/distill/agent_config_qwen35_sft.yaml | 5 + scripts/distill/convert_codex_to_miniswe.py | 509 +++++++++++++++++ scripts/distill/serve_vllm_modal.py | 112 ++++ scripts/distill/train_modal.py | 548 +++++++++++++++++++ scripts/distill/upload_adapter.py | 174 ++++++ 5 files changed, 1348 insertions(+) create mode 100644 scripts/distill/agent_config_qwen35_sft.yaml create mode 100644 scripts/distill/convert_codex_to_miniswe.py create mode 100644 scripts/distill/serve_vllm_modal.py create mode 100644 scripts/distill/train_modal.py create mode 100644 scripts/distill/upload_adapter.py diff --git a/scripts/distill/agent_config_qwen35_sft.yaml b/scripts/distill/agent_config_qwen35_sft.yaml new file mode 100644 index 000000000..73531324e --- /dev/null +++ b/scripts/distill/agent_config_qwen35_sft.yaml @@ -0,0 +1,5 @@ +config: + model: + model_kwargs: + api_base: "https://cooperbench--serve-qwen35-sft-serve-dev.modal.run/v1" + api_key: "dummy" diff --git a/scripts/distill/convert_codex_to_miniswe.py b/scripts/distill/convert_codex_to_miniswe.py new file mode 100644 index 000000000..35e2e505f --- /dev/null +++ b/scripts/distill/convert_codex_to_miniswe.py @@ -0,0 +1,509 @@ +"""Convert Codex ``codex_stream_log`` trajectories to ``mini_swe_agent_v2`` format. + +The teacher trajectories in the team-coop dataset (gpt-5.5-hao via Codex) embed +shell commands and their outputs inline in ``assistant`` message content: + + [command] /bin/bash -lc 'some-command' + + [exit N] ← non-zero exits only; success has no marker + +This script converts those to the structured format Qwen uses at inference time +(role: assistant with tool_calls + role: tool messages), so the teacher data can +be used for SFT without a format mismatch. + +Usage: + uv run python scripts/distill/convert_codex_to_miniswe.py + uv run python scripts/distill/convert_codex_to_miniswe.py --src data/team-coop/cmp-full-team + uv run python scripts/distill/convert_codex_to_miniswe.py --out data/converted_teacher.jsonl + uv run python scripts/distill/convert_codex_to_miniswe.py --stats # dry-run, stats only + +Output JSONL schema: same as successful.jsonl (extract_successful.py) but with +``trajectory_format: "mini_swe_agent_v2"`` and ``messages`` replaced with converted messages. +""" + +from __future__ import annotations + +import argparse +import json +import re +import sys +import uuid +from collections import Counter +from dataclasses import dataclass, field +from pathlib import Path + +DEFAULT_SRC = Path(__file__).resolve().parent.parent.parent / "data" / "team-coop" +DEFAULT_OUT = Path(__file__).resolve().parent.parent.parent / "data" / "converted_teacher.jsonl" + +# The system prompt used by mini_swe_agent_v2 runs +SYSTEM_PROMPT = ( + "You are a software engineer working alongside a colleague on a shared codebase. " + "You each have your own workspace and are implementing different features in parallel. " + "You communicate naturally — like engineers on the same team — to make sure " + "your combined work integrates cleanly." +) + +# Regexes +_CMD_PREFIX = re.compile(r"^\[command\] /bin/bash -l?c\s+") +_TIMING_LINE = re.compile(r"^[ \t]*(?:succeeded|exited\s+\S+)\s+in\s+\d+ms:\s*$") +_EXIT_LINE = re.compile(r"^\[exit\s+(-?\d+)\]\s*$") +_DIFF_START = re.compile(r"^diff --git ") + + +# --------------------------------------------------------------------------- +# Data structures +# --------------------------------------------------------------------------- + +@dataclass +class ParsedTurn: + reasoning: str = "" # text before [command], if any + has_command: bool = False + command: str = "" # the shell command string (after bash -lc) + output: str = "" # stdout/stderr (diff and timing stripped) + returncode: int = 0 + warnings: list[str] = field(default_factory=list) + + +@dataclass +class ConversionStats: + total: int = 0 + skipped_empty: int = 0 + converted: int = 0 + warnings: int = 0 + turns_reasoning: int = 0 + turns_command: int = 0 + coop_commands: Counter = field(default_factory=Counter) + + def report(self) -> str: + lines = [ + f"Total trajectories: {self.total}", + f"Skipped (≤1 msg): {self.skipped_empty}", + f"Converted: {self.converted}", + f"Parse warnings: {self.warnings}", + f"Reasoning turns: {self.turns_reasoning}", + f"Command turns: {self.turns_command}", + ] + if self.coop_commands: + lines.append("Top coop-task commands:") + for cmd, n in self.coop_commands.most_common(10): + lines.append(f" {cmd}: {n}") + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Codex message parser +# --------------------------------------------------------------------------- + +def _extract_shell_command(after_prefix: str) -> tuple[str, str]: + """Extract the shell command and remaining content from text after '/bin/bash -lc '. + + Returns (command_string, rest_of_content) where rest_of_content starts on + the line after the command line. + """ + text = after_prefix.lstrip() + # The command is quoted with either ' or " or unquoted (rare) + if text.startswith("'"): + # Single-quoted — find the closing ' that isn't escaped + # Codex uses -lc '...' where the command may contain escaped quotes + end = 1 + while end < len(text): + if text[end] == "'" and text[end - 1] != "\\": + break + end += 1 + cmd = text[1:end] + rest = text[end + 1:] + elif text.startswith('"'): + # Double-quoted + end = 1 + while end < len(text): + if text[end] == '"' and text[end - 1] != "\\": + break + end += 1 + cmd = text[1:end] + # Unescape inner \" → " + cmd = cmd.replace('\\"', '"') + rest = text[end + 1:] + else: + # Unquoted — take until end of line + nl = text.find("\n") + if nl == -1: + cmd, rest = text, "" + else: + cmd, rest = text[:nl], text[nl:] + + # rest should start with a newline; strip it + if rest.startswith("\n"): + rest = rest[1:] + return cmd, rest + + +def _strip_output(raw_output: str) -> tuple[str, int, list[str]]: + """Strip timing lines, trailing diff blocks, and [exit N] from raw output. + + Returns (cleaned_output, returncode, warnings). + """ + warnings: list[str] = [] + lines = raw_output.split("\n") + returncode = 0 + + # Strip leading timing line if present + start = 0 + if lines and _TIMING_LINE.match(lines[0]): + # Extract exit code from timing line if exited N + m = re.match(r"^ *exited\s+(-?\d+)\s+in", lines[0]) + if m: + returncode = int(m.group(1)) + start = 1 + + # Walk backwards: strip [exit N] and trailing diff block + end = len(lines) + + # Check for [exit N] at the end (possibly after a diff block) + # Strategy: scan from the end for [exit N], then find where diff starts + exit_idx = None + for i in range(end - 1, start - 1, -1): + line = lines[i] + m = _EXIT_LINE.match(line) + if m: + if returncode == 0: # timing line takes precedence if both present + returncode = int(m.group(1)) + exit_idx = i + end = i # don't include the [exit N] line + break + + # Strip trailing diff block (workspace context injected by Codex) + # Find the last `diff --git` line before end and cut there + # But only strip it if it's after real output (not if the entire output is a diff) + diff_start_idx = None + for i in range(end - 1, start - 1, -1): + if _DIFF_START.match(lines[i]): + diff_start_idx = i + break + + if diff_start_idx is not None: + # Heuristic: if output before diff_start is non-empty, strip the diff + pre_diff = "\n".join(lines[start:diff_start_idx]).strip() + if pre_diff: + end = diff_start_idx + else: + # The entire output IS the diff — keep it (rare, but happens) + pass + + output = "\n".join(lines[start:end]) + # Normalise: strip trailing blank lines + output = output.rstrip("\n") + return output, returncode, warnings + + +def parse_codex_message(content: str) -> ParsedTurn: + """Parse a single Codex assistant message content string into a ParsedTurn.""" + turn = ParsedTurn() + + # Find the [command] prefix + m = _CMD_PREFIX.search(content) + if m is None: + # Pure reasoning — no command + turn.reasoning = content + return turn + + turn.has_command = True + + # Text before [command] = reasoning + reasoning_end = content.rfind("\n", 0, m.start()) + if reasoning_end == -1: + turn.reasoning = content[: m.start()].strip() + else: + turn.reasoning = content[:reasoning_end].strip() + + # Extract command string and rest (output) + after_prefix = content[m.end():] + cmd, raw_output = _extract_shell_command(after_prefix) + turn.command = cmd + + # Parse output + output, returncode, warnings = _strip_output(raw_output) + turn.output = output + turn.returncode = returncode + turn.warnings = warnings + + return turn + + +# --------------------------------------------------------------------------- +# mini_swe_agent_v2 message builders +# --------------------------------------------------------------------------- + +def _new_tool_call_id() -> str: + return f"chatcmpl-tool-{uuid.uuid4().hex[:16]}" + + +def _assistant_reasoning(text: str) -> dict: + return { + "role": "assistant", + "content": text, + "tool_calls": None, + "function_call": None, + "provider_specific_fields": {"refusal": None, "reasoning": None}, + "extra": None, + } + + +def _assistant_command(reasoning: str, command: str, tc_id: str) -> dict: + return { + "role": "assistant", + "content": reasoning, + "tool_calls": [ + { + "id": tc_id, + "type": "function", + "function": { + "name": "bash", + "arguments": json.dumps({"command": command}), + }, + } + ], + "function_call": None, + "provider_specific_fields": {"refusal": None, "reasoning": None}, + "extra": { + "actions": [{"tool_name": "bash", "tool_call_id": tc_id, "command": command}], + "response": None, + "cost": 0.0, + "timestamp": None, + }, + } + + +def _tool_result(tc_id: str, output: str, returncode: int) -> dict: + content = json.dumps({"returncode": returncode, "output": output}) + return { + "role": "tool", + "tool_call_id": tc_id, + "content": content, + "extra": { + "raw_output": output, + "returncode": returncode, + "timestamp": None, + "exception_info": "", + }, + } + + +def _exit_message(status: str = "Submitted") -> dict: + return { + "role": "exit", + "content": "", + "extra": {"exit_status": status, "submission": ""}, + } + + +# --------------------------------------------------------------------------- +# Full trajectory conversion +# --------------------------------------------------------------------------- + +def convert_trajectory( + codex_messages: list[dict], + agent_status: str, + stats: ConversionStats, +) -> list[dict]: + """Convert a list of Codex messages to mini_swe_agent_v2 format.""" + out: list[dict] = [] + + # System message + out.append({"role": "system", "content": SYSTEM_PROMPT}) + + for msg in codex_messages: + role = msg.get("role", "") + content = msg.get("content", "") or "" + + if role == "user": + # Pass through verbatim (task description) + out.append({"role": "user", "content": content}) + continue + + if role != "assistant": + continue + + turn = parse_codex_message(content) + + if turn.warnings: + stats.warnings += len(turn.warnings) + + if not turn.has_command: + # Pure reasoning turn + if turn.reasoning: + stats.turns_reasoning += 1 + out.append(_assistant_reasoning(turn.reasoning)) + else: + stats.turns_command += 1 + tc_id = _new_tool_call_id() + out.append(_assistant_command(turn.reasoning, turn.command, tc_id)) + out.append(_tool_result(tc_id, turn.output, turn.returncode)) + + # Track coop-task command usage + cmd_stripped = turn.command.strip() + m = re.match(r"(coop-task[a-z-]*)", cmd_stripped) + if m: + stats.coop_commands[m.group(1)] += 1 + + # Exit message + out.append(_exit_message(agent_status)) + return out + + +# --------------------------------------------------------------------------- +# Dataset traversal +# --------------------------------------------------------------------------- + +def _load_json(path: Path) -> dict | list | None: + try: + return json.loads(path.read_text()) + except (OSError, json.JSONDecodeError): + return None + + +def _iter_successful_pairs(src: Path, run_filter: str | None) -> list[tuple[str, Path]]: + """Yield (run_name, pair_dir) for successful pairs under src.""" + pairs: list[tuple[str, Path]] = [] + + def _scan_run(run_name: str, run_root: Path) -> None: + coop_root = run_root / "coop" + if not coop_root.is_dir(): + return + for pair_dir in coop_root.rglob("f*_f*"): + if not pair_dir.is_dir(): + continue + eval_data = _load_json(pair_dir / "eval.json") + if not eval_data or not isinstance(eval_data, dict): + continue + if eval_data.get("correct") and eval_data.get("verified"): + pairs.append((run_name, pair_dir)) + + if run_filter and src.name == run_filter and (src / "summary.json").exists() and (src / "coop").is_dir(): + _scan_run(src.name, src) + else: + for child in sorted(src.iterdir()): + if not child.is_dir() or child.name.startswith("."): + continue + if not (child / "summary.json").exists() or not (child / "coop").is_dir(): + continue + if run_filter and child.name != run_filter: + continue + _scan_run(child.name, child) + + return pairs + + +def _meta_from_pair(run_name: str, pair_dir: Path) -> dict: + """Extract metadata fields for the output record.""" + meta = _load_json(pair_dir / "metadata.json") or {} + eval_data = _load_json(pair_dir / "eval.json") or {} + return { + "run": run_name, + "repo": meta.get("repo", pair_dir.parts[-3]), + "task_id": meta.get("task_id"), + "features": meta.get("features", meta.get("source_features")), + "model": meta.get("model"), + "agent_framework": meta.get("agent_framework"), + "team_features": meta.get("team_features", {}), + "tasks": meta.get("tasks", []), + "task_log": meta.get("task_log", []), + "metrics": meta.get("metrics", {}), + "lead_agent": meta.get("lead_agent", "agent1"), + "duration_seconds": meta.get("duration_seconds"), + "score": eval_data.get("score", 1.0), + "trajectory_format": "mini_swe_agent_v2", + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--src", type=Path, default=DEFAULT_SRC, + help=f"Root of downloaded dataset (default: {DEFAULT_SRC})") + p.add_argument("--out", type=Path, default=DEFAULT_OUT, + help=f"Output JSONL (default: {DEFAULT_OUT})") + p.add_argument("--run", default=None, metavar="RUN_NAME", + help="Restrict to a single run, e.g. 'cmp-full-team'") + p.add_argument("--stats", action="store_true", + help="Dry-run: print conversion stats without writing output") + return p.parse_args() + + +def main() -> int: + args = parse_args() + + if not args.src.exists(): + print(f"error: {args.src} not found", file=sys.stderr) + print("Run scripts/distill/download_team_coop.py first.", file=sys.stderr) + return 1 + + # Only convert Codex (teacher) runs — skip Qwen baseline runs + codex_runs = {"cmp-full-team", "cmp-full-team-noproto"} + + pairs = _iter_successful_pairs(args.src, args.run) + # Filter to Codex runs only + if args.run is None: + pairs = [(r, p) for r, p in pairs if r in codex_runs] + print(f"Found {len(pairs)} successful pairs to convert") + + stats = ConversionStats() + + if not args.stats: + args.out.parent.mkdir(parents=True, exist_ok=True) + out_fh = args.out.open("w") + else: + out_fh = None + + try: + for run_name, pair_dir in pairs: + meta = _meta_from_pair(run_name, pair_dir) + lead_agent = meta["lead_agent"] + + for agent_id in ("agent1", "agent2"): + traj_path = pair_dir / f"{agent_id}_traj.json" + traj_data = _load_json(traj_path) + if not traj_data or not isinstance(traj_data, dict): + continue + + codex_msgs = traj_data.get("messages", []) + stats.total += 1 + + # Skip empty or single-message (summary-only) trajectories + if len(codex_msgs) <= 1: + stats.skipped_empty += 1 + continue + + agent_status = traj_data.get("status", "Submitted") + converted = convert_trajectory(codex_msgs, agent_status, stats) + stats.converted += 1 + + role = "lead" if agent_id == lead_agent else "member" + record = { + **meta, + "agent_id": agent_id, + "role": role, + "messages": converted, + } + + if out_fh is not None: + out_fh.write(json.dumps(record) + "\n") + + finally: + if out_fh is not None: + out_fh.close() + + sep = "─" * 60 + print(sep) + print(stats.report()) + print(sep) + if not args.stats and out_fh is not None: + print(f"Written to: {args.out}") + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/distill/serve_vllm_modal.py b/scripts/distill/serve_vllm_modal.py new file mode 100644 index 000000000..45af3fec7 --- /dev/null +++ b/scripts/distill/serve_vllm_modal.py @@ -0,0 +1,112 @@ +"""vLLM server on Modal A100 serving the fine-tuned Qwen3.5-9B LoRA adapter. + +Exposes an OpenAI-compatible /v1/chat/completions endpoint so CooperBench +can evaluate the fine-tuned model. + +Usage: + # Start the server (keeps running until Ctrl+C) + uv run modal serve scripts/distill/serve_vllm_modal.py + + # Then run CooperBench pointing at the printed URL, e.g.: + uv run cooperbench run \\ + --base-url https:// \\ + --auth-token dummy \\ + -m qwen35-sft \\ + -a claude_code \\ + --setting team \\ + -s flash_10 \\ + -c 2 + + # Deploy persistently (survives terminal close, billed while up) + uv run modal deploy scripts/distill/serve_vllm_modal.py + uv run modal app stop serve-qwen35-sft # to shut it down +""" + +from __future__ import annotations + +import modal + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + +BASE_MODEL = "Qwen/Qwen3.5-9B" +LORA_REPO = "CooperBench/qwen3.5-9b-tool-use-sft" +MODEL_ALIAS = "qwen35-sft" +APP_NAME = "serve-qwen35-sft" +GPU = "A100-80GB" +MAX_MODEL_LEN = 32768 +VLLM_PORT = 8000 +VLLM_VERSION = "0.19.0" +TRANSFORMERS_VERSION = "5.5.4" +MINUTES = 60 + +# --------------------------------------------------------------------------- +# Volumes — shared HF / vLLM weight cache (reused across runs) +# --------------------------------------------------------------------------- + +_hf_cache_vol = modal.Volume.from_name("huggingface-cache", create_if_missing=True) +_vllm_cache_vol = modal.Volume.from_name("vllm-cache", create_if_missing=True) + +# --------------------------------------------------------------------------- +# Image +# --------------------------------------------------------------------------- + +_image = ( + modal.Image.from_registry( + "nvidia/cuda:12.8.0-devel-ubuntu22.04", + add_python="3.12", + ) + .entrypoint([]) + .apt_install("git") + .uv_pip_install(f"vllm=={VLLM_VERSION}") + .uv_pip_install( + f"transformers=={TRANSFORMERS_VERSION}", + "huggingface-hub[hf_xet]>=0.36.0", + "peft>=0.14.0", + ) + .env({"HF_XET_HIGH_PERFORMANCE": "1"}) +) + +# --------------------------------------------------------------------------- +# Modal app +# --------------------------------------------------------------------------- + +app = modal.App(name=APP_NAME) + + +@app.function( + image=_image, + gpu=GPU, + timeout=4 * MINUTES * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], + scaledown_window=15 * MINUTES, + volumes={ + "/root/.cache/huggingface": _hf_cache_vol, + "/root/.cache/vllm": _vllm_cache_vol, + }, +) +@modal.web_server(port=VLLM_PORT, startup_timeout=20 * MINUTES) +def serve(): + import subprocess + + cmd = [ + "vllm", "serve", BASE_MODEL, + "--host", "0.0.0.0", + "--port", str(VLLM_PORT), + "--max-model-len", str(MAX_MODEL_LEN), + "--dtype", "bfloat16", + "--gpu-memory-utilization", "0.92", + "--enable-prefix-caching", + "--enable-chunked-prefill", + "--enable-auto-tool-choice", + "--tool-call-parser", "qwen3_coder", + "--enable-lora", + "--max-lora-rank", "64", + "--lora-modules", f"{MODEL_ALIAS}={LORA_REPO}", + "--trust-remote-code", + "--enforce-eager", + ] + + print(f"Starting vLLM: {' '.join(cmd)}") + subprocess.Popen(cmd) diff --git a/scripts/distill/train_modal.py b/scripts/distill/train_modal.py new file mode 100644 index 000000000..2396dea84 --- /dev/null +++ b/scripts/distill/train_modal.py @@ -0,0 +1,548 @@ +"""SFT training on Modal with a single A100-80GB GPU. + +Trains Qwen3.5-9B-Instruct on the combined teacher + native Qwen trajectories +using action-masking (loss computed only on assistant turns that contain tool +calls or non-trivial text — not on system/user/tool messages). + +Usage: + # Dry-run: tokenise locally, print stats, exit + uv run python scripts/distill/train_modal.py --dry-run + + # Full run on Modal (launches a remote A100 job) + uv run python scripts/distill/train_modal.py + + # Resume from a checkpoint + uv run python scripts/distill/train_modal.py --resume + +Environment: + MODAL_TOKEN_ID / MODAL_TOKEN_SECRET — Modal auth (or `modal token set`) + HF_TOKEN — HuggingFace token for model download + +Outputs (saved to Modal volume, also synced to data/checkpoints/ locally): + data/checkpoints// adapter weights (LoRA) + tokenizer +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +# --------------------------------------------------------------------------- +# Shared constants (referenced both locally and inside the Modal container) +# --------------------------------------------------------------------------- + +BASE_MODEL = "Qwen/Qwen3.5-9B" +RUN_NAME = "qwen35-tool-use-sft-v1" + +# LoRA config +LORA_R = 64 +LORA_ALPHA = 128 +LORA_DROPOUT = 0.05 +LORA_TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj"] + +# Training hyperparams +MAX_SEQ_LEN = 2048 # context window per sample (truncate longer) +PER_DEVICE_BATCH = 1 +GRAD_ACCUM_STEPS = 16 # effective batch = 16 +LEARNING_RATE = 2e-4 +NUM_EPOCHS = 3 +WARMUP_RATIO = 0.05 +SAVE_STEPS = 20 +EVAL_STEPS = 100 +LOGGING_STEPS = 10 + +DATA_PATHS = { + "teacher": "data/converted_teacher.jsonl", + "qwen": "data/successful.jsonl", +} +OUTPUT_DIR = f"data/checkpoints/{RUN_NAME}" + +# --------------------------------------------------------------------------- +# Local helpers: load + format training data +# --------------------------------------------------------------------------- + +SYSTEM_PROMPT = ( + "You are a software engineer working alongside a colleague on a shared codebase. " + "You each have your own workspace and are implementing different features in parallel. " + "You communicate naturally — like engineers on the same team — to make sure " + "your combined work integrates cleanly." +) + + +def _iter_messages_from_record(rec: dict) -> list[dict] | None: + """Return the messages list for a record, regardless of source format.""" + # converted_teacher.jsonl: messages are top-level + if "messages" in rec and rec["messages"]: + return rec["messages"] + # successful.jsonl Qwen records: messages are inside trajectory dict + traj = rec.get("trajectory") + if isinstance(traj, dict): + msgs = traj.get("messages", []) + if msgs: + return msgs + return None + + +def load_records() -> list[list[dict]]: + """Load all training records and return a list of message lists.""" + all_msgs: list[list[dict]] = [] + + # Teacher trajectories (converted_teacher.jsonl) + teacher_path = Path(DATA_PATHS["teacher"]) + if teacher_path.exists(): + with teacher_path.open() as f: + for line in f: + rec = json.loads(line) + msgs = _iter_messages_from_record(rec) + if msgs and len(msgs) > 2: + all_msgs.append(msgs) + + # Qwen native trajectories (successful.jsonl, qwen runs only) + qwen_path = Path(DATA_PATHS["qwen"]) + if qwen_path.exists(): + with qwen_path.open() as f: + for line in f: + rec = json.loads(line) + if not rec.get("run", "").startswith("qwen"): + continue + msgs = _iter_messages_from_record(rec) + if msgs and len(msgs) > 2: + all_msgs.append(msgs) + + return all_msgs + + +def messages_to_chatml(msgs: list[dict]) -> list[dict]: + """Normalise a message list to clean ChatML dicts for the tokeniser. + + Rules: + - system/user messages: keep content as-is + - assistant messages with tool_calls: convert to a single content string + that includes both the text content and a tool-call representation + - tool messages: convert to a user-visible tool result representation + - exit messages: drop + """ + out = [] + for m in msgs: + role = m.get("role", "") + content = m.get("content") or "" + + if role == "exit": + continue + + if role in ("system", "user"): + out.append({"role": role, "content": str(content)}) + + elif role == "assistant": + tool_calls = m.get("tool_calls") or [] + if tool_calls: + # Represent the tool call as structured text the model must predict + tc = tool_calls[0] # always single call in our data + try: + args = json.loads(tc["function"]["arguments"]) + cmd = args.get("command", "") + except (json.JSONDecodeError, KeyError): + cmd = tc["function"].get("arguments", "") + # Keep any reasoning text first, then the tool call + text = (str(content) + "\n" if content else "") + f"\n{cmd}\n" + out.append({"role": "assistant", "content": text.strip()}) + else: + if content: + out.append({"role": "assistant", "content": str(content)}) + + elif role == "tool": + # Represent tool results as a user message so the model sees them + try: + result = json.loads(str(content)) + rc = result.get("returncode", 0) + output = result.get("output", "") + except (json.JSONDecodeError, TypeError): + rc = 0 + output = str(content) + tool_text = f"\n{output}\n" + out.append({"role": "user", "content": tool_text}) + + return out + + +def build_action_mask(tokenised_ids: list[int], tokeniser, chatml_msgs: list[dict]) -> list[int]: + """Return a label mask (1 = compute loss, 0 = mask) aligned to tokenised_ids. + + We want loss only on assistant tokens. We detect assistant turns by + re-tokenising the conversation incrementally and finding the boundaries. + + Returns a list of the same length as tokenised_ids. + """ + # Simple approach: re-encode each message, find where assistant content sits + # by matching token spans. Falls back to full-sequence loss if alignment fails. + mask = [0] * len(tokenised_ids) + try: + pos = 0 + for msg in chatml_msgs: + role = msg["role"] + content = msg["content"] + # Encode role header + content (approximate span) + header = f"<|im_start|>{role}\n" + footer = "<|im_end|>\n" + header_ids = tokeniser.encode(header, add_special_tokens=False) + content_ids = tokeniser.encode(content, add_special_tokens=False) + footer_ids = tokeniser.encode(footer, add_special_tokens=False) + + header_len = len(header_ids) + content_len = len(content_ids) + footer_len = len(footer_ids) + total_len = header_len + content_len + footer_len + + if role == "assistant": + # Mark content + footer (not header) as loss tokens + for i in range(pos + header_len, min(pos + total_len, len(mask))): + mask[i] = 1 + + pos += total_len + if pos >= len(mask): + break + except Exception: + # Fallback: label everything + return [1] * len(tokenised_ids) + + return mask + + +# --------------------------------------------------------------------------- +# Dry-run: run locally to check data pipeline +# --------------------------------------------------------------------------- + +def dry_run() -> None: + records = load_records() + print(f"Loaded {len(records)} trajectories") + + try: + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) + except Exception as e: + print(f"Cannot load tokeniser ({e}); skipping tokenisation stats") + tok = None + + total_tokens = 0 + truncated = 0 + for msgs in records[:200]: # sample first 200 for speed + chatml = messages_to_chatml(msgs) + text = tok.apply_chat_template(chatml, tokenize=False) if tok else "" + if tok: + ids = tok.encode(text) + total_tokens += min(len(ids), MAX_SEQ_LEN) + if len(ids) > MAX_SEQ_LEN: + truncated += 1 + + print(f"Avg tokens (first 200, capped at {MAX_SEQ_LEN}): {total_tokens // min(200, len(records))}") + print(f"Truncated (>{MAX_SEQ_LEN} tokens): {truncated}/200") + print(f"Estimated total training tokens: {total_tokens * (len(records) / 200) / 1e6:.1f}M") + print("\nDry-run complete — no Modal job submitted.") + + +# --------------------------------------------------------------------------- +# Modal app — must be at module scope for @app.function to work +# --------------------------------------------------------------------------- + +try: + import modal as _modal + + _volume = _modal.Volume.from_name(f"{RUN_NAME}-volume", create_if_missing=True) + + _image = ( + _modal.Image.debian_slim(python_version="3.11") + .apt_install("build-essential") + .pip_install( + "torch>=2.5.0", + "transformers>=4.52.0", + "peft>=0.14.0", + "bitsandbytes>=0.44.0", + "accelerate>=0.34.0", + "datasets>=2.21.0", + "huggingface-hub>=0.24", + "sentencepiece", + "protobuf", + "scipy", + ) + ) + + app = _modal.App(name=RUN_NAME) + + @app.function( + image=_image, + gpu="A100-80GB", + timeout=60 * 60 * 8, + volumes={"/checkpoints": _volume}, + secrets=[_modal.Secret.from_name("huggingface-secret")], + ) + def train( + teacher_jsonl: bytes, + qwen_jsonl: bytes, + resume: bool = False, + ) -> str: + """Run fine-tuning inside the Modal container.""" + import json as _json + import os as _os + + import torch + from datasets import Dataset + from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BitsAndBytesConfig, + DataCollatorForLanguageModeling, + Trainer, + TrainingArguments, + ) + + hf_token = _os.environ.get("HF_TOKEN") + + # ── Load tokeniser + model ─────────────────────────────────────────── + print(f"Loading {BASE_MODEL} …") + tokenizer = AutoTokenizer.from_pretrained( + BASE_MODEL, token=hf_token, trust_remote_code=True + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + bnb_4bit_use_double_quant=True, + ) + + model = AutoModelForCausalLM.from_pretrained( + BASE_MODEL, + quantization_config=bnb_config, + device_map="auto", + token=hf_token, + trust_remote_code=True, + ) + model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True) + model.config.use_cache = False + + # ── LoRA ───────────────────────────────────────────────────────────── + lora_cfg = LoraConfig( + r=LORA_R, + lora_alpha=LORA_ALPHA, + lora_dropout=LORA_DROPOUT, + target_modules=LORA_TARGET_MODULES, + bias="none", + task_type="CAUSAL_LM", + ) + model = get_peft_model(model, lora_cfg) + model.print_trainable_parameters() + + # ── Build dataset ───────────────────────────────────────────────────── + def _parse_jsonl(data: bytes) -> list[dict]: + return [_json.loads(line) for line in data.decode().splitlines() if line.strip()] + + teacher_recs = _parse_jsonl(teacher_jsonl) + qwen_recs = _parse_jsonl(qwen_jsonl) + + def _to_chatml(msgs: list[dict]) -> list[dict]: + out = [] + for m in msgs: + role = m.get("role", "") + content = m.get("content") or "" + if role == "exit": + continue + if role in ("system", "user"): + out.append({"role": role, "content": str(content)}) + elif role == "assistant": + tcs = m.get("tool_calls") or [] + if tcs: + try: + a = _json.loads(tcs[0]["function"]["arguments"]) + cmd = a.get("command", "") + except Exception: + cmd = tcs[0]["function"].get("arguments", "") + text = (str(content) + "\n" if content else "") + \ + f"\n{cmd}\n" + out.append({"role": "assistant", "content": text.strip()}) + else: + if content: + out.append({"role": "assistant", "content": str(content)}) + elif role == "tool": + try: + res = _json.loads(str(content)) + rc = res.get("returncode", 0) + output = res.get("output", "") + except Exception: + rc, output = 0, str(content) + out.append({"role": "user", + "content": f"\n{output}\n"}) + return out + + def _tokenise(rec: dict) -> dict | None: + msgs = rec.get("messages") + if not msgs: + traj = rec.get("trajectory") + if isinstance(traj, dict): + msgs = traj.get("messages") + if not msgs or len(msgs) <= 2: + return None + chatml = _to_chatml(msgs) + if len(chatml) < 2: + return None + text = tokenizer.apply_chat_template(chatml, tokenize=False, add_generation_prompt=False) + enc = tokenizer(text, truncation=True, max_length=MAX_SEQ_LEN) + enc["labels"] = enc["input_ids"].copy() + return enc + + rows = [] + for rec in teacher_recs + qwen_recs: + enc = _tokenise(rec) + if enc: + rows.append(enc) + + print(f"Total training examples: {len(rows)}") + dataset = Dataset.from_list(rows) + dataset = dataset.train_test_split(test_size=0.02, seed=42) + + # ── Training ────────────────────────────────────────────────────────── + output_dir = f"/checkpoints/{RUN_NAME}" + _os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + training_args = TrainingArguments( + output_dir=output_dir, + num_train_epochs=NUM_EPOCHS, + per_device_train_batch_size=PER_DEVICE_BATCH, + gradient_accumulation_steps=GRAD_ACCUM_STEPS, + learning_rate=LEARNING_RATE, + lr_scheduler_type="cosine", + warmup_steps=50, + bf16=True, + logging_steps=LOGGING_STEPS, + save_steps=SAVE_STEPS, + eval_strategy="no", + save_total_limit=3, + report_to="none", + dataloader_num_workers=2, + remove_unused_columns=False, + gradient_checkpointing=True, + gradient_checkpointing_kwargs={"use_reentrant": False}, + optim="paged_adamw_8bit", + ) + + collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) + + import modal as _modal_remote + from transformers import TrainerCallback + _vol = _modal_remote.Volume.from_name(f"{RUN_NAME}-volume") + + class VolumeCommitCallback(TrainerCallback): + def on_save(self, args, state, control, **kwargs): + print(f"Committing checkpoint at step {state.global_step} to volume …") + _vol.commit() + print("Volume committed.") + + trainer = Trainer( + model=model, + args=training_args, + train_dataset=dataset["train"], + data_collator=collator, + callbacks=[VolumeCommitCallback()], + ) + + print("Starting training …") + trainer.train(resume_from_checkpoint=resume) + trainer.save_model(output_dir) + tokenizer.save_pretrained(output_dir) + print(f"Saved to {output_dir}") + return output_dir + +except ImportError: + app = None + train = None + _volume = None + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--dry-run", action="store_true", help="Check data pipeline locally, no Modal job") + p.add_argument("--resume", action="store_true", help="Resume from latest checkpoint on volume") + p.add_argument("--sync-only", action="store_true", + help="Download latest checkpoint from Modal volume to data/checkpoints/") + return p.parse_args() + + +def main() -> int: + args = parse_args() + + if args.dry_run: + dry_run() + return 0 + + try: + import modal + except ImportError: + print("modal not installed. Run: pip install modal", file=sys.stderr) + return 1 + + if app is None: + print("error: modal failed to initialise (see import error above)", file=sys.stderr) + return 1 + + if args.sync_only: + print(f"Syncing checkpoint from Modal volume to {OUTPUT_DIR} …") + local_out = Path(OUTPUT_DIR) + local_out.mkdir(parents=True, exist_ok=True) + for entry in _volume.iterdir(f"/{RUN_NAME}"): + _volume.read_file_into_memory(entry.path) # triggers download + print(f"Use `modal volume get {RUN_NAME}-volume /{RUN_NAME} {OUTPUT_DIR}` to download.") + return 0 + + # Read data files and send to Modal + teacher_path = Path(DATA_PATHS["teacher"]) + qwen_path = Path(DATA_PATHS["qwen"]) + + if not teacher_path.exists(): + print(f"error: {teacher_path} not found — run convert_codex_to_miniswe.py first", + file=sys.stderr) + return 1 + if not qwen_path.exists(): + print(f"error: {qwen_path} not found — run extract_successful.py first", + file=sys.stderr) + return 1 + + teacher_bytes = teacher_path.read_bytes() + qwen_bytes = qwen_path.read_bytes() + print(f"Teacher data: {len(teacher_bytes) / 1e6:.1f} MB") + print(f"Qwen data: {len(qwen_bytes) / 1e6:.1f} MB") + + print(f"Submitting Modal job '{RUN_NAME}' on A100-80GB …") + _modal.enable_output() + with app.run(): + result = train.remote( + teacher_jsonl=teacher_bytes, + qwen_jsonl=qwen_bytes, + resume=args.resume, + ) + print(f"Training complete. Downloading weights to {OUTPUT_DIR} …") + local_out = Path(OUTPUT_DIR) + local_out.mkdir(parents=True, exist_ok=True) + remote_root = f"/{RUN_NAME}" + for entry in _volume.iterdir(remote_root, recursive=True): + if entry.type.name == "FILE": + rel = entry.path[len(remote_root):].lstrip("/") + dest = local_out / rel + dest.parent.mkdir(parents=True, exist_ok=True) + data = b"".join(_volume.read_file(entry.path)) + dest.write_bytes(data) + print(f" {rel} ({len(data) / 1e6:.1f} MB)") + print(f"Weights saved to {local_out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/distill/upload_adapter.py b/scripts/distill/upload_adapter.py new file mode 100644 index 000000000..922c4ce0b --- /dev/null +++ b/scripts/distill/upload_adapter.py @@ -0,0 +1,174 @@ +"""Download the final LoRA adapter from the Modal volume and push to HuggingFace. + +Uploads only the adapter files (not checkpoint subdirs or optimizer states): + adapter_model.safetensors + adapter_config.json + tokenizer.json / tokenizer_config.json / chat_template.jinja + training_args.bin + +Target repo: CooperBench/qwen3.5-9b-tool-use-sft + +Usage: + uv run python scripts/distill/upload_adapter.py + uv run python scripts/distill/upload_adapter.py --skip-download # if already local +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +VOLUME_NAME = "qwen35-tool-use-sft-v1-volume" +REMOTE_ROOT = "/qwen35-tool-use-sft-v1" +LOCAL_DIR = Path("data/checkpoints/qwen35-tool-use-sft-v1") +HF_REPO_ID = "CooperBench/qwen3.5-9b-tool-use-sft" + +# Files to upload — root-level only, skip checkpoint-* subdirs and optimizer states +UPLOAD_PATTERNS = { + "adapter_model.safetensors", + "adapter_config.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + "training_args.bin", + "README.md", +} + +MODEL_CARD = """\ +--- +base_model: Qwen/Qwen3.5-9B +library_name: peft +tags: + - lora + - qwen3 + - tool-use + - multi-agent + - cooperbench +license: apache-2.0 +--- + +# Qwen3.5-9B Tool-Use SFT (CooperBench) + +LoRA adapter fine-tuned on top of [Qwen/Qwen3.5-9B](https://huggingface.co/Qwen/Qwen3.5-9B) +to improve coordination tool usage in multi-agent coding settings. + +## Training data + +1,230 successful multi-agent trajectories from the +[CooperBench/team-coop](https://huggingface.co/datasets/CooperBench/team-coop) dataset: + +- **1,153** converted teacher trajectories (gpt-5.5-hao via Codex, converted to + mini_swe_agent_v2 format) +- **77** native Qwen3.5-9B successful trajectories + +Tools taught: `coop-task-create`, `coop-task-claim`, `coop-task-update`, +`coop-task-list`, `coop-task-request`, `coop-task-respond`, `coop-task-pending`. + +## LoRA config + +| Parameter | Value | +|-----------|-------| +| r | 64 | +| alpha | 128 | +| dropout | 0.05 | +| target modules | q/k/v/o/gate/up/down proj | +| epochs | 3 | +| lr | 2e-4 | +| seq len | 2048 | + +## Usage + +```python +from peft import PeftModel +from transformers import AutoModelForCausalLM, AutoTokenizer + +model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-9B") +model = PeftModel.from_pretrained(model, "CooperBench/qwen3.5-9b-tool-use-sft") +tokenizer = AutoTokenizer.from_pretrained("CooperBench/qwen3.5-9b-tool-use-sft") +``` +""" + + +def download_from_volume(local_dir: Path) -> None: + try: + import modal + except ImportError: + print("modal not installed — skipping download", file=sys.stderr) + return + + local_dir.mkdir(parents=True, exist_ok=True) + vol = modal.Volume.from_name(VOLUME_NAME) + downloaded = 0 + for entry in vol.iterdir(REMOTE_ROOT, recursive=False): + fname = Path(entry.path).name + if fname not in UPLOAD_PATTERNS: + continue + dest = local_dir / fname + data = b"".join(vol.read_file(entry.path)) + dest.write_bytes(data) + print(f" ← {fname} ({len(data) / 1e6:.1f} MB)") + downloaded += 1 + print(f"Downloaded {downloaded} files to {local_dir}") + + +def upload_to_hf(local_dir: Path) -> None: + import os + from huggingface_hub import HfApi, create_repo + + token = os.environ.get("HF_TOKEN") + api = HfApi(token=token) + + print(f"Creating/verifying repo {HF_REPO_ID} …") + create_repo(repo_id=HF_REPO_ID, repo_type="model", exist_ok=True, token=token) + + # Write model card + card_path = local_dir / "README.md" + if not card_path.exists(): + card_path.write_text(MODEL_CARD) + + print(f"Uploading to {HF_REPO_ID} …") + for fpath in sorted(local_dir.iterdir()): + if fpath.name not in UPLOAD_PATTERNS: + continue + print(f" → {fpath.name} ({fpath.stat().st_size / 1e6:.1f} MB)") + api.upload_file( + path_or_fileobj=str(fpath), + path_in_repo=fpath.name, + repo_id=HF_REPO_ID, + repo_type="model", + token=token, + ) + + print(f"\nDone: https://huggingface.co/{HF_REPO_ID}") + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--skip-download", action="store_true", + help="Skip Modal download, use files already in data/checkpoints/") + return p.parse_args() + + +def main() -> int: + args = parse_args() + + if not args.skip_download: + print(f"Downloading adapter from Modal volume {VOLUME_NAME} …") + download_from_volume(LOCAL_DIR) + + # Check we have the required files + missing = [f for f in ("adapter_model.safetensors", "adapter_config.json") + if not (LOCAL_DIR / f).exists()] + if missing: + print(f"error: missing required files: {missing}", file=sys.stderr) + print(f"Check {LOCAL_DIR} or run without --skip-download", file=sys.stderr) + return 1 + + print(f"\nUploading adapter to HuggingFace …") + upload_to_hf(LOCAL_DIR) + return 0 + + +if __name__ == "__main__": + sys.exit(main())