import os
import sys
import argparse

parser = argparse.ArgumentParser(description='wrapper')
parser.add_argument('--job_data_root', default="/root/Documents/GroupLung/Datasets/det_data/job_data_det_train", type=str)
args = parser.parse_args()

def run():
    import torch
    assert torch.cuda.is_available()
    print('num_gpus: ',torch.cuda.device_count()) # 正常
    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)