fix(data): preserve LMDB label availability across batches#5839
Open
njzjz-bot wants to merge 1 commit into
Open
fix(data): preserve LMDB label availability across batches#5839njzjz-bot wants to merge 1 commit into
njzjz-bot wants to merge 1 commit into
Conversation
Contributor
📝 WalkthroughWalkthroughLMDB batching and validation now group frames by atom count and scalar ChangesLMDB availability-aware batching
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant LmdbDataReader
participant SameNlocBatchSampler
participant collate_lmdb_frames
participant FullValidator
LmdbDataReader->>SameNlocBatchSampler: group frames by nloc and find signature
SameNlocBatchSampler->>collate_lmdb_frames: emit homogeneous batches
collate_lmdb_frames->>FullValidator: provide scalar-consistent find flags
FullValidator->>FullValidator: evaluate grouped validation frames
Possibly related PRs
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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/dpmodel/utils/lmdb_data.py`:
- Around line 1195-1210: Update the sampling flow around
group_indices_by_find_signature and _expand_indices_by_blocks to allocate each
block’s global target across all (nloc, signature) groups before expansion,
using remainder-preserving integer allocation so per-group rounding preserves
the exact block total. Pass each group’s allocated target into expansion, and
update the estimator logic around _block_total_actual to use the same allocation
rather than independently recomputing targets per signature.
In `@source/tests/pt_expt/test_lmdb_training.py`:
- Around line 202-213: Update the test around DataRequirementItem registration
to realize or capture the lazy sampler iterator before calling
ds.add_data_requirements(). Assert that registration replaces the previous
iterator, then obtain subsequent batches and verify they use the newly
registered requirement signatures.
In `@source/tests/pt/test_lmdb_dataloader.py`:
- Around line 388-423: Extend the test around EnergyStdLoss to cover the force
branch using start_pref_f and limit_pref_f. Invoke the loss with a zero-force
model and force/find_force labels, asserting zero loss when find_force is 0.0
and positive loss when it is 1.0, alongside the existing energy assertions.
- Around line 407-408: Update the zero_energy_model closure in the loop so it
binds the current batch value at definition time, avoiding a late-bound
loop-variable closure flagged by Ruff B023 while preserving the existing energy
tensor behavior. Run ruff check . to verify the lint issue is resolved.
🪄 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: 2bd6db4e-0f72-43dd-8c28-134427b8cdbe
📒 Files selected for processing (7)
deepmd/dpmodel/utils/lmdb_data.pydeepmd/pt/train/validation.pydeepmd/pt/utils/lmdb_dataset.pydeepmd/pt_expt/utils/lmdb_dataset.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt/test_validation.pysource/tests/pt_expt/test_lmdb_training.py
Partition LMDB frames by atom count and scalar find-flag signature before training, statistics, auto-probability, and distributed batching. Allocate auto-probability expansion targets exactly across homogeneous groups so repeated rounding cannot over-sample a block. Evaluate full LMDB validation in the same homogeneous groups so default-filled optional labels do not bias metrics. Cover standard PyTorch energy and force masking, iterator resets, distributed lengths, statistics loaders, pt_expt batches, and partial-label validation. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
njzjz-bot
force-pushed
the
fix/lmdb-collate-find-flags-5636
branch
from
July 17, 2026 01:33
2fa9dae to
a5d7da6
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_lmdb_data.py (1)
752-753: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
batch.countto explicitly count occurrences.Currently,
index in batchevaluates to a boolean. This correctly counts the total frames here because the mock reader's batch size is explicitly set to 1. However, if the batch size were ever increased to pack multiple identical indices, this expression would only count the number of batches containing the index rather than the total occurrences. Usingbatch.count(index)is more robust and semantically clearer.♻️ Proposed refactor
batches = list(sampler) - counts = [sum(index in batch for batch in batches) for index in (0, 1)] + counts = [sum(batch.count(index) for batch in batches) for index in (0, 1)]🤖 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/common/dpmodel/test_lmdb_data.py` around lines 752 - 753, Update the counts comprehension in the sampler test to use batch.count(index) instead of membership testing, so it sums all occurrences of each index across batches rather than only counting batches that contain it.
🤖 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/common/dpmodel/test_lmdb_data.py`:
- Around line 752-753: Update the counts comprehension in the sampler test to
use batch.count(index) instead of membership testing, so it sums all occurrences
of each index across batches rather than only counting batches that contain it.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2bf6ab22-4ac4-4374-82bd-0d46f30282c7
📒 Files selected for processing (8)
deepmd/dpmodel/utils/lmdb_data.pydeepmd/pt/train/validation.pydeepmd/pt/utils/lmdb_dataset.pydeepmd/pt_expt/utils/lmdb_dataset.pysource/tests/common/dpmodel/test_lmdb_data.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt/test_validation.pysource/tests/pt_expt/test_lmdb_training.py
🚧 Files skipped from review as they are similar to previous changes (7)
- deepmd/pt_expt/utils/lmdb_dataset.py
- deepmd/pt/train/validation.py
- deepmd/pt/utils/lmdb_dataset.py
- source/tests/pt_expt/test_lmdb_training.py
- source/tests/pt/test_validation.py
- source/tests/pt/test_lmdb_dataloader.py
- deepmd/dpmodel/utils/lmdb_data.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5839 +/- ##
==========================================
- Coverage 78.58% 78.33% -0.26%
==========================================
Files 1050 1050
Lines 120637 120751 +114
Branches 4356 4357 +1
==========================================
- Hits 94801 94586 -215
- Misses 24278 24593 +315
- Partials 1558 1572 +14 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Member
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #5636
Summary
find_*signatureSimply stacking
find_*values would not be safe here. Current PyTorch and dpmodel losses reduce a whole batch before applying scalar availability flags, and display/validation paths convert those flags to booleans. A vector flag would broadcast incorrectly, produce vector-valued losses, or fail boolean conversion. Homogeneous batching is the smaller compatible fix.Why existing tests missed this
The direct collator test used
find_energy=1in every frame and did not assert the flag value. All LMDB factories stored energy and force in every record, so same-nloc batching was also label-homogeneous by accident. Reader consistency tests inspect individual frames before collation, and loss tests construct homogeneous scalar flags without connecting them to LMDB batches. Auto-probability tests also used one availability signature, so they could not expose independent per-signature rounding that changed a block target of 3 into 4 sampled frames.Full-validation coverage only varied atom count.
LmdbTestDatatherefore always observed globally present labels, and no test combined real labels with default-filled frames or checked order-independent optional-label metrics.Validation
find_*flag is zeroruff format .ruff check .Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit