chore: import upstream snapshot with attribution

scripts/benchmark_takehome.py

@@ -0,0 +1,408 @@
+#!/usr/bin/env python3
+"""
+Benchmark single agent vs swarm on the Anthropic Performance Take-Home.
+
+Usage:
+    BENCHMARK_TIMEOUT=5 python scripts/benchmark_takehome.py single
+    BENCHMARK_TIMEOUT=10 python scripts/benchmark_takehome.py swarm
+    BENCHMARK_TIMEOUT=10 python scripts/benchmark_takehome.py both
+"""
+
+import socket
+import json
+import os
+import sys
+import time
+import select
+import shutil
+import subprocess
+import threading
+from pathlib import Path
+
+DEBUG_SOCKET = f"/run/user/{os.getuid()}/jcode-debug.sock"
+TAKEHOME_SOURCE = os.environ.get(
+    "TAKEHOME_SOURCE", str(Path.home() / "original_performance_takehome")
+)
+BENCHMARK_DIR = "/tmp/takehome-benchmark"
+TIMEOUT_MINUTES = int(os.environ.get('BENCHMARK_TIMEOUT', '10'))
+BASELINE = 147734
+
+
+def send_cmd(cmd: str, session_id: str = None, timeout: float = 300) -> tuple:
+    """Send a debug command and get response."""
+    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    sock.connect(DEBUG_SOCKET)
+    sock.setblocking(False)
+
+    req = {"type": "debug_command", "id": 1, "command": cmd}
+    if session_id:
+        req["session_id"] = session_id
+
+    sock.send((json.dumps(req) + '\n').encode())
+
+    start = time.time()
+    data = b""
+    while time.time() - start < timeout:
+        ready, _, _ = select.select([sock], [], [], 1.0)
+        if ready:
+            try:
+                chunk = sock.recv(65536)
+                if not chunk:
+                    break
+                data += chunk
+                if b'\n' in data:
+                    break
+            except BlockingIOError:
+                continue
+
+    sock.close()
+
+    if not data:
+        return False, "", "Timeout"
+
+    try:
+        resp = json.loads(data.decode().strip())
+        return resp.get('ok', False), resp.get('output', ''), resp.get('error', '')
+    except json.JSONDecodeError as e:
+        return False, "", f"JSON error: {e}"
+
+
+def create_session(working_dir: str) -> tuple:
+    """Create a session and return (session_id, friendly_name)."""
+    ok, output, err = send_cmd(f"create_session:{working_dir}", timeout=120)
+    if not ok:
+        raise RuntimeError(f"Failed to create session: {err}")
+    data = json.loads(output)
+    return data['session_id'], data.get('friendly_name', data['session_id'][:12])
+
+
+def destroy_session(session_id: str):
+    """Destroy a session."""
+    send_cmd(f"destroy_session:{session_id}")
+
+
+def setup_workspace(name: str) -> str:
+    """Create a clean copy of the take-home."""
+    workspace = os.path.join(BENCHMARK_DIR, name)
+    if os.path.exists(workspace):
+        shutil.rmtree(workspace)
+    shutil.copytree(TAKEHOME_SOURCE, workspace)
+    return workspace
+
+
+def get_cycles(workspace: str) -> int:
+    """Run tests and return cycle count."""
+    try:
+        result = subprocess.run(
+            ["python", "tests/submission_tests.py", "-v"],
+            cwd=workspace,
+            capture_output=True,
+            text=True,
+            timeout=120
+        )
+        for line in (result.stdout + result.stderr).split('\n'):
+            if 'CYCLES:' in line:
+                return int(line.split('CYCLES:')[1].strip())
+    except Exception as e:
+        print(f"Error getting cycles: {e}")
+    return BASELINE
+
+
+def make_single_prompt(workspace: str) -> str:
+    return f"""You are optimizing a VLIW SIMD kernel for Anthropic's performance take-home.
+
+IMPORTANT: You MUST work in this directory: {workspace}
+All file paths should be relative to or within this directory.
+
+Goal: Reduce the cycle count from 147,734 to as low as possible.
+
+Key files (in {workspace}):
+- problem.py: Defines the Machine, instruction set (VLIW bundles, vector ops, VLEN=16)
+- perf_takehome.py: Contains KernelBuilder.build_kernel() - this is what you optimize
+- tests/submission_tests.py: Run to verify correctness and see cycle count
+
+DO NOT modify tests/ folder.
+
+Key optimizations to try:
+1. VLIW parallelism - pack independent operations into single bundles
+2. Vector operations - use VLEN=16 to process 16 elements at once
+3. Reduce memory access latency - batch loads/stores
+4. Optimize the hash function - it runs many times per element
+
+Start by reading {workspace}/problem.py to understand the machine, then optimize build_kernel().
+After each change, run `cd {workspace} && python tests/submission_tests.py` to check correctness and cycles.
+
+Work efficiently - focus on the highest-impact optimizations first."""
+
+
+def run_single_agent() -> dict:
+    """Run a single agent benchmark using async messaging."""
+    print("\n" + "=" * 60)
+    print(f"SINGLE AGENT BENCHMARK (timeout: {TIMEOUT_MINUTES}m)")
+    print("=" * 60)
+
+    workspace = setup_workspace("single")
+    print(f"Workspace: {workspace}")
+
+    start_time = time.time()
+    session_id = None
+    best_cycles = BASELINE
+
+    try:
+        session_id, name = create_session(workspace)
+        print(f"Session: {name}")
+
+        # Initial cycles
+        cycles = get_cycles(workspace)
+        print(f"Baseline: {cycles} cycles")
+
+        # Send optimization task asynchronously
+        print("\nStarting optimization (async)...")
+        prompt = make_single_prompt(workspace)
+
+        # Use message_async to start the job
+        ok, output, err = send_cmd(f"message_async:{prompt}", session_id, timeout=30)
+        if not ok:
+            print(f"Failed to start async job: {err}")
+            return {"approach": "single", "cycles": BASELINE, "time_seconds": 0, "error": err}
+
+        job_data = json.loads(output)
+        job_id = job_data.get("job_id")
+        print(f"Job started: {job_id}")
+
+        timeout_seconds = TIMEOUT_MINUTES * 60
+        last_cycles = BASELINE
+        check_interval = 30  # Check every 30 seconds
+
+        while time.time() - start_time < timeout_seconds:
+            elapsed = time.time() - start_time
+
+            # Check job status
+            ok, status_output, _ = send_cmd(f"job_status:{job_id}", session_id, timeout=10)
+            if ok:
+                try:
+                    status = json.loads(status_output)
+                    job_status = status.get("status", "unknown")
+                    if job_status in ["completed", "failed"]:
+                        print(f"\n[{elapsed/60:.1f}m] Job {job_status}")
+                        break
+                except:
+                    pass
+
+            # Check current cycles in workspace
+            cycles = get_cycles(workspace)
+            if cycles < best_cycles:
+                best_cycles = cycles
+                print(f"[{elapsed/60:.1f}m] NEW BEST: {cycles} cycles ({BASELINE/cycles:.2f}x speedup)")
+            elif cycles != last_cycles:
+                print(f"[{elapsed/60:.1f}m] Cycles: {cycles}")
+            last_cycles = cycles
+
+            time.sleep(check_interval)
+
+        # Final check
+        cycles = get_cycles(workspace)
+        if cycles < best_cycles:
+            best_cycles = cycles
+
+        elapsed = time.time() - start_time
+        print(f"\nFinal: {best_cycles} cycles in {elapsed/60:.1f}m ({BASELINE/best_cycles:.2f}x)")
+
+        return {
+            "approach": "single",
+            "cycles": best_cycles,
+            "time_seconds": elapsed,
+            "workspace": workspace
+        }
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return {
+            "approach": "single",
+            "cycles": BASELINE,
+            "time_seconds": time.time() - start_time,
+            "error": str(e)
+        }
+    finally:
+        if session_id:
+            destroy_session(session_id)
+
+
+def run_swarm(n_agents: int = 2) -> dict:
+    """Run autonomous swarm benchmark using swarm_message_async.
+
+    This uses the full swarm capability where ONE agent becomes coordinator,
+    creates a plan, and spawns subagents automatically.
+    """
+    print("\n" + "=" * 60)
+    print(f"AUTONOMOUS SWARM BENCHMARK (timeout: {TIMEOUT_MINUTES}m)")
+    print("=" * 60)
+
+    workspace = setup_workspace("swarm")
+    print(f"Workspace: {workspace}")
+
+    start_time = time.time()
+    session_id = None
+    best_cycles = BASELINE
+
+    try:
+        # Create ONE session - it becomes coordinator and spawns agents
+        session_id, name = create_session(workspace)
+        print(f"Coordinator: {name}")
+
+        baseline = get_cycles(workspace)
+        print(f"Baseline: {baseline} cycles")
+
+        # Use swarm_message_async - this will:
+        # 1. Plan subtasks automatically
+        # 2. Spawn subagents to work in parallel
+        # 3. Integrate results
+        prompt = f"""Optimize the VLIW SIMD kernel in {workspace}/perf_takehome.py to minimize cycle count.
+
+Current baseline: 147,734 cycles. Goal: as low as possible.
+
+The problem:
+- {workspace}/problem.py defines the machine (VLEN=16 vectors, VLIW bundles, slot limits)
+- {workspace}/perf_takehome.py has build_kernel() which needs optimization
+- Run `cd {workspace} && python tests/submission_tests.py` to verify correctness and check cycles
+
+Key optimization strategies:
+1. Vectorization - use VLEN=16 to process 16 elements at once
+2. VLIW packing - bundle independent operations together
+3. Reduce memory latency - batch loads/stores
+4. Optimize hash function - it runs many times per element
+
+Break this into parallel subtasks and spawn agents to work on different optimizations.
+DO NOT modify tests/ folder."""
+
+        print("\nStarting autonomous swarm (swarm_message_async)...")
+        ok, output, err = send_cmd(f"swarm_message_async:{prompt}", session_id, timeout=30)
+        if not ok:
+            print(f"Failed to start swarm: {err}")
+            return {"approach": "swarm", "cycles": BASELINE, "time_seconds": 0, "error": err}
+
+        job_data = json.loads(output)
+        job_id = job_data.get("job_id")
+        print(f"Swarm job started: {job_id}")
+
+        timeout_seconds = TIMEOUT_MINUTES * 60
+        last_cycles = BASELINE
+        check_interval = 30
+
+        while time.time() - start_time < timeout_seconds:
+            elapsed = time.time() - start_time
+
+            # Check job status
+            ok, status_output, _ = send_cmd(f"job_status:{job_id}", session_id, timeout=10)
+            if ok:
+                try:
+                    status = json.loads(status_output)
+                    job_status = status.get("status", "unknown")
+                    if job_status == "completed":
+                        print(f"\n[{elapsed/60:.1f}m] Swarm completed!")
+                        break
+                    elif job_status == "failed":
+                        print(f"\n[{elapsed/60:.1f}m] Swarm failed: {status.get('error', 'unknown')}")
+                        break
+                except:
+                    pass
+
+            # Check swarm members (to see how many agents were spawned)
+            ok, swarm_output, _ = send_cmd("swarm:members", session_id, timeout=10)
+            if ok and elapsed < 60:  # Only print once early on
+                try:
+                    print(f"[{elapsed/60:.1f}m] Swarm: {swarm_output[:100]}...")
+                except:
+                    pass
+
+            # Check current cycles
+            cycles = get_cycles(workspace)
+            if cycles < best_cycles:
+                best_cycles = cycles
+                print(f"[{elapsed/60:.1f}m] NEW BEST: {cycles} cycles ({BASELINE/cycles:.2f}x)")
+            elif cycles != last_cycles:
+                print(f"[{elapsed/60:.1f}m] Cycles: {cycles}")
+            last_cycles = cycles
+
+            time.sleep(check_interval)
+
+        # Final check
+        cycles = get_cycles(workspace)
+        if cycles < best_cycles:
+            best_cycles = cycles
+
+        elapsed = time.time() - start_time
+        print(f"\nFinal: {best_cycles} cycles in {elapsed/60:.1f}m ({BASELINE/best_cycles:.2f}x)")
+
+        return {
+            "approach": "swarm",
+            "cycles": best_cycles,
+            "time_seconds": elapsed,
+            "workspace": workspace
+        }
+
+    except Exception as e:
+        print(f"Error: {e}")
+        return {
+            "approach": "swarm",
+            "cycles": BASELINE,
+            "time_seconds": time.time() - start_time,
+            "error": str(e)
+        }
+    finally:
+        if session_id:
+            destroy_session(session_id)
+
+
+def print_results(results: dict):
+    """Print comparison table."""
+    print("\n" + "=" * 60)
+    print("RESULTS")
+    print("=" * 60)
+    print(f"{'Approach':<15} {'Cycles':<12} {'Time':<10} {'Speedup':<10}")
+    print("-" * 60)
+
+    for name, data in results.items():
+        cycles = data['cycles']
+        time_m = data['time_seconds'] / 60
+        speedup = BASELINE / cycles
+        print(f"{name:<15} {cycles:<12} {time_m:<10.1f}m {speedup:<10.2f}x")
+
+    if len(results) > 1:
+        winner = min(results.items(), key=lambda x: x[1]['cycles'])
+        print(f"\nWinner: {winner[0]} ({winner[1]['cycles']} cycles)")
+
+
+def main():
+    if len(sys.argv) < 2:
+        print(__doc__)
+        sys.exit(1)
+
+    mode = sys.argv[1].lower()
+    os.makedirs(BENCHMARK_DIR, exist_ok=True)
+
+    print(f"Benchmark timeout: {TIMEOUT_MINUTES} minutes per approach")
+    print(f"Set BENCHMARK_TIMEOUT env var to change (e.g., BENCHMARK_TIMEOUT=30)")
+
+    if mode == "single":
+        r = run_single_agent()
+        print_results({"single": r})
+
+    elif mode == "swarm":
+        r = run_swarm()
+        print_results({"swarm": r})
+
+    elif mode == "both":
+        results = {}
+        results['single'] = run_single_agent()
+        results['swarm'] = run_swarm()
+        print_results(results)
+
+    else:
+        print(f"Unknown mode: {mode}")
+        print(__doc__)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()