Multi data collection script

multi_data_collection_script.py

@@ -0,0 +1,166 @@
+import os
+import shutil
+import subprocess
+import argparse
+from pathlib import Path
+from datetime import datetime
+import time
+import json
+import tqdm
+from analyse_results import extract_result, get_immediate_subdirectories, analyze_json_file
+import glob
+
+
+"""
+This script is intended to run the evaluation script multiple times and then automatically aggregate the 
+successful logs into a subfolder for each run, based on the success marked in the experiment folder. Then 
+at the end it will aggregate everything into a json file, ready for training. 
+"""
+
+
+def extract_result_single_agent(folder_path):
+    folder_name = os.path.basename(folder_path)
+    json_files = glob.glob(os.path.join(folder_path, "*.json"))
+    if not json_files:
+        print(f"No JSON files found in {folder_name}")
+        return None
+    else:
+        outcome = False
+        for json_file in json_files:
+            outcome = analyze_json_file(json_file)
+            if outcome:
+                return True
+        return False
+
+
+def identify_success_folders(download_dir, num_agents):
+    folders = get_immediate_subdirectories(download_dir)
+    
+    if num_agents == 1:
+        extract_result_fn = extract_result_single_agent
+    else:
+        extract_result_fn = extract_result
+    
+    total = 0
+    successful = 0
+    successful_exp_list = []
+    
+    for folder_path in tqdm.tqdm(folders):
+        folder_name = os.path.basename(folder_path)
+
+        try: 
+            total += 1
+            success = int(extract_result_fn(folder_path))
+            successful += success
+            if success:
+                successful_exp_list.append(folder_path)
+        except Exception as e:
+            print(f"Error processing {folder_name}: {e}")
+    
+    return successful_exp_list
+
+
+def run_data_collection(args):
+    # Set up directories
+    LOGS_DIR = Path("logs")
+    SUCCESSFUL_DIR = Path(f"successful_run_logs_{datetime.now().strftime('%Y-%m-%d')}")
+    FULL_RUN_LOGS_DIR = Path(f"full_run_logs_{datetime.now().strftime('%Y-%m-%d')}")
+    EXPERIMENTS_DIR = Path("experiments")
+    BOTS_DIR = Path("bots")
+    LOGS_DIR.mkdir(exist_ok=True)
+    SUCCESSFUL_DIR.mkdir(exist_ok=True)
+    FULL_RUN_LOGS_DIR.mkdir(exist_ok=True)
+
+    # Parse tasks and repetitions
+    TASKS_TO_RUN = []
+    for task_spec in args.tasks:
+        parts = task_spec.split(':')
+        if len(parts) == 2:
+            task_path, repeats = parts[0], int(parts[1])
+            TASKS_TO_RUN.append((task_path, repeats))
+        else:
+            print(f"Warning: Invalid task specification '{task_spec}', expected format 'path:repeats'")
+            
+    # First clear anything named Andy_ or Jill_ from the bots/ folder
+    for bot_dir in BOTS_DIR.glob("*"):
+        if bot_dir.name.startswith(("Andy_", "Jill_", "agent_")):
+            shutil.rmtree(bot_dir)
+
+    run_counter = 1
+    for task_path, repeats in TASKS_TO_RUN:
+        for rep in range(repeats):
+            run_id = f"run_{run_counter:03d}"
+            print(f"\n Starting {task_path} (rep {rep + 1}/{repeats}) -> {run_id}")
+
+            # Track start time to locate experiment folder
+            before = set(EXPERIMENTS_DIR.glob("*"))
+
+            # Run evaluation
+            subprocess.run([
+                "python", args.eval_script,
+                "--api", args.api,
+                "--model", args.model,
+                "--task_path", task_path,
+                "--num_agents", str(args.num_agents),
+                "--num_parallel", str(args.num_parallel)
+            ], check=True)
+
+            # Wait for experiment folder to appear
+            time.sleep(20)  # avoid race condition
+            after = set(EXPERIMENTS_DIR.glob("*"))
+            new_experiments = list(after - before)
+            assert len(new_experiments) == 1, f"Expected one new experiment folder, found {len(new_experiments)}"
+            experiment_dir = new_experiments[0]
+
+            print(f"Found experiment folder: {experiment_dir}")
+            
+            # Identify successful experiments
+            successful_exp_list = identify_success_folders(experiment_dir, args.num_agents)
+            
+            # Save successful logs and results
+            success_output_dir = SUCCESSFUL_DIR / run_id
+            success_output_dir.mkdir(parents=True, exist_ok=True)
+            # Identify the ones that are successful
+            for exp_path in successful_exp_list:
+                exp_name = os.path.basename(exp_path)
+                # For each agent, find and copy their logs for this successful experiment
+                for bot_dir in BOTS_DIR.glob("*"):
+                    if bot_dir.name.startswith(("Andy_", "Jill_", "agent_")):
+                        agent_logs_dir = bot_dir / "logs"
+                        if agent_logs_dir.exists():
+                            # Look for the experiment directory directly under logs
+                            exp_dir = agent_logs_dir / exp_name
+                            if exp_dir.exists():
+                                # Add timestamp for uniqueness
+                                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
+                                dest_dir = success_output_dir / f"{bot_dir.name}_{timestamp}_{exp_name}"
+                                shutil.copytree(exp_dir, dest_dir)
+                                print(f"Copied successful log directory: {exp_dir} -> {dest_dir}")
+            
+            # Move full logs to the full logs dir, aka anything named Jill_ or Andy_
+            full_logs_dir = FULL_RUN_LOGS_DIR / run_id
+            full_logs_dir.mkdir(parents=True, exist_ok=True)
+            for bot_dir in BOTS_DIR.glob("*"):
+                if bot_dir.name.startswith(("Andy_", "Jill_", "agent_")):
+                    # bot_dir is already the full path, no need for agent_dir
+                    dest_dir = full_logs_dir / bot_dir.name
+                    if bot_dir.exists():
+                        shutil.copytree(bot_dir, dest_dir, dirs_exist_ok=True)
+                        print(f"Copied full agent directory for {bot_dir.name} to {dest_dir}")
+
+            run_counter += 1
+
+    print("\nAll evaluations done and successful runs saved.")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run multiple evaluations and collect successful logs")
+    parser.add_argument("--eval_script", default="evaluation_script.py", help="Path to evaluation script")
+    parser.add_argument("--api", default="vllm", help="API to use")
+    parser.add_argument("--model", default="meta-llama/Meta-Llama-3-8B-Instruct", help="Model to use")
+    parser.add_argument("--num_agents", type=int, default=2, help="Number of agents")
+    parser.add_argument("--num_parallel", type=int, default=2, help="Number of parallel runs")
+    parser.add_argument("--tasks", nargs="+", default=["tasks/crafting_tasks/test_tasks/tasks_2_agents.json:2"], 
+                        help="Tasks to run in format 'path:repeats'")
+    
+    args = parser.parse_args()
+    run_data_collection(args)