| 1040 | # checkpoint. File operations are cheap to make atomic, that's why. |
| 1041 | |
| 1042 | def _on_commit_callback(array_names): # Runs after writing ckpt is done. |
| 1043 | with gfile.GFile(f"{path}-CUR", "w") as f: |
| 1044 | f.write(curr) |
| 1045 | |
| 1046 | last = "" |
| 1047 | if gfile.exists(f"{path}-LAST"): |
| 1048 | with gfile.GFile(f"{path}-LAST", "r") as f: |
| 1049 | last = f.read() |
| 1050 | |
| 1051 | gfile.rename(f"{path}-CUR", f"{path}-LAST", overwrite=True) |
| 1052 | |
| 1053 | if last.endswith("-tmp"): |
| 1054 | # If pre-emption happens here, some old checkpoints may not be |
| 1055 | # deleted. |
| 1056 | multiprocessing.pool.ThreadPool().map( |
| 1057 | gfile.rmtree, |
| 1058 | [f"{path}-{last}/{name}" for name in array_names]) |
| 1059 | gfile.rmtree(f"{path}-{last}") |
| 1060 | |
| 1061 | # NOTE: The jax checkpoint manager automatically waits for the previous save |
| 1062 | # to be finished before writing again, so we don't need to do it here. |