reactagent/prepare_task.py

""" Prepare a benchmark folder for a task. """

import os
import subprocess
import sys
import json

benchmarks_dir = os.path.dirname(os.path.realpath(__file__)) + "/benchmarks"

def get_task_info(task):
    """Get research problem and benchmark folder name for task"""
    research_problem = None
    benchmark_folder_name= None

    # Retrieve task from benchmarks
    tasks = json.load(open(os.path.join(benchmarks_dir, "tasks.json")))
    if task in tasks:
        research_problem = tasks[task].get("research_problem", None)
        benchmark_folder_name = tasks[task].get("benchmark_folder_name", None)

    elif task in os.listdir(benchmarks_dir) and os.path.isdir(os.path.join(benchmarks_dir, task, "env")):
        # default benchmarks
        benchmark_folder_name = task 
    
    else:
        raise ValueError(f"task {task} not supported in benchmarks")

    if research_problem is None:
        research_problem_file = os.path.join(benchmarks_dir, benchmark_folder_name, "scripts", "research_problem.txt")
        if os.path.exists(research_problem_file):
            # Load default research problem from file
            with open(research_problem_file, "r") as f:
                research_problem = f.read()

    return benchmark_folder_name, research_problem


def prepare_task(benchmark_dir, python="python"):
    """ Run prepare.py in the scripts folder of the benchmark if it exists and has not been run yet. """
    if os.path.exists(os.path.join(benchmark_dir, "scripts", "prepare.py")) and not os.path.exists(os.path.join(benchmark_dir, "scripts", "prepared")):
        print("Running prepare.py ...")
        p = subprocess.run([python, "prepare.py"], cwd=os.path.join(benchmark_dir,"scripts"))
        if p.returncode != 0:
            print("prepare.py failed")
            sys.exit(1)
        else:
            with open(os.path.join(benchmark_dir, "scripts", "prepared"), "w") as f:
                f.write("success")
        print("prepare.py finished")
    else:
        print("prepare.py not found or already prepared")

if __name__ == "__main__":

    task = sys.argv[1]
    if len(sys.argv) > 2:
        python = sys.argv[2]
    else:
        python = "python"
    benchmark_name, _ = get_task_info(task)
    benchmark_dir = os.path.join(benchmarks_dir, benchmark_name)
    prepare_task(benchmark_dir, python=python)