"int(os.environ["LOCAL_RANK"])").
Below is an expository example of a training script that checkpoints
on each epoch, hence the worst-case progress lost on failure is one
full epoch worth of training.
def main():
        args = parse_args(sys.argv[1:])
        state = load_checkpoint(args.checkpoint_path)
        initialize(state)
    # torch.distributed.run ensures that this will work
    # by exporting all the env vars needed to initialize the process group
    torch.distributed.init_process_group(backend=args.backend)

    for i in range(state.epoch, state.total_num_epochs)
         for batch in iter(state.dataset)
             train(batch, state.model)

         state.epoch += 1
         save_checkpoint(state)

For concrete examples of torchelastic-compliant train scripts, visit
our examples page.