Repository: google/flax

Issue:
flax.training.checkpoint.save_checkpoint with keep=1 leads to "file already exists" error
I'm using Jax latest, Tensorflow latest, Jaxlib 0.1.59.

Let say I run a python code which does a save_checkpoint with keep=1 once. Then, I rerun the python code again, it will give me a "tensorflow.python.framework.errors_impl.AlreadyExistsError: file already exists". This is really annoying because this means that I have to manually delete the checkpoints if I want to reuse the same python code which happens a lot when debugging. 

I think that this happens because the extra files are only deleted at the end of save_checkpoint.

>  File "/localscratch/jolicoea.63359842.0/1/ScoreSDEMore/run_lib.py", line 469, in evaluate
>     checkpoints.save_checkpoint(
>   File "/localscratch/jolicoea.63359842.0/1/env/lib/python3.8/site-packages/flax/training/checkpoints.py", line 99, in save_checkpoint
>     gfile.rename(ckpt_tmp_path, ckpt_path)
>   File "/localscratch/jolicoea.63359842.0/1/env/lib/python3.8/site-packages/tensorflow/python/lib/io/file_io.py", line 548, in rename_v2
>     _pywrap_file_io.RenameFile(
> tensorflow.python.framework.errors_impl.AlreadyExistsError: file already exists


Discussion:
I'm not quite sure how to resolve this. We definitely want to avoid the case where we accidentally overwrite a checkpoint.
For debugging purposes wouldn't it be better to not checkpoint at all if you later want to discard these files anyway?
There could be an option called "overwrite" which defaults to False.
We could add an overwrite option. But I think the current implementation has more issues when you reuse the same checkpoint directory. It might also store let's say at step 10 while you also have a checkpoint at step 100 leading to the new checkpoint at step 10 to be removed.

I'll try to add an overwrite option while fixing the issue of storing older checkpoints than are already stored

Resolving patch:
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -298,4 +298,16 @@ def __call__(self, x):
         return x
   """
   def __init__(self, scope_name):
-    super().__init__(f'Duplicate use of scope name: "{scope_name}"')
\ No newline at end of file
+    super().__init__(f'Duplicate use of scope name: "{scope_name}"')
+
+
+class InvalidCheckpointError(FlaxError):
+  """
+  A checkpoint cannot be stored in a directory that already has
+  a checkpoint at the current or a later step.
+
+  You can pass `overwrite=True` to disable this behavior and
+  overwrite existing checkpoints in the target directory.
+  """
+  def __init__(self, path, step):
+    super().__init__(f'Trying to save an outdated checkpoint at step: "{step}" and path: "{path}".')
\ No newline at end of file
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
--- a/flax/training/checkpoints.py
+++ b/flax/training/checkpoints.py
@@ -25,6 +25,7 @@
 
 from absl import logging
 from flax import core
+from flax import errors
 from flax import serialization
 from tensorflow.io import gfile
 
@@ -73,7 +74,8 @@ def save_checkpoint(ckpt_dir,
                     target,
                     step,
                     prefix='checkpoint_',
-                    keep=1):
+                    keep=1,
+                    overwrite=False):
   """Save a checkpoint of the model.
 
   Attempts to be pre-emption safe by writing to temporary before
@@ -85,7 +87,8 @@ def save_checkpoint(ckpt_dir,
     step: int or float: training step number or other metric number.
     prefix: str: checkpoint file name prefix.
     keep: number of past checkpoint files to keep.
-
+    overwrite: overwrite existing checkpoint files if a checkpoint
+      at the current or a later step already exits (default: False).
   Returns:
     Filename of saved checkpoint.
   """
@@ -94,16 +97,38 @@ def save_checkpoint(ckpt_dir,
   ckpt_tmp_path = _checkpoint_path(ckpt_dir, 'tmp', prefix)
   ckpt_path = _checkpoint_path(ckpt_dir, step, prefix)
   gfile.makedirs(os.path.dirname(ckpt_path))
+  base_path = os.path.join(ckpt_dir, prefix)
+  checkpoint_files = gfile.glob(base_path + '*')
+
+  if ckpt_path in checkpoint_files:
+    if not overwrite:
+      raise errors.InvalidCheckpointError(ckpt_path, step)
+  else:
+    checkpoint_files.append(ckpt_path)
+
+  checkpoint_files = natural_sort(checkpoint_files)
+  if ckpt_path != checkpoint_files[-1]:
+    if not overwrite:
+      raise errors.InvalidCheckpointError(ckpt_path, step)
+
   with gfile.GFile(ckpt_tmp_path, 'wb') as fp:
     fp.write(serialization.to_bytes(target))
 
   # Rename once serialization and writing finished.
-  gfile.rename(ckpt_tmp_path, ckpt_path)
+  gfile.rename(ckpt_tmp_path, ckpt_path, overwrite=overwrite)
   logging.info('Saved checkpoint at %s', ckpt_path)
+  print(ckpt_path)
+
+  # Remove newer checkpoints
+  if overwrite:
+    ind = checkpoint_files.index(ckpt_path) + 1
+    newer_ckpts = checkpoint_files[ind:]
+    checkpoint_files = checkpoint_files[:ind]
+    for path in newer_ckpts:
+      logging.info('Removing checkpoint at %s', path)
+      gfile.remove(path)
 
   # Remove old checkpoint files.
-  base_path = os.path.join(ckpt_dir, f'{prefix}')
-  checkpoint_files = natural_sort(gfile.glob(base_path + '*'))
   if len(checkpoint_files) > keep:
     old_ckpts = checkpoint_files[:-keep]
     for path in old_ckpts:


Test patch:
diff --git a/tests/checkpoints_test.py b/tests/checkpoints_test.py
--- a/tests/checkpoints_test.py
+++ b/tests/checkpoints_test.py
@@ -21,6 +21,7 @@
 from absl.testing import absltest
 import flax
 from flax import core
+from flax import errors
 from flax.training import checkpoints
 import jax
 from jax import numpy as jnp
@@ -156,6 +157,32 @@ def test_save_restore_checkpoints(self):
       checkpoints.restore_checkpoint(
           tmp_dir, test_object0, step=5, prefix='test_')
 
+  def test_overwrite_checkpoints(self):
+    tmp_dir = self.create_tempdir().full_path
+    test_object0 = {'a': np.array([0, 0, 0], np.int32)}
+    test_object = {'a': np.array([1, 2, 3], np.int32)}
+
+    checkpoints.save_checkpoint(
+        tmp_dir, test_object0, 0, keep=1)
+    with self.assertRaises(errors.InvalidCheckpointError):
+      checkpoints.save_checkpoint(
+          tmp_dir, test_object, 0, keep=1)
+    checkpoints.save_checkpoint(
+          tmp_dir, test_object, 0, keep=1, overwrite=True)
+    new_object = checkpoints.restore_checkpoint(tmp_dir, test_object0)
+    jtu.check_eq(new_object, test_object)
+    checkpoints.save_checkpoint(
+          tmp_dir, test_object0, 2, keep=1, overwrite=True)
+    new_object = checkpoints.restore_checkpoint(tmp_dir, test_object)
+    jtu.check_eq(new_object, test_object0)
+    with self.assertRaises(errors.InvalidCheckpointError):
+      checkpoints.save_checkpoint(
+            tmp_dir, test_object, 1, keep=1)
+    checkpoints.save_checkpoint(
+          tmp_dir, test_object, 1, keep=1, overwrite=True)
+    new_object = checkpoints.restore_checkpoint(tmp_dir, test_object0)
+    jtu.check_eq(new_object, test_object)
+
   def test_save_restore_checkpoints_w_float_steps(self):
     tmp_dir = self.create_tempdir().full_path
     test_object0 = {'a': np.array([0, 0, 0], np.int32),
@@ -174,20 +201,14 @@ def test_save_restore_checkpoints_w_float_steps(self):
     jtu.check_eq(new_object, test_object1)
     checkpoints.save_checkpoint(
         tmp_dir, test_object1, 2.0, prefix='test_', keep=1)
-    checkpoints.save_checkpoint(
-        tmp_dir, test_object2, 1.0, prefix='test_', keep=1)
-    new_object = checkpoints.restore_checkpoint(
-        tmp_dir, test_object0, prefix='test_')
-    jtu.check_eq(new_object, test_object1)
+    with self.assertRaises(errors.InvalidCheckpointError):
+      checkpoints.save_checkpoint(
+          tmp_dir, test_object2, 1.0, prefix='test_', keep=1)
     checkpoints.save_checkpoint(
         tmp_dir, test_object2, 3.0, prefix='test_', keep=2)
-    checkpoints.save_checkpoint(
-        tmp_dir, test_object1, -1.0, prefix='test_', keep=2)
-    new_object = checkpoints.restore_checkpoint(
-        tmp_dir, test_object0, prefix='test_')
     self.assertIn('test_3.0', os.listdir(tmp_dir))
     self.assertIn('test_2.0', os.listdir(tmp_dir))
-    jtu.check_eq(new_object, test_object2)
+    jtu.check_eq(new_object, test_object1)
 
   def test_save_restore_checkpoints_target_none(self):
     tmp_dir = self.create_tempdir().full_path