From b6928d8388c2d1089fa5aef65ff4405067d071ed Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 23:33:41 +0800
Subject: [PATCH 1/6] feat(pt_expt): honor training.save_dir via
 TorchCheckpointManager

Introduce a pt_expt-owned checkpoint manager on top of the shared
CheckpointStore layout so periodic paths, atomic publication, latest
aliases, retention, and restart resolution share one filesystem owner.

Trainer code now supplies payloads only; chief-rank atomic torch.save
plus os.replace keeps an interrupted write from replacing the previous
latest checkpoint. Refs #5817.
---
 deepmd/dpmodel/train/__init__.py              |   2 +
 deepmd/dpmodel/train/checkpoint.py            | 147 ++++++++++++++++-
 deepmd/pt_expt/entrypoints/main.py            |  24 +--
 deepmd/pt_expt/train/checkpoint.py            | 152 ++++++++++++++++++
 deepmd/pt_expt/train/training.py              |  27 ++--
 .../common/dpmodel/test_train_checkpoint.py   |  54 +++++++
 .../tests/pt_expt/test_checkpoint_manager.py  | 124 ++++++++++++++
 7 files changed, 496 insertions(+), 34 deletions(-)
 create mode 100644 deepmd/pt_expt/train/checkpoint.py
 create mode 100644 source/tests/pt_expt/test_checkpoint_manager.py

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..5c69c1b5be 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,8 @@ 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.
 
         Parameters
         ----------
@@ -181,12 +238,17 @@ 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]] = []
         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:
                 continue
             if current_step is not None and step > current_step:
                 path.unlink(missing_ok=True)
@@ -198,9 +260,86 @@ def prune(self, current: Path) -> None:
         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:
+                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 when ``spec`` is ``latest`` or a bare directory that
+        should follow that store's pointer and alias conventions.
+    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():
+        if store is not None:
+            resolved = store.resolve_latest()
+            if resolved is not None and resolved.is_file():
+                return resolved
+        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..667f5c13b5 100644
--- a/deepmd/pt_expt/entrypoints/main.py
+++ b/deepmd/pt_expt/entrypoints/main.py
@@ -42,6 +42,9 @@
 from deepmd.utils.stat_file import (
     StatFileSpec,
 )
+from deepmd.dpmodel.train import (
+    resolve_checkpoint_path,
+)
 from deepmd.utils.summary import SummaryPrinter as BaseSummaryPrinter
 
 log = logging.getLogger(__name__)
@@ -889,24 +892,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..934e66b868
--- /dev/null
+++ b/deepmd/pt_expt/train/checkpoint.py
@@ -0,0 +1,152 @@
+# 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 logging
+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,
+    )
+
+log = logging.getLogger(__name__)
+
+__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..9fc88b125e 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,13 +2031,16 @@ 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,
             rank=self.rank,
         )
+        self.ckpt_store = self.checkpoint_manager.store
+        self.ema_ckpt_store = self.checkpoint_manager.ema_store
 
         # Learning rate -------------------------------------------------------
         self.lr_schedule = make_learning_rate_schedule(
@@ -2959,18 +2964,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 +3068,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..4894876b7f 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,56 @@ 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_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
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..73eb8cb365
--- /dev/null
+++ b/source/tests/pt_expt/test_checkpoint_manager.py
@@ -0,0 +1,124 @@
+# 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()

From a235559add92d919ae22beedaa5d8815ac4417ec Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Oct 2026 15:34:54 +0000
Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 deepmd/pt_expt/entrypoints/main.py              | 4 +---
 source/tests/pt_expt/test_checkpoint_manager.py | 8 ++++++--
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py
index 667f5c13b5..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,
@@ -42,9 +43,6 @@
 from deepmd.utils.stat_file import (
     StatFileSpec,
 )
-from deepmd.dpmodel.train import (
-    resolve_checkpoint_path,
-)
 from deepmd.utils.summary import SummaryPrinter as BaseSummaryPrinter
 
 log = logging.getLogger(__name__)
diff --git a/source/tests/pt_expt/test_checkpoint_manager.py b/source/tests/pt_expt/test_checkpoint_manager.py
index 73eb8cb365..d6c5f9b843 100644
--- a/source/tests/pt_expt/test_checkpoint_manager.py
+++ b/source/tests/pt_expt/test_checkpoint_manager.py
@@ -1,9 +1,13 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Tests for the pt_expt Torch checkpoint manager."""
 
-from __future__ import annotations
+from __future__ import (
+    annotations,
+)
 
-from pathlib import Path
+from pathlib import (
+    Path,
+)
 
 import pytest
 import torch

From 92bdbbe28536d8cd2a4f76271b311c7cc3e44336 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Fri, 9 Oct 2026 23:47:00 +0800
Subject: [PATCH 3/6] fix(train): resolve directory checkpoints against that
 directory

resolve_checkpoint_path(dir, store=...) previously returned
store.resolve_latest() whenever store was set, so resolving run B's
directory while holding store A yielded A's checkpoint.

Use store only for the literal "latest" handle; directories always read
their own checkpoint pointer / prefix alias. Drop unused trainer
aliases of the manager stores. Add a two-directory regression test.
---
 deepmd/dpmodel/train/checkpoint.py            | 12 +++----
 deepmd/pt_expt/train/training.py              |  2 --
 .../common/dpmodel/test_train_checkpoint.py   | 32 +++++++++++++++++++
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/deepmd/dpmodel/train/checkpoint.py b/deepmd/dpmodel/train/checkpoint.py
index 5c69c1b5be..0bbbed1365 100644
--- a/deepmd/dpmodel/train/checkpoint.py
+++ b/deepmd/dpmodel/train/checkpoint.py
@@ -283,8 +283,9 @@ def resolve_checkpoint_path(
         suffix), a directory holding a ``checkpoint`` pointer / prefix alias,
         or the literal ``latest`` when ``store`` is provided.
     store : CheckpointStore, optional
-        Store consulted when ``spec`` is ``latest`` or a bare directory that
-        should follow that store's pointer and alias conventions.
+        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.
 
@@ -316,10 +317,9 @@ def resolve_checkpoint_path(
         return path
 
     if path.is_dir():
-        if store is not None:
-            resolved = store.resolve_latest()
-            if resolved is not None and resolved.is_file():
-                return resolved
+        # 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())
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index 9fc88b125e..2820ee108b 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -2039,8 +2039,6 @@ def initialize_statistics(
             ema_prefix=self.ema_save_ckpt,
             rank=self.rank,
         )
-        self.ckpt_store = self.checkpoint_manager.store
-        self.ema_ckpt_store = self.checkpoint_manager.ema_store
 
         # Learning rate -------------------------------------------------------
         self.lr_schedule = make_learning_rate_schedule(
diff --git a/source/tests/common/dpmodel/test_train_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py
index 4894876b7f..225ed2faed 100644
--- a/source/tests/common/dpmodel/test_train_checkpoint.py
+++ b/source/tests/common/dpmodel/test_train_checkpoint.py
@@ -288,3 +288,35 @@ def test_resolve_latest_prefers_pointer_then_alias(tmp_path: Path) -> None:
     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
+

From ee83c3c1f5e46b3dcf33831675c28aa1e24bd412 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 9 Oct 2026 15:48:04 +0000
Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 source/tests/common/dpmodel/test_train_checkpoint.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/source/tests/common/dpmodel/test_train_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py
index 225ed2faed..67fed193c0 100644
--- a/source/tests/common/dpmodel/test_train_checkpoint.py
+++ b/source/tests/common/dpmodel/test_train_checkpoint.py
@@ -289,6 +289,7 @@ def test_resolve_latest_prefers_pointer_then_alias(tmp_path: Path) -> None:
     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.
 
@@ -319,4 +320,3 @@ def test_resolve_directory_ignores_foreign_store(tmp_path: Path) -> None:
     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
-

From 5c3a6df74286f634b26f38d32f96b91f8b153b93 Mon Sep 17 00:00:00 2001
From: OutisLi-Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 00:41:13 +0800
Subject: [PATCH 5/6] chore(train): clarify retention except and drop unused
 logger

Document why OSError from path.resolve is ignored during retention, and
remove an unused module logger from TorchCheckpointManager.
---
 deepmd/dpmodel/train/checkpoint.py | 2 ++
 deepmd/pt_expt/train/checkpoint.py | 3 ---
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/deepmd/dpmodel/train/checkpoint.py b/deepmd/dpmodel/train/checkpoint.py
index 0bbbed1365..dcc282714b 100644
--- a/deepmd/dpmodel/train/checkpoint.py
+++ b/deepmd/dpmodel/train/checkpoint.py
@@ -264,6 +264,8 @@ def prune(self, current: Path) -> None:
                 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)
 
diff --git a/deepmd/pt_expt/train/checkpoint.py b/deepmd/pt_expt/train/checkpoint.py
index 934e66b868..0619524f48 100644
--- a/deepmd/pt_expt/train/checkpoint.py
+++ b/deepmd/pt_expt/train/checkpoint.py
@@ -11,7 +11,6 @@
     annotations,
 )
 
-import logging
 import os
 from pathlib import (
     Path,
@@ -34,8 +33,6 @@
         Mapping,
     )
 
-log = logging.getLogger(__name__)
-
 __all__ = ["TorchCheckpointManager"]
 
 

From 7b578dc805db5af2f967fdf24fdd7ecbe2d85a00 Mon Sep 17 00:00:00 2001
From: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com>
Date: Sat, 10 Oct 2026 10:43:21 +0800
Subject: [PATCH 6/6] fix(train): count protected latest alias in retention
 budget

Prune skipped every protected path from the deletable set but only
budgeted the fresh checkpoint. PT prunes before publish, so the previous
latest-alias target was a second protected in-store file that silently
exceeded max_keep. Count all protected in-window numbered checkpoints
toward the window while still refusing to delete them.
---
 deepmd/dpmodel/train/checkpoint.py            | 14 ++++++++---
 .../common/dpmodel/test_train_checkpoint.py   | 25 +++++++++++++++++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/deepmd/dpmodel/train/checkpoint.py b/deepmd/dpmodel/train/checkpoint.py
index dcc282714b..110688339f 100644
--- a/deepmd/dpmodel/train/checkpoint.py
+++ b/deepmd/dpmodel/train/checkpoint.py
@@ -226,7 +226,9 @@ def prune(self, current: Path) -> None:
         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 and the target of the latest
-        alias are never removed, nor are unrelated files.
+        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
         ----------
@@ -240,6 +242,11 @@ def prune(self, current: Path) -> None:
         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():
@@ -249,15 +256,14 @@ def prune(self, current: Path) -> None:
             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:
diff --git a/source/tests/common/dpmodel/test_train_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py
index 67fed193c0..2330b5aed3 100644
--- a/source/tests/common/dpmodel/test_train_checkpoint.py
+++ b/source/tests/common/dpmodel/test_train_checkpoint.py
@@ -273,6 +273,31 @@ def test_prune_never_deletes_the_latest_alias_target(tmp_path: Path) -> None:
     _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",