Fix checkpoint callback state keys

src/autocast/callbacks/checkpoint.py

@@ -95,6 +95,20 @@ def on_fit_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None
         super().on_fit_start(trainer, pl_module)
         self._maybe_resolve_fractional_train_steps(trainer)
 
+    @property
+    def state_key(self) -> str:
+        return self._generate_state_key(
+            monitor=self.monitor,
+            mode=self.mode,
+            every_n_train_steps=self._every_n_train_steps,
+            every_n_epochs=self._every_n_epochs,
+            train_time_interval=self._train_time_interval,
+            start_after_fraction=self.start_after_fraction,
+            stop_after_fraction=self.stop_after_fraction,
+            every_n_train_steps_fraction=self.every_n_train_steps_fraction,
+            monitor_optional=self.monitor_optional,
+        )
+
     def _maybe_resolve_fractional_train_steps(self, trainer: L.Trainer) -> None:
         if self.every_n_train_steps_fraction is None or self._resolved_fractional_steps:
             return

tests/scripts/test_training.py

@@ -11,6 +11,7 @@
 import torch
 from conftest import get_optimizer_config
 from hydra import compose, initialize_config_dir
+from hydra.utils import instantiate
 from lightning.pytorch.callbacks import ModelCheckpoint, Timer
 from matplotlib import pyplot as plt
 from omegaconf import DictConfig, OmegaConf, open_dict
@@ -326,6 +327,38 @@ def test_progress_model_checkpoint_disables_default_epoch_trigger():
     assert callback._every_n_epochs == 0
 
 
+def test_progress_model_checkpoint_state_key_includes_progress_window():
+    callbacks = [
+        ProgressModelCheckpoint(
+            monitor="val_multicoverage",
+            monitor_optional=True,
+            stop_after_fraction=0.25,
+            mode="min",
+            save_top_k=1,
+            filename="best-pre-{epoch:04d}",
+        ),
+        ProgressModelCheckpoint(
+            monitor="val_multicoverage",
+            monitor_optional=True,
+            start_after_fraction=0.25,
+            mode="min",
+            save_top_k=1,
+            filename="best-from0p25-{epoch:04d}",
+        ),
+        ProgressModelCheckpoint(
+            monitor="val_multicoverage",
+            monitor_optional=True,
+            start_after_fraction=0.5,
+            mode="min",
+            save_top_k=1,
+            filename="best-from0p50-{epoch:04d}",
+        ),
+    ]
+
+    state_keys = [callback.state_key for callback in callbacks]
+    assert len(state_keys) == len(set(state_keys))
+
+
 @pytest.mark.parametrize(
     "trigger_kwargs",
     [
@@ -558,6 +591,19 @@ def test_default_trainer_config_tracks_coverage_winkler_and_plots(config_dir: st
     )
 
 
+def test_default_trainer_config_instantiates_callbacks(config_dir: str):
+    trainer_cfg = OmegaConf.load(Path(config_dir) / "trainer" / "default.yaml")
+
+    trainer = instantiate(trainer_cfg)
+    progress_state_keys = [
+        callback.state_key
+        for callback in trainer.callbacks
+        if isinstance(callback, ProgressModelCheckpoint)
+    ]
+
+    assert len(progress_state_keys) == len(set(progress_state_keys))
+
+
 def test_epd_config_forward_smoke(config_dir: str, toy_batch: Batch, dummy_datamodule):
     model_cfg = _load_config(
         config_dir,