run_homeragents_neo.py

import argparse
import json
import os
import shutil
from utils_autogen.agent_interact import main
from pathlib import Path
from generate_run_config import main as generate_run_config
from ruamel.yaml import YAML
import debugpy
from utils.constants import *
yaml = YAML()
yaml.preserve_quotes = True  # Preserve the original quoting style
yaml.width = 4096
import asyncio



model_name_short = 'gpt-41' 

repo_root_dir = os.path.dirname(
    os.path.abspath(__file__)
)


models_config_filepath = Path(f"{repo_root_dir}/configs/models.yaml")

with open(models_config_filepath, "r") as file:
    models_config = yaml.load(file)
MODEL_NAMES = {}
for model_short in models_config["model"]:
    MODEL_NAMES[model_short] = models_config["model"][model_short]["name"]
    
# Create the parser
parser = argparse.ArgumentParser()

# Add arguments
parser.add_argument("--generate-config", action="store_true", help="Generate configuration file")
parser.add_argument("--debug", action="store_true", help="Debug mode")
parser.add_argument("--force-new", action="store_true", help="Force new run to remove existing output directory")
parser.add_argument("--tag", default="synthesis", help="Experiment tag")
# Parse the arguments
args = parser.parse_args()

tag = args.tag
exp_id = f"{tag}_{model_name_short}"

if args.debug:
    print("Waiting for debugger to attach...")
    debugpy.listen(("0.0.0.0", 5678))
    debugpy.wait_for_client()
    print("Debugger is attached.")
    
    

# if removing previous run, then regenerate the config
if args.force_new:
    args.generate_config = True

# check if the config is already generated, if not, generate new
output_exp_dir = Path(f"{repo_root_dir}/exp/{exp_id}")
if not os.path.exists(output_exp_dir) or args.generate_config or args.force_new:
    print ('Config not found, generating new ... ')
    generate_run_config(tag, model_name_short, prompt_file='configs/prompts.json', use_thinking_tokens=False)
else:
    print ('Config found successfully, skipping generation ... ') 
    
if args.force_new:
    try:
        print (f'removing previous runs under model {model_name_short} tag {tag}')
        # perform clean-up of previous runs of the same experiment tag
        model_name = MODEL_NAMES[model_name_short]
        # first list all task folders within tasks
        task_list = os.listdir('tasks')
        for i, task_del in enumerate(task_list):
            if not '-' in task_del:
                continue
            subtask_list = os.listdir(f'tasks/{task_del}/outputs')
            for subtask_del in subtask_list:
                print ('running task folder', task_del, i, 'out of', len(task_list))
                print ('running subtask', subtask_del)
                output_subtask_dir = f'tasks/{task_del}/outputs/{subtask_del}'
                folder_list = os.listdir(output_subtask_dir)
                folder_name_to_clean = f'{model_name}_{tag}'
                folder_path_to_clean = f'{output_subtask_dir}/{folder_name_to_clean}'
                if folder_name_to_clean in folder_list:
                    print(f"Cleaning up {output_subtask_dir}")
                    shutil.rmtree(folder_path_to_clean)
                else:
                    print(f"Skipping {output_subtask_dir}")
    except Exception as e:
        print (f'Error while removing previous runs: {e}, exception handled, continue with removing and running new tasks ... ')
      
exp_config_path = Path(f"{repo_root_dir}/exp/{exp_id}/{CONFIG}/{exp_id}.yaml")
exp_config = yaml.load(exp_config_path)
        
max_iter = exp_config[ENV][MAX_ITER]

model_name = MODEL_NAMES[model_name_short]

generate_folder_name = 'subtasks_neo'
task_list = list(sorted(os.listdir('tasks')))


async def run_all_tasks():
    for i, task in enumerate(task_list):
        if not '-' in task:
            continue
        print ('running task', task, i, 'out of', len(task_list))
        subtask_list = [list(sorted(os.listdir(f'tasks/{task}/subtasks')))[0]] 
        for subtask in subtask_list:
            print ('running task folder', task, i, 'out of', len(task_list))
            print ('running subtask', subtask)
            config_file = f'tasks/{task}/subtasks/{subtask}'
            subtask_name = subtask.split('.')[0]
            

            if not args.force_new:
                # skip tasks already generated
                if task in generated_tasks:
                    print (f'skipping task {task} as it is already generated')
                    continue
        
            max_retries = 3
            retry_count = 0
            while retry_count < max_retries:
                try:
                    await main(docker_name='officebench',
                        container_name=f'officebench-debug-{tag}-{model_name_short}',
                        dockerfile_path='./docker/Dockerfile',
                        model_name=model_name,
                        task_dir=f'tasks/{task}',
                        config_file=config_file,
                        task=None,
                        tag=tag,
                        max_iter=max_iter,
                        mode='force_new',
                        exp_config=exp_config,
                        memory_dir=None,
                        generate_folder_name=generate_folder_name)
                    print(f"Task {task} with subtask {subtask} completed successfully.")
                    break  # Exit the retry loop if successful
                except Exception as e:
                    print(f"Error occurred: {e}. Retrying {retry_count + 1}/{max_retries}...")
                    retry_count += 1
                    if retry_count == max_retries:
                        print("Max retries reached. Skipping this task.")
                

if __name__ == "__main__":
    asyncio.run(run_all_tasks())