fix(jax): write the checkpoint pointer beside save_ckpt (#5726)

## Problem

Fixes #5678. JAX training writes checkpoint directories and the stable
`.jax` link relative to `save_ckpt` (which may include a directory), but
always wrote the `checkpoint` pointer file to the current working
directory with a value that still carried the directory prefix, e.g.
`runs/water/model.ckpt.jax`. The freeze entrypoint looks for the pointer
inside the folder it is given and resolves the pointer's value relative
to that folder (`checkpoint_folder / pointer`). So for `save_ckpt =
runs/water/model.ckpt`, the pointer was written to `./checkpoint` (not
`runs/water/checkpoint`) and, even if relocated, its value would have
double-prefixed to `runs/water/runs/water/model.ckpt.jax`. Passing
`runs/water` to freeze or restart-style tooling could not find or
resolve the checkpoint, even though the matching checkpoint directory
and `.jax` link were written there.

## Fix

Write the pointer into `Path(save_ckpt).parent` and store a value
relative to that directory (the basename only). For the default bare
`save_ckpt` (parent is `.`) the pointer stays in the CWD with the same
value, so existing behavior is unchanged; only directory-valued
`save_ckpt` is affected.

## Test

Adds `source/tests/jax/test_checkpoint_pointer.py`, which drives
`_save_checkpoint` with the checkpoint I/O mocked. The directory case
asserts the pointer lands beside the checkpoint (`subdir/checkpoint`)
with a basename value (`model.ckpt.jax`) and not in the CWD — this fails
on master — and a bare-name control asserts the pointer stays in the
CWD. The trainer's pointer writing previously had no coverage; the
existing freeze test hand-wrote a correct pointer and never exercised
the writer.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Bug Fixes**
- Fixed checkpoint “pointer” file placement to be written alongside the
configured checkpoint save directory/path rather than always in the
current working directory.
- Updated pointer contents to reference the correct checkpoint name,
improving downstream checkpoint resolution during freeze/restart.

- **Tests**
- Added unit tests validating pointer location and contents for both
nested save paths and bare checkpoint filenames.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/jax/train/trainer.py

@@ -841,8 +841,14 @@ def _save_checkpoint(self, step: int) -> None:
         log.info(f"Trained model has been saved to: {ckpt_path!s}")
         _link_checkpoint(ckpt_path, Path(f"{self.save_ckpt}.jax"))
         self._cleanup_old_checkpoints()
-        with open("checkpoint", "w") as fp:
-            fp.write(f"{self.save_ckpt}.jax")
+        # Write the pointer next to the checkpoint prefix, with a value relative
+        # to that directory (basename only). The freeze entrypoint looks for the
+        # pointer inside the folder it is given and resolves the value relative
+        # to it, so a directory-valued save_ckpt would otherwise be unresolvable.
+        ckpt_dir = Path(self.save_ckpt).parent
+        ckpt_dir.mkdir(parents=True, exist_ok=True)
+        with open(ckpt_dir / "checkpoint", "w") as fp:
+            fp.write(f"{Path(self.save_ckpt).name}.jax")
 
     def _save_full_validation_checkpoint(
         self,

source/tests/jax/test_checkpoint_pointer.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Test the JAX trainer writes its checkpoint pointer beside save_ckpt.
+
+JAX training writes checkpoint directories and the stable ``.jax`` link relative
+to ``save_ckpt`` (which may include a directory), but used to always write the
+``checkpoint`` pointer file to the current working directory with a value that
+still carried the directory prefix. The freeze entrypoint expects the pointer
+inside the folder it is given and resolves its value relative to that folder, so
+a directory-valued ``save_ckpt`` broke freeze/restart tooling.
+"""
+
+import os
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+from unittest.mock import (
+    patch,
+)
+
+from deepmd.jax.train.trainer import (
+    DPTrainer,
+)
+
+
+class TestCheckpointPointer(unittest.TestCase):
+    def setUp(self) -> None:
+        self.tmpdir = tempfile.TemporaryDirectory()
+        self.cwd = os.getcwd()
+        os.chdir(self.tmpdir.name)
+
+    def tearDown(self) -> None:
+        os.chdir(self.cwd)
+        self.tmpdir.cleanup()
+
+    def _save(self, save_ckpt: str) -> None:
+        trainer = DPTrainer.__new__(DPTrainer)
+        trainer.save_ckpt = save_ckpt
+        with (
+            patch.object(DPTrainer, "_write_checkpoint"),
+            patch.object(DPTrainer, "_cleanup_old_checkpoints"),
+            patch("deepmd.jax.train.trainer._link_checkpoint"),
+        ):
+            trainer._save_checkpoint(1)
+
+    def test_pointer_beside_ckpt_for_subdir(self) -> None:
+        # save_ckpt with a directory: pointer must land in that directory with a
+        # value relative to it (basename only), so freeze(subdir) resolves it.
+        self._save("subdir/model.ckpt")
+        pointer = Path("subdir") / "checkpoint"
+        self.assertTrue(pointer.is_file())
+        self.assertEqual(pointer.read_text(), "model.ckpt.jax")
+        self.assertFalse(Path("checkpoint").exists())
+
+    def test_pointer_in_cwd_for_bare_name(self) -> None:
+        # default/bare save_ckpt: pointer stays in the CWD (parent is ".").
+        self._save("model.ckpt")
+        pointer = Path("checkpoint")
+        self.assertTrue(pointer.is_file())
+        self.assertEqual(pointer.read_text(), "model.ckpt.jax")
+
+
+if __name__ == "__main__":
+    unittest.main()