fix(pt): allow concurrent stat cache reads#5773
Conversation
Add explicit read-only access while preserving update mode by default.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughTraining now supports ChangesStatistics cache mode handling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingConfig
participant PTEntrypoint
participant StatisticsCache
participant DescriptorStats
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->>DescriptorStats: load_or_compute_stats with cache and data
DescriptorStats->>HDF5: validate required items
HDF5-->>DescriptorStats: restore stats or report missing items
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 (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 78.58% 78.34% -0.24%
==========================================
Files 1050 1050
Lines 120637 120656 +19
Branches 4356 4354 -2
==========================================
- Hits 94801 94528 -273
- Misses 24278 24568 +290
- Partials 1558 1560 +2 ☔ 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
Keep update-mode fallback coverage while asserting that read-only caches reject missing observed-type statistics.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
High — cannot reuse a regular energy-model stat cache produced by update
The new read-only preflight unconditionally requires bias_atom_<key> and std_atom_<key> for every supplied output key. However, the real PT atomic-model path calls:
keys=list(self.atomic_output_def().keys())atomic_output_def() includes runtime-derived outputs in addition to fitting outputs, including mask. A standard energy model therefore has at least:
["energy", "mask"]The update path filters keys against labels in the training data. Since mask has no label, the cache only gets:
bias_atom_energy
std_atom_energy
not bias_atom_mask / std_atom_mask.
Reusing that same model and cache with stat_file_mode: read then fails before sampling with FileNotFoundError for the missing mask entries. That breaks the primary workflow this PR enables: one process writes the cache with update, then multiple training processes share it with read.
Please make BaseAtomicModel.change_out_bias() pass only the outputs that need persistent bias/std statistics—for example:
keys=self.bias_keys
# or: list(self.fitting_output_def().keys())This keeps derived outputs such as mask out of the read-cache completeness check.
Please also add an end-to-end regression test that:
- writes a directory/HDF5 stat cache using a real PT energy model in
updatemode; - reloads it with the same model in
readmode; and - verifies loading succeeds without invoking the sampler.
The current tests call compute_output_stats() directly with keys=["energy"], so they bypass the atomic_output_def() path that injects mask.
— OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.6-terra)
|
Addressed in d62b215:
Validation:
|
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
High — read mode still samples data when descriptor statistics are missing
_prepare_stat_file_path() verifies only that the cache root exists. Descriptor statistics are stored one level deeper under the type-map group and the descriptor-specific EnvMatStatSe.get_hash() group. If a cache exists but that hash group is absent—for example, because it was generated for different descriptor parameters or is incomplete—the descriptor's compute_input_stats() treats the path as a cache miss, calls the lazy sampler, and then tries to save the result. The save is rejected only at mkdir() with ValueError("Cannot mkdir to read-only path").
That violates the advertised read-mode semantics: an incomplete cache should fail before sampling, rather than performing an expensive data pass and surfacing a generic write error. Please make the descriptor-stat loading path distinguish a read-only miss and raise a clear FileNotFoundError before invoking merged() (preferably centrally in EnvMatStat.load_or_compute_stats, while also validating the required datasets).
Please add a regression test using a read-only cache that contains the type-map root but omits the descriptor-hash statistics group. Its sampler should fail if called, and the expected result should be FileNotFoundError with no cache mutation.
— OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.6-terra)
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed the latest head (2626abf). The prior read-only cache issues are addressed:
- output-stat validation now scopes required entries to persistent bias keys, excluding derived outputs such as
mask; - missing or partial descriptor-stat groups fail before the lazy sampler runs;
- the update → read path is covered by end-to-end HDF5 regression tests.
The implementation and tests look good to me.
— OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.6-terra)
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
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 on-disk statistics cache access, including rank-aware handling in multi-process training.readmode, missing or incomplete caches now fail fast with clearer errors; required cache items are validated before use.stat_file_mode: "update".