fix: publish checkpoints atomically

- 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>

astrai/serialization/checkpoint.py

@@ -1,7 +1,11 @@
 """Model checkpoint serialization helpers."""
 
+import hashlib
 import io
 import json
+import os
+import shutil
+import tempfile
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -16,6 +20,8 @@ from astrai.parallel.setup import get_rank
 _META_FILE = "meta.json"
 _CONFIG_FILE = "config.json"
 _WEIGHTS_FILE = "model.safetensors"
+_MANIFEST_FILE = "manifest.json"
+_CHECKPOINT_FORMAT_VERSION = 1
 
 
 def save_safetensors(state_dict: dict, path: Union[str, Path]):
@@ -79,6 +85,85 @@ def load_torch(path: Union[str, Path], broadcast: bool = False) -> Any:
     return torch.load(buf, map_location="cpu", weights_only=False)
 
 
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as file:
+        for chunk in iter(lambda: file.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _sync_file(path: Path):
+    with path.open("rb") as file:
+        os.fsync(file.fileno())
+
+
+def _sync_directory(path: Path):
+    flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)
+    fd = os.open(path, flags)
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def _checkpoint_manifest(
+    save_path: Path,
+    state_dict: Dict[str, Any],
+    meta: Dict[str, Any],
+) -> Dict[str, Any]:
+    files = {}
+    for path in sorted(save_path.iterdir()):
+        if path.is_file() and path.name != _MANIFEST_FILE:
+            files[path.name] = {
+                "size": path.stat().st_size,
+                "sha256": _sha256(path),
+            }
+    return {
+        "format_version": _CHECKPOINT_FORMAT_VERSION,
+        "created_at": meta["timestamp"],
+        "optimizer_step": meta.get("optimizer_step"),
+        "policy_version": meta.get("policy_version"),
+        "tensors": sorted(state_dict),
+        "files": files,
+    }
+
+
+def _validate_manifest(
+    save_path: Path,
+    verify_checksums: bool = False,
+    broadcast: bool = False,
+) -> dict:
+    def validate() -> dict:
+        manifest = load_json(save_path / _MANIFEST_FILE)
+        if manifest.get("format_version") != _CHECKPOINT_FORMAT_VERSION:
+            raise ValueError(
+                "Unsupported checkpoint format version: "
+                f"{manifest.get('format_version')}"
+            )
+
+        files = manifest.get("files")
+        if not isinstance(files, dict):
+            raise ValueError("Checkpoint manifest has no file table")
+        for required in (_META_FILE, _CONFIG_FILE, _WEIGHTS_FILE):
+            if required not in files:
+                raise ValueError(f"Checkpoint manifest is missing {required}")
+
+        for name, descriptor in files.items():
+            if Path(name).name != name or not isinstance(descriptor, dict):
+                raise ValueError(f"Invalid checkpoint manifest entry: {name!r}")
+            path = save_path / name
+            if not path.is_file():
+                raise ValueError(f"Checkpoint file is missing: {name}")
+            if path.stat().st_size != descriptor.get("size"):
+                raise ValueError(f"Checkpoint file size mismatch: {name}")
+            if verify_checksums and _sha256(path) != descriptor.get("sha256"):
+                raise ValueError(f"Checkpoint file checksum mismatch: {name}")
+        return manifest
+
+    return _broadcast_load(validate, broadcast)
+
+
 def save_model(config: dict, state_dict: dict, save_directory: str):
     save_path = Path(save_directory)
     save_path.mkdir(parents=True, exist_ok=True)
@@ -151,7 +236,18 @@ class Checkpoint:
 
     def save(self, save_dir: str):
         save_path = Path(save_dir)
-        save_path.mkdir(parents=True, exist_ok=True)
+        save_path.parent.mkdir(parents=True, exist_ok=True)
+        if save_path.exists() and not save_path.is_dir():
+            raise FileExistsError(
+                f"Checkpoint path exists and is not a directory: {save_path}"
+            )
+
+        staging_path = Path(
+            tempfile.mkdtemp(
+                prefix=f".{save_path.name}.tmp-",
+                dir=save_path.parent,
+            )
+        )
 
         meta = {
             "epoch": self.epoch,
@@ -159,16 +255,60 @@ class Checkpoint:
             "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
             **self.meta,
         }
-        save_json(meta, save_path / _META_FILE)
-        save_json(self.config, save_path / _CONFIG_FILE)
-        save_safetensors(self.state_dict, save_path / _WEIGHTS_FILE)
-        for key, value in self.extra.items():
-            save_torch(value, save_path / f"{key}.pt")
+        retired_path: Optional[Path] = None
+        try:
+            save_json(meta, staging_path / _META_FILE)
+            save_json(self.config, staging_path / _CONFIG_FILE)
+            save_safetensors(self.state_dict, staging_path / _WEIGHTS_FILE)
+            for key, value in self.extra.items():
+                save_torch(value, staging_path / f"{key}.pt")
+
+            manifest = _checkpoint_manifest(staging_path, self.state_dict, meta)
+            save_json(manifest, staging_path / _MANIFEST_FILE)
+            for path in staging_path.iterdir():
+                if path.is_file():
+                    _sync_file(path)
+            _sync_directory(staging_path)
+
+            # Re-publishing an existing step (re-runs into the same output
+            # directory) atomically retires the old payload first; a crash
+            # between the two renames leaves the previous checkpoint hidden
+            # under the retired name instead of a partial directory.
+            retired_path = None
+            if save_path.exists():
+                retired_path = Path(
+                    tempfile.mkdtemp(
+                        prefix=f".{save_path.name}.retired-",
+                        dir=save_path.parent,
+                    )
+                )
+                retired_path.rmdir()
+                os.replace(save_path, retired_path)
+            os.replace(staging_path, save_path)
+            _sync_directory(save_path.parent)
+        except BaseException:
+            shutil.rmtree(staging_path, ignore_errors=True)
+            raise
+        finally:
+            if retired_path is not None:
+                shutil.rmtree(retired_path, ignore_errors=True)
 
     @classmethod
-    def load(cls, save_dir: str, broadcast: bool = False) -> "Checkpoint":
+    def load(
+        cls,
+        save_dir: str,
+        broadcast: bool = False,
+        verify_checksums: bool = False,
+    ) -> "Checkpoint":
         save_path = Path(save_dir)
 
+        if (save_path / _MANIFEST_FILE).exists():
+            _validate_manifest(
+                save_path,
+                verify_checksums=verify_checksums,
+                broadcast=broadcast,
+            )
+
         meta = load_json(save_path / _META_FILE, broadcast)
         config = load_json(save_path / _CONFIG_FILE, broadcast)
         state_dict = load_state_dict(save_path / _WEIGHTS_FILE, broadcast=broadcast)
@@ -188,13 +328,22 @@ class Checkpoint:
         )
 
     @classmethod
-    def load_any(cls, save_dir: str, broadcast: bool = False) -> Optional["Checkpoint"]:
+    def load_any(
+        cls,
+        save_dir: str,
+        broadcast: bool = False,
+        verify_checksums: bool = False,
+    ) -> Optional["Checkpoint"]:
         save_path = Path(save_dir)
         meta_path = save_path / _META_FILE
         weights_path = save_path / _WEIGHTS_FILE
 
         if meta_path.exists():
-            return cls.load(save_dir, broadcast=broadcast)
+            return cls.load(
+                save_dir,
+                broadcast=broadcast,
+                verify_checksums=verify_checksums,
+            )
 
         weights_path = save_path / _WEIGHTS_FILE
         index_path = save_path / "model.safetensors.index.json"

astrai/trainer/train_callback.py

@@ -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

docs/developer/architecture.md

@@ -259,8 +259,8 @@ classDiagram
             +dict meta
             +dict config
             +save(save_dir)
-            +load(save_dir, broadcast) Checkpoint
-            +load_any(save_dir, broadcast) Optional[Checkpoint]
+            +load(save_dir, broadcast, verify_checksums) Checkpoint
+            +load_any(save_dir, broadcast, verify_checksums) Optional[Checkpoint]
         }
     }
 

docs/developer/docker-training.md

@@ -183,8 +183,14 @@ config.json
 model.safetensors
 optimizer.pt
 scheduler.pt
+manifest.json
 ```
 
+New checkpoints write `manifest.json` after every payload file, sync the complete
+staging directory, and then atomically rename that directory into place. Legacy
+checkpoints without a manifest remain resumable when the original required files
+are complete.
+
 `start` resumes the latest complete checkpoint and ignores partial writes. If no
 complete checkpoint exists, `/models/base/config.json` and
 `/models/base/model.safetensors` are required. `stop` sends `SIGTERM`; the

docs/guides/distributed.md

@@ -165,7 +165,9 @@ Checkpoints are saved by **rank-0 only**. The flow:
    - `ddp`: `model.module.state_dict()`
    - `fsdp`: `unshard()` → `full_tensor()` → `reshard()` (collective on all ranks, result kept only on rank-0)
 3. Non-rank-0 ranks get `None` — the save is skipped.
-4. Rank-0 writes `meta.json`, `config.json`, `model.safetensors`, and optional `{key}.pt` (optimizer/scheduler state).
+4. Rank-0 writes metadata, weights, optional optimizer/scheduler state, and a
+   checksum manifest to a hidden sibling directory, then atomically renames the
+   complete checkpoint into place.
 
 > **FSDP note**: Even though only rank-0 saves, all ranks must participate in `unwrap_model` because `unshard()` and `full_tensor()` are collective operations. The barriers in `checkpoint_context` keep all ranks in lockstep.
 

docs/guides/training.md

@@ -196,11 +196,19 @@ Callback wraps each `DecoderBlock.forward` with `torch.utils.checkpoint.checkpoi
 
 ```
 Checkpoint(state_dict, epoch, consumed_samples, extra, meta, config)
-  ├── save(save_dir)    meta.json (epoch/consumed_samples/timestamp) + config.json (model config) + model.safetensors + optional {key}.pt (optimizer.pt, scheduler.pt)
-  └── load(save_dir, broadcast=False)    loads from local disk; set broadcast=True to broadcast metadata from rank-0
+  ├── save(save_dir)    atomically publishes manifest.json + metadata + weights + optional {key}.pt
+  └── load(save_dir, broadcast=False, verify_checksums=False)    loads locally or broadcasts from rank-0
 ```
 
-`Checkpoint.save()` writes whenever it is called. During training, `CheckpointCallback` uses the executor checkpoint context so only rank 0 receives a state dict and calls `save()`.
+`Checkpoint.save()` writes to a hidden sibling staging directory, records file
+sizes and SHA-256 checksums in `manifest.json`, flushes the files, and atomically
+renames the completed directory into place. Published checkpoint directories are
+immutable: saving to an existing non-empty path raises `FileExistsError`. Legacy
+checkpoints without a manifest remain loadable. Pass `verify_checksums=True` when
+loading to hash every published file.
+
+During training, `CheckpointCallback` uses the executor checkpoint context so
+only rank 0 receives a state dict and calls `save()`.
 
 Optimizer/scheduler state persisted by default via `Checkpoint.extra`.  
 Model config (`context.model_config`) saved into `config.json` during training via `CheckpointCallback`.

tests/data/test_checkpoint.py

@@ -1,6 +1,8 @@
+import json
 import os
 import tempfile
 
+import pytest
 import torch
 import torch.distributed as dist
 from torch.optim import AdamW
@@ -8,6 +10,7 @@ from torch.optim.lr_scheduler import CosineAnnealingLR
 
 from astrai.parallel.setup import get_rank, spawn_parallel_fn
 from astrai.serialization import Checkpoint
+from astrai.serialization import checkpoint as checkpoint_module
 
 
 def test_single_process():
@@ -65,6 +68,103 @@ def test_checkpoint_with_extra():
         assert "state" in loaded.extra["optimizer"]
 
 
+def test_checkpoint_is_atomically_published_with_manifest(tmp_path, monkeypatch):
+    target = tmp_path / "epoch_1_step_9"
+    checkpoint = Checkpoint(
+        state_dict={"weight": torch.arange(4)},
+        epoch=1,
+        consumed_samples=36,
+        meta={"optimizer_step": 9, "policy_version": 3},
+        config={"hidden_size": 4},
+    )
+    original_save = checkpoint_module.save_safetensors
+
+    def observe_staging(state_dict, path):
+        assert not target.exists()
+        assert path.parent.name.startswith(f".{target.name}.tmp-")
+        original_save(state_dict, path)
+
+    monkeypatch.setattr(checkpoint_module, "save_safetensors", observe_staging)
+    checkpoint.save(target)
+
+    assert target.is_dir()
+    assert not list(tmp_path.glob(f".{target.name}.tmp-*"))
+    manifest = json.loads((target / "manifest.json").read_text())
+    assert manifest["format_version"] == 1
+    assert manifest["optimizer_step"] == 9
+    assert manifest["policy_version"] == 3
+    assert manifest["tensors"] == ["weight"]
+    assert set(manifest["files"]) == {
+        "config.json",
+        "meta.json",
+        "model.safetensors",
+    }
+    assert len(manifest["files"]["model.safetensors"]["sha256"]) == 64
+    assert Checkpoint.load(target, verify_checksums=True).consumed_samples == 36
+
+
+def test_checkpoint_failure_never_publishes_partial_directory(tmp_path, monkeypatch):
+    target = tmp_path / "epoch_0_step_1"
+    checkpoint = Checkpoint(state_dict={"weight": torch.ones(2)})
+
+    def fail_save(*args, **kwargs):
+        assert not target.exists()
+        raise RuntimeError("injected write failure")
+
+    monkeypatch.setattr(checkpoint_module, "save_safetensors", fail_save)
+    with pytest.raises(RuntimeError, match="injected write failure"):
+        checkpoint.save(target)
+
+    assert not target.exists()
+    assert not list(tmp_path.glob(f".{target.name}.tmp-*"))
+
+
+def test_checkpoint_republish_atomically_replaces_published_directory(tmp_path):
+    target = tmp_path / "epoch_0_step_1"
+    Checkpoint(state_dict={"weight": torch.ones(2)}).save(target)
+    replacement = Checkpoint(
+        state_dict={"weight": torch.zeros(3)},
+        meta={"optimizer_step": 1},
+        config={"hidden_size": 3},
+    )
+    replacement.save(target)
+
+    assert target.is_dir()
+    assert not list(tmp_path.glob(f".{target.name}.tmp-*"))
+    assert not list(tmp_path.glob(f".{target.name}.retired-*"))
+    loaded = Checkpoint.load(target, verify_checksums=True)
+    torch.testing.assert_close(loaded.state_dict["weight"], torch.zeros(3))
+
+
+def test_checkpoint_checksum_verification_detects_same_size_corruption(tmp_path):
+    target = tmp_path / "epoch_0_step_1"
+    Checkpoint(state_dict={"weight": torch.ones(2)}).save(target)
+    meta_path = target / "meta.json"
+    corrupted = meta_path.read_bytes()
+    replacement = b"X" if corrupted[0:1] != b"X" else b"Y"
+    meta_path.write_bytes(replacement + corrupted[1:])
+
+    with pytest.raises(ValueError, match="checksum mismatch: meta.json"):
+        Checkpoint.load(target, verify_checksums=True)
+
+
+def test_checkpoint_load_accepts_legacy_directory_without_manifest(tmp_path):
+    target = tmp_path / "legacy"
+    target.mkdir()
+    checkpoint_module.save_json(
+        {"epoch": 2, "consumed_samples": 20}, target / "meta.json"
+    )
+    checkpoint_module.save_json({"hidden_size": 2}, target / "config.json")
+    checkpoint_module.save_safetensors(
+        {"weight": torch.ones(2)}, target / "model.safetensors"
+    )
+
+    loaded = Checkpoint.load(target, verify_checksums=True)
+
+    assert loaded.epoch == 2
+    assert loaded.consumed_samples == 20
+
+
 def simple_training():
     model = torch.nn.Linear(10, 5)
     optimizer = AdamW(model.parameters(), lr=1e-3)

tests/trainer/test_callbacks.py

@@ -170,3 +170,7 @@ def test_checkpoint_captures_completed_optimizer_step(
     )
     assert checkpoint.extra["optimizer"]["state"]
     assert checkpoint.extra["scheduler"]["last_epoch"] == 1
+    assert checkpoint.meta["optimizer_step"] == 1
+    assert (
+        Path(base_test_env["test_dir"]) / "epoch_0_step_1" / "metric.jsonl"
+    ).is_file()