diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index 831b669..a18fdcb 100644 --- a/astrai/trainer/train_callback.py +++ b/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 diff --git a/tests/trainer/test_callbacks.py b/tests/trainer/test_callbacks.py index ad97881..3194755 100644 --- a/tests/trainer/test_callbacks.py +++ b/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