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/serialization/checkpoint.py
+158 -9
View File
@@ -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
+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