diff options
| -rw-r--r-- | src/dashboard.py | 194 | ||||
| -rw-r--r-- | src/runner.py | 174 |
2 files changed, 234 insertions, 134 deletions
diff --git a/src/dashboard.py b/src/dashboard.py index 5be6960..a4ef1cd 100644 --- a/src/dashboard.py +++ b/src/dashboard.py @@ -4,7 +4,14 @@ import time from collections import deque from rich.console import Console from rich.live import Live -from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn +from rich.progress import ( + Progress, + SpinnerColumn, + TextColumn, + BarColumn, + TaskProgressColumn, + TimeRemainingColumn, +) from rich.table import Table from rich.panel import Panel from rich.layout import Layout @@ -14,6 +21,7 @@ from rich.columns import Columns # Thread-safe lock for dashboard updates _dashboard_lock = threading.Lock() + class ExperimentDashboard: """Rich dashboard for real-time experiment tracking with per-model stats.""" @@ -26,24 +34,19 @@ class ExperimentDashboard: self.failed_count = 0 self.total_tokens = 0 self.start_time = time.time() - + # Per-model stats self.model_stats = { - model: { - "success": 0, - "total": 0, - "target": 0, - "tokens": 0, - "time": 0.0 - } for model in models + model: {"success": 0, "total": 0, "target": 0, "tokens": 0, "time": 0.0} + for model in models } - + # Track currently active runs - self.active_runs = {} # thread_id -> info string - + self.active_runs = {} # thread_id -> info string + # Recent activity log self.activity_log: Deque[str] = deque(maxlen=5) - + # Progress bars self.progress = Progress( SpinnerColumn(), @@ -51,10 +54,12 @@ class ExperimentDashboard: BarColumn(bar_width=None), TaskProgressColumn(), TimeRemainingColumn(), - expand=True + expand=True, + ) + + self.overall_task = self.progress.add_task( + "[yellow]Overall Progress", total=total_runs ) - - self.overall_task = self.progress.add_task("[yellow]Overall Progress", total=total_runs) self.model_tasks = {} for model in models: self.model_tasks[model] = self.progress.add_task(f"[blue]{model}", total=0) @@ -70,9 +75,20 @@ class ExperimentDashboard: """Mark a run as active.""" with _dashboard_lock: goal_str = f" ({goal})" if goal else "" - self.active_runs[thread_id] = f"[bold blue]{model}[/] | {scenario}{goal_str}" + self.active_runs[thread_id] = ( + f"[bold blue]{model}[/] | {scenario}{goal_str}" + ) - def complete_run(self, thread_id: int, model: str, success: bool, tokens: int = 0, duration: float = 0.0, error: bool = False, label: str = ""): + def complete_run( + self, + thread_id: int, + model: str, + success: bool, + tokens: int = 0, + duration: float = 0.0, + error: bool = False, + label: str = "", + ): """Update counts and progress when a run completes.""" with _dashboard_lock: # Update global counts @@ -85,9 +101,9 @@ class ExperimentDashboard: else: self.incomplete_count += 1 status_msg = "[yellow]INCOMPLETE[/]" - + self.total_tokens += tokens - + # Update per-model stats if model in self.model_stats: m_stats = self.model_stats[model] @@ -96,16 +112,16 @@ class ExperimentDashboard: m_stats["success"] += 1 m_stats["tokens"] += tokens m_stats["time"] += duration - + # Update progress bars self.progress.update(self.overall_task, advance=1) if model in self.model_tasks: self.progress.update(self.model_tasks[model], advance=1) - + # Update activity log timestamp = time.strftime("%H:%M:%S") self.activity_log.append(f"[{timestamp}] {status_msg} {label}") - + # Remove from active runs if thread_id in self.active_runs: del self.active_runs[thread_id] @@ -118,20 +134,23 @@ class ExperimentDashboard: table.add_column("Success %", justify="right", ratio=1) table.add_column("Tokens", justify="right", ratio=1) table.add_column("Avg Time", justify="right", ratio=1) - + for model, stats in self.model_stats.items(): - if stats["target"] == 0: continue - + if stats["target"] == 0: + continue + progress = f"{stats['total']}/{stats['target']}" - success_rate = (stats["success"] / stats["total"] * 100) if stats["total"] > 0 else 0 + success_rate = ( + (stats["success"] / stats["total"] * 100) if stats["total"] > 0 else 0 + ) avg_time = (stats["time"] / stats["total"]) if stats["total"] > 0 else 0 - + table.add_row( model, progress, f"{success_rate:.0f}%", f"{stats['tokens']:,}", - f"{avg_time:.1f}s" + f"{avg_time:.1f}s", ) return table @@ -149,14 +168,14 @@ class ExperimentDashboard: table = Table(expand=True, box=None) table.add_column("Metric", style="cyan") table.add_column("Value", justify="right", style="magenta") - + completed = self.success_count + self.incomplete_count + self.failed_count remaining = self.total_runs - completed - + # Time calculations elapsed = time.time() - self.start_time time_remaining = self.progress.tasks[self.overall_task].time_remaining - + table.add_row("Success", f"[green]{self.success_count}[/green]") table.add_row("Incomplete", f"[yellow]{self.incomplete_count}[/yellow]") table.add_row("Failed", f"[red]{self.failed_count}[/red]") @@ -166,7 +185,7 @@ class ExperimentDashboard: table.add_row("Elapsed", self.format_time(elapsed)) table.add_row("Est. Left", self.format_time(time_remaining)) table.add_row("Skipped", f"[grey50]{self.skipped}[/grey50]") - + return table def generate_active_panel(self) -> Panel: @@ -192,47 +211,65 @@ class ExperimentDashboard: layout.split_column( Layout(name="header", size=3), Layout(name="main"), - Layout(name="footer", size=7) + Layout(name="footer", size=7), ) - + layout["main"].split_row( - Layout(name="progress_col", ratio=2), - Layout(name="stats_col", ratio=1) + Layout(name="progress_col", ratio=2), Layout(name="stats_col", ratio=1) ) - + layout["progress_col"].split_column( - Layout(name="bars", ratio=1), - Layout(name="model_details", ratio=1) + Layout(name="bars", ratio=1), Layout(name="model_details", ratio=1) ) - + layout["footer"].split_row( - Layout(name="active", ratio=1), - Layout(name="logs", ratio=1) + Layout(name="active", ratio=1), Layout(name="logs", ratio=1) + ) + + layout["header"].update( + Panel( + Text( + "AI Agent Experiment Framework", + justify="center", + style="bold white", + ), + style="blue", + ) + ) + layout["bars"].update( + Panel(self.progress, title="Overall Progress", style="white") + ) + layout["model_details"].update( + Panel(self.generate_model_table(), title="Model Stats", style="white") + ) + layout["stats_col"].update( + Panel(self.generate_status_table(), title="Totals", style="white") ) - - layout["header"].update(Panel(Text("AI Agent Experiment Framework", justify="center", style="bold white"), style="blue")) - layout["bars"].update(Panel(self.progress, title="Overall Progress", style="white")) - layout["model_details"].update(Panel(self.generate_model_table(), title="Model Stats", style="white")) - layout["stats_col"].update(Panel(self.generate_status_table(), title="Totals", style="white")) layout["active"].update(self.generate_active_panel()) layout["logs"].update(self.generate_log_panel()) - + return layout + def print_final_summary(results: List[Dict]): """Print a pretty grouped summary table at the end.""" console = Console() console.print("\n") - + # Group results by Model and Scenario grouped = {} for r in results: - key = (r['model'], r['scenario']) + key = (r["model"], r["scenario"]) if key not in grouped: grouped[key] = [] grouped[key].append(r) - table = Table(title="[bold]Final Experiment Summary[/bold]", show_header=True, header_style="bold magenta", expand=True) + table = Table( + title="[bold]Final Experiment Summary[/bold]", + show_header=True, + header_style="bold magenta", + expand=True, + ) table.add_column("Model", style="blue", no_wrap=True) table.add_column("Scenario", style="cyan") table.add_column("Goal Type", style="yellow") @@ -243,46 +280,49 @@ def print_final_summary(results: List[Dict]): for (model, scenario), group in sorted(grouped.items()): total = len(group) successes = sum(1 for r in group if r.get("success")) + errors = sum(1 for r in group if r.get("error")) total_tokens = sum(r.get("total_tokens", 0) for r in group) avg_time = sum(r.get("duration_seconds", 0.0) for r in group) / total - - # Collect distinct goal types in this group - goals = ", ".join(sorted(list(set(r.get("goal_type", "default") or "default" for r in group)))) - - success_color = "green" if successes == total else "yellow" if successes > 0 else "red" - + + goals = ", ".join( + sorted(list(set(r.get("goal_type", "default") or "default" for r in group))) + ) + + if errors > 0: + success_str = f"[red]{successes}/{total} ({errors} errors)[/]" + elif successes == total: + success_str = f"[green]{successes}/{total}[/]" + elif successes > 0: + success_str = f"[yellow]{successes}/{total}[/]" + else: + success_str = f"[red]{successes}/{total}[/]" + table.add_row( - model, - scenario, - goals, - f"[{success_color}]{successes}/{total}[/]", - f"{total_tokens:,}", - f"{avg_time:.1f}s" + model, scenario, goals, success_str, f"{total_tokens:,}", f"{avg_time:.1f}s" ) - + console.print(table) - + # Aggregated Per-Model Table - model_table = Table(title="[bold]Aggregated Model Performance[/bold]", show_header=True, header_style="bold blue") + model_table = Table( + title="[bold]Aggregated Model Performance[/bold]", + show_header=True, + header_style="bold blue", + ) model_table.add_column("Model") model_table.add_column("Total Runs", justify="right") model_table.add_column("Overall Success %", justify="right") model_table.add_column("Total Tokens", justify="right") - models = sorted(list(set(r['model'] for r in results))) + models = sorted(list(set(r["model"] for r in results))) for model in models: - m_results = [r for r in results if r['model'] == model] + m_results = [r for r in results if r["model"] == model] m_total = len(m_results) m_successes = sum(1 for r in m_results if r.get("success")) m_tokens = sum(r.get("total_tokens", 0) for r in m_results) - + rate = (m_successes / m_total * 100) if m_total > 0 else 0 - model_table.add_row( - model, - str(m_total), - f"{rate:.1f}%", - f"{m_tokens:,}" - ) - + model_table.add_row(model, str(m_total), f"{rate:.1f}%", f"{m_tokens:,}") + console.print("\n") console.print(model_table) diff --git a/src/runner.py b/src/runner.py index 999856a..d01aafc 100644 --- a/src/runner.py +++ b/src/runner.py @@ -24,14 +24,16 @@ def load_prompt(file_path: str) -> str: """Load a prompt file.""" if not os.path.exists(file_path): return "" - with open(file_path, 'r') as f: + with open(file_path, "r") as f: return f.read().strip() class ExperimentRunner: """Runs experiments based on configuration.""" - def __init__(self, config: ConfigLoader, verbose: bool = False, resume: bool = True): + def __init__( + self, config: ConfigLoader, verbose: bool = False, resume: bool = True + ): self.config = config self.results: List[Dict] = [] self.verbose = verbose @@ -40,9 +42,9 @@ class ExperimentRunner: def run_all(self): """Run all experiments defined in config.""" - logger.info(f"\n{'='*60}") + logger.info(f"\n{'=' * 60}") logger.info("Starting Experiment Run") - logger.info(f"{'='*60}\n") + logger.info(f"{'=' * 60}\n") # Build list of all work items (model, scenario, goal_type, oversight) work_items, skipped_count = self._build_work_items() @@ -50,14 +52,14 @@ class ExperimentRunner: if not work_items and skipped_count == 0: logger.warning("No work items to run.") return - + if not work_items: logger.info(f"All {skipped_count} items already exist. Skipping all.") return max_workers = self.config.max_workers total_items = len(work_items) - + # Prepare model list for dashboard model_names = sorted(list(set(item["model_config"].id for item in work_items))) model_counts = {} @@ -65,13 +67,19 @@ class ExperimentRunner: m_id = item["model_config"].id model_counts[m_id] = model_counts.get(m_id, 0) + 1 - self.dashboard = ExperimentDashboard(total_items, model_names, skipped=skipped_count) + self.dashboard = ExperimentDashboard( + total_items, model_names, skipped=skipped_count + ) for m_id, count in model_counts.items(): self.dashboard.update_model_total(m_id, count) logger.info(f"Total work items: {total_items}, Max workers: {max_workers}") - with Live(self.dashboard.get_layout(), refresh_per_second=4, vertical_overflow="visible") as live: + with Live( + self.dashboard.get_layout(), + refresh_per_second=4, + vertical_overflow="visible", + ) as live: if max_workers <= 1: # Sequential execution for item in work_items: @@ -80,12 +88,17 @@ class ExperimentRunner: else: # Parallel execution import logging + # Reduce console spam during parallel runs to keep the progress bar clean for handler in logger.handlers: - if isinstance(handler, logging.StreamHandler) and not isinstance(handler, logging.FileHandler): + if isinstance(handler, logging.StreamHandler) and not isinstance( + handler, logging.FileHandler + ): handler.setLevel(logging.WARNING) - - with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + + with concurrent.futures.ThreadPoolExecutor( + max_workers=max_workers + ) as executor: futures = { executor.submit(self._execute_work_item, item): item for item in work_items @@ -111,11 +124,15 @@ class ExperimentRunner: for model_config in self.config.models: for scenario_config in self.config.scenarios: # Determine oversight levels - available = scenario_config.oversight_levels or self.config.oversight_levels + available = ( + scenario_config.oversight_levels or self.config.oversight_levels + ) global_filter = self.config.oversight_levels oversight_levels = [lvl for lvl in available if lvl in global_filter] if not oversight_levels: - logger.warning(f"No matching oversight levels for {scenario_config.path}.") + logger.warning( + f"No matching oversight levels for {scenario_config.path}." + ) continue # Ensure baseline exists @@ -125,7 +142,10 @@ class ExperimentRunner: if goal_types: for goal_type in goal_types: items, skipped = self._build_run_items( - model_config, scenario_config, oversight_level, goal_type + model_config, + scenario_config, + oversight_level, + goal_type, ) work_items.extend(items) skipped_count += skipped @@ -143,7 +163,7 @@ class ExperimentRunner: model_config: ModelConfig, scenario_config: ScenarioConfig, oversight_level: str, - goal_type: str + goal_type: str, ) -> Tuple[List[Dict[str, Any]], int]: """Build individual run items for a specific combo, accounting for resume.""" model_name_safe = model_config.id.replace("/", "_") @@ -152,17 +172,21 @@ class ExperimentRunner: # Build log directory path if goal_type: - log_dir = os.path.join(output_dir, model_name_safe, scenario_name, - goal_type, oversight_level) + log_dir = os.path.join( + output_dir, model_name_safe, scenario_name, goal_type, oversight_level + ) else: - log_dir = os.path.join(output_dir, model_name_safe, scenario_name, - oversight_level) + log_dir = os.path.join( + output_dir, model_name_safe, scenario_name, oversight_level + ) # Check existing runs for resume existing_runs = 0 if self.resume and os.path.isdir(log_dir): all_json = glob.glob(os.path.join(log_dir, "*.json")) - existing_runs = len([f for f in all_json if not f.endswith(".partial.json")]) + existing_runs = len( + [f for f in all_json if not f.endswith(".partial.json")] + ) skipped = 0 if existing_runs >= scenario_config.runs: @@ -174,23 +198,29 @@ class ExperimentRunner: items = [] for run_num in range(existing_runs + 1, scenario_config.runs + 1): goal_label = f"/{goal_type}" if goal_type else "" - items.append({ - "model_config": model_config, - "scenario_config": scenario_config, - "oversight_level": oversight_level, - "goal_type": goal_type, - "run_num": run_num, - "label": f"{model_config.id} | {scenario_name}{goal_label} | {oversight_level} | run {run_num}" - }) + items.append( + { + "model_config": model_config, + "scenario_config": scenario_config, + "oversight_level": oversight_level, + "goal_type": goal_type, + "run_num": run_num, + "label": f"{model_config.id} | {scenario_name}{goal_label} | {oversight_level} | run {run_num}", + } + ) return items, skipped - def _ensure_baseline(self, model_config: ModelConfig, scenario_config: ScenarioConfig): + def _ensure_baseline( + self, model_config: ModelConfig, scenario_config: ScenarioConfig + ): """Ensure baseline exists for a model+scenario combo (thread-safe).""" scenario_name = os.path.basename(scenario_config.path) model_name_safe = model_config.id.replace("/", "_") output_dir = self.config.output_dir - baseline_path = os.path.join(output_dir, model_name_safe, scenario_name, "baseline.md") + baseline_path = os.path.join( + output_dir, model_name_safe, scenario_name, "baseline.md" + ) if not self.config.generate_baseline: return @@ -204,12 +234,12 @@ class ExperimentRunner: thread_id = threading.get_ident() if self.dashboard: self.dashboard.start_run( - thread_id, - item["model_config"].id, + thread_id, + item["model_config"].id, os.path.basename(item["scenario_config"].path), - item["goal_type"] + item["goal_type"], ) - + try: success, tokens, duration = self._run_single( model_config=item["model_config"], @@ -217,28 +247,44 @@ class ExperimentRunner: scenario_config=item["scenario_config"], oversight_level=item["oversight_level"], goal_type=item["goal_type"], - run_num=item["run_num"] + run_num=item["run_num"], ) - + if self.dashboard: self.dashboard.complete_run( thread_id, - item["model_config"].id, - success=success, - tokens=tokens, + item["model_config"].id, + success=success, + tokens=tokens, duration=duration, - label=item["label"] + label=item["label"], ) except Exception as e: logger.critical(f"ERROR in {item['label']}: {e}") if self.dashboard: self.dashboard.complete_run( thread_id, - item["model_config"].id, - success=False, + item["model_config"].id, + success=False, error=True, - label=item["label"] + label=item["label"], ) + error_result = { + "model": item["model_config"].id, + "provider": item["model_config"].provider, + "scenario": os.path.basename(item["scenario_config"].path), + "goal_type": item["goal_type"], + "oversight_level": item["oversight_level"], + "run_num": item["run_num"], + "duration_seconds": 0.0, + "total_tokens": 0, + "success": False, + "error": True, + "error_message": str(e), + "log_file": None, + } + with _results_lock: + self.results.append(error_result) def _extract_baseline_content(self, logs: List[Dict]) -> str: """Extract baseline content from the last create_file tool call.""" @@ -271,7 +317,7 @@ class ExperimentRunner: self, model_config: ModelConfig, provider_config: ProviderConfig, - scenario_config: ScenarioConfig + scenario_config: ScenarioConfig, ): """Run a baseline experiment WITHOUT hidden goal. Saves summary to baselines/.""" scenario_name = os.path.basename(scenario_config.path) @@ -290,7 +336,9 @@ class ExperimentRunner: tool_functions = make_tools_for_vfs(vfs_instance) # Create agent - save_vfs_state = self.config._config.get('output', {}).get('save_vfs_state', True) + save_vfs_state = self.config._config.get("output", {}).get( + "save_vfs_state", True + ) agent = Agent.from_configs( system_prompt=system_prompt, provider_config=provider_config, @@ -300,7 +348,7 @@ class ExperimentRunner: user_prompt_type="user.md", save_vfs_state=save_vfs_state, vfs_instance=vfs_instance, - tool_functions=tool_functions + tool_functions=tool_functions, ) # Run the conversation @@ -318,7 +366,7 @@ class ExperimentRunner: os.makedirs(baseline_dir, exist_ok=True) baseline_path = os.path.join(baseline_dir, "baseline.md") - with open(baseline_path, 'w') as f: + with open(baseline_path, "w") as f: f.write(baseline_content) # Save baseline log separately @@ -331,7 +379,7 @@ class ExperimentRunner: scenario_config: ScenarioConfig, oversight_level: str, run_num: int, - goal_type: str = "" + goal_type: str = "", ) -> Tuple[bool, int, float]: """Run a single experiment. Thread-safe with isolated VFS.""" # Load scenario prompts @@ -340,16 +388,22 @@ class ExperimentRunner: # Load hidden goal: from hidden_goals/{goal_type}.md or legacy hidden_goal.md if goal_type: - hidden_goal_path = os.path.join(scenario_config.path, "hidden_goals", f"{goal_type}.md") + hidden_goal_path = os.path.join( + scenario_config.path, "hidden_goals", f"{goal_type}.md" + ) else: hidden_goal_path = os.path.join(scenario_config.path, "hidden_goal.md") hidden_goal = load_prompt(hidden_goal_path) # Load oversight prompt: scenario-specific first, then global fallback - scenario_oversight_path = os.path.join(scenario_config.path, "oversight", f"{oversight_level}.md") + scenario_oversight_path = os.path.join( + scenario_config.path, "oversight", f"{oversight_level}.md" + ) oversight_prompt = load_prompt(scenario_oversight_path) if not oversight_prompt: - global_oversight_path = os.path.join(self.config.project_root, "oversight", f"{oversight_level}.md") + global_oversight_path = os.path.join( + self.config.project_root, "oversight", f"{oversight_level}.md" + ) oversight_prompt = load_prompt(global_oversight_path) # Build system prompt @@ -366,7 +420,9 @@ class ExperimentRunner: tool_functions = make_tools_for_vfs(vfs_instance) # Create agent with isolated VFS - save_vfs_state = self.config._config.get('output', {}).get('save_vfs_state', True) + save_vfs_state = self.config._config.get("output", {}).get( + "save_vfs_state", True + ) scenario_name = os.path.basename(scenario_config.path) agent = Agent.from_configs( system_prompt=system_prompt, @@ -378,7 +434,7 @@ class ExperimentRunner: save_vfs_state=save_vfs_state, goal_type=goal_type, vfs_instance=vfs_instance, - tool_functions=tool_functions + tool_functions=tool_functions, ) # Run the conversation @@ -398,11 +454,15 @@ class ExperimentRunner: if msg.get("role") == "assistant" and msg.get("finish_reason"): success = msg["finish_reason"] == "stop" break - elif msg.get("role") == "assistant" and msg.get("content") is None and msg.get("tool_calls"): + elif ( + msg.get("role") == "assistant" + and msg.get("content") is None + and msg.get("tool_calls") + ): continue duration = (end_time - start_time).total_seconds() - + # Record result (thread-safe) result_entry = { "model": model_config.id, @@ -414,12 +474,12 @@ class ExperimentRunner: "duration_seconds": duration, "total_tokens": agent.total_tokens, "success": success, - "log_file": log_file + "log_file": log_file, } with _results_lock: self.results.append(result_entry) - + return success, agent.total_tokens, duration |