diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 1463dad928..bb1748a52a 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -4,6 +4,7 @@ from .checkpoint import ( CheckpointStore, build_checkpoint_stores, + resolve_checkpoint_path, resolve_keep_ckpt_count, ) from .data import ( @@ -67,6 +68,7 @@ "iter_training_task_configs", "make_task_maps", "print_data_summaries", + "resolve_checkpoint_path", "resolve_keep_ckpt_count", "resolve_step_schedule", ] diff --git a/deepmd/dpmodel/train/checkpoint.py b/deepmd/dpmodel/train/checkpoint.py index 5c3925c26f..110688339f 100644 --- a/deepmd/dpmodel/train/checkpoint.py +++ b/deepmd/dpmodel/train/checkpoint.py @@ -12,6 +12,7 @@ ) import logging +import os from math import ( ceil, ) @@ -37,7 +38,24 @@ log = logging.getLogger(__name__) -__all__ = ["CheckpointStore", "build_checkpoint_stores", "resolve_keep_ckpt_count"] +__all__ = [ + "CheckpointStore", + "build_checkpoint_stores", + "resolve_checkpoint_path", + "resolve_keep_ckpt_count", +] + + +def _atomic_write_text(path: Path, text: str) -> None: + """Replace ``path`` with ``text`` via a temporary sibling and ``os.replace``.""" + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + temporary.write_text(text) + os.replace(temporary, path) + except BaseException: + temporary.unlink(missing_ok=True) + raise class CheckpointStore: @@ -49,6 +67,12 @@ class CheckpointStore: Publishing a checkpoint points the prefix-named files at it, so a consumer that only knows the prefix always reaches the newest checkpoint. + Nested and absolute ``prefix`` values keep their directory for the + latest-checkpoint alias and the optional pointer file's default location + semantics; only the file name of ``prefix`` seeds numbered files under + ``save_dir``. When ``save_dir`` is unset, numbered files share the prefix + directory, which is the historical prefix-relative layout. + Parameters ---------- prefix : str or Path @@ -84,6 +108,11 @@ def __init__( self.suffix = suffix self.pointer_file = Path(pointer_file) if pointer_file is not None else None + @property + def alias_path(self) -> Path: + """Path of the prefix-named latest-checkpoint alias for this store.""" + return Path(f"{self.prefix}{self.suffix}") + def prepare(self) -> None: """Create the directories receiving the checkpoints and the symlinks.""" self.directory.mkdir(parents=True, exist_ok=True) @@ -159,7 +188,34 @@ def publish(self, path: Path) -> None: self.prefix.parent.mkdir(parents=True, exist_ok=True) symlink_prefix_files(str(path.with_suffix("")), str(self.prefix)) if self.pointer_file is not None: - self.pointer_file.write_text(str(path)) + _atomic_write_text(self.pointer_file, str(path)) + + def resolve_latest(self) -> Path | None: + """Return the checkpoint currently published by this store, if any. + + Prefers the pointer file when this store owns one, then the + prefix-named alias. Relative pointer entries are resolved against the + process working directory, matching how they were recorded. + """ + if self.pointer_file is not None and self.pointer_file.is_file(): + recorded = Path(self.pointer_file.read_text().strip()) + if recorded.is_file(): + return recorded + alias = self.alias_path + if alias.exists(): + return alias.resolve() if alias.is_symlink() else alias + return None + + def _protected_targets(self, current: Path) -> set[Path]: + """Paths retention must not delete: the fresh write and the latest alias.""" + protected: set[Path] = set() + for candidate in (current, self.alias_path): + try: + if candidate.exists(): + protected.add(candidate.resolve()) + except OSError: + continue + return protected def prune(self, current: Path) -> None: """Drop the checkpoints made obsolete by a fresh one. @@ -169,7 +225,10 @@ def prune(self, current: Path) -> None: them in place would let the retention window discard the freshly written checkpoint instead, so a rerun in a finished directory would keep no result at all. The window then retains the newest ``max_keep`` - checkpoints. The checkpoint just written is never removed. + checkpoints. The checkpoint just written and the target of the latest + alias are never removed, nor are unrelated files. Protected in-store + checkpoints still occupy slots of that window, so prune-before-publish + (PT) and publish-then-prune (pt_expt) both honor ``max_keep``. Parameters ---------- @@ -181,26 +240,114 @@ def prune(self, current: Path) -> None: if self.max_keep < 1: return current_step = self.step_of(current) if self.holds(current) else None + protected = self._protected_targets(current) retained: list[tuple[int, Path]] = [] + # Protected numbered files are skipped from deletion but still count + # toward ``max_keep``. Counting only ``current`` would under-budget + # when prune runs before publish and the previous latest alias target + # is a distinct in-store checkpoint. + occupied = 0 for path in self.directory.glob(f"*{self.suffix}"): step = self.step_of(path) if step is None or path.is_symlink(): continue - if current_step is not None and path.name == current.name: + try: + resolved = path.resolve() + except OSError: + continue + if resolved in protected: + if current_step is None or step <= current_step: + occupied += 1 continue if current_step is not None and step > current_step: path.unlink(missing_ok=True) else: retained.append((step, path)) retained.sort(key=lambda item: (item[0], item[1].name)) - # The current checkpoint occupies one slot of the window when this - # store holds it. - occupied = 1 if current_step is not None else 0 excess = max(0, len(retained) + occupied - self.max_keep) for _, path in retained[:excess]: + try: + if path.resolve() in protected: + continue + except OSError: + # Unresolvable path cannot be in ``protected``; fall through + # to unlink with missing_ok. + pass path.unlink(missing_ok=True) +def resolve_checkpoint_path( + spec: str | Path, + *, + store: CheckpointStore | None = None, + suffix: str = ".pt", +) -> Path: + """Resolve an explicit checkpoint, a prefix, or a latest-checkpoint handle. + + Parameters + ---------- + spec : str or Path + An existing checkpoint file, a checkpoint prefix (with or without the + suffix), a directory holding a ``checkpoint`` pointer / prefix alias, + or the literal ``latest`` when ``store`` is provided. + store : CheckpointStore, optional + Store consulted only when ``spec`` is the literal ``latest``. Directory + specs always read that directory's own pointer / prefix alias, matching + freeze-without-store behavior, so a foreign store cannot hijack them. + suffix : str, optional + Suffix appended when ``spec`` looks like a prefix rather than a file. + + Returns + ------- + Path + Path of an existing checkpoint file. + + Raises + ------ + FileNotFoundError + When no checkpoint can be resolved from ``spec``. + """ + text = str(spec).strip() + if text == "latest": + if store is None: + raise FileNotFoundError( + "Cannot resolve checkpoint 'latest' without a checkpoint store." + ) + resolved = store.resolve_latest() + if resolved is None: + raise FileNotFoundError( + "No latest checkpoint is published for the configured store." + ) + return resolved + + path = Path(text) + if path.is_file(): + return path + + if path.is_dir(): + # Directories always resolve against their own contents. ``store`` is + # reserved for the literal ``latest`` handle above; consulting it here + # would ignore ``dir`` and return another run's checkpoint. + pointer = path / "checkpoint" + if pointer.is_file(): + recorded = Path(pointer.read_text().strip()) + candidate = recorded if recorded.is_absolute() else path / recorded + if candidate.is_file(): + return candidate + if recorded.is_file(): + return recorded + for name in ("model.ckpt" + suffix, "model.ckpt.pt"): + alias = path / name + if alias.exists(): + return alias.resolve() if alias.is_symlink() else alias + raise FileNotFoundError(f"Cannot find a checkpoint in directory {path}.") + + suffixed = path if path.name.endswith(suffix) else Path(f"{path}{suffix}") + if suffixed.exists(): + return suffixed.resolve() if suffixed.is_symlink() else suffixed + raise FileNotFoundError(f"Cannot find checkpoint {spec!r}.") + + def resolve_keep_ckpt_count( ckpt_keep_ratio: float | None, num_steps: int, save_freq: int ) -> int | None: diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 7a381b4938..7a87936175 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -21,6 +21,7 @@ iter_training_task_configs, make_task_maps, print_data_summaries, + resolve_checkpoint_path, ) from deepmd.dpmodel.utils.lmdb_data import ( is_lmdb, @@ -889,24 +890,9 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None: output=FLAGS.output, ) elif FLAGS.command == "freeze": - if Path(FLAGS.checkpoint_folder).is_dir(): - checkpoint_path = Path(FLAGS.checkpoint_folder) - # pt_expt training saves a symlink "model.ckpt.pt" → latest ckpt - default_ckpt = checkpoint_path / "model.ckpt.pt" - if default_ckpt.exists(): - FLAGS.model = str(default_ckpt) - else: - raise FileNotFoundError( - f"Cannot find checkpoint in '{checkpoint_path}'. " - "Expected 'model.ckpt.pt' (created by pt_expt training)." - ) - else: - model_path = Path(FLAGS.checkpoint_folder) - if not model_path.exists(): - raise FileNotFoundError( - f"Checkpoint path '{model_path}' does not exist." - ) - FLAGS.model = str(model_path) + # Directories resolve through the checkpoint pointer / latest alias so + # a save_dir layout still freezes the published periodic checkpoint. + FLAGS.model = str(resolve_checkpoint_path(FLAGS.checkpoint_folder)) # Suffix defaulting lives in freeze(): the correct suffix depends on # the RESOLVED lower kind (native-spin models force 'graph'), which # is only known after the checkpoint's model is built there. diff --git a/deepmd/pt_expt/train/checkpoint.py b/deepmd/pt_expt/train/checkpoint.py new file mode 100644 index 0000000000..0619524f48 --- /dev/null +++ b/deepmd/pt_expt/train/checkpoint.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Torch checkpoint publication for the pt_expt training backend. + +Path layout and retention live in :mod:`deepmd.dpmodel.train.checkpoint`. This +module owns the Torch-specific serialization boundary: atomic ``torch.save``, +rank-gated filesystem side effects, and restart path resolution against the +configured stores. +""" + +from __future__ import ( + annotations, +) + +import os +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import torch + +from deepmd.dpmodel.train import ( + CheckpointStore, + build_checkpoint_stores, + resolve_checkpoint_path, +) + +if TYPE_CHECKING: + from collections.abc import ( + Mapping, + ) + +__all__ = ["TorchCheckpointManager"] + + +def _atomic_torch_save(payload: Mapping[str, Any], path: Path) -> None: + """Serialize ``payload`` to ``path`` without replacing it until complete.""" + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") + try: + torch.save(dict(payload), temporary) + os.replace(temporary, path) + except BaseException: + temporary.unlink(missing_ok=True) + raise + + +class TorchCheckpointManager: + """Filesystem owner for pt_expt checkpoint families. + + The trainer assembles checkpoint payloads (including any collective gather + required under sharding). This manager resolves numbered paths, writes them + atomically, updates latest aliases, and applies retention. Only the chief + rank performs filesystem mutations; other ranks are no-ops for write, + commit, and preparation side effects. + """ + + def __init__( + self, + store: CheckpointStore, + ema_store: CheckpointStore, + *, + rank: int = 0, + ) -> None: + self.store = store + self.ema_store = ema_store + self.rank = int(rank) + + @classmethod + def from_training_params( + cls, + training_params: Mapping[str, Any], + *, + num_steps: int, + ema_prefix: str | Path, + rank: int = 0, + ) -> TorchCheckpointManager: + """Build the manager from a normalized ``training`` section.""" + store, ema_store = build_checkpoint_stores( + training_params, + num_steps=num_steps, + ema_prefix=ema_prefix, + rank=rank, + ) + return cls(store, ema_store, rank=rank) + + @property + def is_chief(self) -> bool: + """Whether this rank owns checkpoint filesystem side effects.""" + return self.rank == 0 + + def store_for(self, *, ema: bool = False) -> CheckpointStore: + """Return the regular or EMA checkpoint store.""" + return self.ema_store if ema else self.store + + def path_for(self, step: int, *, ema: bool = False) -> Path: + """Return the numbered path for a periodic checkpoint step.""" + return self.store_for(ema=ema).path_for(step) + + def write(self, path: Path, payload: Mapping[str, Any]) -> None: + """Atomically serialize ``payload`` to ``path`` on the chief rank.""" + if not self.is_chief: + return + _atomic_torch_save(payload, Path(path)) + + def commit(self, path: Path, *, ema: bool = False) -> None: + """Publish the latest alias and apply retention after a successful write.""" + if not self.is_chief: + return + store = self.store_for(ema=ema) + store.publish(Path(path)) + store.prune(Path(path)) + + def save_step( + self, + step: int, + payload: Mapping[str, Any], + *, + ema: bool = False, + ) -> Path: + """Write, publish and prune a periodic checkpoint for ``step``. + + Returns + ------- + Path + The numbered path used for this step. Non-chief ranks still return + the path even though they perform no filesystem work, so callers + can log a single location. + """ + path = self.path_for(step, ema=ema) + self.write(path, payload) + self.commit(path, ema=ema) + return path + + def save_at(self, path: Path, payload: Mapping[str, Any]) -> None: + """Write a checkpoint to an explicit path without publishing or pruning. + + Validation-best checkpoints use an independent namespace and retention + policy owned by the validator; they must not move the periodic latest + alias. + """ + self.write(Path(path), payload) + + def resolve_restart(self, spec: str | Path, *, ema: bool = False) -> Path: + """Resolve an explicit or latest restart target for one checkpoint family.""" + return resolve_checkpoint_path(spec, store=self.store_for(ema=ema)) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..2820ee108b 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -56,7 +56,6 @@ TrainingTask, TrainingTaskCollection, TrainStepResult, - build_checkpoint_stores, change_model_out_bias, change_model_out_bias_by_task, resolve_step_schedule, @@ -111,6 +110,9 @@ from deepmd.pt_expt.model.graph_lower import ( model_uses_graph_lower, ) +from deepmd.pt_expt.train.checkpoint import ( + TorchCheckpointManager, +) from deepmd.pt_expt.train.ema import ( EMA_CHECKPOINT_KEY, ModelEMA, @@ -2029,8 +2031,9 @@ def initialize_statistics( # Checkpoint layout ---------------------------------------------------- # num_steps is final here, so a retention ratio can be converted into an - # absolute keep count once. - self.ckpt_store, self.ema_ckpt_store = build_checkpoint_stores( + # absolute keep count once. The manager owns path resolution, atomic + # publication, latest aliases, and retention for each family. + self.checkpoint_manager = TorchCheckpointManager.from_training_params( training_params, num_steps=self.num_steps, ema_prefix=self.ema_save_ckpt, @@ -2959,18 +2962,15 @@ def save_checkpoint(self, step: int) -> None: # Abort before writing if any gradient norm since the previous # checkpoint was non-finite, so a diverged interval is not persisted. self.nonfinite_grad_guard.raise_if_nonfinite(self.wrapper.named_parameters) - ckpt_path = self.ckpt_store.path_for(step) + ckpt_path = self.checkpoint_manager.path_for(step) self._save_checkpoint_to_path(ckpt_path, step=step) + self.checkpoint_manager.commit(ckpt_path) if self.rank == 0: - self.ckpt_store.publish(ckpt_path) - self.ckpt_store.prune(ckpt_path) log.info(f"Saved model to {ckpt_path}") if self.model_ema is not None: - ema_path = self.ema_ckpt_store.path_for(step) + ema_path = self.checkpoint_manager.path_for(step, ema=True) self._save_checkpoint_to_path(ema_path, step=step, use_ema_weights=True) - if self.rank == 0: - self.ema_ckpt_store.publish(ema_path) - self.ema_ckpt_store.prune(ema_path) + self.checkpoint_manager.commit(ema_path, ema=True) def _save_full_validation_checkpoint( self, @@ -3066,8 +3066,9 @@ def _write_checkpoint( state["optimizer"] = optim_state if include_ema_state and self.model_ema is not None: state[EMA_CHECKPOINT_KEY] = self.model_ema.state_dict() - ckpt_path.parent.mkdir(parents=True, exist_ok=True) - torch.save(state, ckpt_path) + # Payload assembly may need every rank under sharding; filesystem + # publication is chief-only and atomic inside the manager. + self.checkpoint_manager.write(ckpt_path, state) def _collect_checkpoint_states( self, diff --git a/source/tests/common/dpmodel/test_train_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py index c697a2e8c7..2330b5aed3 100644 --- a/source/tests/common/dpmodel/test_train_checkpoint.py +++ b/source/tests/common/dpmodel/test_train_checkpoint.py @@ -11,6 +11,7 @@ from deepmd.dpmodel.train import ( CheckpointStore, build_checkpoint_stores, + resolve_checkpoint_path, resolve_keep_ckpt_count, ) @@ -234,3 +235,113 @@ def test_ema_store_accepts_an_explicit_retention_override( assert store.max_keep == 7 assert ema_store.max_keep == 2 + + +def test_nested_and_absolute_prefix_resolution(tmp_path: Path) -> None: + """Only the prefix file name seeds numbered files under save_dir.""" + nested = CheckpointStore(tmp_path / "run" / "nested" / "model.ckpt") + assert nested.path_for(1) == tmp_path / "run" / "nested" / "model.ckpt-1.pt" + assert nested.alias_path == tmp_path / "run" / "nested" / "model.ckpt.pt" + + relocated = CheckpointStore( + tmp_path / "run" / "nested" / "model.ckpt", + save_dir=tmp_path / "models", + ) + assert relocated.path_for(1) == tmp_path / "models" / "model.ckpt-1.pt" + assert relocated.alias_path == tmp_path / "run" / "nested" / "model.ckpt.pt" + + absolute = CheckpointStore( + tmp_path / "abs" / "model.ckpt", + save_dir=tmp_path / "models", + ) + assert absolute.path_for(2) == tmp_path / "models" / "model.ckpt-2.pt" + assert absolute.alias_path == tmp_path / "abs" / "model.ckpt.pt" + + +def test_prune_never_deletes_the_latest_alias_target(tmp_path: Path) -> None: + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=1) + older = _write(store.path_for(1)) + newer = _write(store.path_for(2)) + store.publish(older) + + # Retention is asked to keep only the fresh write, but the published latest + # still names the older file and must survive. + store.prune(newer) + + assert older.exists() + assert newer.exists() + _assert_prefix_alias(tmp_path / "model.ckpt.pt", older, "model.ckpt-1.pt") + + +def test_prune_before_publish_counts_previous_latest_in_budget( + tmp_path: Path, +) -> None: + """PT writes prune before publish; the prior alias target still counts. + + Skipping every protected path without budgeting it would leave + ``max_keep + 1`` files when the fresh write and the previous latest are + distinct. Publish does not prune again, so the completed save must already + honor the configured window. + """ + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=2) + for step in (2, 3, 4): + _write(store.path_for(step)) + store.publish(store.path_for(3)) + + store.prune(store.path_for(4)) + + assert not store.path_for(2).exists() + assert store.path_for(3).exists() + assert store.path_for(4).exists() + _assert_prefix_alias( + tmp_path / "model.ckpt.pt", store.path_for(3), "model.ckpt-3.pt" + ) + + +def test_resolve_latest_prefers_pointer_then_alias(tmp_path: Path) -> None: + store = CheckpointStore( + tmp_path / "model.ckpt", + save_dir=tmp_path / "ckpts", + pointer_file=tmp_path / "checkpoint", + ) + store.prepare() + path = _write(store.path_for(4)) + store.publish(path) + + assert store.resolve_latest() == path + + assert resolve_checkpoint_path("latest", store=store) == path + assert resolve_checkpoint_path(tmp_path, store=store) == path + assert resolve_checkpoint_path(tmp_path / "model.ckpt") == path + + +def test_resolve_directory_ignores_foreign_store(tmp_path: Path) -> None: + """A directory spec must read THAT directory, even when a store is passed. + + ``store`` is only for the literal ``latest`` handle. Consulting it for a + directory would return another run's checkpoint (e.g. resolving run B's + directory while holding store A). + """ + run_a = tmp_path / "run_a" + run_b = tmp_path / "run_b" + store_a = CheckpointStore( + run_a / "model.ckpt", + save_dir=run_a / "ckpts", + pointer_file=run_a / "checkpoint", + ) + store_b = CheckpointStore( + run_b / "model.ckpt", + save_dir=run_b / "ckpts", + pointer_file=run_b / "checkpoint", + ) + store_a.prepare() + store_b.prepare() + path_a = _write(store_a.path_for(1)) + path_b = _write(store_b.path_for(2)) + store_a.publish(path_a) + store_b.publish(path_b) + + assert resolve_checkpoint_path("latest", store=store_a) == path_a + assert resolve_checkpoint_path(run_b, store=store_a) == path_b + assert resolve_checkpoint_path(run_b) == path_b + assert resolve_checkpoint_path(run_a, store=store_b) == path_a diff --git a/source/tests/pt_expt/test_checkpoint_manager.py b/source/tests/pt_expt/test_checkpoint_manager.py new file mode 100644 index 0000000000..d6c5f9b843 --- /dev/null +++ b/source/tests/pt_expt/test_checkpoint_manager.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the pt_expt Torch checkpoint manager.""" + +from __future__ import ( + annotations, +) + +from pathlib import ( + Path, +) + +import pytest +import torch + +from deepmd.pt_expt.train.checkpoint import ( + TorchCheckpointManager, +) + + +def test_save_dir_holds_periodic_files_and_cwd_aliases( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + manager = TorchCheckpointManager.from_training_params( + { + "save_ckpt": "model.ckpt", + "save_dir": "models", + "max_ckpt_keep": 2, + "save_freq": 1, + }, + num_steps=4, + ema_prefix="model.ckpt_ema", + rank=0, + ) + + for step in (1, 2, 3): + manager.save_step(step, {"model": {"step": step}, "optimizer": {}}) + + assert sorted(path.name for path in Path("models").glob("model.ckpt-*.pt")) == [ + "model.ckpt-2.pt", + "model.ckpt-3.pt", + ] + assert Path("model.ckpt.pt").resolve() == (tmp_path / "models" / "model.ckpt-3.pt") + assert Path("checkpoint").read_text() == str(Path("models") / "model.ckpt-3.pt") + assert not (Path("models") / "model.ckpt.pt").exists() + + +def test_interrupted_write_leaves_previous_latest_intact( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + manager = TorchCheckpointManager.from_training_params( + {"save_ckpt": "model.ckpt", "save_dir": "models", "max_ckpt_keep": 5}, + num_steps=2, + ema_prefix="model.ckpt_ema", + ) + first = manager.save_step(1, {"model": {"step": 1}}) + assert torch.load(first, weights_only=True)["model"]["step"] == 1 + + real_save = torch.save + + def boom(obj, f, *args, **kwargs): + if hasattr(f, "name") and str(f.name).endswith(".tmp"): + raise RuntimeError("disk full") + return real_save(obj, f, *args, **kwargs) + + monkeypatch.setattr(torch, "save", boom) + with pytest.raises(RuntimeError, match="disk full"): + manager.save_step(2, {"model": {"step": 2}}) + + assert first.exists() + assert Path("model.ckpt.pt").resolve() == first.resolve() + assert Path("checkpoint").read_text() == str(first) + assert not (tmp_path / "models" / "model.ckpt-2.pt").exists() + assert list(Path("models").glob(".model.ckpt-2.pt.*.tmp")) == [] + + +def test_non_chief_rank_performs_no_filesystem_mutations( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + manager = TorchCheckpointManager.from_training_params( + {"save_ckpt": "model.ckpt", "save_dir": "models"}, + num_steps=1, + ema_prefix="model.ckpt_ema", + rank=1, + ) + path = manager.save_step(1, {"model": {}}) + assert path == Path("models") / "model.ckpt-1.pt" + assert not path.exists() + assert not Path("model.ckpt.pt").exists() + assert not Path("checkpoint").exists() + + +def test_resolve_restart_from_explicit_and_latest( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + manager = TorchCheckpointManager.from_training_params( + {"save_ckpt": "nested/model.ckpt", "save_dir": "models"}, + num_steps=1, + ema_prefix="nested/model.ckpt_ema", + ) + written = manager.save_step(7, {"model": {"ok": True}}) + + assert manager.resolve_restart(written) == written + assert manager.resolve_restart("latest") == written + assert manager.resolve_restart(".") == written + assert manager.resolve_restart("nested/model.ckpt").resolve() == written.resolve() + with pytest.raises(FileNotFoundError): + manager.resolve_restart("typo.ckpt") + + +def test_validation_save_does_not_move_latest_alias( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + manager = TorchCheckpointManager.from_training_params( + {"save_ckpt": "model.ckpt", "save_dir": "models"}, + num_steps=1, + ema_prefix="model.ckpt_ema", + ) + periodic = manager.save_step(1, {"model": {"kind": "periodic"}}) + best = tmp_path / "best" / "model.ckpt-1.pt" + manager.save_at(best, {"model": {"kind": "best"}}) + + assert torch.load(best, weights_only=True)["model"]["kind"] == "best" + assert Path("model.ckpt.pt").resolve() == periodic.resolve()