diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..403ee8b --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "SWE-bench_Pro-os"] + path = SWE-bench_Pro-os + url = https://github.com/scaleapi/SWE-bench_Pro-os diff --git a/SWE-bench_Pro-os b/SWE-bench_Pro-os new file mode 160000 index 0000000..ca10a60 --- /dev/null +++ b/SWE-bench_Pro-os @@ -0,0 +1 @@ +Subproject commit ca10a60a5fcae51e6948ffe1485d4153d421e6c5 diff --git a/configs/agents/mini_swe_swebench_pro.yaml b/configs/agents/mini_swe_swebench_pro.yaml new file mode 100644 index 0000000..9155c3a --- /dev/null +++ b/configs/agents/mini_swe_swebench_pro.yaml @@ -0,0 +1,30 @@ +step_limit: 250 +cost_limit: 0.0 +cwd: "/app" +timeout: 180 +environment_class: modal +action_regex: '```mswea_bash_command\s*\n(.*?)\n```' +system_template: | + You are a software engineer working in a Linux container. + The repository is checked out at /app and all dependencies are installed. + Fix the reported issue by modifying the source files. + Every response must contain exactly one bash command in this format: + + ```mswea_bash_command + your_command_here + ``` + + When you are done and confident the issue is fixed, submit by running: + ```mswea_bash_command + echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT + ``` +instance_template: | + Please solve this task: {{task}} + + The repository is at /app. Use bash commands to explore the code, edit files, and run tests. + When the issue is fixed, submit with: + ```mswea_bash_command + echo COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT + ``` +extra_model_kwargs: + drop_params: true diff --git a/configs/labeling/laguna_xs2_swebench_pro_test_labeler.yaml b/configs/labeling/laguna_xs2_swebench_pro_test_labeler.yaml new file mode 100644 index 0000000..7b05185 --- /dev/null +++ b/configs/labeling/laguna_xs2_swebench_pro_test_labeler.yaml @@ -0,0 +1,12 @@ +# Label 16 test trajectories from Laguna-XS.2 on SWE-bench Pro +# Usage: +# sbatch slurm/swebench_pro_label.sh --config configs/labeling/laguna_xs2_swebench_pro_test_labeler.yaml + +trajectory_dir: generations/swebench_pro/laguna_xs2_test +output_dir: labels/swebench_pro/laguna_xs2_test +scripts_dir: SWE-bench_Pro-os +modal_app_name: program-probes-labeler-pro +sandbox_timeout: 3600 +eval_timeout: 600 +resume: true +n_workers: 16 diff --git a/configs/runs/laguna_xs2_swebench_pro_full.yaml b/configs/runs/laguna_xs2_swebench_pro_full.yaml new file mode 100644 index 0000000..f00826f --- /dev/null +++ b/configs/runs/laguna_xs2_swebench_pro_full.yaml @@ -0,0 +1,44 @@ +# Full run: Laguna-XS.2 on all 731 SWE-bench Pro instances, 10 runs each, 4xA100 thin +# +# Submit as 5-shard array job: +# sbatch --array=0-4 slurm/swebench_pro_run_thin.sh \ +# --run-config configs/runs/laguna_xs2_swebench_pro_full.yaml + +# --- Model / vLLM --- +model_id: poolside/Laguna-XS.2 +served_model_name: laguna-xs2 +tensor_parallel_size: 4 +max_model_len: 262144 +dtype: bfloat16 +gpu_memory_utilization: 0.9 +vllm_port: 18000 +vllm_startup_timeout_s: 3600 +vllm_log_path: logs/vllm_laguna_xs2_swebench_pro_full.log +vllm_extra_args: + - "--reasoning-parser" + - "poolside_v1" + - "--enable-auto-tool-choice" + - "--tool-call-parser" + - "poolside_v1" + - "--default-chat-template-kwargs" + - '{"enable_thinking": true}' + +# --- Generation (thinking mode, coding) --- +temperature: 0.7 +top_p: 0.95 +top_k: 20 + +# --- Agent --- +step_limit: 100 +command_timeout: 180 + +# --- Execution --- +n_runs: 10 # 10 trajectories per instance (pass@10) +n_instances: -1 # all instances in this shard +n_workers: 20 # 20 parallel Modal containers per shard (5 shards × 20 = 100 total) +modal_app_name: program-probes-swebench-pro +modal_timeout: 1800 # 30 min per sandbox +output_dir: generations/swebench_pro/laguna_xs2_full + +# --- SWE-bench Pro --- +swe_bench_pro_scripts_dir: SWE-bench_Pro-os diff --git a/configs/runs/laguna_xs2_swebench_pro_test.yaml b/configs/runs/laguna_xs2_swebench_pro_test.yaml new file mode 100644 index 0000000..e8701d9 --- /dev/null +++ b/configs/runs/laguna_xs2_swebench_pro_test.yaml @@ -0,0 +1,44 @@ +# Test run: Laguna-XS.2 on 16 SWE-bench Pro instances, 4xA100 thin +# +# Submit with: +# sbatch --gpus=4 slurm/swebench_pro_run.sh \ +# --run-config configs/runs/laguna_xs2_swebench_pro_test.yaml + +# --- Model / vLLM --- +model_id: poolside/Laguna-XS.2 +served_model_name: laguna-xs2 +tensor_parallel_size: 4 +max_model_len: 262144 +dtype: bfloat16 +gpu_memory_utilization: 0.9 +vllm_port: 18000 +vllm_startup_timeout_s: 900 +vllm_log_path: logs/vllm_laguna_xs2_swebench_pro_test.log +vllm_extra_args: + - "--reasoning-parser" + - "poolside_v1" + - "--enable-auto-tool-choice" + - "--tool-call-parser" + - "poolside_v1" + - "--default-chat-template-kwargs" + - '{"enable_thinking": true}' + +# --- Generation (thinking mode, coding) --- +temperature: 0.7 +top_p: 0.95 +top_k: 20 + +# --- Agent --- +step_limit: 100 +command_timeout: 180 # seconds per bash command in Modal sandbox + +# --- Execution --- +n_runs: 1 # one attempt per instance +n_instances: 16 # stop after 16 unique instances +n_workers: 16 # 16 parallel agent threads; vLLM batches their requests +modal_app_name: program-probes-swebench-pro +modal_timeout: 1800 # seconds for the full sandbox lifecycle +output_dir: generations/swebench_pro/laguna_xs2_test + +# --- SWE-bench Pro --- +swe_bench_pro_scripts_dir: SWE-bench_Pro-os # local clone of github.com/scaleapi/SWE-bench_Pro-os diff --git a/configs/runs/qwen36_27b_swebench_pro_full.yaml b/configs/runs/qwen36_27b_swebench_pro_full.yaml new file mode 100644 index 0000000..4fb7412 --- /dev/null +++ b/configs/runs/qwen36_27b_swebench_pro_full.yaml @@ -0,0 +1,41 @@ +# Full run: Qwen3.6-27B on all 731 SWE-bench Pro instances, 8xA100 80GB +# +# Submit as 4-shard array job: +# sbatch --array=0-3 -t 02:00:00 slurm/swebench_pro_run.sh \ +# --run-config configs/runs/qwen36_27b_swebench_pro_full.yaml + +# --- Model / vLLM --- +model_id: Qwen/Qwen3.6-27B +served_model_name: qwen36-27b +tensor_parallel_size: 8 +max_model_len: 262144 +dtype: bfloat16 +gpu_memory_utilization: 0.9 +vllm_port: 18000 +vllm_startup_timeout_s: 900 +vllm_log_path: logs/vllm_qwen36_27b_swebench_pro_full.log +vllm_extra_args: + - "--reasoning-parser" + - "qwen3" + - "--enable-auto-tool-choice" + - "--tool-call-parser" + - "qwen3_coder" + +# --- Generation --- +temperature: 1.0 +top_p: 0.95 + +# --- Agent --- +step_limit: 100 +command_timeout: 180 + +# --- Execution --- +n_runs: 1 +n_instances: -1 # all instances in this shard +n_workers: 8 +modal_app_name: program-probes-swebench-pro +modal_timeout: 3600 +output_dir: generations/swebench_pro/qwen36_27b_full + +# --- SWE-bench Pro --- +swe_bench_pro_scripts_dir: SWE-bench_Pro-os diff --git a/configs/runs/qwen36_27b_swebench_pro_test.yaml b/configs/runs/qwen36_27b_swebench_pro_test.yaml new file mode 100644 index 0000000..8f2da65 --- /dev/null +++ b/configs/runs/qwen36_27b_swebench_pro_test.yaml @@ -0,0 +1,41 @@ +# Test run: Qwen3.6-27B on 16 SWE-bench Pro instances, 8xA100 80GB +# +# Submit with: +# sbatch slurm/swebench_pro_run.sh --run-config configs/runs/qwen36_27b_swebench_pro_test.yaml + +# --- Model / vLLM --- +model_id: Qwen/Qwen3.6-27B +served_model_name: qwen36-27b +tensor_parallel_size: 8 +max_model_len: 262144 +dtype: bfloat16 +gpu_memory_utilization: 0.9 +vllm_port: 18000 +vllm_startup_timeout_s: 900 +vllm_log_path: logs/vllm_qwen36_27b_swebench_pro.log +vllm_extra_args: + - "--reasoning-parser" + - "qwen3" + - "--enable-auto-tool-choice" + - "--tool-call-parser" + - "qwen3_coder" + +# --- Generation --- +temperature: 1.0 +top_p: 0.95 +# max_new_tokens: not set — no per-call cap, lets thinking models use the full remaining context + +# --- Agent --- +step_limit: 100 +command_timeout: 180 # seconds per bash command in Modal sandbox + +# --- Execution --- +n_runs: 1 # one attempt per instance +n_instances: 16 # stop after 16 unique instances +n_workers: 16 # 16 parallel agent threads; vLLM batches their requests +modal_app_name: program-probes-swebench-pro +modal_timeout: 1800 # seconds for the full sandbox lifecycle +output_dir: generations/swebench_pro/qwen36_27b_test + +# --- SWE-bench Pro --- +swe_bench_pro_scripts_dir: SWE-bench_Pro-os # local clone of github.com/scaleapi/SWE-bench_Pro-os diff --git a/configs/tasks/swe_bench_pro.yaml b/configs/tasks/swe_bench_pro.yaml new file mode 100644 index 0000000..f3caf54 --- /dev/null +++ b/configs/tasks/swe_bench_pro.yaml @@ -0,0 +1,3 @@ +dataset: ScaleAI/SWE-bench_Pro +adapter: swe_bench_pro +execution_timeout: 300 diff --git a/run_labeler_pro.py b/run_labeler_pro.py new file mode 100644 index 0000000..cec8ff5 --- /dev/null +++ b/run_labeler_pro.py @@ -0,0 +1,106 @@ +"""Run the SWE-bench Pro labeler on generated agent trajectories. + +Replays each edit step of agent trajectories in Modal sandboxes, computing +probe labels (compiles, test_results) at every intermediate code state using +the Pro eval harness (entryscript.sh → output.json). + +Usage: + uv run python run_labeler_pro.py --config configs/labeling/laguna_xs2_swebench_pro_test_labeler.yaml + +The config YAML fields are documented in SwebenchProLabelerConfig (src/configs.py). +""" + +import argparse +import json +import random +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +from src.configs import SwebenchProLabelerConfig, load_config +from src.labeling.swebench_pro_labeler import label_pro_trajectory, load_pro_instance + + +def _process( + tf: Path, + iid: str, + cfg: SwebenchProLabelerConfig, + out_dir: Path, + scripts_dir: Path, + jitter: float = 0.0, +) -> str: + label_path = out_dir / f"{tf.stem}_labels.json" + if cfg.resume and label_path.exists(): + return f"{tf.name}: skipped (labels exist)" + if jitter > 0: + time.sleep(random.uniform(0, jitter)) + instance = load_pro_instance(iid) + label_pro_trajectory( + tf, + instance=instance, + scripts_dir=scripts_dir, + output_path=label_path, + modal_app_name=cfg.modal_app_name, + sandbox_timeout=cfg.sandbox_timeout, + eval_timeout=cfg.eval_timeout, + ) + return f"{tf.name}: done" + + +def main() -> None: + parser = argparse.ArgumentParser(description="Label SWE-bench Pro agent trajectories") + parser.add_argument( + "--config", required=True, + help="Path to SwebenchProLabelerConfig YAML (see configs/labeling/)", + ) + args = parser.parse_args() + + cfg: SwebenchProLabelerConfig = load_config(args.config, SwebenchProLabelerConfig) + + traj_dir = Path(cfg.trajectory_dir) + out_dir = Path(cfg.output_dir) + scripts_dir = Path(cfg.scripts_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + traj_files = sorted( + f for f in traj_dir.glob("*.json") + if not f.name.startswith("index") and "_labels" not in f.name + ) + if cfg.single: + traj_files = [f for f in traj_files if f.name == cfg.single] + + # Extract instance IDs: strip _run suffix if present + def _iid(tf: Path) -> str: + stem = tf.stem + # e.g. django__django-12345_run01 → django__django-12345 + if "_run" in stem: + stem = stem[: stem.rfind("_run")] + return stem + + if cfg.instances: + traj_files = [f for f in traj_files if _iid(f) in cfg.instances] + + print(f"Labeling {len(traj_files)} trajectories from {traj_dir} → {out_dir}", flush=True) + + jitter = 2.0 if cfg.n_workers > 1 else 0.0 + + if cfg.n_workers == 1: + for tf in traj_files: + iid = _iid(tf) + result = _process(tf, iid, cfg, out_dir, scripts_dir, jitter=0.0) + print(result, flush=True) + else: + with ThreadPoolExecutor(max_workers=cfg.n_workers) as pool: + futures = { + pool.submit(_process, tf, _iid(tf), cfg, out_dir, scripts_dir, jitter): tf.name + for tf in traj_files + } + for fut in as_completed(futures): + try: + print(fut.result(), flush=True) + except Exception as exc: + print(f"{futures[fut]}: ERROR — {exc}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/run_probe.py b/run_probe.py index 0f0a88b..3a94dfe 100644 --- a/run_probe.py +++ b/run_probe.py @@ -1,7 +1,7 @@ import argparse import json from src.configs import ModelConfig, load_config -from src.probe import run_sweep, run_final, create_sweep +from src.probe import run_sweep, run_final, run_eval, create_sweep def main(): @@ -50,6 +50,12 @@ def main(): final_p.add_argument("--pos-weight", type=float, default=1.0, help="Positive class weight multiplier (only used with --loss weighted_cross_entropy)") + eval_p = subparsers.add_parser("eval", help="Evaluate pre-trained weights on n_eval_bins bins (no retraining)") + eval_p.add_argument("--weights-run-id", default=None, + help="Run ID whose weights.pt to load (defaults to --run-id)") + eval_p.add_argument("--output-run-id", required=True, + help="Run ID under which to save results.pt") + args = parser.parse_args() if args.n_eval_bins is not None and args.n_eval_bins != args.n_bins and args.n_bins != 1: parser.error("--n-eval-bins requires --n-bins 1 or --n-eval-bins equal to --n-bins") @@ -81,6 +87,20 @@ def main(): fixed_params=args.fixed, shuffle_labels=args.shuffle_labels, ) + elif args.mode == "eval": + run_eval( + weights_run_id=args.weights_run_id or args.run_id, + output_run_id=args.output_run_id, + probe_name=args.probe, + probe_layers=model_cfg.probe_layers, + seed=args.seed, + cache_dir=args.cache_dir, + cache_run_id=cache_run_id, + results_dir=args.results_dir, + probe_arch=args.probe_arch, + n_eval_bins=args.n_eval_bins or 10, + eval_bin_axis=args.eval_bin_axis, + ) else: run_final( run_id=args.run_id, diff --git a/run_swebench_pro_agent.py b/run_swebench_pro_agent.py new file mode 100644 index 0000000..8a33205 --- /dev/null +++ b/run_swebench_pro_agent.py @@ -0,0 +1,348 @@ +"""Run mini-SWE-agent on SWE-bench Pro instances with Modal sandboxes. + +Each instance gets a fresh Modal sandbox loaded with its per-instance Docker +image from jefzda/sweap-images (the repo is already at /app). A single vLLM +server stays up for the full job; all parallel agent threads share it. + +Prerequisite: clone the official harness repo next to this project (or point +--scripts-dir at it): + git clone https://github.com/scaleapi/SWE-bench_Pro-os + +Usage (single shard): + export MODAL_TOKEN_ID="..." + export MODAL_TOKEN_SECRET="..." + + uv run python run_swebench_pro_agent.py \\ + --run-config configs/runs/qwen36_27b_swebench_pro_test.yaml + +Usage (array job — driven by swebench_pro_run.sh): + uv run python run_swebench_pro_agent.py \\ + --run-config configs/runs/qwen36_27b_swebench_pro_test.yaml \\ + --shard-rank 2 --num-shards 8 +""" + +from __future__ import annotations + +import argparse +import re +import threading +import urllib.request +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +import wandb + +from src.agents.mini_swe import MiniSweAgentAdapter +from src.agents.swe_bench_pro_environment import SWEBenchProModalEnvironment +from src.agents.vllm_server import VllmServer +from src.configs import SWEBenchRunConfig, load_config +from src.tasks.swe_bench_pro import SWEBenchProAdapter + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--run-config", required=True, help="Path to SWEBenchRunConfig YAML") + parser.add_argument("--shard-rank", type=int, default=0) + parser.add_argument("--num-shards", type=int, default=1) + parser.add_argument("--resume", action="store_true", default=True) + parser.add_argument("--no-resume", dest="resume", action="store_false") + return parser.parse_args() + + +def _out_path(output_dir: Path, instance_id: str, run_idx: int) -> Path: + if run_idx == 0: + return output_dir / f"{instance_id}.json" + return output_dir / f"{instance_id}_run{run_idx:02d}.json" + + +def _extract_agent_metrics(agent: MiniSweAgentAdapter) -> dict: + messages = agent.last_messages or [] + result = agent.last_result or {} + + prompt_tokens = 0 + completion_tokens = 0 + n_steps = 0 + timestamps: list[float] = [] + + for msg in messages: + extra = msg.get("extra") or {} + if "actions" not in extra: + continue + n_steps += 1 + usage = (extra.get("response") or {}).get("usage") or {} + prompt_tokens += usage.get("prompt_tokens", 0) + completion_tokens += usage.get("completion_tokens", 0) + ts = extra.get("timestamp") + if ts is not None: + timestamps.append(ts) + + duration_s = (max(timestamps) - min(timestamps)) if len(timestamps) >= 2 else 0.0 + + return { + "agent/n_steps": n_steps, + "agent/exit_status": result.get("exit_status", "unknown"), + "agent/prompt_tokens": prompt_tokens, + "agent/completion_tokens": completion_tokens, + "agent/total_tokens": prompt_tokens + completion_tokens, + "agent/tool_calls": len(agent.command_history), + "agent/duration_s": duration_s, + } + + +_VLLM_METRICS = { + "vllm:num_requests_running": "vllm/requests_running", + "vllm:num_requests_waiting": "vllm/requests_waiting", + "vllm:gpu_cache_usage_perc": "vllm/gpu_cache_usage_perc", + "vllm:prompt_tokens_total": "vllm/prompt_tokens_total", + "vllm:generation_tokens_total": "vllm/generation_tokens_total", +} + + +def _scrape_vllm_metrics(base_url: str) -> dict[str, float]: + try: + url = f"{base_url.rstrip('/')}/metrics" + with urllib.request.urlopen(url, timeout=5) as resp: + text = resp.read().decode() + except Exception: + return {} + + out: dict[str, float] = {} + for line in text.splitlines(): + if line.startswith("#"): + continue + m = re.match(r"^([\w:]+)(?:\{[^}]*\})?\s+([\d.e+\-]+)", line) + if m and m.group(1) in _VLLM_METRICS: + try: + out[_VLLM_METRICS[m.group(1)]] = float(m.group(2)) + except ValueError: + pass + return out + + +def _vllm_metrics_poller(base_url: str, stop: threading.Event, interval: int = 30) -> None: + while not stop.wait(interval): + metrics = _scrape_vllm_metrics(base_url) + if metrics: + wandb.log(metrics) + + +def _run_one( + instance: dict, + run_idx: int, + *, + cfg: SWEBenchRunConfig, + server: VllmServer, + output_dir: Path, +) -> dict: + instance_id = instance["instance_id"] + label = f"{instance_id}[run{run_idx}]" + out_path = _out_path(output_dir, instance_id, run_idx) + + env = SWEBenchProModalEnvironment( + instance, + scripts_dir=cfg.swe_bench_pro_scripts_dir, + timeout=cfg.modal_timeout, + app_name=cfg.modal_app_name, + ) + agent = MiniSweAgentAdapter( + model_name=server.model_name, + base_url=server.base_url, + api_key="EMPTY", + agent_config=cfg.to_pro_agent_config(), + generation_config=cfg.to_generation_config(), + environment=env, + ) + + outcome = None + patch = "" + error = None + eval_log = "" + agent_metrics: dict = {} + + try: + task = ( + f"Repository: {instance.get('repo', '')}\n\n" + f"Issue:\n{instance['problem_statement']}" + ) + agent.run_task(task) + agent_metrics = _extract_agent_metrics(agent) + patch = env.get_patch() + outcome, eval_log = env.evaluate(patch=patch) + print( + f"[{label}] outcome={'PASS' if outcome else 'FAIL'} " + f"patch={len(patch)}chars steps={agent_metrics.get('agent/n_steps', '?')} " + f"exit={agent_metrics.get('agent/exit_status', '?')}", + flush=True, + ) + except Exception as exc: + error = str(exc) + print(f"[{label}] ERROR: {exc}", flush=True) + finally: + env.close() + + if error is None: + agent.save_trajectory( + out_path, + tokenizer_name=cfg.model_id, + metadata={ + "instance_id": instance_id, + "run_idx": run_idx, + "repo": instance.get("repo", ""), + "base_commit": instance.get("base_commit", ""), + "outcome": outcome, + "eval_log": eval_log[-5000:] if eval_log else "", + "patch": patch, + "run_config": cfg.model_dump(), + }, + ) + + return { + "instance_id": instance_id, + "run_idx": run_idx, + "outcome": outcome, + "patch_len": len(patch), + **agent_metrics, + **({"error": error} if error else {}), + } + + +def main() -> None: + args = _parse_args() + cfg = load_config(args.run_config, SWEBenchRunConfig) + assert isinstance(cfg, SWEBenchRunConfig) + + output_dir = Path(cfg.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + from src.configs import TaskConfig + task_cfg = TaskConfig(dataset="ScaleAI/SWE-bench_Pro", adapter="swe_bench_pro") + adapter = SWEBenchProAdapter() + all_instances = adapter.load_dataset(task_cfg) + + instances = [inst for i, inst in enumerate(all_instances) if i % args.num_shards == args.shard_rank] + print(f"[swe-bench-pro] shard {args.shard_rank}/{args.num_shards}: {len(instances)} instances", flush=True) + + jobs = [ + (inst, run_idx) + for inst in instances + for run_idx in range(cfg.n_runs) + ] + + if args.resume: + n_before = len(jobs) + jobs = [ + (inst, run_idx) for inst, run_idx in jobs + if not _out_path(output_dir, inst["instance_id"], run_idx).exists() + ] + print(f"[swe-bench-pro] resume: skipped {n_before - len(jobs)}, {len(jobs)} remaining", flush=True) + + if cfg.n_instances > 0: + seen: set[str] = set() + limited = [] + for inst, run_idx in jobs: + seen.add(inst["instance_id"]) + if len(seen) <= cfg.n_instances: + limited.append((inst, run_idx)) + jobs = limited + + if not jobs: + print("[swe-bench-pro] nothing to do.", flush=True) + return + + print(f"[swe-bench-pro] {len(jobs)} jobs, {cfg.n_workers} workers", flush=True) + + run_name = f"{Path(args.run_config).stem}_shard{args.shard_rank}of{args.num_shards}" + wandb.init( + project="program-probes", + name=run_name, + config={**cfg.model_dump(), "shard_rank": args.shard_rank, "num_shards": args.num_shards}, + resume="allow", + ) + + vllm_cfg = cfg.to_vllm_server_config() + if vllm_cfg.log_path: + stem = vllm_cfg.log_path.replace(".log", "") + vllm_cfg = vllm_cfg.model_copy(update={"log_path": f"{stem}_shard{args.shard_rank}.log"}) + model_cfg = cfg.to_model_config() + + n_total = len(jobs) + n_done = n_passed = n_failed = n_errored = 0 + + with VllmServer(model_cfg, vllm_cfg) as server: + print(f"[swe-bench-pro] vLLM ready at {server.base_url} ({server.model_name})", flush=True) + + stop_poller = threading.Event() + poller = threading.Thread( + target=_vllm_metrics_poller, + args=(server.base_url, stop_poller), + daemon=True, + ) + poller.start() + + try: + with ThreadPoolExecutor(max_workers=cfg.n_workers) as pool: + futures = { + pool.submit( + _run_one, + inst, + run_idx, + cfg=cfg, + server=server, + output_dir=output_dir, + ): (inst["instance_id"], run_idx) + for inst, run_idx in jobs + } + for future in as_completed(futures): + instance_id, run_idx = futures[future] + try: + summary = future.result() + except Exception as exc: + print(f"[swe-bench-pro] unhandled exception for {instance_id}[run{run_idx}]: {exc}", flush=True) + summary = {"instance_id": instance_id, "run_idx": run_idx, "outcome": None, "patch_len": 0, "error": str(exc)} + + n_done += 1 + errored = "error" in summary + passed = (not errored) and bool(summary.get("outcome")) + if errored: + n_errored += 1 + elif passed: + n_passed += 1 + else: + n_failed += 1 + + wandb.log({ + "instance/outcome": 1 if passed else 0, + "instance/patch_len": summary.get("patch_len", 0), + "instance/errored": int(errored), + "instance/n_steps": summary.get("agent/n_steps", 0), + "instance/prompt_tokens": summary.get("agent/prompt_tokens", 0), + "instance/completion_tokens": summary.get("agent/completion_tokens", 0), + "instance/total_tokens": summary.get("agent/total_tokens", 0), + "instance/tool_calls": summary.get("agent/tool_calls", 0), + "instance/duration_s": summary.get("agent/duration_s", 0), + "instance/exit_status": summary.get("agent/exit_status", "unknown"), + "progress/completed": n_done, + "progress/passed": n_passed, + "progress/failed": n_failed, + "progress/errored": n_errored, + "progress/pass_rate": n_passed / n_done, + "progress/remaining": n_total - n_done, + }) + finally: + stop_poller.set() + poller.join(timeout=5) + + wandb.summary.update({ + "total": n_total, + "passed": n_passed, + "failed": n_failed, + "errored": n_errored, + "pass_rate": n_passed / n_done if n_done else 0.0, + }) + wandb.finish() + + print(f"\n[swe-bench-pro] done.", flush=True) + + +if __name__ == "__main__": + main() diff --git a/slurm/swebench_pro_label.sh b/slurm/swebench_pro_label.sh new file mode 100644 index 0000000..4b8882e --- /dev/null +++ b/slurm/swebench_pro_label.sh @@ -0,0 +1,16 @@ +#!/bin/bash +# SLURM launcher for SWE-bench Pro trajectory labeling. +# +# Usage: +# sbatch slurm/swebench_pro_label.sh --config configs/labeling/laguna_xs2_swebench_pro_test_labeler.yaml +# +#SBATCH -J pp-pro-label +#SBATCH -p berzelius-cpu +#SBATCH -t 12:00:00 +#SBATCH -o logs/pro_labeler_%A_%a.out +#SBATCH -e logs/pro_labeler_%A_%a.err + +set -euo pipefail +mkdir -p logs + +uv run python run_labeler_pro.py "$@" diff --git a/slurm/swebench_pro_run.sh b/slurm/swebench_pro_run.sh new file mode 100644 index 0000000..2996ccb --- /dev/null +++ b/slurm/swebench_pro_run.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# SWE-bench Pro agent run — SLURM launcher. +# +# Single shard: +# sbatch slurm/swebench_pro_run.sh --run-config configs/runs/qwen36_27b_swebench_pro_test.yaml +# +# Array job (4 shards): +# sbatch --array=0-3 slurm/swebench_pro_run.sh \ +# --run-config configs/runs/qwen36_27b_swebench_pro_full.yaml +# +# Prerequisite: initialise the SWE-bench_Pro-os submodule (one-time): +# git submodule update --init SWE-bench_Pro-os +# +#SBATCH -J pp-swebench-pro +#SBATCH -p berzelius +#SBATCH --gpus=8 +#SBATCH -C fat +#SBATCH -t 12:00:00 +#SBATCH -o logs/swebench_pro_%A_%a.out +#SBATCH -e logs/swebench_pro_%A_%a.err + +set -euo pipefail +mkdir -p logs + +module load buildenv-gcccuda/12.4.1-gcc13.3.0 +unset CPATH +export LIBRARY_PATH="/usr/local/cuda/lib64:${LIBRARY_PATH:-}" +export CUDA_HOME=/usr/local/cuda +export PATH="/usr/local/cuda/bin:$PATH" +export SSL_CERT_FILE=/etc/pki/tls/cert.pem + +RANK=${SLURM_ARRAY_TASK_ID:-0} +NUM_SHARDS=${NUM_SHARDS:-${SLURM_ARRAY_TASK_COUNT:-1}} + +uv run python run_swebench_pro_agent.py \ + --shard-rank "$RANK" \ + --num-shards "$NUM_SHARDS" \ + --resume \ + "$@" diff --git a/slurm/swebench_pro_run_thin.sh b/slurm/swebench_pro_run_thin.sh new file mode 100644 index 0000000..e4a41d1 --- /dev/null +++ b/slurm/swebench_pro_run_thin.sh @@ -0,0 +1,36 @@ +#!/bin/bash +# SWE-bench Pro agent run — thin A100 nodes (4 GPUs, no fat constraint). +# +# Single shard: +# sbatch --gpus=4 slurm/swebench_pro_run_thin.sh \ +# --run-config configs/runs/laguna_xs2_swebench_pro_test.yaml +# +# Array job (5 shards): +# sbatch --array=0-4 slurm/swebench_pro_run_thin.sh \ +# --run-config configs/runs/laguna_xs2_swebench_pro_full.yaml +# +#SBATCH -J pp-swebench-pro +#SBATCH -p berzelius +#SBATCH --gpus=4 +#SBATCH -t 32:00:00 +#SBATCH -o logs/swebench_pro_%A_%a.out +#SBATCH -e logs/swebench_pro_%A_%a.err + +set -euo pipefail +mkdir -p logs + +module load buildenv-gcccuda/12.4.1-gcc13.3.0 +unset CPATH +export LIBRARY_PATH="/usr/local/cuda/lib64:${LIBRARY_PATH:-}" +export CUDA_HOME=/usr/local/cuda +export PATH="/usr/local/cuda/bin:$PATH" +export SSL_CERT_FILE=/etc/pki/tls/cert.pem + +RANK=${SLURM_ARRAY_TASK_ID:-0} +NUM_SHARDS=${NUM_SHARDS:-${SLURM_ARRAY_TASK_COUNT:-1}} + +uv run python run_swebench_pro_agent.py \ + --shard-rank "$RANK" \ + --num-shards "$NUM_SHARDS" \ + --resume \ + "$@" diff --git a/src/agents/swe_bench_pro_environment.py b/src/agents/swe_bench_pro_environment.py new file mode 100644 index 0000000..b9f150a --- /dev/null +++ b/src/agents/swe_bench_pro_environment.py @@ -0,0 +1,179 @@ +from __future__ import annotations + +import ast +import json +from pathlib import Path +from typing import Any + +from src.agents.modal_environment import ModalSandboxEnvironment + + +def _get_image_uri(instance: dict) -> str: + """Return the Docker Hub URI for a SWE-bench Pro instance. + + The tag is stored directly in the dataset as `dockerhub_tag`. + """ + tag = instance["dockerhub_tag"] + return f"jefzda/sweap-images:{tag}" + + +def _extract_env_exports(dockerfile_text: str) -> str: + """Convert Dockerfile ENV lines to shell export statements.""" + lines = [] + for line in dockerfile_text.splitlines(): + stripped = line.strip() + if stripped.startswith("ENV "): + lines.append(stripped.replace("ENV ", "export ", 1)) + return "\n".join(lines) + + +def _build_entry_script(instance: dict, scripts_dir: Path) -> str: + """Build the entryscript.sh content for a Pro instance. + + Mirrors create_entryscript() from the official SWE-bench_Pro-os harness + (https://github.com/scaleapi/SWE-bench_Pro-os/blob/main/swe_bench_pro_eval.py). + """ + instance_id = instance["instance_id"] + base_commit = instance["base_commit"] + before_repo_set_cmd = instance["before_repo_set_cmd"].strip().split("\n")[-1] + selected_files = ",".join(ast.literal_eval(instance["selected_test_files_to_run"])) + + base_df = scripts_dir / "dockerfiles" / "base_dockerfile" / instance_id / "Dockerfile" + inst_df = scripts_dir / "dockerfiles" / "instance_dockerfile" / instance_id / "Dockerfile" + + env_exports = "" + for df_path in (base_df, inst_df): + if df_path.exists(): + env_exports += _extract_env_exports(df_path.read_text()) + "\n" + + return f"""\ +#!/bin/bash +set -uo pipefail +{env_exports} +cd /app +git reset --hard {base_commit} +git checkout {base_commit} +git apply -v /workspace/patch.diff +{before_repo_set_cmd} +bash /workspace/run_script.sh {selected_files} > /workspace/stdout.log 2> /workspace/stderr.log +python /workspace/parser.py /workspace/stdout.log /workspace/stderr.log /workspace/output.json +""" + + +def _grade(output_json: dict, instance: dict) -> bool: + """Return True iff all fail_to_pass and pass_to_pass tests are PASSED.""" + passed = {t["name"] for t in output_json.get("tests", []) if t["status"] == "PASSED"} + f2p = set(ast.literal_eval(instance["fail_to_pass"])) + p2p = set(ast.literal_eval(instance["pass_to_pass"])) + return (f2p | p2p) <= passed + + +class SWEBenchProModalEnvironment(ModalSandboxEnvironment): + """ModalSandboxEnvironment pre-loaded with a SWE-bench Pro Docker image. + + `scripts_dir` must point to a local clone of + https://github.com/scaleapi/SWE-bench_Pro-os so that per-instance + run_scripts/ and dockerfiles/ are available at eval time. + """ + + def __init__( + self, + instance: dict, + *, + scripts_dir: str | Path, + timeout: int = 180, + app_name: str = "program-probes-swebench-pro", + ): + try: + import modal + except ImportError as exc: + raise ImportError("modal is required for SWEBenchProModalEnvironment.") from exc + + image_uri = _get_image_uri(instance) + image = modal.Image.from_registry(image_uri) + super().__init__(app_name=app_name, image=image, timeout=timeout) + self.instance = instance + self.scripts_dir = Path(scripts_dir) + + def execute(self, action: dict[str, Any], cwd: str = "") -> dict[str, Any]: + result = super().execute(action, cwd=cwd) + diff = self._capture_diff() + result["diff"] = diff + return result + + def _capture_diff(self) -> str: + try: + process = self.sandbox.exec("bash", "-lc", "cd /app && git diff HEAD -- .", timeout=30) + diff = process.stdout.read() + if hasattr(process, "wait"): + process.wait() + return diff or "" + except Exception: + return "" + + def get_patch(self) -> str: + return self._capture_diff() + + def _write_sandbox_file(self, remote_path: str, content: str) -> None: + """Write a text file directly into the sandbox via sandbox.open().""" + with self.sandbox.open(remote_path, "w") as f: + f.write(content) + + def evaluate(self, patch: str = "", _eval_script: str = "") -> tuple[bool, str]: + """Run the Pro eval harness and return (resolved, log). + + Uploads patch.diff, run_script.sh, parser.py, and entryscript.sh into + /workspace/ inside the sandbox, then executes entryscript.sh and reads + back /workspace/output.json to grade the result. + + Pass `patch` explicitly to avoid a redundant get_patch() RPC when the + caller already holds the diff. If omitted, get_patch() is called here. + `_eval_script` is accepted for interface compatibility but ignored — the + entry script is generated from the instance fields and local run_scripts/. + """ + instance_id = self.instance["instance_id"] + run_scripts_dir = self.scripts_dir / "run_scripts" / instance_id + + run_script_path = run_scripts_dir / "run_script.sh" + parser_path = run_scripts_dir / "parser.py" + for p in (run_script_path, parser_path): + if not p.exists(): + print(f"[swe-pro-eval] missing required file: {p}", flush=True) + return False, f"missing {p}" + + if not patch: + patch = self.get_patch() + if not patch.strip(): + print("[swe-pro-eval] no patch to evaluate", flush=True) + return False, "no patch" + + entry_script = _build_entry_script(self.instance, self.scripts_dir) + + try: + self.sandbox.exec("bash", "-lc", "mkdir -p /workspace", timeout=30).wait() + + self._write_sandbox_file("/workspace/patch.diff", patch) + self._write_sandbox_file("/workspace/run_script.sh", run_script_path.read_text()) + self._write_sandbox_file("/workspace/parser.py", parser_path.read_text()) + self._write_sandbox_file("/workspace/entryscript.sh", entry_script) + + process = self.sandbox.exec( + "bash", "/workspace/entryscript.sh", timeout=self.timeout + ) + stdout = process.stdout.read() + stderr = process.stderr.read() if getattr(process, "stderr", None) else "" + if hasattr(process, "wait"): + process.wait() + log = f"{stdout}\n{stderr}".strip() + print(f"[swe-pro-eval output]\n{log[-3000:]}", flush=True) + + with self.sandbox.open("/workspace/output.json", "r") as f: + output_json = json.load(f) + + resolved = _grade(output_json, self.instance) + print(f"[swe-pro-eval] resolved={resolved}", flush=True) + return resolved, log + + except Exception as exc: + print(f"[swe-pro-eval] exception during evaluation: {exc}", flush=True) + return False, str(exc) diff --git a/src/configs.py b/src/configs.py index b0cd03e..85f2510 100644 --- a/src/configs.py +++ b/src/configs.py @@ -120,6 +120,9 @@ class SWEBenchRunConfig(BaseModel): modal_timeout: int = 300 # seconds for the Modal sandbox lifecycle output_dir: str = "generations/swebench" + # SWE-bench Pro only: path to a local clone of github.com/scaleapi/SWE-bench_Pro-os + swe_bench_pro_scripts_dir: str = "SWE-bench_Pro-os" + def to_model_config(self) -> "ModelConfig": return ModelConfig( model_id=self.model_id, @@ -156,13 +159,16 @@ def to_generation_config(self) -> "GenerationConfig": return GenerationConfig(**kwargs) def to_agent_config(self) -> "MiniSweAgentConfig": + return self._load_agent_config("configs/agents/mini_swe_swebench.yaml") + + def to_pro_agent_config(self) -> "MiniSweAgentConfig": + return self._load_agent_config("configs/agents/mini_swe_swebench_pro.yaml") + + def _load_agent_config(self, config_path: str) -> "MiniSweAgentConfig": import yaml as _yaml from pathlib import Path as _Path - _swe_cfg_path = _Path(__file__).parent.parent / "configs/agents/mini_swe_swebench.yaml" - if _swe_cfg_path.exists(): - base = _yaml.safe_load(_swe_cfg_path.read_text()) - else: - base = {} + _cfg_path = _Path(__file__).parent.parent / config_path + base = _yaml.safe_load(_cfg_path.read_text()) if _cfg_path.exists() else {} base["step_limit"] = self.step_limit base["timeout"] = self.command_timeout return MiniSweAgentConfig.model_validate(base) @@ -187,6 +193,12 @@ class SwebenchLabelerConfig(BaseModel): n_workers: int = 1 # parallel Modal sandboxes +class SwebenchProLabelerConfig(SwebenchLabelerConfig): + """Extends SwebenchLabelerConfig with the SWE-bench Pro scripts directory.""" + + scripts_dir: str = "SWE-bench_Pro-os" + + def load_config(path: str, model: type[BaseModel]) -> BaseModel: import yaml with open(path) as f: diff --git a/src/export_swebench_dashboard.py b/src/export_swebench_dashboard.py index 896e484..b648378 100644 --- a/src/export_swebench_dashboard.py +++ b/src/export_swebench_dashboard.py @@ -267,7 +267,7 @@ def _load_labels_for(f: Path) -> dict: (run_dir / "stats.json").write_text(json.dumps(stats_out)) # Per-bin majority baselines - cache_run = output_run_id or run_id + cache_run = run_id majority_baselines_per_bin = { p: _per_bin_majority_baseline_from_cache(Path(cache_dir) / cache_run, p, n_bins, eval_bin_axis=eval_bin_axis) for p in probe_names diff --git a/src/labeling/swebench_pro_labeler.py b/src/labeling/swebench_pro_labeler.py new file mode 100644 index 0000000..52da45d --- /dev/null +++ b/src/labeling/swebench_pro_labeler.py @@ -0,0 +1,290 @@ +"""Label SWE-bench Pro agent trajectories by replaying cumulative diffs in Modal sandboxes. + +For each trajectory, identifies edit steps (commands that changed files via their +recorded cumulative diff), then applies each diff from a clean git state inside a +per-instance Modal sandbox and evaluates the code using the Pro eval harness. + +Use ``run_labeler_pro.py`` as the CLI entry point. +""" + +from __future__ import annotations + +import json +import os +import tempfile +from pathlib import Path +from typing import Any, Dict, List, Optional + +from src.agents.swe_bench_pro_environment import _build_entry_script, _grade, _get_image_uri + + +_COMPILE_ERROR_PATTERNS = ( + "SyntaxError", + "ImportError", + "ModuleNotFoundError", + "ERROR collecting", + "import file mismatch", +) + + +def _compiles_from_log(log: Optional[str]) -> Optional[bool]: + if log is None: + return None + for pattern in _COMPILE_ERROR_PATTERNS: + if pattern in log: + return False + return True + + +def _create_pro_sandbox(instance: Dict, app_name: str, timeout: int) -> Any: + """Create a Modal sandbox pre-loaded with the instance's Pro Docker image.""" + import modal + + app = modal.App.lookup(app_name, create_if_missing=True) + image_uri = _get_image_uri(instance) + image = modal.Image.from_registry(image_uri) + return modal.Sandbox.create(app=app, image=image, timeout=timeout) + + +def _exec_in_sandbox(sandbox: Any, command: str, timeout: int = 180) -> Dict: + """Execute a bash command in *sandbox* and return ``{returncode, output}``.""" + process = sandbox.exec("bash", "-lc", command, timeout=timeout) + stdout = process.stdout.read() + stderr = process.stderr.read() if getattr(process, "stderr", None) is not None else "" + if hasattr(process, "wait"): + process.wait() + returncode = getattr(process, "returncode", 0) + output = stdout + if stderr: + output = f"{stdout}\n{stderr}" if stdout else stderr + return {"returncode": returncode, "output": output} + + +def _write_sandbox_file(sandbox: Any, remote_path: str, content: str) -> None: + with sandbox.open(remote_path, "w") as f: + f.write(content) + + +def _apply_patch_in_sandbox(sandbox: Any, patch: str, base_commit: str) -> bool: + """Reset /app to base_commit and apply *patch*. Returns True on success.""" + _exec_in_sandbox( + sandbox, + f"cd /app && git reset --hard {base_commit} && git checkout {base_commit}", + timeout=30, + ) + if not patch.strip(): + return True + + with tempfile.NamedTemporaryFile(mode="w", suffix=".patch", delete=False) as f: + f.write(patch) + patch_local = f.name + + try: + with open(patch_local) as pf: + content = pf.read() + upload = f"cat > /tmp/edit.patch << 'PATCHEOF'\n{content}\nPATCHEOF" + r = _exec_in_sandbox(sandbox, upload, timeout=30) + if r["returncode"] != 0: + return False + r = _exec_in_sandbox(sandbox, "cd /app && git apply /tmp/edit.patch", timeout=30) + return r["returncode"] == 0 + finally: + os.unlink(patch_local) + + +def _run_pro_eval( + sandbox: Any, + instance: Dict, + scripts_dir: Path, + patch: str, + eval_timeout: int = 600, +) -> tuple[Optional[bool], Optional[bool], Optional[Dict]]: + """Apply *patch* and run the Pro eval harness. + + Returns (resolved, compiles, output_json). All values are None on error. + """ + base_commit = instance["base_commit"] + instance_id = instance["instance_id"] + + run_scripts_dir = scripts_dir / "run_scripts" / instance_id + run_script_path = run_scripts_dir / "run_script.sh" + parser_path = run_scripts_dir / "parser.py" + + for p in (run_script_path, parser_path): + if not p.exists(): + print(f"[pro-labeler] missing required file: {p}", flush=True) + return None, None, None + + ok = _apply_patch_in_sandbox(sandbox, patch, base_commit) + if not ok: + print(f"[pro-labeler] {instance_id}: git apply failed", flush=True) + return None, None, None + + entry_script = _build_entry_script(instance, scripts_dir) + + try: + _exec_in_sandbox(sandbox, "mkdir -p /workspace", timeout=30) + _write_sandbox_file(sandbox, "/workspace/patch.diff", patch) + _write_sandbox_file(sandbox, "/workspace/run_script.sh", run_script_path.read_text()) + _write_sandbox_file(sandbox, "/workspace/parser.py", parser_path.read_text()) + _write_sandbox_file(sandbox, "/workspace/entryscript.sh", entry_script) + + process = sandbox.exec("bash", "/workspace/entryscript.sh", timeout=eval_timeout) + stdout = process.stdout.read() + stderr = process.stderr.read() if getattr(process, "stderr", None) else "" + if hasattr(process, "wait"): + process.wait() + log = f"{stdout}\n{stderr}".strip() + print(f"[pro-labeler] eval output (last 1000 chars):\n{log[-1000:]}", flush=True) + + with sandbox.open("/workspace/output.json", "r") as f: + output_json = json.load(f) + + resolved = _grade(output_json, instance) + compiles = _compiles_from_log(log) + return resolved, compiles, output_json + + except Exception as exc: + print(f"[pro-labeler] {instance_id}: eval exception: {exc}", flush=True) + return None, None, None + + +def label_pro_trajectory( + trajectory_path: Path, + *, + instance: Dict, + scripts_dir: Path, + output_path: Path, + modal_app_name: str = "program-probes-labeler-pro", + sandbox_timeout: int = 3600, + eval_timeout: int = 600, +) -> Dict: + """Label a single Pro trajectory by replaying cumulative diffs in a Modal sandbox. + + Mirrors the structure of swebench_labeler.label_trajectory but uses the Pro + eval harness (entryscript.sh → output.json → _grade()) instead of the pytest + log parser. + """ + with open(trajectory_path) as f: + traj = json.load(f) + + command_history = traj.get("command_history", []) + if not command_history: + print(f"[pro-labeler] {trajectory_path.name}: no command history", flush=True) + result = { + "instance_id": instance["instance_id"], + "trajectory_path": str(trajectory_path), + "edits": [], + } + _write_labels(result, output_path) + return result + + # Identify edit steps — first command in each group that shares a cumulative diff + edit_indices: List[int] = [] + last_diff = "" + for i, entry in enumerate(command_history): + d = entry.get("diff", "") + if d and d != last_diff: + edit_indices.append(i) + last_diff = d + + if not edit_indices: + print(f"[pro-labeler] {trajectory_path.name}: no edits", flush=True) + result = { + "instance_id": instance["instance_id"], + "trajectory_path": str(trajectory_path), + "edits": [], + } + _write_labels(result, output_path) + return result + + print( + f"[pro-labeler] {trajectory_path.name}: {len(edit_indices)} edits at cmds {edit_indices}", + flush=True, + ) + + sandbox = _create_pro_sandbox(instance, modal_app_name, sandbox_timeout) + + edits: List[Dict] = [] + try: + # Baseline: clean checkout (no patch) + print("[pro-labeler] clean checkout: running baseline eval ...", flush=True) + resolved, compiles, output_json = _run_pro_eval( + sandbox, instance, scripts_dir, patch="", eval_timeout=eval_timeout + ) + print( + f"[pro-labeler] clean checkout: resolved={resolved}, compiles={compiles}", + flush=True, + ) + edits.append({ + "cmd_idx": -1, + "compiles": compiles, + "test_results": _to_test_results(resolved, output_json), + }) + + for cmd_idx in edit_indices: + cumulative_diff = command_history[cmd_idx].get("diff", "") + print(f"[pro-labeler] cmd[{cmd_idx}]: running eval ...", flush=True) + + resolved, compiles, output_json = _run_pro_eval( + sandbox, instance, scripts_dir, + patch=cumulative_diff, eval_timeout=eval_timeout, + ) + if resolved is None and compiles is None: + edits.append({ + "cmd_idx": cmd_idx, + "compiles": None, + "test_results": None, + "apply_error": True, + }) + print(f"[pro-labeler] cmd[{cmd_idx}]: eval failed", flush=True) + continue + + edits.append({ + "cmd_idx": cmd_idx, + "compiles": compiles, + "test_results": _to_test_results(resolved, output_json), + }) + print( + f"[pro-labeler] cmd[{cmd_idx}]: resolved={resolved}, compiles={compiles}", + flush=True, + ) + finally: + sandbox.terminate() + + result = { + "instance_id": instance["instance_id"], + "trajectory_path": str(trajectory_path), + "edits": edits, + } + _write_labels(result, output_path) + print(f"[pro-labeler] wrote {output_path}", flush=True) + return result + + +def _to_test_results(resolved: Optional[bool], output_json: Optional[Dict]) -> Optional[Dict]: + """Convert Pro output.json to the same shape as swebench_labeler test_results.""" + if output_json is None: + return None + tests = output_json.get("tests", []) + passed = [t["name"] for t in tests if t.get("status") == "PASSED"] + failed = [t["name"] for t in tests if t.get("status") == "FAILED"] + error = [t["name"] for t in tests if t.get("status") not in ("PASSED", "FAILED")] + return {"passed": passed, "failed": failed, "error": error, "resolved": resolved} + + +def _write_labels(result: Dict, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + json.dump(result, f, indent=2) + + +def load_pro_instance(instance_id: str) -> Dict: + """Load a single SWE-bench Pro instance from HuggingFace.""" + from datasets import load_dataset + + data = load_dataset("ScaleAI/SWE-bench_Pro", split="test") + matches = [inst for inst in data if inst["instance_id"] == instance_id] + if not matches: + raise ValueError(f"Instance {instance_id!r} not found in SWE-bench Pro") + return dict(matches[0]) diff --git a/src/probe.py b/src/probe.py index ed880ee..cb69642 100644 --- a/src/probe.py +++ b/src/probe.py @@ -576,3 +576,136 @@ def log_fn(metrics: dict) -> None: torch.save(all_results, out / "results.pt") torch.save(all_weights, out / "weights.pt") return {layer: [vars(r) for r in rs] for layer, rs in all_results.items()} + + +def run_eval( + weights_run_id: str, + output_run_id: str, + probe_name: str, + probe_layers: list[int], + seed: int = 42, + cache_dir: str = "cache", + cache_run_id: str | None = None, + results_dir: str = "results", + probe_arch: str = "linear", + n_eval_bins: int = 10, + eval_bin_axis: str = "position", +) -> None: + """Evaluate pre-trained probe weights on n_eval_bins bins without retraining. + + Loads weights from results///weights.pt, runs forward + passes on the test split of the cache, and saves results to + results///results.pt. + """ + weights_path = Path(results_dir) / weights_run_id / probe_name / "weights.pt" + all_weights = torch.load(weights_path, weights_only=False) + + cache_base = Path(cache_dir) / (cache_run_id or weights_run_id) / probe_name + all_results: dict[int, list[ProbeResult]] = {} + + for layer_idx in probe_layers: + cache_path = cache_base / f"layer_{layer_idx}.pt" + data = torch.load(str(cache_path), weights_only=False) + H = data["H"] + y = data["y"] + rel_pos = data["rel_pos"] + step_idx = data.get("step_idx") + sample_ids = data["sample_id"] + group_ids = data["group_id"] + + train_groups, val_groups, test_groups = _split_groups(group_ids, seed) + test_samples = {sid for sid, gid in zip(sample_ids, group_ids) if gid in test_groups} + val_samples = {sid for sid, gid in zip(sample_ids, group_ids) if gid in val_groups} + + layer_weights = all_weights[layer_idx] + # pooled probe: single entry at bin 0 + bin0 = layer_weights[0] + mean = bin0["mean"] + hidden_dim = H.shape[1] + model = _build_probe(probe_arch, hidden_dim) + model.load_state_dict(bin0["state_dict"]) + model.eval() + + # Per-trajectory max step for step_relative binning + sample_max_step: dict[str, int] = {} + if step_idx is not None: + for sid, s in zip(sample_ids, step_idx.tolist()): + if s > sample_max_step.get(sid, 0): + sample_max_step[sid] = s + + valid = y >= 0 + results = [] + + for eb in range(n_eval_bins): + if eval_bin_axis == "step_relative": + if step_idx is None: + raise ValueError("eval_bin_axis='step_relative' requires step_idx in cache") + bin_mask = torch.tensor([ + valid[i].item() + and _bin_index(step_idx[i].item() / max(sample_max_step.get(sample_ids[i], 1), 1), n_eval_bins) == eb + for i in range(len(y)) + ]) + else: # position + bin_mask = torch.tensor([ + valid[i].item() + and _bin_index(rel_pos[i].item(), n_eval_bins) == eb + for i in range(len(y)) + ]) + + if not bin_mask.any(): + continue + + mask_idx = bin_mask.nonzero(as_tuple=True)[0] + bin_samples = [sample_ids[i] for i in mask_idx.tolist()] + bin_H = H[bin_mask].float() - mean + bin_y = y[bin_mask] + + val_mask = torch.tensor([s in val_samples for s in bin_samples]) + test_mask = torch.tensor([s in test_samples for s in bin_samples]) + if val_mask.sum() == 0 or test_mask.sum() == 0: + continue + + H_val_e, y_val_e = bin_H[val_mask], bin_y[val_mask] + H_test_e, y_test_e = bin_H[test_mask], bin_y[test_mask] + + with torch.no_grad(): + val_probs = torch.softmax(model(H_val_e), dim=1)[:, 1].numpy() + test_probs = torch.softmax(model(H_test_e), dim=1)[:, 1].numpy() + + val_preds_np = (val_probs >= 0.5).astype(int) + test_preds_np = (test_probs >= 0.5).astype(int) + val_labels_np = y_val_e.numpy() + test_labels_np = y_test_e.numpy() + + val_m = _clf_metrics(val_probs, val_preds_np, val_labels_np) + test_m = _clf_metrics(test_probs, test_preds_np, test_labels_np) + + n_test = H_test_e.shape[0] + n_val = H_val_e.shape[0] + # n_train: total train tokens in this eval bin (informational) + train_mask = torch.tensor([s not in val_samples and s not in test_samples for s in bin_samples]) + n_train = int(train_mask.sum().item()) + + results.append(ProbeResult( + layer=layer_idx, + bin_idx=eb, + val_acc=float((val_preds_np == val_labels_np).mean()), + val_f1=val_m["f1"], val_precision=val_m["precision"], + val_recall=val_m["recall"], val_auc=val_m["auc"], + val_ece=val_m["ece"], val_brier=val_m["brier"], + test_acc=float((test_preds_np == test_labels_np).mean()), + test_f1=test_m["f1"], test_precision=test_m["precision"], + test_recall=test_m["recall"], test_auc=test_m["auc"], + test_ece=test_m["ece"], test_brier=test_m["brier"], + n_train=n_train, n_val=n_val, n_test=n_test, + n_pos_test=int((test_labels_np == 1).sum()), + n_epochs=0, + )) + + all_results[layer_idx] = results + print(f" [eval] layer {layer_idx}: {len(results)} bins") + + out = Path(results_dir) / output_run_id / probe_name + out.mkdir(parents=True, exist_ok=True) + torch.save(all_results, out / "results.pt") + print(f"[eval] saved to {out / 'results.pt'}") diff --git a/src/tasks/swe_bench_pro.py b/src/tasks/swe_bench_pro.py new file mode 100644 index 0000000..0e70462 --- /dev/null +++ b/src/tasks/swe_bench_pro.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from src.tasks.base import TaskAdapter, ChatPrompt +from src.configs import TaskConfig + +_SYSTEM_TEMPLATE = ( + "You are a software engineer working in a Linux container. " + "The repository is checked out at /app and all dependencies are installed. " + "Fix the reported issue by modifying the source files. " + "Every response must contain exactly one bash command in this format:\n\n" + "```mswea_bash_command\n" + "your_command_here\n" + "```\n\n" + "When you are done, submit by running a command whose first output line is:\n" + "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" +) + + +class SWEBenchProAdapter(TaskAdapter): + def load_dataset(self, task_config: TaskConfig) -> list[dict]: + from datasets import load_dataset as hf_load_dataset + dataset_name = task_config.dataset or "ScaleAI/SWE-bench_Pro" + data = hf_load_dataset(dataset_name, split="test") + return list(data) + + def format_prompt(self, sample: dict) -> ChatPrompt: + problem = sample["problem_statement"] + repo = sample.get("repo", "the repository") + user_content = ( + f"Repository: {repo}\n\n" + f"Issue:\n{problem}\n\n" + "The repository is checked out at /app. " + "Fix the issue by modifying the source files. " + "When done, submit by running a command whose first output line is:\n" + "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" + ) + return ChatPrompt(user_content=user_content, system_content=_SYSTEM_TEMPLATE) + + def check_correct(self, generated: str, sample: dict) -> bool: + return False + + def group_id(self, sample: dict) -> str: + return sample["instance_id"] + + def sample_id(self, sample: dict) -> str: + return sample["instance_id"] diff --git a/tests/test_swebench_pro_evaluate_modal.py b/tests/test_swebench_pro_evaluate_modal.py new file mode 100644 index 0000000..b14a297 --- /dev/null +++ b/tests/test_swebench_pro_evaluate_modal.py @@ -0,0 +1,77 @@ +"""Integration test for SWEBenchProModalEnvironment.evaluate(). + +Validates that the grading logic returns False before the ground-truth patch is +applied and True after. Requires Modal credentials, network access, and the SWE-bench_Pro-os submodule +initialised (git submodule update --init SWE-bench_Pro-os). + +Run from the project root: + uv run python tests/test_swebench_pro_evaluate_modal.py \\ + --scripts-dir SWE-bench_Pro-os \\ + [--instance-id ...] +""" + +import argparse + +# Python instance with a relatively fast test suite. +DEFAULT_INSTANCE_ID = "instance_ansible__ansible-0ea40e09d1b35bcb69ff4d9cecf3d0defa4b36e8-v30a923fb5c164d6cd18280c02422f75e611e8fb2" + + +def _load_instance(instance_id: str) -> dict: + from datasets import load_dataset + data = load_dataset("ScaleAI/SWE-bench_Pro", split="test") + matches = [inst for inst in data if inst["instance_id"] == instance_id] + if not matches: + raise ValueError(f"Instance {instance_id!r} not found in SWE-bench Pro") + return dict(matches[0]) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--instance-id", default=DEFAULT_INSTANCE_ID) + parser.add_argument( + "--scripts-dir", + default="SWE-bench_Pro-os", + help="Path to local clone of github.com/scaleapi/SWE-bench_Pro-os", + ) + parser.add_argument("--modal-app-name", default="program-probes-swebench-pro-eval-test") + parser.add_argument("--modal-timeout", type=int, default=600) + args = parser.parse_args() + + from src.agents.swe_bench_pro_environment import SWEBenchProModalEnvironment + + print(f"Loading instance {args.instance_id!r}...") + instance = _load_instance(args.instance_id) + patch = instance.get("patch", "") + print(f" repo: {instance['repo']}") + print(f" language: {instance.get('repo_language', '?')}") + print(f" patch lines: {len(patch.splitlines())}") + print(f" docker tag: {instance['dockerhub_tag']}") + + env = SWEBenchProModalEnvironment( + instance, + scripts_dir=args.scripts_dir, + timeout=args.modal_timeout, + app_name=args.modal_app_name, + ) + + with env: + print("\n--- smoke test ---") + smoke = env.execute({"command": "cd /app && git log --oneline -1"}) + print(smoke["output"]) + + print("\n--- applying ground-truth patch ---") + apply = env.execute({"command": f"cd /app && git apply - << 'PATCH_EOF'\n{patch}\nPATCH_EOF"}) + print(f" rc={apply['returncode']} {apply['output'][:200]}") + assert apply["returncode"] == 0, f"Patch did not apply cleanly: {apply['output']}" + + print("\n--- evaluate AFTER patch (expect True) ---") + applied_patch = env.get_patch() + result_after, log = env.evaluate(patch=applied_patch) + print(f"evaluate() = {result_after}") + assert result_after is True, f"Expected True after patch, got {result_after}\nLog:\n{log[-2000:]}" + + print("\nAll assertions passed.") + + +if __name__ == "__main__": + main()