fix(pt): allow concurrent stat cache reads#5773
Conversation
Add explicit read-only access while preserving update mode by default.
📝 WalkthroughWalkthroughTraining now supports ChangesStatistics cache mode handling
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingConfig
participant PTEntrypoint
participant StatisticsCache
participant StatisticsValidation
participant HDF5
TrainingConfig->>PTEntrypoint: provide stat_file and stat_file_mode
PTEntrypoint->>StatisticsCache: prepare cache on rank 0
StatisticsCache->>HDF5: open in read or update mode
HDF5-->>StatisticsCache: return cache path
StatisticsCache->>StatisticsValidation: require expected stat items
StatisticsValidation-->>StatisticsCache: restore entries or raise missing-item error
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 (1)
source/tests/pt/test_stat_file_mode.py (1)
85-97: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding tests for
_prepare_stat_file_pathedge cases.The tests cover the main scenarios well. Two edge cases in
_prepare_stat_file_paththat lack direct test coverage:
- Invalid
stat_file_mode(lines 130-134 inmain.py) — raisesValueErrorwith a descriptive messagestat_file=Nonewithstat_file_mode="read"(lines 135-138 inmain.py) — raisesValueErrorThese are caught upstream by
argcheck.pyvalidation, so the risk is low, but direct unit tests would protect against regressions if the function is called from new call sites.🤖 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 `@source/tests/pt/test_stat_file_mode.py` around lines 85 - 97, Extend test_stat_file_mode_configuration_validation with direct unit coverage for _prepare_stat_file_path: assert invalid stat_file_mode raises ValueError with a descriptive message, and assert stat_file=None with stat_file_mode="read" also raises ValueError. Use the existing configuration/test fixtures and preserve the current normalization assertions.
🤖 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 `@source/tests/pt/test_stat_file_mode.py`:
- Around line 85-97: Extend test_stat_file_mode_configuration_validation with
direct unit coverage for _prepare_stat_file_path: assert invalid stat_file_mode
raises ValueError with a descriptive message, and assert stat_file=None with
stat_file_mode="read" also raises ValueError. Use the existing
configuration/test fixtures and preserve the current normalization assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 9de43e18-1401-4328-97ca-38177d11aaaa
📒 Files selected for processing (5)
deepmd/pt/entrypoints/main.pydeepmd/utils/argcheck.pydeepmd/utils/path.pyexamples/water/dpa4/input.jsonsource/tests/pt/test_stat_file_mode.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Opening a completed HDF5 statistics cache with mode r is the right direction, but the current implementation does not yet provide the documented load-only semantics.
Cache completeness depends on model-specific type-map suffixes, descriptor hashes, and required statistic keys, so checking only the top-level path in the entrypoint cannot establish that the cache is complete. I recommend making the statistics-loading layer distinguish these cases explicitly: load an existing item; compute and save a missing item only when the path is writable; and fail immediately with the missing statistic name when the path is read-only. This avoids sampling data before an eventual generic write error.
The cache-opening helper is also infrastructure rather than entrypoint logic. A focused module such as deepmd/utils/stat_file.py would keep creation/opening policy separate from main.py, while the generic read-only guards should remain in deepmd/utils/path.py and be covered in source/tests/common/test_path.py.
The CodeRabbit edge-test suggestion is reasonable but secondary. In particular, its statement that both cases are caught upstream is not correct for stat_file_mode="read" without stat_file.
Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5773 +/- ##
==========================================
- Coverage 79.70% 79.57% -0.14%
==========================================
Files 1020 1020
Lines 116359 116376 +17
Branches 4305 4303 -2
==========================================
- Hits 92742 92603 -139
- Misses 22076 22228 +152
- Partials 1541 1545 +4 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/dpmodel/utils/stat.py (1)
91-103: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winMake
observed_typeoptional in read-only restores
_collect_and_set_observed_typeskips_save_observed_type_to_filewheneverpreset_observed_typeis provided, so caches created through that path can legitimately lackobserved_type. In read mode, the unconditional_require_stat_file_items(..., ["observed_type"])turns that into a hardFileNotFoundErroreven though the rest of the cache is usable. Relax this guard or fall back to recomputing when the entry is missing.🤖 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/dpmodel/utils/stat.py` around lines 91 - 103, The read-only restore path should tolerate caches without an observed_type entry. Update _restore_observed_type_from_file to avoid unconditionally requiring observed_type, return the stored values when the file exists, and return None when it is absent so callers can recompute them.
🤖 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.
Outside diff comments:
In `@deepmd/dpmodel/utils/stat.py`:
- Around line 91-103: The read-only restore path should tolerate caches without
an observed_type entry. Update _restore_observed_type_from_file to avoid
unconditionally requiring observed_type, return the stored values when the file
exists, and return None when it is absent so callers can recompute them.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 96803061-6577-4f51-9f78-1e6cc4e48250
📒 Files selected for processing (8)
deepmd/dpmodel/fitting/general_fitting.pydeepmd/dpmodel/utils/stat.pydeepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/task/fitting.pydeepmd/pt/utils/stat.pydeepmd/utils/argcheck.pysource/tests/common/test_path.pysource/tests/pt/test_stat_file_mode.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/utils/argcheck.py
What
training.stat_file_modeoption withupdate(default) and read-onlyreadmodes.Why
Multiple PyTorch training processes may need to consume the same completed statistics cache concurrently. Opening the cache in update mode permits writes and can cause HDF5 contention or accidental mutation. A read-only mode lets workers share an existing cache safely while preserving the current writable behavior for cache creation and updates.
Impact
stat_file_modekeeps the existingupdatebehavior.stat_file_mode: "read"requires a complete existing cache and rejects writes or directory creation.Checks
Summary by CodeRabbit
stat_file_mode(update/read) to control how the statistics cache is used.stat_fileis missing, the cache is incomplete, or required statistics items are absent.stat_file_mode: "update".stat_file_modevalidation and read-only HDF5 behavior.