ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: publish checkpoints atomically

Browse Source 
- Write checkpoint payloads to a hidden sibling staging directory, add a versioned checksum manifest, fsync the completed payload, and publish it with an atomic rename
- Republishing an existing step retires the old payload under a hidden sibling name before the atomic rename, so re-runs into the same output directory replace the previous checkpoint instead of raising FileExistsError
- Keep legacy checkpoints loadable, add optional checksum verification, and align metric flushing with checkpoint publication

Co-authored-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
This commit is contained in:
ViperEkura
2026-09-02 15:29:22 +08:00
committed by 0z5a
co-authored by 0z5a
parent 01bcd0d105
commit 1fad50d847
8 changed files with 292 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/trainer/train_callback.py
+8 -4
View File
@@ -153,8 +153,6 @@ class CheckpointCallback(TrainCallback):
self.last_ckpt_step = context.optimizer_step
def _save_checkpoint(self, context: TrainContext):
self.last_ckpt_step = context.optimizer_step
with context.executor.checkpoint_context(context.model) as state_dict:
if state_dict is not None:
save_path = os.path.join(
@@ -162,7 +160,10 @@ class CheckpointCallback(TrainCallback):
f"epoch_{context.epoch}_step_{context.optimizer_step}",
)
extra = self.save_extra_fn(context)
meta = context.config.to_dict()
meta = {
**context.config.to_dict(),
"optimizer_step": context.optimizer_step,
}
context.checkpoint = Checkpoint(
state_dict=state_dict,
epoch=context.epoch,
@@ -172,6 +173,7 @@ class CheckpointCallback(TrainCallback):
meta=meta,
)
context.checkpoint.save(save_path)
self.last_ckpt_step = context.optimizer_step
def after_optimizer_step(self, context: TrainContext):
if context.optimizer_step - self.last_ckpt_step >= self.interval:
@@ -182,7 +184,8 @@ class CheckpointCallback(TrainCallback):
self._save_checkpoint(context)
def on_error(self, context: TrainContext):
self._save_checkpoint(context)
if context.optimizer_step != self.last_ckpt_step:
self._save_checkpoint(context)
@staticmethod
def save_extra(context: TrainContext) -> dict:
@@ -361,6 +364,7 @@ class MetricCallback(TrainCallback):
step_metrics = [m for m in self.metrics if m != "val_loss"]
self._append("step", context, **self._metrics(context, step_metrics))
def after_optimizer_step(self, context):
if context.optimizer_step - self.last_log_flush_step >= self.save_interval:
self._flush(context.epoch, context.optimizer_step)
self.last_log_flush_step = context.optimizer_step