Spaces:
Running
Running
| """ 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) |