openadapt-ml 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

openadapt_ml/benchmarks/runner.py

@@ -1,418 +0,0 @@
-"""Evaluation runner for benchmarks.
-
-This module provides functions to run agents on benchmarks and collect results.
-
-Example:
-    from openadapt_ml.benchmarks import WAAAdapter, PolicyAgent, evaluate_agent_on_benchmark
-
-    adapter = WAAAdapter(waa_repo_path="/path/to/WAA")
-    agent = PolicyAgent(policy)
-    results = evaluate_agent_on_benchmark(agent, adapter, max_steps=50)
-
-    print(f"Success rate: {sum(r.success for r in results) / len(results):.1%}")
-"""
-
-from __future__ import annotations
-
-import logging
-import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from dataclasses import dataclass
-from typing import Callable
-
-from openadapt_ml.benchmarks.agent import BenchmarkAgent
-from openadapt_ml.benchmarks.base import (
-    BenchmarkAdapter,
-    BenchmarkAction,
-    BenchmarkObservation,
-    BenchmarkResult,
-    BenchmarkTask,
-)
-from openadapt_ml.benchmarks.data_collection import ExecutionTraceCollector
-from openadapt_ml.benchmarks.live_tracker import LiveEvaluationTracker
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class EvaluationConfig:
-    """Configuration for benchmark evaluation.
-
-    Attributes:
-        max_steps: Maximum steps per task.
-        parallel: Number of parallel workers (if supported).
-        save_trajectories: Whether to save full trajectories in results.
-        verbose: Whether to print progress.
-        on_step: Optional callback called after each step.
-        on_task_complete: Optional callback called after each task.
-        save_execution_traces: Whether to save execution traces for viewer.
-        model_id: Model identifier for execution traces.
-        output_dir: Output directory for benchmark results.
-        run_name: Name for this evaluation run.
-        enable_live_tracking: Whether to enable live evaluation progress tracking.
-        live_tracking_file: Path to live tracking JSON file.
-    """
-
-    max_steps: int = 50
-    parallel: int = 1
-    save_trajectories: bool = True
-    verbose: bool = True
-    on_step: Callable[[BenchmarkObservation, BenchmarkAction, int], None] | None = None
-    on_task_complete: Callable[[BenchmarkResult], None] | None = None
-    save_execution_traces: bool = True
-    model_id: str = "unknown"
-    output_dir: str = "benchmark_results"
-    run_name: str | None = None
-    enable_live_tracking: bool = True
-    live_tracking_file: str = "benchmark_live.json"
-
-
-def evaluate_agent_on_benchmark(
-    agent: BenchmarkAgent,
-    adapter: BenchmarkAdapter,
-    task_ids: list[str] | None = None,
-    max_steps: int = 50,
-    parallel: int = 1,
-    config: EvaluationConfig | None = None,
-) -> list[BenchmarkResult]:
-    """Run agent on benchmark tasks and collect results.
-
-    Args:
-        agent: Agent to evaluate.
-        adapter: Benchmark adapter.
-        task_ids: Specific tasks to run (None = all tasks).
-        max_steps: Maximum steps per task (overridden by config if provided).
-        parallel: Number of parallel workers (overridden by config if provided).
-        config: Full evaluation configuration.
-
-    Returns:
-        List of BenchmarkResult for each task.
-    """
-    if config is None:
-        config = EvaluationConfig(max_steps=max_steps, parallel=parallel)
-
-    # Load tasks
-    if task_ids is not None:
-        tasks = [adapter.load_task(tid) for tid in task_ids]
-    else:
-        tasks = adapter.list_tasks()
-
-    if config.verbose:
-        logger.info(f"Evaluating {len(tasks)} tasks on {adapter.name}")
-
-    # Initialize execution trace collector if enabled
-    trace_collector = None
-    if config.save_execution_traces:
-        trace_collector = ExecutionTraceCollector(
-            benchmark_name=adapter.name,
-            run_name=config.run_name,
-            model_id=config.model_id,
-            output_dir=config.output_dir,
-        )
-        if config.verbose:
-            logger.info(f"Saving execution traces to: {trace_collector.run_dir}")
-
-    # Initialize live evaluation tracker if enabled
-    live_tracker = None
-    if config.enable_live_tracking:
-        live_tracker = LiveEvaluationTracker(
-            output_file=config.live_tracking_file,
-            total_tasks=len(tasks),
-        )
-        if config.verbose:
-            logger.info(f"Live tracking enabled: {config.live_tracking_file}")
-
-    # Run evaluation
-    if config.parallel > 1 and adapter.supports_parallel:
-        results = _evaluate_parallel(agent, adapter, tasks, config, trace_collector, live_tracker)
-    else:
-        results = _evaluate_sequential(agent, adapter, tasks, config, trace_collector, live_tracker)
-
-    # Save summary if trace collection is enabled
-    if trace_collector is not None:
-        trace_collector.save_summary(results)
-
-    # Mark live tracking as complete
-    if live_tracker is not None:
-        live_tracker.finish()
-
-    # Log summary
-    if config.verbose:
-        success_count = sum(1 for r in results if r.success)
-        success_rate = success_count / len(results) if results else 0
-        avg_steps = sum(r.num_steps for r in results) / len(results) if results else 0
-        logger.info(
-            f"Evaluation complete: {success_count}/{len(results)} "
-            f"({success_rate:.1%}) success, {avg_steps:.1f} avg steps"
-        )
-
-    return results
-
-
-def _evaluate_sequential(
-    agent: BenchmarkAgent,
-    adapter: BenchmarkAdapter,
-    tasks: list[BenchmarkTask],
-    config: EvaluationConfig,
-    trace_collector: ExecutionTraceCollector | None = None,
-    live_tracker: LiveEvaluationTracker | None = None,
-) -> list[BenchmarkResult]:
-    """Run evaluation sequentially.
-
-    Args:
-        agent: Agent to evaluate.
-        adapter: Benchmark adapter.
-        tasks: Tasks to evaluate.
-        config: Evaluation configuration.
-        trace_collector: Optional trace collector for saving execution data.
-        live_tracker: Optional live evaluation tracker.
-
-    Returns:
-        List of results.
-    """
-    results = []
-    for i, task in enumerate(tasks):
-        if config.verbose:
-            logger.info(f"Task {i + 1}/{len(tasks)}: {task.task_id}")
-
-        result = _run_single_task(agent, adapter, task, config, trace_collector, live_tracker)
-        results.append(result)
-
-        if config.on_task_complete:
-            config.on_task_complete(result)
-
-    return results
-
-
-def _evaluate_parallel(
-    agent: BenchmarkAgent,
-    adapter: BenchmarkAdapter,
-    tasks: list[BenchmarkTask],
-    config: EvaluationConfig,
-    trace_collector: ExecutionTraceCollector | None = None,
-    live_tracker: LiveEvaluationTracker | None = None,
-) -> list[BenchmarkResult]:
-    """Run evaluation in parallel.
-
-    Note: This requires the adapter to support parallel execution
-    (e.g., via multiple VM instances).
-
-    Args:
-        agent: Agent to evaluate.
-        adapter: Benchmark adapter.
-        tasks: Tasks to evaluate.
-        config: Evaluation configuration.
-        trace_collector: Optional trace collector for saving execution data.
-        live_tracker: Optional live evaluation tracker.
-
-    Returns:
-        List of results.
-    """
-    results = []
-
-    with ThreadPoolExecutor(max_workers=config.parallel) as executor:
-        # Submit all tasks
-        future_to_task = {
-            executor.submit(_run_single_task, agent, adapter, task, config, trace_collector, live_tracker): task
-            for task in tasks
-        }
-
-        # Collect results as they complete
-        for future in as_completed(future_to_task):
-            task = future_to_task[future]
-            try:
-                result = future.result()
-                results.append(result)
-
-                if config.on_task_complete:
-                    config.on_task_complete(result)
-
-                if config.verbose:
-                    status = "SUCCESS" if result.success else "FAIL"
-                    logger.info(f"Task {task.task_id}: {status}")
-
-            except Exception as e:
-                logger.error(f"Task {task.task_id} failed with error: {e}")
-                results.append(
-                    BenchmarkResult(
-                        task_id=task.task_id,
-                        success=False,
-                        score=0.0,
-                        error=str(e),
-                    )
-                )
-
-    return results
-
-
-def _run_single_task(
-    agent: BenchmarkAgent,
-    adapter: BenchmarkAdapter,
-    task: BenchmarkTask,
-    config: EvaluationConfig,
-    trace_collector: ExecutionTraceCollector | None = None,
-    live_tracker: LiveEvaluationTracker | None = None,
-) -> BenchmarkResult:
-    """Run a single task and return result.
-
-    Args:
-        agent: Agent to evaluate.
-        adapter: Benchmark adapter.
-        task: Task to run.
-        config: Evaluation configuration.
-        trace_collector: Optional trace collector for saving execution data.
-        live_tracker: Optional live evaluation tracker.
-
-    Returns:
-        BenchmarkResult.
-    """
-    start_time = time.perf_counter()
-    history: list[tuple[BenchmarkObservation, BenchmarkAction]] = []
-
-    # Start trace collection if enabled
-    if trace_collector is not None:
-        trace_collector.start_task(task)
-
-    # Start live tracking if enabled
-    if live_tracker is not None:
-        live_tracker.start_task(task)
-
-    try:
-        # Reset agent and environment
-        agent.reset()
-        obs = adapter.reset(task)
-
-        done = False
-        steps = 0
-        max_steps = task.time_limit_steps or config.max_steps
-
-        while not done and steps < max_steps:
-            # Get action from agent
-            action = agent.act(obs, task, history if config.save_trajectories else None)
-
-            # Extract reasoning if available from PolicyAgent
-            reasoning = None
-            if hasattr(action, "raw_action") and action.raw_action:
-                reasoning = action.raw_action.get("thought")
-
-            # Record step in trace collector
-            if trace_collector is not None:
-                trace_collector.record_step(steps, obs, action, reasoning)
-
-            # Record step in live tracker
-            if live_tracker is not None:
-                live_tracker.record_step(steps, obs, action, reasoning)
-
-            # Record step in history
-            if config.save_trajectories:
-                history.append((obs, action))
-
-            if config.on_step:
-                config.on_step(obs, action, steps)
-
-            # Check for terminal action
-            if action.type == "done":
-                done = True
-                break
-
-            # Execute action
-            obs, done, info = adapter.step(action)
-            steps += 1
-
-        # Evaluate result
-        result = adapter.evaluate(task)
-
-        # Update result with trajectory info
-        result.steps = history if config.save_trajectories else []
-        result.num_steps = steps
-        result.total_time_seconds = time.perf_counter() - start_time
-
-        # Finish trace collection if enabled
-        if trace_collector is not None:
-            trace_collector.finish_task(result)
-
-        # Finish live tracking if enabled
-        if live_tracker is not None:
-            live_tracker.finish_task(result)
-
-        return result
-
-    except Exception as e:
-        logger.error(f"Error running task {task.task_id}: {e}")
-        result = BenchmarkResult(
-            task_id=task.task_id,
-            success=False,
-            score=0.0,
-            steps=history if config.save_trajectories else [],
-            num_steps=len(history),
-            error=str(e),
-            total_time_seconds=time.perf_counter() - start_time,
-        )
-
-        # Finish trace collection even on error
-        if trace_collector is not None:
-            trace_collector.finish_task(result)
-
-        return result
-
-
-def compute_metrics(results: list[BenchmarkResult]) -> dict:
-    """Compute aggregate metrics from evaluation results.
-
-    Args:
-        results: List of BenchmarkResult from evaluation.
-
-    Returns:
-        Dict with aggregate metrics.
-    """
-    if not results:
-        return {
-            "num_tasks": 0,
-            "success_rate": 0.0,
-            "avg_score": 0.0,
-            "avg_steps": 0.0,
-            "avg_time_seconds": 0.0,
-        }
-
-    num_tasks = len(results)
-    success_count = sum(1 for r in results if r.success)
-    total_score = sum(r.score for r in results)
-    total_steps = sum(r.num_steps for r in results)
-    total_time = sum(r.total_time_seconds for r in results)
-
-    return {
-        "num_tasks": num_tasks,
-        "success_rate": success_count / num_tasks,
-        "avg_score": total_score / num_tasks,
-        "avg_steps": total_steps / num_tasks,
-        "avg_time_seconds": total_time / num_tasks,
-        "success_count": success_count,
-        "fail_count": num_tasks - success_count,
-    }
-
-
-def compute_domain_metrics(
-    results: list[BenchmarkResult], tasks: list[BenchmarkTask]
-) -> dict[str, dict]:
-    """Compute per-domain metrics.
-
-    Args:
-        results: List of BenchmarkResult.
-        tasks: List of BenchmarkTask (to get domain info).
-
-    Returns:
-        Dict mapping domain to metrics dict.
-    """
-    # Build task_id -> domain mapping
-    task_domains = {t.task_id: t.domain for t in tasks}
-
-    # Group results by domain
-    domain_results: dict[str, list[BenchmarkResult]] = {}
-    for result in results:
-        domain = task_domains.get(result.task_id, "unknown")
-        if domain not in domain_results:
-            domain_results[domain] = []
-        domain_results[domain].append(result)
-
-    # Compute metrics per domain
-    return {domain: compute_metrics(res) for domain, res in domain_results.items()}