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>

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