import os
import datetime
import shutil

from torch.distributed import barrier, init_process_group, destroy_process_group


init_process_group(backend="nccl", timeout=datetime.timedelta(seconds=12 * 60 * 60))

from_fp = "/mnt/round-surf/checkpoints/2024-04-15_22-54-10/last_ckpt.pt"
to_fp = "/tmp/13b.pt"

# Check if the destination exists
if os.path.exists(to_fp):
    print("deleting destination")
    # Delete if it exists
    os.remove(to_fp)

barrier()

print("copy", from_fp, "to", to_fp)
# Make the copy
shutil.copyfile(from_fp, to_fp)

barrier()

print("done!")

destroy_process_group()