Upload smoke_test.py

smoke_test.py

@@ -0,0 +1,371 @@
+"""
+Cascade Smoke Test — Self-Contained HF Job Script
+
+Tests the cascade agent end-to-end on one SWE-bench instance.
+Installs miniconda if needed, sets up environment, runs agent,
+verifies patch.
+
+Usage via hf_jobs:
+  operation: run
+  script: "https://huggingface.co/narcolepticchicken/agent-cost-optimizer/resolve/main/smoke_test.py"
+  dependencies: ["huggingface_hub", "datasets", "trackio"]
+  hardware: a10g-largex2
+  timeout: 4h
+  env:
+    INSTANCE_ID: "django__django-14315"
+"""
+
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+import time
+import traceback
+from datetime import datetime
+from pathlib import Path
+
+# ============================================================
+# Bootstrap: ensure conda is available
+# ============================================================
+
+def ensure_conda():
+    """Install miniconda if not present."""
+    for path in [os.path.expanduser("~/miniconda3/bin/conda"), "/opt/conda/bin/conda"]:
+        if os.path.exists(path):
+            return path
+    
+    print("📦 Installing Miniconda...")
+    rc, out, err = subprocess_run(
+        "wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && "
+        "bash /tmp/miniconda.sh -b -p $HOME/miniconda3",
+        timeout=120
+    )
+    if rc != 0:
+        print(f"Conda install failed: {err}")
+        return None
+    
+    conda_path = os.path.expanduser("~/miniconda3/bin/conda")
+    # Add to PATH for this session
+    os.environ["PATH"] = os.path.expanduser("~/miniconda3/bin:") + os.environ.get("PATH", "")
+    return conda_path
+
+
+def subprocess_run(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 conda_run(conda, subcmd, timeout=300):
+    """Run a conda command with timeout."""
+    return subprocess_run(f"{conda} {subcmd}", timeout=timeout)
+
+
+# ============================================================
+# Agent
+# ============================================================
+
+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(instance, repo_dir, conda, env_name):
+    """Run T1 → T2 cascade agent."""
+    from huggingface_hub import InferenceClient
+    
+    T1 = "meta-llama/Llama-3.1-8B-Instruct"
+    T2 = "meta-llama/Llama-3.3-70B-Instruct"
+    
+    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. 
+Use pytest to verify. Be thorough but efficient."""
+
+    messages = [
+        {"role": "system", "content": system},
+        {"role": "user", "content": f"PROBLEM:\n{problem}\n\nStart by exploring the repository structure."}
+    ]
+    
+    tiers = [("T1", T1, 30), ("T2", T2, 30)]
+    
+    for tier_name, model_id, max_turns in tiers:
+        print(f"\n[{tier_name}] {model_id}")
+        client = 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} tok, {len(text)} ch")
+            
+            patch = extract_patch(text)
+            if patch:
+                print(f"  ✅ PATCH ({len(patch)} ch)")
+                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()
+                cmd = cmd.replace("pytest", f"{conda} run -n {env_name} python -m pytest") if "pytest" in cmd else cmd
+                print(f"  $ {cmd[:120]}")
+                rc, stdout, stderr = subprocess_run(cmd, cwd=str(repo_dir), timeout=60)
+                output = (stdout + stderr)[:1500]
+                if rc != 0:
+                    output += f" [EXIT:{rc}]"
+                messages.append({"role": "user", "content": f"<output>\n{output}\n</output>"})
+            
+            if "<submit>" in text:
+                break
+    
+    return {"patch": None, "tier": None, "turns": 0, "input_tokens": 0, "output_tokens": 0}
+
+
+def verify_patch(instance, model_patch, repo_dir, conda, env_name):
+    """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
+    subprocess_run(f"cd {repo_dir} && git checkout -f {base_commit}", timeout=30)
+    subprocess_run(f"cd {repo_dir} && git clean -fd", timeout=30)
+    
+    # Apply model patch
+    (Path(repo_dir) / "_aco.patch").write_text(model_patch)
+    rc, out, err = subprocess_run(f"cd {repo_dir} && git apply --check _aco.patch", timeout=10)
+    if rc != 0:
+        return {"resolved": False, "error": f"patch --check: {err[:200]}"}
+    rc, out, err = subprocess_run(f"cd {repo_dir} && git apply _aco.patch", timeout=10)
+    if rc != 0:
+        return {"resolved": False, "error": f"patch apply: {err[:200]}"}
+    
+    # Apply test_patch
+    (Path(repo_dir) / "_aco_test.patch").write_text(test_patch)
+    subprocess_run(f"cd {repo_dir} && git apply --check _aco_test.patch && git apply _aco_test.patch || git apply --reject _aco_test.patch", timeout=10)
+    
+    # Run F2P tests
+    f2p_str = ' '.join(f2p[:10])
+    cmd = f"cd {repo_dir} && {conda} run -n {env_name} python -m pytest -v --tb=short -x {f2p_str}"
+    print(f"  F2P: pytest {' '.join(f2p[:3])}...")
+    rc, out, err = subprocess_run(cmd, timeout=300)
+    
+    if rc == 0:
+        p2p = instance.get("PASS_TO_PASS", [])
+        if p2p:
+            p2p_str = ' '.join(p2p[:10])
+            cmd2 = f"cd {repo_dir} && {conda} run -n {env_name} python -m pytest -v --tb=short -x {p2p_str}"
+            rc2, out2, err2 = subprocess_run(cmd2, timeout=300)
+            if rc2 != 0:
+                return {"resolved": False, "error": f"P2P: {(out2+err2)[:200]}"}
+        return {"resolved": True, "test_output": (out + err)[:500]}
+    
+    failures = len(re.findall(r'FAILED', out + err))
+    return {"resolved": False, "error": f"{failures} F2P failures", "test_output": (out + err)[:500]}
+
+
+def setup_environment(conda, instance, repo_dir, env_name):
+    """Create conda environment for the repo."""
+    env_commit = instance.get("environment_setup_commit", "")
+    
+    if env_commit:
+        subprocess_run(f"cd {repo_dir} && git fetch origin {env_commit}", timeout=60)
+        subprocess_run(f"cd {repo_dir} && git checkout {env_commit}", timeout=30)
+    
+    # Find environment.yml
+    env_yml = None
+    for c in ["environment.yml", "dev/environment.yml", ".github/environment.yml",
+               "ci/environment.yml"]:
+        if (Path(repo_dir) / c).exists():
+            env_yml = c
+            break
+    
+    if env_yml:
+        print(f"  Using {env_yml}")
+        rc, out, err = conda_run(conda, f"env create -f {repo_dir}/{env_yml} -n {env_name} --quiet", timeout=600)
+    else:
+        print(f"  Creating basic python=3.10 env")
+        rc, out, err = conda_run(conda, f"create -n {env_name} python=3.10 pip -y --quiet", timeout=300)
+    
+    if rc != 0:
+        print(f"  ⚠️ env creation failed: {err[:200]}")
+        # Try with just pip
+        rc2, out2, err2 = conda_run(conda, f"create -n {env_name} python=3.10 pip -y --quiet", timeout=300)
+        if rc2 != 0:
+            return False, f"conda env: {err[:200]}"
+    
+    # Install repo
+    base_commit = instance["base_commit"]
+    subprocess_run(f"cd {repo_dir} && git fetch origin {base_commit}", timeout=60)
+    subprocess_run(f"cd {repo_dir} && git checkout {base_commit}", timeout=30)
+    rc, out, err = conda_run(conda, f"run -n {env_name} pip install -e . --quiet", 
+                             cwd=str(repo_dir), timeout=300)
+    if rc != 0:
+        print(f"  ⚠️ pip install: {err[:200]}")
+        # Try without -e
+        rc2, out2, err2 = conda_run(conda, f"run -n {env_name} pip install . --quiet",
+                                     cwd=str(repo_dir), timeout=300)
+    
+    return True, ""
+
+
+def main():
+    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"
+    
+    print("="*60)
+    print(f"🚀 CASCADE SMOKE TEST — {datetime.now().isoformat()}")
+    print("="*60)
+    print(f"  Instance: {INSTANCE_ID}")
+    print(f"  T1: {T1_MODEL}")
+    print(f"  T2: {T2_MODEL}")
+    
+    # Ensure conda
+    conda = ensure_conda()
+    if not conda:
+        print("❌ Cannot install conda")
+        sys.exit(1)
+    print(f"\n✅ Conda: {conda}")
+    
+    # Load instance
+    print(f"\n[1/5] Loading instance...")
+    import datasets
+    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_ID} not found!")
+        sys.exit(1)
+    
+    print(f"  Repo: {instance['repo']}")
+    print(f"  Base: {instance['base_commit'][:12]}")
+    print(f"  F2P: {len(instance.get('FAIL_TO_PASS', []))} tests")
+    print(f"  P2P: {len(instance.get('PASS_TO_PASS', []))} tests")
+    
+    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 {instance['repo']}...")
+        t0 = time.time()
+        url = f"https://github.com/{instance['repo']}.git"
+        rc, out, err = subprocess_run(f"git clone --depth 50 {url} {repo_dir}", timeout=300)
+        if rc != 0:
+            rc, out, err = subprocess_run(f"git clone {url} {repo_dir}", timeout=600)
+        clone_t = time.time() - t0
+        if rc != 0:
+            print(f"❌ Clone: {err[:200]}")
+            sys.exit(1)
+        print(f"  Done ({clone_t:.0f}s)")
+        
+        # Environment
+        print(f"\n[3/5] Setting up conda env '{env_name}'...")
+        t0 = time.time()
+        ok, err = setup_environment(conda, instance, repo_dir, env_name)
+        env_t = time.time() - t0
+        if not ok:
+            print(f"❌ Env setup: {err}")
+            sys.exit(1)
+        print(f"  Done ({env_t:.0f}s)")
+        
+        # Agent
+        print(f"\n[4/5] Running cascade agent...")
+        t0 = time.time()
+        agent = run_cascade(instance, repo_dir, conda, env_name)
+        agent_t = time.time() - t0
+        
+        if not agent["patch"]:
+            print(f"❌ No patch!")
+            sys.exit(1)
+        
+        print(f"  ✅ {agent['tier']}, {agent['turns']} turns, {agent['input_tokens']}+{agent['output_tokens']} tokens")
+        
+        # Verify
+        print(f"\n[5/5] Verifying...")
+        t0 = time.time()
+        verify = verify_patch(instance, agent["patch"], repo_dir, conda, env_name)
+        verify_t = time.time() - t0
+        
+        # Result
+        status = "✅ RESOLVED" if verify["resolved"] else "❌ FAILED"
+        print(f"\n{'='*60}")
+        print(f"{status}")
+        print(f"{'='*60}")
+        print(f"  Times: clone={clone_t:.0f}s env={env_t:.0f}s agent={agent_t:.0f}s verify={verify_t:.0f}s")
+        if not verify["resolved"]:
+            print(f"  Error: {verify.get('error', 'unknown')[:300]}")
+        
+        # Save
+        result = {
+            "instance_id": INSTANCE_ID,
+            "repo": instance["repo"],
+            "resolved": verify["resolved"],
+            "tier": agent["tier"],
+            "turns": agent["turns"],
+            "input_tokens": agent["input_tokens"],
+            "output_tokens": agent["output_tokens"],
+            "times": {"clone": clone_t, "env": env_t, "agent": agent_t, "verify": verify_t},
+            "error": verify.get("error"),
+            "patch_preview": agent["patch"][:500],
+        }
+        
+        with open("smoke_result.json", "w") as f:
+            json.dump(result, f, indent=2)
+        print(f"\nSaved: smoke_result.json")
+        
+        # Cleanup
+        conda_run(conda, f"env remove -n {env_name} -y --quiet", timeout=30)
+    
+    return 0 if verify["resolved"] else 1
+
+
+if __name__ == "__main__":
+    try:
+        sys.exit(main())
+    except Exception as e:
+        print(f"💥 {e}")
+        traceback.print_exc()
+        sys.exit(1)