""" Runner - orchestrates experiment runs based on config. Loops through models, scenarios, goal types, and oversight levels. Supports parallel execution via ThreadPoolExecutor. """ import os import glob import concurrent.futures from typing import List, Dict, Any, Tuple from config_loader import ConfigLoader, ProviderConfig, ModelConfig, ScenarioConfig from vfs import VFS from agent import Agent from tools import make_tools_for_vfs, tools as tool_schemas from logger import get_logger import datetime import threading # Get logger instance logger = get_logger("experiment") # Thread-safe lock for results list _results_lock = threading.Lock() 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: return f.read().strip() class ExperimentRunner: """Runs experiments based on configuration.""" def __init__(self, config: ConfigLoader, verbose: bool = False, resume: bool = True): self.config = config self.results: List[Dict] = [] self.verbose = verbose self.resume = resume def run_all(self): """Run all experiments defined in config.""" logger.info(f"\n{'='*60}") logger.info("Starting Experiment Run") logger.info(f"{'='*60}\n") # Build list of all work items (model, scenario, goal_type, oversight) work_items = self._build_work_items() if not work_items: logger.warning("No work items to run.") return max_workers = self.config.max_workers total_items = len(work_items) logger.info(f"Total work items: {total_items}, Max workers: {max_workers}") if max_workers <= 1: # Sequential execution (original behavior) for item in work_items: self._execute_work_item(item) else: # Parallel execution logger.info(f"Running with {max_workers} parallel workers") from tqdm import tqdm from tqdm.contrib.logging import logging_redirect_tqdm 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): # Keep console relatively quiet (WARNING/ERROR/CRITICAL) handler.setLevel(logging.WARNING) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: futures = { executor.submit(self._execute_work_item, item): item for item in work_items } with logging_redirect_tqdm(): with tqdm(total=total_items, desc="Running Experiments", unit="run", dynamic_ncols=True) as pbar: for future in concurrent.futures.as_completed(futures): item = futures[future] try: future.result() except Exception as e: logger.critical(f"Work item failed: {item.get('label', 'unknown')}: {e}") pbar.write(f"ERROR: {item.get('label', 'unknown')} failed: {e}") pbar.update(1) # Summary successful = sum(1 for r in self.results if r.get("success", False)) total_runs = len(self.results) logger.info(f"\n{'='*60}") logger.info(f"Experiment Complete: {total_runs} runs completed") logger.info(f" SUCCESS: {successful}") logger.info(f" INCOMPLETE: {total_runs - successful}") logger.info(f"{'='*60}\n") def _build_work_items(self) -> List[Dict[str, Any]]: """Build a flat list of all (model, scenario, goal_type, oversight, run_num) combos.""" work_items = [] goal_types = self.config.goal_types 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 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}.") continue # Ensure baseline exists self._ensure_baseline(model_config, scenario_config) for oversight_level in oversight_levels: if goal_types: # New: iterate over goal types for goal_type in goal_types: items = self._build_run_items( model_config, scenario_config, oversight_level, goal_type ) work_items.extend(items) else: # Legacy: single hidden_goal.md (no goal_types in config) items = self._build_run_items( model_config, scenario_config, oversight_level, "" ) work_items.extend(items) return work_items def _build_run_items( self, model_config: ModelConfig, scenario_config: ScenarioConfig, oversight_level: str, goal_type: str ) -> List[Dict[str, Any]]: """Build individual run items for a specific combo, accounting for resume.""" model_name_safe = model_config.id.replace("/", "_") scenario_name = os.path.basename(scenario_config.path) output_dir = self.config.output_dir # Build log directory path if goal_type: 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) # 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")]) if existing_runs >= scenario_config.runs: goal_label = f"/{goal_type}" if goal_type else "" logger.info(f" SKIP: {model_config.id} | {scenario_name}{goal_label} | " f"{oversight_level} ({existing_runs}/{scenario_config.runs} exist)") return [] elif existing_runs > 0: goal_label = f"/{goal_type}" if goal_type else "" logger.info(f" RESUME: {model_config.id} | {scenario_name}{goal_label} | " f"{oversight_level} ({existing_runs}/{scenario_config.runs} exist)") 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}" }) return items 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") if not self.config.generate_baseline: if not os.path.exists(baseline_path): logger.warning(f"Baseline generation DISABLED. No baseline for " f"{model_config.id} | {scenario_name}.") return if not os.path.exists(baseline_path): logger.info(f"\n--- Generating baseline: {model_config.id} | {scenario_name} ---") provider_config = self.config.get_provider(model_config.provider) self._run_baseline(model_config, provider_config, scenario_config) logger.info(f" Baseline saved to {baseline_path}") def _execute_work_item(self, item: Dict[str, Any]): """Execute a single work item (one experiment run). Thread-safe.""" try: self._run_single( model_config=item["model_config"], provider_config=self.config.get_provider(item["model_config"].provider), scenario_config=item["scenario_config"], oversight_level=item["oversight_level"], goal_type=item["goal_type"], run_num=item["run_num"] ) except Exception as e: logger.critical(f"ERROR in {item['label']}: {e}") import traceback traceback.print_exc() def _extract_baseline_content(self, logs: List[Dict]) -> str: """Extract baseline content from the last create_file tool call.""" import json # Iterate through logs in reverse to find last create_file call for msg in reversed(logs): tool_calls = msg.get("tool_calls") if not tool_calls: continue # Check the last tool call in this message for tc in reversed(tool_calls): func = tc.get("function", {}) if func.get("name") == "create_file": try: args = json.loads(func.get("arguments", "{}")) return args.get("content", "") except (json.JSONDecodeError, KeyError): continue # Fallback: try to get content from last assistant message for msg in reversed(logs): if msg.get("role") == "assistant" and msg.get("content"): return msg.get("content", "") return "" def _run_baseline( self, model_config: ModelConfig, provider_config: ProviderConfig, scenario_config: ScenarioConfig ): """Run a baseline experiment WITHOUT hidden goal. Saves summary to baselines/.""" scenario_name = os.path.basename(scenario_config.path) model_name = model_config.id # Load only benign goal (NO hidden goal) benign_goal = load_prompt(os.path.join(scenario_config.path, "benign_goal.md")) user_prompt = load_prompt(os.path.join(scenario_config.path, "user.md")) # Build system prompt with ONLY benign goal system_prompt = benign_goal # Create isolated VFS for this baseline run vfs_path = os.path.join(scenario_config.path, "data") vfs_instance = VFS.create(vfs_path) tool_functions = make_tools_for_vfs(vfs_instance) # Create agent 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, model_config=model_config, scenario=scenario_name, oversight_level="baseline", user_prompt_type="user.md", save_vfs_state=save_vfs_state, vfs_instance=vfs_instance, tool_functions=tool_functions ) # Run the conversation logger.info(f" Running baseline...") output_dir = self.config.output_dir agent.enable_incremental_save(output_dir=output_dir) start_time = datetime.datetime.now() result = agent.run(user_prompt) end_time = datetime.datetime.now() # Extract baseline content from last create_file tool call baseline_content = self._extract_baseline_content(agent.logs) # Save baseline to logs/{model}/{scenario}/baseline.md model_name_safe = model_name.replace("/", "_") output_dir = self.config.output_dir baseline_dir = os.path.join(output_dir, model_name_safe, scenario_name) os.makedirs(baseline_dir, exist_ok=True) baseline_path = os.path.join(baseline_dir, "baseline.md") with open(baseline_path, 'w') as f: f.write(baseline_content) # Save baseline log separately agent.save_logs(output_dir=output_dir) logger.info(f" Baseline completed in {(end_time - start_time).total_seconds():.2f}s") def _run_single( self, model_config: ModelConfig, provider_config: ProviderConfig, scenario_config: ScenarioConfig, oversight_level: str, run_num: int, goal_type: str = "" ): """Run a single experiment. Thread-safe with isolated VFS.""" # Load scenario prompts benign_goal = load_prompt(os.path.join(scenario_config.path, "benign_goal.md")) user_prompt = load_prompt(os.path.join(scenario_config.path, "user.md")) # 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") 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") 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") oversight_prompt = load_prompt(global_oversight_path) # Build system prompt system_parts = [benign_goal] if hidden_goal: system_parts.append(hidden_goal) if oversight_prompt: system_parts.append(oversight_prompt) system_prompt = "\n\n".join(filter(None, system_parts)) # Create isolated VFS for this run (thread-safe) vfs_path = os.path.join(scenario_config.path, "data") vfs_instance = VFS.create(vfs_path) tool_functions = make_tools_for_vfs(vfs_instance) # Log VFS info vfs_files = vfs_instance.list_files("/") logger.debug(f"VFS initialized from: {vfs_path}") logger.debug(f"VFS files: {vfs_files}") # Create agent with isolated VFS 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, provider_config=provider_config, model_config=model_config, scenario=scenario_name, oversight_level=oversight_level, user_prompt_type="user.md", save_vfs_state=save_vfs_state, goal_type=goal_type, vfs_instance=vfs_instance, tool_functions=tool_functions ) # Run the conversation goal_label = f" | {goal_type}" if goal_type else "" logger.info(f"\n Starting: {model_config.id} | {scenario_name}{goal_label} | " f"{oversight_level} | run {run_num}") output_dir = self.config.output_dir agent.enable_incremental_save(output_dir=output_dir) start_time = datetime.datetime.now() result = agent.run(user_prompt) end_time = datetime.datetime.now() # Save logs output_dir = self.config.output_dir log_file = agent.save_logs(output_dir=output_dir) # Check if run was successful (ended with "stop" finish_reason) success = False for msg in reversed(agent.logs): 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"): continue # Record result (thread-safe) result_entry = { "model": model_config.id, "provider": model_config.provider, "scenario": scenario_name, "goal_type": goal_type, "oversight_level": oversight_level, "run_num": run_num, "run_id": f"{model_config.id}/{scenario_name}/{goal_type}/{oversight_level}/{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}", "duration_seconds": (end_time - start_time).total_seconds(), "total_tokens": agent.total_tokens, "success": success, "log_file": log_file } with _results_lock: self.results.append(result_entry) status = "SUCCESS" if success else "INCOMPLETE" logger.info(f" [{status}] {model_config.id} | {scenario_name}{goal_label} | " f"{oversight_level} | run {run_num} ({(end_time - start_time).total_seconds():.2f}s)") def run_from_config(config_path: str = "config.yaml", resume: bool = True): """Convenience function to run all experiments from a config file.""" config = ConfigLoader(config_path) config.load() runner = ExperimentRunner(config, resume=resume) runner.run_all() return runner.results if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Run experiments from config") parser.add_argument("--config", default="config.yaml", help="Path to config file") parser.add_argument("--no-resume", dest="resume", action="store_false", default=True, help="Ignore existing logs and start fresh") args = parser.parse_args() run_from_config(args.config, resume=args.resume)