From 1fad50d8476abe7d4b4cb527eb6e174c511fd409 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 2 Sep 2026 12:40:34 +0800 Subject: [PATCH] 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 --- astrai/serialization/checkpoint.py | 167 +++++++++++++++++++++++++++-- astrai/trainer/train_callback.py | 12 ++- docs/developer/architecture.md | 4 +- docs/developer/docker-training.md | 6 ++ docs/guides/distributed.md | 4 +- docs/guides/training.md | 14 ++- tests/data/test_checkpoint.py | 100 +++++++++++++++++ tests/trainer/test_callbacks.py | 4 + 8 files changed, 292 insertions(+), 19 deletions(-) diff --git a/astrai/serialization/checkpoint.py b/astrai/serialization/checkpoint.py index d72c699..70eecfb 100644 --- a/astrai/serialization/checkpoint.py +++ b/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" diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index a391837..ba5c1ef 100644 --- a/astrai/trainer/train_callback.py +++ b/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 diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 791d331..eb2b181 100644 --- a/docs/developer/architecture.md +++ b/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] } } diff --git a/docs/developer/docker-training.md b/docs/developer/docker-training.md index 0049ae8..991ee53 100644 --- a/docs/developer/docker-training.md +++ b/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 diff --git a/docs/guides/distributed.md b/docs/guides/distributed.md index 95adb10..063567d 100644 --- a/docs/guides/distributed.md +++ b/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. diff --git a/docs/guides/training.md b/docs/guides/training.md index 6e4a64a..e90ff7a 100644 --- a/docs/guides/training.md +++ b/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`. diff --git a/tests/data/test_checkpoint.py b/tests/data/test_checkpoint.py index 9d8c858..f19abbc 100644 --- a/tests/data/test_checkpoint.py +++ b/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) diff --git a/tests/trainer/test_callbacks.py b/tests/trainer/test_callbacks.py index de283b1..85613ae 100644 --- a/tests/trainer/test_callbacks.py +++ b/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()