run_experiment.py

#!/usr/bin/env python3
"""
Run a complete experiment: agent + tests for selected benchmarks.
Reads experiment config from YAML file, orchestrates runs, saves results.
"""

import yaml
import json
import shutil
import sys
import argparse
from concurrent.futures import ProcessPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

from run_agent import run_agent_for_benchmark, process_agent_config
from run_benchmarks import run_tests_for_benchmark
from generate_prompts import generate_prompt, load_benchmark_config


def load_experiment_config(config_path: Path) -> dict:
    """Load and validate experiment configuration."""
    if not config_path.exists():
        print(f"Error: Experiment config not found: {config_path}")
        sys.exit(1)

    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Validate required fields
    required_fields = ['name', 'agent', 'benchmarks']
    for field in required_fields:
        if field not in config:
            print(f"Error: Missing required field '{field}' in experiment config")
            sys.exit(1)

    # Set defaults
    config.setdefault('timeout', 3600)
    config.setdefault('description', '')
    config.setdefault('parallel', 1)
    config.setdefault('template_version', None)  # None = use existing prompt.md
    config.setdefault('iterations', 1)  # For Ralph Wiggum experiments
    config.setdefault('test_each_iteration', False)  # Test after each iteration vs once at end

    return config


def load_agents_config(agents_path: Path) -> dict:
    """Load agent configurations from agents.yml."""
    if not agents_path.exists():
        print(f"Error: agents.yml not found: {agents_path}")
        sys.exit(1)

    with open(agents_path) as f:
        return yaml.safe_load(f)


def validate_experiment(config: dict, agents: dict, dataset_dir: Path) -> bool:
    """Validate experiment config against available agents and benchmarks."""
    # Check agent exists
    if config['agent'] not in agents:
        print(f"Error: Unknown agent '{config['agent']}'")
        print(f"Available agents: {', '.join(agents.keys())}")
        return False

    # Check benchmarks exist
    for benchmark in config['benchmarks']:
        benchmark_dir = dataset_dir / benchmark
        if not benchmark_dir.exists():
            print(f"Error: Unknown benchmark '{benchmark}'")
            print(f"Available benchmarks: {', '.join(d.name for d in dataset_dir.iterdir() if d.is_dir())}")
            return False

    return True


def create_experiment_dir(base_dir: Path, experiment_name: str) -> Path:
    """Create timestamped experiment output directory."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    exp_dir = base_dir / f"{timestamp}_{experiment_name}"
    exp_dir.mkdir(parents=True, exist_ok=True)
    return exp_dir


def save_config_snapshot(exp_dir: Path, experiment_config: dict, agents_config: dict):
    """Save copies of configs used for reproducibility."""
    with open(exp_dir / "experiment.yml", 'w') as f:
        yaml.dump(experiment_config, f, default_flow_style=False)

    with open(exp_dir / "agents.yml", 'w') as f:
        yaml.dump(agents_config, f, default_flow_style=False)


def get_agent_config(agents: dict, agent_name: str, model_override: Optional[str] = None) -> dict:
    """Get agent config, applying model override if specified, then process templates."""
    agent_config = agents[agent_name].copy()

    if model_override:
        agent_config['model'] = model_override

    # Process template placeholders (cli_version, model, api_key_env)
    return process_agent_config(agent_config)


def initialize_workspace(benchmark_dir: Path, output_dir: Path, benchmark_name: str) -> Path:
    """
    Initialize workspace from scratch by cloning src and creating dst with git.
    Returns path to the initialized workspace.
    """
    import subprocess

    dest_workspace = output_dir / "workspace"
    dest_workspace.mkdir(parents=True, exist_ok=True)

    # Load metadata for provenance info
    metadata_file = benchmark_dir / "metadata.yml"
    source_repo = None
    source_commit = None
    source_subdir = None
    if metadata_file.exists():
        with open(metadata_file) as f:
            metadata = yaml.safe_load(f)
            provenance = metadata.get("provenance", {})
            source_repo = provenance.get("source_repo")
            source_commit = provenance.get("source_commit")
            source_subdir = provenance.get("source_subdir")  # Optional subdirectory

    # Clone src/ from remote
    temp_src = dest_workspace / "src"
    if source_repo and source_commit:
        print(f"  [{benchmark_name}] Cloning {source_repo}...")
        result = subprocess.run(
            ["git", "clone", "--recurse-submodules", source_repo, str(temp_src)],
            capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"  [{benchmark_name}] Clone failed: {result.stderr}")
            raise RuntimeError(f"Failed to clone {source_repo}")

        # Checkout specific commit
        result = subprocess.run(
            ["git", "checkout", source_commit],
            cwd=temp_src, capture_output=True, text=True
        )
        if result.returncode != 0:
            print(f"  [{benchmark_name}] Checkout failed: {result.stderr}")
            raise RuntimeError(f"Failed to checkout {source_commit}")

        # Initialize submodules at this commit
        subprocess.run(
            ["git", "submodule", "update", "--init", "--recursive"],
            cwd=temp_src, capture_output=True
        )
        print(f"  [{benchmark_name}] Cloned at {source_commit[:8]}")
    else:
        # Fallback: copy from workspace/src if it exists
        workspace_src = benchmark_dir / "workspace" / "src"
        if workspace_src.exists():
            shutil.copytree(workspace_src, temp_src)
            print(f"  [{benchmark_name}] Copied workspace/src (no provenance)")
        else:
            raise RuntimeError(f"No source_repo in metadata and no workspace/src")

    # Create dst/ with initialized git repo
    dst_dir = dest_workspace / "dst"
    dst_dir.mkdir(parents=True, exist_ok=True)
    subprocess.run(["git", "init"], cwd=dst_dir, capture_output=True)
    subprocess.run(
        ["git", "config", "user.email", "agent@benchmark.local"],
        cwd=dst_dir, capture_output=True
    )
    subprocess.run(
        ["git", "config", "user.name", "Benchmark Agent"],
        cwd=dst_dir, capture_output=True
    )

    # Generate prompt.md from benchmark.yml
    benchmark_config = load_benchmark_config(benchmark_dir)
    if benchmark_config:
        prompt_content = generate_prompt(benchmark_config, "v1")
        with open(dest_workspace / "prompt.md", 'w') as f:
            f.write(prompt_content)
    else:
        print(f"  [{benchmark_name}] Warning: Could not generate prompt.md")

    return dest_workspace


def run_single_benchmark(
    benchmark_name: str,
    dataset_dir: Path,
    exp_dir: Path,
    agent_config: dict,
    timeout: int,
    template_version: str = None,
    iterations: int = 1,
    test_each_iteration: bool = False
) -> dict:
    """
    Run agent and tests for a single benchmark.
    This function is designed to be called in parallel.
    Workspace is copied to experiment folder for isolation.

    Args:
        template_version: Template version for prompt generation (None = use existing prompt.md)
        iterations: Number of iterations for Ralph Wiggum experiments (default: 1)
        test_each_iteration: If True, run tests after each iteration (for accuracy curves).
                            If False (default), run all iterations then test once at end.
    """
    print(f"\n{'=' * 70}")
    print(f"[START] Benchmark: {benchmark_name}")
    print("=" * 70)

    # Create benchmark output directory
    benchmark_output_dir = exp_dir / benchmark_name
    benchmark_output_dir.mkdir(exist_ok=True)

    # Get benchmark source directory
    benchmark_src_dir = dataset_dir / benchmark_name

    # Initialize workspace from scratch (clone src, create git-enabled dst)
    workspace_dir = initialize_workspace(benchmark_src_dir, benchmark_output_dir, benchmark_name)

    # Initialize result record
    result = {
        "agent": agent_config['name'],
        "cli_version": agent_config.get('cli_version', 'unknown'),
        "model": agent_config.get('model', 'unknown'),
        "benchmark": benchmark_name,
        "template_version": template_version or "existing",
        "iterations": iterations,
        "test_each_iteration": test_each_iteration,
    }

    if test_each_iteration:
        # MODE 2: Test after each iteration (for accuracy curves)
        iteration_results = []
        total_time = 0

        # Clear the agent log file before starting iterations (so append mode works correctly)
        agent_log_file = benchmark_output_dir / "agent_log.jsonl"
        if agent_log_file.exists():
            agent_log_file.unlink()

        for i in range(1, iterations + 1):
            print(f"  === Iteration {i}/{iterations} ===")

            # Run agent for single iteration (fresh start each time)
            agent_result = run_agent_for_benchmark(
                benchmark_name=benchmark_name,
                benchmark_dir=benchmark_src_dir,
                agent_config=agent_config,
                output_dir=benchmark_output_dir,
                timeout=timeout,
                workspace_dir=workspace_dir,
                template_version=template_version,
                iterations=1  # Single iteration at a time
            )

            iter_time = agent_result.get("elapsed_seconds", 0)
            total_time += iter_time

            # Run tests
            test_result = run_tests_for_benchmark(
                benchmark_name=benchmark_name,
                benchmark_dir=benchmark_src_dir,
                output_dir=benchmark_output_dir,
                test_target="dst",
                workspace_dir=workspace_dir
            )

            # CRITICAL: Delete tests after each run to keep them hidden from agent
            tests_dir = workspace_dir / "tests"
            if tests_dir.exists():
                shutil.rmtree(tests_dir)

            # Record iteration result
            iter_result = {
                "iteration": i,
                "agent_time": iter_time,
                "build_success": test_result.get("build_success", False),
                "tests_passed": test_result.get("tests_passed", 0),
                "tests_total": test_result.get("tests_total", 0),
                "pass_rate": test_result.get("pass_rate", 0.0)
            }
            iteration_results.append(iter_result)

            # Print iteration summary
            build_status = "✅" if iter_result["build_success"] else "❌"
            print(f"  Iteration {i}: {build_status} {iter_result['tests_passed']}/{iter_result['tests_total']} ({iter_result['pass_rate']:.1%}) in {iter_time:.1f}s")

        # Store iteration results
        result["iteration_results"] = iteration_results
        result["time_seconds"] = total_time
        result["iterations_completed"] = iterations

        # Final results from last iteration
        final = iteration_results[-1] if iteration_results else {}
        result["build_success"] = final.get("build_success", False)
        result["tests_passed"] = final.get("tests_passed", 0)
        result["tests_total"] = final.get("tests_total", 0)
        result["pass_rate"] = final.get("pass_rate", 0.0)

        # Print benchmark summary
        build_status = "✅" if result["build_success"] else "❌"
        print(f"\n[DONE] {benchmark_name}: Final {build_status} {result['tests_passed']}/{result['tests_total']} ({result['pass_rate']:.1%}) in {total_time:.1f}s")

    else:
        # MODE 1: Run all iterations, then test once at end (existing behavior)
        iterations_str = f" x{iterations} iterations" if iterations > 1 else ""
        print(f"  [{benchmark_name}] Running Agent{iterations_str}...")
        agent_result = run_agent_for_benchmark(
            benchmark_name=benchmark_name,
            benchmark_dir=benchmark_src_dir,
            agent_config=agent_config,
            output_dir=benchmark_output_dir,
            timeout=timeout,
            workspace_dir=workspace_dir,
            template_version=template_version,
            iterations=iterations
        )

        result["time_seconds"] = agent_result.get("elapsed_seconds", 0)
        result["commits"] = agent_result.get("commits", 0)
        result["termination_reason"] = agent_result.get("termination_reason", "unknown")
        result["iterations_completed"] = agent_result.get("iterations_completed", 1)

        # Run tests (uses isolated workspace)
        print(f"  [{benchmark_name}] Running Tests...")
        test_result = run_tests_for_benchmark(
            benchmark_name=benchmark_name,
            benchmark_dir=benchmark_src_dir,
            output_dir=benchmark_output_dir,
            test_target="dst",
            workspace_dir=workspace_dir
        )

        result["build_success"] = test_result.get("build_success", False)
        result["tests_passed"] = test_result.get("tests_passed", 0)
        result["tests_total"] = test_result.get("tests_total", 0)
        result["pass_rate"] = test_result.get("pass_rate", 0.0)

        # Print benchmark summary
        build_status = "✅" if result["build_success"] else "❌"
        print(f"\n[DONE] {benchmark_name}: {build_status} {result['tests_passed']}/{result['tests_total']} ({result['pass_rate']:.1%}) in {result['time_seconds']}s")

    return result


def run_experiment(config_path: Path):
    """Main experiment runner."""
    # Setup paths
    root_dir = Path(__file__).parent
    dataset_dir = root_dir / "dataset"
    agents_path = root_dir / "agents.yml"
    experiments_dir = root_dir / "experiments"

    # Load configs
    print("Loading experiment config...")
    experiment_config = load_experiment_config(config_path)
    agents_config = load_agents_config(agents_path)

    # Validate
    print("Validating experiment...")
    if not validate_experiment(experiment_config, agents_config, dataset_dir):
        sys.exit(1)

    # Get agent config with optional model override
    agent_config = get_agent_config(
        agents_config,
        experiment_config['agent'],
        experiment_config.get('model_override')
    )

    # Create experiment directory
    exp_dir = create_experiment_dir(experiments_dir, experiment_config['name'])
    print(f"Experiment directory: {exp_dir}")

    # Save config snapshots
    save_config_snapshot(exp_dir, experiment_config, agents_config)

    # Print experiment info
    print("=" * 70)
    print("MCode Experiment Runner")
    print("=" * 70)
    print(f"Name: {experiment_config['name']}")
    print(f"Description: {experiment_config.get('description', 'N/A')}")
    print(f"Agent: {agent_config['name']}")
    print(f"CLI Version: {agent_config.get('cli_version', 'unknown')}")
    print(f"Model: {agent_config.get('model', 'unknown')}")
    parallel = experiment_config['parallel']
    template_version = experiment_config.get('template_version')
    iterations = experiment_config.get('iterations', 1)
    test_each_iteration = experiment_config.get('test_each_iteration', False)
    print(f"Benchmarks: {', '.join(experiment_config['benchmarks'])}")
    print(f"Timeout: {experiment_config['timeout']}s per iteration")
    print(f"Iterations: {iterations}")
    print(f"Test Each Iteration: {test_each_iteration}")
    print(f"Parallel: {parallel}")
    print(f"Template: {template_version or 'existing prompt.md'}")
    print("=" * 70)

    # Results collector
    results = []
    benchmarks = experiment_config['benchmarks']

    if parallel <= 1:
        # Sequential execution
        for benchmark_name in benchmarks:
            result = run_single_benchmark(
                benchmark_name=benchmark_name,
                dataset_dir=dataset_dir,
                exp_dir=exp_dir,
                agent_config=agent_config,
                timeout=experiment_config['timeout'],
                template_version=template_version,
                iterations=iterations,
                test_each_iteration=test_each_iteration
            )
            results.append(result)
    else:
        # Parallel execution
        print(f"\nRunning {len(benchmarks)} benchmarks with {parallel} workers...")
        with ProcessPoolExecutor(max_workers=parallel) as executor:
            # Submit all benchmark jobs
            future_to_benchmark = {
                executor.submit(
                    run_single_benchmark,
                    benchmark_name=benchmark_name,
                    dataset_dir=dataset_dir,
                    exp_dir=exp_dir,
                    agent_config=agent_config,
                    timeout=experiment_config['timeout'],
                    template_version=template_version,
                    iterations=iterations,
                    test_each_iteration=test_each_iteration
                ): benchmark_name
                for benchmark_name in benchmarks
            }

            # Collect results as they complete
            for future in as_completed(future_to_benchmark):
                benchmark_name = future_to_benchmark[future]
                try:
                    result = future.result()
                    results.append(result)
                except Exception as e:
                    print(f"[ERROR] {benchmark_name} failed with exception: {e}")
                    results.append({
                        "agent": agent_config['name'],
                        "cli_version": agent_config.get('cli_version', 'unknown'),
                        "model": agent_config.get('model', 'unknown'),
                        "benchmark": benchmark_name,
                        "time_seconds": 0,
                        "commits": 0,
                        "termination_reason": "exception",
                        "build_success": False,
                        "tests_passed": 0,
                        "tests_total": 0,
                        "pass_rate": 0.0,
                        "error": str(e)
                    })

    # Generate summary.jsonl
    print(f"\n{'=' * 70}")
    print("Generating summary...")
    summary_file = exp_dir / "summary.jsonl"
    with open(summary_file, 'w') as f:
        for r in results:
            f.write(json.dumps(r) + "\n")

    # Print final summary
    print(f"\n{'=' * 70}")
    print("EXPERIMENT SUMMARY")
    print("=" * 70)
    for r in results:
        build_status = "✅" if r["build_success"] else "❌"
        print(f"  {r['benchmark']}: {build_status} {r['tests_passed']}/{r['tests_total']} ({r['pass_rate']:.1%})")
    print("=" * 70)

    print(f"\nExperiment complete!")
    print(f"Results saved to: {exp_dir}")
    print(f"Summary: {summary_file}")


def main():
    parser = argparse.ArgumentParser(description='Run MCode experiment')
    parser.add_argument('config', type=str, help='Path to experiment config YAML file')
    args = parser.parse_args()

    config_path = Path(args.config)
    run_experiment(config_path)


if __name__ == "__main__":
    main()