fix(jax): write the checkpoint pointer beside save_ckpt

[wanghan-iapcm]

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.

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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: d713d61a-c160-4cbd-8c29-c2e37d5ca021

📥 Commits

Reviewing files that changed from the base of the PR and between bc2095b and 767425c.

📒 Files selected for processing (1)

• source/tests/jax/test_checkpoint_pointer.py

🚧 Files skipped from review as they are similar to previous changes (1)

• source/tests/jax/test_checkpoint_pointer.py

---

📝 Walkthrough

Walkthrough

DPTrainer._save_checkpoint now writes the checkpoint pointer beside the configured save_ckpt prefix and stores its basename. Tests cover subdirectory and bare-name checkpoint paths.

Changes

Checkpoint pointer fix

Layer / File(s)	Summary
Write pointer relative to save_ckpt directory deepmd/jax/train/trainer.py	_save_checkpoint creates the save_ckpt directory when needed and writes the relative checkpoint pointer there.
Pointer placement tests source/tests/jax/test_checkpoint_pointer.py	Tests isolate the working directory and verify pointer location and basename content for subdirectory and bare-name paths.

Estimated code review effort: 2 (Simple) | ~10 minutes

Possibly related issues

• [Code scan] Write JAX checkpoint pointers beside save_ckpt #5678: Directly addresses the checkpoint pointer placement behavior implemented here.

Possibly related PRs

• deepmodeling/deepmd-kit#5603: Touches the same JAX checkpoint-saving and pointer-linking logic.

Suggested reviewers: njzjz, iProzd

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main fix: writing the JAX checkpoint pointer beside save_ckpt.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.53%. Comparing base (0acd0e5) to head (767425c).
⚠️ Report is 3 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5726      +/-   ##
==========================================
- Coverage   79.62%   79.53%   -0.09%     
==========================================
  Files        1014     1015       +1     
  Lines      115533   115740     +207     
  Branches     4276     4274       -2     
==========================================
+ Hits        91995    92056      +61     
- Misses      21994    22139     +145     
- Partials     1544     1545       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.