fix: always leave a checkpoint when training is interrupted

- save an emergency checkpoint even when the signal is handled before the first optimizer step, where optimizer_step == last_ckpt_step used to skip the save entirely (CI race on slow cold-start runners)
- track a saved-this-session flag so interrupted runs always have at least one checkpoint, while normal zero-step completion keeps skipping the save

astrai/trainer/train_callback.py

@@ -175,6 +175,7 @@ class CheckpointCallback(TrainCallback):
         self.weight_only = weight_only
         self.save_extra_fn = save_extra_fn or CheckpointCallback.save_extra
         self.last_ckpt_step = None
+        self._saved = False
 
     def on_train_begin(self, context: TrainContext):
         self.last_ckpt_step = context.optimizer_step
@@ -205,6 +206,7 @@ class CheckpointCallback(TrainCallback):
                 context.checkpoint.save(save_path)
                 _copy_tokenizer_files(context.param_path, save_path)
         self.last_ckpt_step = context.optimizer_step
+        self._saved = True
 
     def after_optimizer_step(self, context: TrainContext):
         if context.optimizer_step - self.last_ckpt_step >= self.interval:
@@ -215,7 +217,11 @@ class CheckpointCallback(TrainCallback):
             self._save_checkpoint(context)
 
     def on_error(self, context: TrainContext):
-        if context.optimizer_step != self.last_ckpt_step:
+        # An interrupted run must always leave at least one checkpoint
+        # behind: on a slow start the signal can be handled before the
+        # first optimizer step, where optimizer_step == last_ckpt_step
+        # and the change-based guard alone would skip the save entirely.
+        if not self._saved or context.optimizer_step != self.last_ckpt_step:
             self._save_checkpoint(context)
 
     @staticmethod

tests/trainer/test_callbacks.py

@@ -10,7 +10,7 @@ from astrai.trainer.train_callback import (
     _copy_tokenizer_files,
 )
 from astrai.trainer.trainer import Trainer
-from tests.helpers import RandomTokenDataset
+from tests.helpers import RandomTokenDataset, load_checkpoint_meta
 
 
 def test_gradient_checkpointing_enable_disable(test_model):
@@ -213,3 +213,31 @@ def test_copy_tokenizer_files_skips_missing_and_none(tmp_path):
     _copy_tokenizer_files(None, str(tmp_path))
     _copy_tokenizer_files(str(tmp_path), str(tmp_path / "out"))
     assert not (tmp_path / "out").exists() or not any((tmp_path / "out").iterdir())
+
+
+def test_interrupt_before_first_step_still_checkpoints(
+    base_test_env, train_config_factory, device
+):
+    """A graceful interrupt handled before the first optimizer step must
+    still leave a checkpoint (regression: optimizer_step == last_ckpt_step
+    skipped the emergency save entirely)."""
+
+    class _ImmediateStop(TrainCallback):
+        def on_epoch_begin(self, context):
+            context.request_stop()
+
+    train_config = train_config_factory(
+        model_fn=lambda: base_test_env["model"],
+        dataset=RandomTokenDataset(length=2),
+        test_dir=base_test_env["test_dir"],
+        device=device,
+        batch_per_device=2,
+        ckpt_interval=1000,
+    )
+
+    trainer = Trainer(train_config)
+    trainer.callbacks.insert(0, _ImmediateStop())
+    trainer.train()
+
+    meta = load_checkpoint_meta(base_test_env["test_dir"])
+    assert meta["optimizer_step"] == 0