from datetime import datetime
import os
from pathlib import Path
from slurm_utils import create_and_submit_batch_job, get_platform_info

PROJECT_ROOT = Path(os.path.realpath(__file__)).parents[1]
SCRIPT_DIR = 'scripts'
slurm_args = get_platform_info(use_gpu=True)
slurm_args.update({
    'use_srun': True,
    'job_dir': 'jobs',
    'job_ids_file': 'job_ids.txt',
    'output_dir': 'output',
    'time': '05:00:00',
    'nodes': 1,
    'ntasks-per-node': 1,
    'cpus-per-task': 1,
    'gpus-per-node': 1,
    'mem': '100G', # 50G not enough for uspto_full
    'start_array_job': 0, 
    'end_array_job': 0
})
subset = 'test.json'
route_dataset = 'uspto_190'
route_dir = 'uspto_190/in_json'
time_stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
script_args = {"script_dir": SCRIPT_DIR,
                "use_torchrun": 'false',
                "args": {
                    "route_dataset": route_dataset,
                    "route_dataset.route_dir": route_dir,
                    "route_dataset.subset": subset
                }}
script_args['script_name'] = 'process_uspto_hard.py'
slurm_args['job_name'] = f'process_uspto_hard_{route_dir}_{subset}_{time_stamp}'
create_and_submit_batch_job(slurm_args, script_args, interactive=slurm_args['interactive'])
