import os
import sys
import argparse

parser = argparse.ArgumentParser(description='wrapper')
parser.add_argument('--job_data_root', default="/data/job_678/job_data_test", type=str)
# parser.add_argument('--local_rank', default=0, type=int)
args = parser.parse_args()

def run():
    import torch
    assert torch.cuda.is_available()
    print('execute testing', flush=True)
    syscommand = fr"export JOB_DATA_ROOT={args.job_data_root} &&" + \
                fr"export TEST=1 &&" + \
                "bash ./test_ddp_k8spod_webui.sh"

    print(fr'wrapper command {syscommand}', flush=True)
    retcode = os.WEXITSTATUS(os.system(syscommand))
    return retcode

if __name__ == "__main__":
    job_data_root = args.job_data_root
    tar_preprocess_dir = os.path.join(job_data_root, 'output', 'preprocess')
    if not os.path.exists(tar_preprocess_dir) or len(os.listdir(tar_preprocess_dir)) == 0:
        print('re-preprocessing...', flush=True)
        from run_preprocess import run as run_preprocess_func
        # run_preprocess
        ret_code = run_preprocess_func()
        if ret_code != 0:
            sys.exit(ret_code)
        print('done...', flush=True)

    print('running inference...', flush=True)
    # run infer
    ret_code = run()
    if ret_code != 0:
        sys.exit(ret_code)
    print('all done', flush=True)