fix(train): scope statistics cache lifecycle#5883
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| self._writable = True | ||
|
|
||
|
|
||
| class _H5StatPath(DPPath): |
| self._writable = True | ||
|
|
||
|
|
||
| class _H5StatPath(DPPath): |
| (path / "value").save_numpy(np.array([1.0])) | ||
| raise RuntimeError("statistics failed") | ||
|
|
||
| assert path is not None |
|
|
||
| def test_hdf5_handle_closes_after_exception(tmp_path: Path) -> None: | ||
| target = tmp_path / "stat.hdf5" | ||
| path = None |
📝 WalkthroughWalkthroughThe PR introduces shared ChangesUnified statistics cache
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
deepmd/pt_expt/train/training.py (1)
1535-1535: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider documenting the two-phase
share_paramscall.The flow is correct (verified): chief recomputes/merges shared stats with
resume=Falseand ties params, values are broadcast to peers, thenresume=Trueis called unconditionally on every rank (including chief, redundantly) to (re-)establish the intra-process parameter tying using the now-synced values — since object-identity sharing can't be broadcast across processes, only tensor values can. This is subtle enough that a short inline comment explaining whyshare_paramsruns twice with differentresumevalues would meaningfully help future readers.Also applies to: 1547-1571
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt_expt/train/training.py` at line 1535, Add a concise inline comment around the two-phase share_params flow near synchronize_model_state, documenting that the chief merges and broadcasts shared-stat values with resume=False, then every rank calls resume=True to restore intra-process parameter identity after cross-process synchronization. Keep the existing call order and behavior unchanged.deepmd/pt/train/training.py (1)
1007-1035: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto the multi-taskzip()formodel_key_prob_map.Static analysis flags
zip(self.model_keys, self.model_prob)(Line 1019) withoutstrict=. Althoughself.model_probis always derived fromself.model_keysand should match in length, the sibling JAX implementation (deepmd/jax/train/trainer.py,_share_model_params) already guards the equivalent call withzip(self.model_keys, model_prob, strict=True). Addingstrict=Truehere costs nothing and fails fast instead of silently truncating/misaligning task↔probability pairs if this invariant is ever broken by a future refactor.🔧 Proposed fix
share_kwargs = { - "model_key_prob_map": dict(zip(self.model_keys, self.model_prob)), + "model_key_prob_map": dict( + zip(self.model_keys, self.model_prob, strict=True) + ), "data_stat_protect": _data_stat_protect[0], }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt/train/training.py` around lines 1007 - 1035, Update the model_key_prob_map construction in the multi-task shared-parameter block to call zip(self.model_keys, self.model_prob, strict=True), matching the invariant enforced by the sibling implementation while preserving the existing dictionary construction.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@deepmd/pt_expt/train/training.py`:
- Line 1535: Add a concise inline comment around the two-phase share_params flow
near synchronize_model_state, documenting that the chief merges and broadcasts
shared-stat values with resume=False, then every rank calls resume=True to
restore intra-process parameter identity after cross-process synchronization.
Keep the existing call order and behavior unchanged.
In `@deepmd/pt/train/training.py`:
- Around line 1007-1035: Update the model_key_prob_map construction in the
multi-task shared-parameter block to call zip(self.model_keys, self.model_prob,
strict=True), matching the invariant enforced by the sibling implementation
while preserving the existing dictionary construction.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: c7b8fe8a-bd3e-4a95-b78a-9606e19e13b6
📒 Files selected for processing (35)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/descriptor/descriptor.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/repflows.pydeepmd/dpmodel/descriptor/repformers.pydeepmd/dpmodel/descriptor/se_e2_a.pydeepmd/dpmodel/descriptor/se_r.pydeepmd/dpmodel/descriptor/se_t.pydeepmd/dpmodel/descriptor/se_t_tebd.pydeepmd/dpmodel/fitting/general_fitting.pydeepmd/dpmodel/train/data.pydeepmd/dpmodel/utils/env_mat_stat.pydeepmd/dpmodel/utils/stat.pydeepmd/jax/common.pydeepmd/jax/train/trainer.pydeepmd/pt/entrypoints/main.pydeepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/task/fitting.pydeepmd/pt/train/training.pydeepmd/pt/utils/stat.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/train/training.pydeepmd/tf2/entrypoints/train.pydeepmd/tf2/train/trainer.pydeepmd/utils/argcheck.pydeepmd/utils/stat_file.pysource/tests/common/dpmodel/test_train_data.pysource/tests/common/stat_file.pysource/tests/common/test_stat_file.pysource/tests/jax/test_model_factory.pysource/tests/jax/test_training.pysource/tests/pt/test_stat_file_mode.pysource/tests/pt_expt/test_entrypoint.pysource/tests/pt_expt/test_training.pysource/tests/tf2/test_training.py
💤 Files with no reviewable changes (1)
- source/tests/pt_expt/test_entrypoint.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5883 +/- ##
==========================================
- Coverage 78.82% 78.61% -0.21%
==========================================
Files 1054 1055 +1
Lines 121769 121987 +218
Branches 4410 4413 +3
==========================================
- Hits 95985 95905 -80
- Misses 24216 24522 +306
+ Partials 1568 1560 -8 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt/train/training.py`:
- Around line 1018-1021: Update the model_key_prob_map construction in the
share_kwargs block to call zip(self.model_keys, self.model_prob) with
strict=True, preserving the existing dictionary conversion while enforcing
equal-length inputs.
In `@deepmd/pt/utils/stat.py`:
- Around line 114-125: Update _restore_from_file to return (None, None) whenever
cached_keys does not include every requested key, allowing compute_output_stats
to recompute statistics for newly added outputs; retain the current dictionaries
only when all keys are cached.
- Line 116: Replace the mutable default value for keys in the relevant function
with a safe immutable or None-based default, initializing ["energy"] inside the
function when no keys are provided. Preserve the existing behavior for callers
that supply keys and ensure Ruff B006 passes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 40a39df4-d3e1-41c3-a1c0-0a004587f8a2
📒 Files selected for processing (35)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/descriptor/descriptor.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/repflows.pydeepmd/dpmodel/descriptor/repformers.pydeepmd/dpmodel/descriptor/se_e2_a.pydeepmd/dpmodel/descriptor/se_r.pydeepmd/dpmodel/descriptor/se_t.pydeepmd/dpmodel/descriptor/se_t_tebd.pydeepmd/dpmodel/fitting/general_fitting.pydeepmd/dpmodel/train/data.pydeepmd/dpmodel/utils/env_mat_stat.pydeepmd/dpmodel/utils/stat.pydeepmd/jax/common.pydeepmd/jax/train/trainer.pydeepmd/pt/entrypoints/main.pydeepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/task/fitting.pydeepmd/pt/train/training.pydeepmd/pt/utils/stat.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/train/training.pydeepmd/tf2/entrypoints/train.pydeepmd/tf2/train/trainer.pydeepmd/utils/argcheck.pydeepmd/utils/stat_file.pysource/tests/common/dpmodel/test_train_data.pysource/tests/common/stat_file.pysource/tests/common/test_stat_file.pysource/tests/jax/test_model_factory.pysource/tests/jax/test_training.pysource/tests/pt/test_stat_file_mode.pysource/tests/pt_expt/test_entrypoint.pysource/tests/pt_expt/test_training.pysource/tests/tf2/test_training.py
💤 Files with no reviewable changes (1)
- source/tests/pt_expt/test_entrypoint.py
🚧 Files skipped from review as they are similar to previous changes (27)
- deepmd/dpmodel/descriptor/se_r.py
- deepmd/dpmodel/train/data.py
- deepmd/dpmodel/descriptor/repformers.py
- source/tests/tf2/test_training.py
- source/tests/common/stat_file.py
- deepmd/dpmodel/atomic_model/base_atomic_model.py
- deepmd/dpmodel/descriptor/se_t.py
- source/tests/common/dpmodel/test_train_data.py
- deepmd/dpmodel/descriptor/se_e2_a.py
- deepmd/dpmodel/utils/env_mat_stat.py
- deepmd/dpmodel/descriptor/descriptor.py
- source/tests/pt_expt/test_training.py
- deepmd/pt/model/task/fitting.py
- source/tests/jax/test_model_factory.py
- deepmd/jax/common.py
- deepmd/dpmodel/descriptor/dpa1.py
- deepmd/jax/train/trainer.py
- deepmd/tf2/entrypoints/train.py
- deepmd/pt/entrypoints/main.py
- source/tests/jax/test_training.py
- deepmd/dpmodel/utils/stat.py
- deepmd/utils/stat_file.py
- deepmd/dpmodel/fitting/general_fitting.py
- source/tests/common/test_stat_file.py
- deepmd/pt_expt/train/training.py
- deepmd/pt_expt/entrypoints/main.py
- source/tests/pt/test_stat_file_mode.py
| share_kwargs = { | ||
| "model_key_prob_map": dict(zip(self.model_keys, self.model_prob)), | ||
| "data_stat_protect": _data_stat_protect[0], | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add strict=True to zip().
zip(self.model_keys, self.model_prob) has no strict= parameter; Ruff (B905) flags this. Since model_prob is always derived from model_keys the lengths should already match, but strict=True documents that invariant and fails fast if it's ever violated.
🔧 Proposed fix
share_kwargs = {
- "model_key_prob_map": dict(zip(self.model_keys, self.model_prob)),
+ "model_key_prob_map": dict(
+ zip(self.model_keys, self.model_prob, strict=True)
+ ),
"data_stat_protect": _data_stat_protect[0],
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| share_kwargs = { | |
| "model_key_prob_map": dict(zip(self.model_keys, self.model_prob)), | |
| "data_stat_protect": _data_stat_protect[0], | |
| } | |
| share_kwargs = { | |
| "model_key_prob_map": dict( | |
| zip(self.model_keys, self.model_prob, strict=True) | |
| ), | |
| "data_stat_protect": _data_stat_protect[0], | |
| } |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 1019-1019: zip() without an explicit strict= parameter
Add explicit value for parameter strict=
(B905)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@deepmd/pt/train/training.py` around lines 1018 - 1021, Update the
model_key_prob_map construction in the share_kwargs block to call
zip(self.model_keys, self.model_prob) with strict=True, preserving the existing
dictionary conversion while enforcing equal-length inputs.
Source: Linters/SAST tools
| def _restore_from_file( | ||
| stat_file_path: DPPath, | ||
| stat_file_path: DPPath | None, | ||
| keys: list[str] = ["energy"], | ||
| ) -> dict | None: | ||
| if stat_file_path is None: | ||
| return None, None | ||
| _require_stat_file_items( | ||
| stat_file_path, | ||
| [item for key in keys for item in (f"bias_atom_{key}", f"std_atom_{key}")], | ||
| ) | ||
| stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys] | ||
| if all(not (ii.is_file()) for ii in stat_files): | ||
| ) -> tuple[dict | None, dict | None]: | ||
| pairs = [(f"bias_atom_{key}", f"std_atom_{key}") for key in keys] | ||
| items = load_paired_items(stat_file_path, pairs) | ||
| if items is None: | ||
| return None, None | ||
| stat_files = [stat_file_path / f"std_atom_{kk}" for kk in keys] | ||
| if all(not (ii.is_file()) for ii in stat_files): | ||
| return None, None | ||
|
|
||
| ret_bias = {} | ||
| ret_std = {} | ||
| for kk in keys: | ||
| fp = stat_file_path / f"bias_atom_{kk}" | ||
| # only read the key that exists | ||
| if fp.is_file(): | ||
| ret_bias[kk] = fp.load_numpy() | ||
| for kk in keys: | ||
| fp = stat_file_path / f"std_atom_{kk}" | ||
| # only read the key that exists | ||
| if fp.is_file(): | ||
| ret_std[kk] = fp.load_numpy() | ||
| cached_keys = [key for key in keys if f"bias_atom_{key}" in items] | ||
| ret_bias = {key: items[f"bias_atom_{key}"] for key in cached_keys} | ||
| ret_std = {key: items[f"std_atom_{key}"] for key in cached_keys} | ||
| return ret_bias, ret_std |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Partial cache coverage is silently accepted instead of triggering recompute.
load_paired_items can legitimately return a dict covering only some of the requested keys in update mode (a fully-uncached output is simply excluded, not treated as "missing"). _restore_from_file doesn't check that cached_keys covers every requested key before returning a non-None result, so compute_output_stats's if bias_atom_e is None: guard (line 308) never fires for this case — statistics for a newly added output key (e.g. adding a property loss to a config that already has a fully-cached energy stat file in update mode) are silently dropped from the returned bias_atom_e/std_atom_e instead of being computed.
🐛 Proposed fix
items = load_paired_items(stat_file_path, pairs)
if items is None:
return None, None
cached_keys = [key for key in keys if f"bias_atom_{key}" in items]
+ if len(cached_keys) != len(keys):
+ # Some requested keys were never cached (e.g. a newly added output);
+ # force the caller to recompute the full set instead of silently
+ # returning stats for only the previously cached keys.
+ return None, None
ret_bias = {key: items[f"bias_atom_{key}"] for key in cached_keys}
ret_std = {key: items[f"std_atom_{key}"] for key in cached_keys}
return ret_bias, ret_std📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _restore_from_file( | |
| stat_file_path: DPPath, | |
| stat_file_path: DPPath | None, | |
| keys: list[str] = ["energy"], | |
| ) -> dict | None: | |
| if stat_file_path is None: | |
| return None, None | |
| _require_stat_file_items( | |
| stat_file_path, | |
| [item for key in keys for item in (f"bias_atom_{key}", f"std_atom_{key}")], | |
| ) | |
| stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys] | |
| if all(not (ii.is_file()) for ii in stat_files): | |
| ) -> tuple[dict | None, dict | None]: | |
| pairs = [(f"bias_atom_{key}", f"std_atom_{key}") for key in keys] | |
| items = load_paired_items(stat_file_path, pairs) | |
| if items is None: | |
| return None, None | |
| stat_files = [stat_file_path / f"std_atom_{kk}" for kk in keys] | |
| if all(not (ii.is_file()) for ii in stat_files): | |
| return None, None | |
| ret_bias = {} | |
| ret_std = {} | |
| for kk in keys: | |
| fp = stat_file_path / f"bias_atom_{kk}" | |
| # only read the key that exists | |
| if fp.is_file(): | |
| ret_bias[kk] = fp.load_numpy() | |
| for kk in keys: | |
| fp = stat_file_path / f"std_atom_{kk}" | |
| # only read the key that exists | |
| if fp.is_file(): | |
| ret_std[kk] = fp.load_numpy() | |
| cached_keys = [key for key in keys if f"bias_atom_{key}" in items] | |
| ret_bias = {key: items[f"bias_atom_{key}"] for key in cached_keys} | |
| ret_std = {key: items[f"std_atom_{key}"] for key in cached_keys} | |
| return ret_bias, ret_std | |
| def _restore_from_file( | |
| stat_file_path: DPPath | None, | |
| keys: list[str] = ["energy"], | |
| ) -> tuple[dict | None, dict | None]: | |
| pairs = [(f"bias_atom_{key}", f"std_atom_{key}") for key in keys] | |
| items = load_paired_items(stat_file_path, pairs) | |
| if items is None: | |
| return None, None | |
| cached_keys = [key for key in keys if f"bias_atom_{key}" in items] | |
| if len(cached_keys) != len(keys): | |
| # Some requested keys were never cached (e.g. a newly added output); | |
| # force the caller to recompute the full set instead of silently | |
| # returning stats for only the previously cached keys. | |
| return None, None | |
| ret_bias = {key: items[f"bias_atom_{key}"] for key in cached_keys} | |
| ret_std = {key: items[f"std_atom_{key}"] for key in cached_keys} | |
| return ret_bias, ret_std |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 116-116: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@deepmd/pt/utils/stat.py` around lines 114 - 125, Update _restore_from_file to
return (None, None) whenever cached_keys does not include every requested key,
allowing compute_output_stats to recompute statistics for newly added outputs;
retain the current dictionaries only when all keys are cached.
| def _restore_from_file( | ||
| stat_file_path: DPPath, | ||
| stat_file_path: DPPath | None, | ||
| keys: list[str] = ["energy"], |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Mutable default argument.
keys: list[str] = ["energy"] uses a mutable default; Ruff (B006) flags this. Not exploitable today since keys is never mutated in-place, but flagged by static analysis and coding guidelines require ruff check . to pass.
🔧 Proposed fix
def _restore_from_file(
stat_file_path: DPPath | None,
- keys: list[str] = ["energy"],
+ keys: list[str] | None = None,
) -> tuple[dict | None, dict | None]:
+ if keys is None:
+ keys = ["energy"]
pairs = [(f"bias_atom_{key}", f"std_atom_{key}") for key in keys]As per coding guidelines: **/*.py: "Install linter and run ruff check . before committing changes or the CI will fail."
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| keys: list[str] = ["energy"], | |
| def _restore_from_file( | |
| stat_file_path: DPPath | None, | |
| keys: list[str] | None = None, | |
| ) -> tuple[dict | None, dict | None]: | |
| if keys is None: | |
| keys = ["energy"] |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 116-116: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@deepmd/pt/utils/stat.py` at line 116, Replace the mutable default value for
keys in the relevant function with a safe immutable or None-based default,
initializing ["energy"] inside the function when no keys are provided. Preserve
the existing behavior for callers that supply keys and ensure Ruff B006 passes.
Sources: Coding guidelines, Linters/SAST tools
Summary
readandupdatemodesRoot cause
Statistics files were opened through
DPH5Path, whose process-widelru_cacheretained liveh5py.Fileobjects for the entire training process. Update-mode cache hits could therefore keep a writable HDF5 handle open and modify file metadata even when no statistics needed to be recomputed. Cache completeness and distributed initialization were also implemented independently by each backend, which allowed partial output pairs and divergent rank-local statistics.Behavior
Existing HDF5 caches in
updatemode are opened read-only and promoted to a writer only after a cache miss. A complete cache hit does not sample training data, acquire a writer handle, or change the file bytes. Every scoped handle is closed when statistics initialization finishes.Output bias and standard-deviation datasets are treated as complete pairs. Replacements use a transient transaction marker, remove stale requested pairs, write only newly computed complete pairs, flush them, and then remove the marker. A successful cache retains the legacy HDF5 layout; an interrupted transaction is recomputed in
updatemode and rejected in strictreadmode.This change intentionally does not modify Paddle, legacy TF1, or HDF5 training-data access.
Validation
git diff --checkpassedPT-expt compile-specific training tests were not completed locally because the Python process aborted inside a generated TorchInductor CPU kernel. The corresponding non-compile paths passed.
Summary by CodeRabbit
New Features
stat_file_mode(read/update) via a sharedStatFileSpecabstraction across training backends.Bug Fixes
Tests