import os
import sys
import argparse

parser = argparse.ArgumentParser(description='wrapper')
parser.add_argument('--job_data_root', default="/data/job_678/job_data_train", 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 training', flush=True)
    syscommand = fr"export JOB_DATA_ROOT={args.job_data_root} &&" + \
                "bash ./train_ddp_k8spod_webui.sh"
                 
    print(fr'wrapper command {syscommand}', flush=True)
    retcode = os.WEXITSTATUS(os.system(syscommand))
    return retcode

if __name__ == "__main__":
    # run training
    ret_code = run()
    if ret_code != 0:
        sys.exit(ret_code)