When unlocking a checkpoint interpret 'missing' as 'already unlocked'

PiperOrigin-RevId: 651915877

checkpoint/orbax/checkpoint/checkpoint_utils.py

@@ -76,7 +76,12 @@ def _unlock_checkpoint(
   """Removes a LOCKED directory to indicate unlocking."""
   if multihost.process_index() == 0:
     logging.info('Unlocking existing step: %d.', step)
-    step_dir = step_name_format.find_step(checkpoint_dir, step).path
+    try:
+      step_dir = step_name_format.find_step(checkpoint_dir, step).path
+    except FileNotFoundError as e:
+      # Checkpoint no longer exists, so there is nothing to unlock.
+      logging.warning('Did not find checkpoint: %s', e)
+      return
     utils.lockdir(step_dir).unlink(missing_ok=True)
 
 

checkpoint/orbax/checkpoint/checkpoint_utils_test.py

@@ -302,6 +302,12 @@ def test_unlock_existing(self):
     self.assertFalse(utils.is_locked(self.directory / str(0)))
     self.assertFalse(utils.is_locked(self.directory / str(1)))
 
+  def test_unlock_deleted(self):
+    # Checkpoint does not exist; `_unlock_checkpoint` returns without raising.
+    checkpoint_utils._unlock_checkpoint(
+        self.directory, step=0, step_name_format=step_lib.standard_name_format()
+    )
+
   @parameterized.parameters(
       (None, None),
       (None, 8),

checkpoint/orbax/checkpoint/path/step.py

@@ -335,7 +335,7 @@ def find_step(self, base_path: epath.PathLike, step: int) -> Metadata:
       return metadata
 
     # Raise detailed error message.
-    raise ValueError(
+    raise FileNotFoundError(
         f'No step path found with name={self.build_name(step)},'
         f' NameFormat={self} for step={step} under {base_path}.'
     )