#!python

if __name__ == '__main__':
    import os, sys
    arguments = sys.argv[1:]
    # file_name
    run_file_name = arguments[0]

    # cuda_visible_devive
    if len(arguments) > 1:
        # 0,1,2,3
        cuda_visible_devive = arguments[1]
    else:
        cuda_visible_devive = None

    # cuda location
    if len(arguments) > 2:
        cuda_loc = arguments[2]
    else:
        cuda_loc = 'localhost'

    os.environ['PARALLEL_TYPE'] = 'ds'

    cuda_ctrl = f' --include {cuda_loc}:{cuda_visible_devive}' if cuda_visible_devive else ''

    command = f'deepspeed{cuda_ctrl} {run_file_name}'

    print(f'run command {command}')
    os.system(command)
