Upload smoke_test_cascade.py

smoke_test_cascade.py

@@ -0,0 +1,407 @@
+"""
+Single-instance smoke test for cascade validation.
+
+This tests the ENTIRE pipeline end-to-end:
+  1. Clone repo
+  2. Set up conda environment
+  3. Run cascade agent (T1 Llama-3.1-8B → T2 Llama-3.3-70B via HF Inference)
+  4. Apply patch + test_patch
+  5. Run FAIL_TO_PASS tests
+
+This is the minimal test to prove the cascade works causally,
+not just correlatively from trace simulation.
+
+Submits to trackio for monitoring.
+
+Usage (hf_jobs):
+  operation: run
+  script: "smoke_test_cascade.py"
+  dependencies: ["huggingface_hub", "datasets", "trackio"]
+  hardware: a10g-largex2 (need GPU for inference + CPU for conda + memory for clone)
+  timeout: 4h
+"""
+
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+import time
+import traceback
+from datetime import datetime
+from pathlib import Path
+
+import datasets
+import huggingface_hub
+import trackio
+
+# ============================================================
+# Trackio setup
+# ============================================================
+trackio.init(
+    project=os.environ.get("TRACKIO_PROJECT", "aco-smoke-test"),
+)
+
+# ============================================================
+# CONFIG
+# ============================================================
+# The easiest SWE-bench instance: django bug with clear fix
+INSTANCE_ID = os.environ.get("INSTANCE_ID", "django__django-14315")
+
+T1_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
+T2_MODEL = "meta-llama/Llama-3.3-70B-Instruct"
+
+trackio.log_params({
+    "instance_id": INSTANCE_ID,
+    "t1_model": T1_MODEL,
+    "t2_model": T2_MODEL,
+    "strategy": "cascade",
+})
+
+
+def run_shell(cmd, cwd=None, timeout=120):
+    result = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True, timeout=timeout, shell=True)
+    return result.returncode, result.stdout, result.stderr
+
+
+def call_model(client, messages, max_tokens=4096):
+    try:
+        completion = client.chat.completions.create(
+            model=client.model,
+            messages=messages,
+            max_tokens=max_tokens,
+            temperature=0.2,
+        )
+        text = completion.choices[0].message.content
+        itok = completion.usage.prompt_tokens if hasattr(completion, 'usage') and completion.usage else 0
+        otok = completion.usage.completion_tokens if hasattr(completion, 'usage') and completion.usage else len(text) // 4
+        return text, itok, otok
+    except Exception as e:
+        return f"[ERROR: {e}]", 0, 0
+
+
+def extract_patch(text):
+    for tag in ['patch', 'diff']:
+        m = re.search(rf'<{tag}>(.*?)</{tag}>', text, re.DOTALL)
+        if m:
+            return m.group(1).strip()
+    for block in ['diff', 'patch']:
+        m = re.search(rf'```{block}\s*\n(.*?)```', text, re.DOTALL)
+        if m:
+            return m.group(1).strip()
+    diff_match = re.search(r'(diff --git a/.*?(?:\n(?:@@|\+\+\+|diff --git|```|</).*)*)', text, re.DOTALL)
+    if diff_match:
+        return diff_match.group(1).strip()
+    return None
+
+
+def run_cascade_agent(instance, repo_dir):
+    """Run T1 then T2. Maximum 30 turns each."""
+    
+    problem = instance.get("problem_statement", "")
+    
+    system = f"""You are fixing a bug in {instance['repo']}. Repository is at {repo_dir}.
+
+Output format:
+- Bash commands: <bash>command</bash>
+- Final patch: <patch>your diff here</patch>
+- Done: <submit>Done</submit>
+
+First explore the codebase to understand the issue, then make a minimal fix and verify it."""
+
+    messages = [
+        {"role": "system", "content": system},
+        {"role": "user", "content": f"PROBLEM:\n{problem}\n\nStart by exploring the repository."}
+    ]
+    
+    tiers = [
+        ("T1", T1_MODEL, 30),
+        ("T2", T2_MODEL, 30),
+    ]
+    
+    for tier_name, model_id, max_turns in tiers:
+        print(f"\n{'='*50}")
+        print(f"[{tier_name}] {model_id}")
+        print(f"{'='*50}")
+        
+        trackio.log({"event": "tier_start", "tier": tier_name, "model": model_id})
+        
+        client = huggingface_hub.InferenceClient(model_id)
+        total_itok = 0
+        total_otok = 0
+        
+        for turn in range(max_turns):
+            text, itok, otok = call_model(client, messages, max_tokens=4096)
+            total_itok += itok
+            total_otok += otok
+            messages.append({"role": "assistant", "content": text})
+            
+            print(f"  Turn {turn+1}: {itok}+{otok} tokens, {len(text)} chars")
+            
+            patch = extract_patch(text)
+            if patch:
+                print(f"  ✅ PATCH FOUND ({len(patch)} chars)")
+                trackio.log({
+                    "event": "patch_found",
+                    "tier": tier_name,
+                    "turn": turn + 1,
+                    "input_tokens": total_itok,
+                    "output_tokens": total_otok,
+                })
+                return {
+                    "patch": patch,
+                    "tier": tier_name,
+                    "turns": turn + 1,
+                    "input_tokens": total_itok,
+                    "output_tokens": total_otok,
+                }
+            
+            cmds = re.findall(r'<bash>(.*?)</bash>', text, re.DOTALL)
+            for cmd in cmds:
+                cmd = cmd.strip()
+                print(f"  $ {cmd[:120]}")
+                rc, stdout, stderr = run_shell(cmd, cwd=str(repo_dir), timeout=30)
+                output = (stdout + stderr)[:1500]
+                if rc != 0:
+                    output += f"\n[EXIT:{rc}]"
+                messages.append({"role": "user", "content": f"<output>\n{output}\n</output>"})
+            
+            if "<submit>" in text:
+                print(f"  [Submit without patch]")
+                break
+        
+        trackio.log({"event": "tier_exhausted", "tier": tier_name, "turns": max_turns})
+    
+    return {"patch": None, "tier": None, "turns": 0, "input_tokens": 0, "output_tokens": 0}
+
+
+def verify_patch(instance, model_patch, repo_dir, env_name=None):
+    """Apply patches, run FAIL_TO_PASS tests."""
+    base_commit = instance.get("base_commit", "")
+    test_patch = instance.get("test_patch", "")
+    f2p = instance.get("FAIL_TO_PASS", [])
+    
+    # Reset repo
+    run_shell(f"cd {repo_dir} && git checkout -f {base_commit}", timeout=30)
+    run_shell(f"cd {repo_dir} && git clean -fd", timeout=30)
+    
+    # Apply model patch
+    patch_file = Path(repo_dir) / "_aco.patch"
+    patch_file.write_text(model_patch)
+    rc, out, err = run_shell(f"cd {repo_dir} && git apply --check _aco.patch", timeout=10)
+    if rc != 0:
+        return {"resolved": False, "error": f"model patch --check: {err[:300]}"}
+    rc, out, err = run_shell(f"cd {repo_dir} && git apply _aco.patch", timeout=10)
+    if rc != 0:
+        return {"resolved": False, "error": f"model patch apply: {err[:300]}"}
+    
+    # Apply test_patch
+    test_file = Path(repo_dir) / "_aco_test.patch"
+    test_file.write_text(test_patch)
+    rc, out, err = run_shell(f"cd {repo_dir} && git apply --check _aco_test.patch", timeout=10)
+    if rc == 0:
+        run_shell(f"cd {repo_dir} && git apply _aco_test.patch", timeout=10)
+    else:
+        run_shell(f"cd {repo_dir} && git apply --reject _aco_test.patch 2>/dev/null; true", timeout=10)
+    
+    # Run F2P tests
+    cmd_prefix = f"conda run -n {env_name} " if env_name else ""
+    f2p_cmd = f"{cmd_prefix}python -m pytest -v --tb=short -x {' '.join(f2p[:10])}"
+    print(f"  F2P: {f2p_cmd[:150]}...")
+    rc, out, err = run_shell(f"cd {repo_dir} && {f2p_cmd}", timeout=300)
+    
+    if rc == 0:
+        # Run P2P regression tests
+        p2p = instance.get("PASS_TO_PASS", [])
+        if p2p:
+            p2p_cmd = f"{cmd_prefix}python -m pytest -v --tb=short -x {' '.join(p2p[:10])}"
+            rc2, out2, err2 = run_shell(f"cd {repo_dir} && {p2p_cmd}", timeout=300)
+            if rc2 != 0:
+                return {"resolved": False, "error": f"P2P regression: {(out2+err2)[:300]}"}
+        return {"resolved": True, "test_output": (out + err)[:500]}
+    
+    return {"resolved": False, "error": f"F2P: {len(re.findall(r'FAILED', out+err))} failures", "test_output": (out + err)[:500]}
+
+
+def main():
+    print(f"🚀 CASCADE SMOKE TEST")
+    print(f"   Instance: {INSTANCE_ID}")
+    print(f"   T1: {T1_MODEL}")
+    print(f"   T2: {T2_MODEL}")
+    print(f"   Time: {datetime.now().isoformat()}")
+    
+    # Load instance
+    print("\n[1/5] Loading SWE-bench_Verified...")
+    ds = datasets.load_dataset("princeton-nlp/SWE-bench_Verified", split="test")
+    instance = None
+    for row in ds:
+        if row["instance_id"] == INSTANCE_ID:
+            instance = dict(row)
+            break
+    
+    if not instance:
+        print(f"❌ Instance {INSTANCE_ID} not found!")
+        trackio.alert("Instance not found", f"{INSTANCE_ID} not in SWE-bench_Verified", level="ERROR")
+        sys.exit(1)
+    
+    print(f"   Repo: {instance['repo']}")
+    print(f"   Base: {instance['base_commit'][:12]}")
+    print(f"   F2P tests: {len(instance.get('FAIL_TO_PASS', []))}")
+    print(f"   P2P tests: {len(instance.get('PASS_TO_PASS', []))}")
+    
+    trackio.log({
+        "event": "instance_loaded",
+        "repo": instance["repo"],
+        "f2p_count": len(instance.get("FAIL_TO_PASS", [])),
+        "p2p_count": len(instance.get("PASS_TO_PASS", [])),
+    })
+    
+    with tempfile.TemporaryDirectory(prefix="aco_smoke_") as tmpdir:
+        repo_dir = Path(tmpdir) / "repo"
+        env_name = f"aco_{INSTANCE_ID.replace('__','_').replace('-','_')[:30]}"
+        
+        # Clone
+        print(f"\n[2/5] Cloning repo...")
+        t0 = time.time()
+        repo = instance["repo"]
+        url = f"https://github.com/{repo}.git"
+        rc, out, err = run_shell(f"git clone --depth 50 {url} {repo_dir}", timeout=300)
+        if rc != 0:
+            rc, out, err = run_shell(f"git clone {url} {repo_dir}", timeout=600)
+        clone_time = time.time() - t0
+        
+        if rc != 0:
+            print(f"❌ Clone failed: {err[:300]}")
+            trackio.alert("Clone failed", err[:200], level="ERROR")
+            sys.exit(1)
+        print(f"   Done ({clone_time:.0f}s)")
+        
+        # Set up conda env
+        print(f"\n[3/5] Setting up conda environment...")
+        t0 = time.time()
+        
+        env_commit = instance.get("environment_setup_commit", "")
+        if env_commit:
+            run_shell(f"cd {repo_dir} && git fetch origin {env_commit}", timeout=60)
+            run_shell(f"cd {repo_dir} && git checkout {env_commit}", timeout=30)
+        
+        env_yml = None
+        for c in ["environment.yml", "dev/environment.yml", ".github/environment.yml",
+                   "ci/environment.yml"]:
+            if (repo_dir / c).exists():
+                env_yml = c
+                break
+        
+        if env_yml:
+            print(f"   Found: {env_yml}")
+            rc, out, err = run_shell(f"cd {repo_dir} && conda env create -f {env_yml} -n {env_name} --quiet", timeout=600)
+        else:
+            print(f"   No environment.yml, creating basic env")
+            rc, out, err = run_shell(f"conda create -n {env_name} python=3.10 pip -y --quiet", timeout=300)
+        
+        if rc != 0:
+            print(f"❌ Conda env creation failed: {err[:300]}")
+            trackio.alert("Conda env failed", err[:300], level="ERROR")
+            sys.exit(1)
+        
+        # Install repo at base_commit
+        base_commit = instance["base_commit"]
+        run_shell(f"cd {repo_dir} && git fetch origin {base_commit}", timeout=60)
+        run_shell(f"cd {repo_dir} && git checkout {base_commit}", timeout=30)
+        rc, out, err = run_shell(f"cd {repo_dir} && conda run -n {env_name} pip install -e . --quiet", timeout=300)
+        if rc != 0:
+            print(f"⚠️ pip install had issues: {err[:200]}")
+        
+        env_time = time.time() - t0
+        print(f"   Done ({env_time:.0f}s)")
+        
+        # Run cascade agent
+        print(f"\n[4/5] Running cascade agent (T1→T2)...")
+        t0 = time.time()
+        agent_result = run_cascade_agent(instance, repo_dir)
+        agent_time = time.time() - t0
+        
+        if not agent_result["patch"]:
+            print(f"❌ No patch produced")
+            trackio.alert("No patch", "Cascade produced no patch", level="ERROR")
+            sys.exit(1)
+        
+        print(f"\n   ✅ Patch: {len(agent_result['patch'])} chars")
+        print(f"   Tier: {agent_result['tier']}")
+        print(f"   Turns: {agent_result['turns']}")
+        print(f"   Tokens: {agent_result['input_tokens']} in + {agent_result['output_tokens']} out")
+        print(f"   Time: {agent_time:.0f}s")
+        
+        trackio.log({
+            "event": "agent_done",
+            "tier": agent_result["tier"],
+            "turns": agent_result["turns"],
+            "input_tokens": agent_result["input_tokens"],
+            "output_tokens": agent_result["output_tokens"],
+            "agent_time": agent_time,
+        })
+        
+        # Verify
+        print(f"\n[5/5] Verifying patch...")
+        t0 = time.time()
+        verify_result = verify_patch(instance, agent_result["patch"], repo_dir, env_name)
+        verify_time = time.time() - t0
+        
+        trackio.log({
+            "event": "verification_done",
+            "resolved": verify_result["resolved"],
+            "verify_time": verify_time,
+        })
+        
+        # Final result
+        print(f"\n{'='*60}")
+        print(f"{'✅ RESOLVED' if verify_result['resolved'] else '❌ FAILED'}")
+        print(f"{'='*60}")
+        print(f"  Instance: {INSTANCE_ID}")
+        print(f"  Tier: {agent_result['tier']}")
+        print(f"  Turns: {agent_result['turns']}")
+        print(f"  Time: clone={clone_time:.0f}s env={env_time:.0f}s agent={agent_time:.0f}s verify={verify_time:.0f}s")
+        
+        if not verify_result["resolved"]:
+            print(f"  Error: {verify_result.get('error', 'unknown')[:300]}")
+            trackio.alert("Not resolved", verify_result.get('error', 'unknown')[:300], level="WARN")
+        else:
+            trackio.alert("✅ Resolved!", f"Instance {INSTANCE_ID} resolved via {agent_result['tier']}", level="INFO")
+        
+        # Save final result
+        final = {
+            "instance_id": INSTANCE_ID,
+            "repo": instance["repo"],
+            "timestamp": datetime.now().isoformat(),
+            "resolved": verify_result["resolved"],
+            "tier": agent_result["tier"],
+            "turns": agent_result["turns"],
+            "input_tokens": agent_result["input_tokens"],
+            "output_tokens": agent_result["output_tokens"],
+            "clone_time": clone_time,
+            "env_time": env_time,
+            "agent_time": agent_time,
+            "verify_time": verify_time,
+            "patch_preview": agent_result["patch"][:500],
+            "error": verify_result.get("error"),
+        }
+        
+        with open("smoke_result.json", "w") as f:
+            json.dump(final, f, indent=2)
+        print(f"\nSaved: smoke_result.json")
+        
+        # Cleanup
+        run_shell(f"conda env remove -n {env_name} -y --quiet", timeout=30)
+    
+    return 0 if verify_result["resolved"] else 1
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except Exception as e:
+        print(f"💥 CRASH: {e}")
+        traceback.print_exc()
+        trackio.alert("Crash", str(e)[:300], level="ERROR")
+        sys.exit(1)