Keep forced checkpoints permanent

Author: dlwh

lib/levanter/src/levanter/checkpoint.py

@@ -521,14 +521,15 @@ def on_step(self, *, tree: PyTree, step: int, force: bool = False):
         my_should_save = force
         my_save_permanent_ckpt = force
 
-        current_every = self._get_current_step_save_interval(step)
-        last_save_time = self._dt_now_injection() - self._last_save_time
-        if current_every is not None and step % current_every == 0:
-            my_should_save = True
-            my_save_permanent_ckpt = True
-        elif self.save_interval and last_save_time >= self.save_interval:
-            my_should_save = True
-            my_save_permanent_ckpt = False
+        if not force:
+            current_every = self._get_current_step_save_interval(step)
+            last_save_time = self._dt_now_injection() - self._last_save_time
+            if current_every is not None and step % current_every == 0:
+                my_should_save = True
+                my_save_permanent_ckpt = True
+            elif self.save_interval and last_save_time >= self.save_interval:
+                my_should_save = True
+                my_save_permanent_ckpt = False
 
         should_save, save_permanent_ckpt = broadcast_one_to_all(
             jnp.array([my_should_save, my_save_permanent_ckpt], dtype=jnp.bool_)

lib/levanter/tests/test_checkpoint.py

@@ -540,6 +540,36 @@ def advance_time(delta_seconds):
         assert _get_checkpoint_steps(tmpdir) == [1, 2]
 
 
+def test_checkpointer_force_save_uses_permanent_path_even_when_time_policy_elapsed():
+    fake_now = datetime.datetime(2021, 1, 1, 0, 0, 0)
+    tick = 10
+
+    def advance_time(delta_seconds):
+        nonlocal fake_now
+        fake_now += timedelta(seconds=delta_seconds)
+
+    with (
+        tempfile.TemporaryDirectory(prefix="checkpoints") as permanent_dir,
+        tempfile.TemporaryDirectory(prefix="temp_checkpoints") as temporary_dir,
+    ):
+        checkpointer = Checkpointer(
+            permanent_dir,
+            timedelta(seconds=tick),
+            [],
+            temporary_base_path=temporary_dir,
+            dt_now_injection=lambda: fake_now,
+        )
+
+        _on_step(checkpointer, 0)
+
+        advance_time(tick)
+        _on_step(checkpointer, 1, force=True)
+        checkpointer.wait_until_finished()
+
+        assert _get_checkpoint_steps(permanent_dir) == [1]
+        assert list(pathlib.Path(temporary_dir).iterdir()) == []
+
+
 def test_load_from_checkpoint_or_initialize():
     In = Axis("in", 2)
     Out = Axis("out", 1)