feat(pt/dpmodel): add lmdb dataloader#5283
Merged
iProzd merged 25 commits intodeepmodeling:masterfrom Apr 7, 2026
Merged
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| def test_getitem_out_of_range(self, lmdb_dir): | ||
| ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) | ||
| with pytest.raises(IndexError): | ||
| ds[999] |
Check notice
Code scanning / CodeQL
Statement has no effect Note test
Contributor
|
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:
📝 WalkthroughWalkthroughAdds comprehensive LMDB support: a framework-agnostic LMDB reader, merge/test utilities, PyTorch LMDB Dataset/collation/samplers, trainer/test and test entrypoint LMDB branches, an example LMDB config, extensive unit tests, and new Changes
Sequence DiagramssequenceDiagram
participant Trainer as Trainer
participant Main as prepare_trainer_input_single
participant Checker as is_lmdb
participant DS as LmdbDataset / DpLoaderSet
participant Reader as LmdbDataReader
participant LMDB as LMDB DB
participant DL as PyTorch DataLoader
Trainer->>Main: request(train_systems, val_systems)
Main->>Checker: is_lmdb(train_systems)
alt LMDB path
Main->>DS: create LmdbDataset(train_path)
Main->>DS: create LmdbDataset(val_path) or DpLoaderSet
else Non-LMDB
Main->>DS: create DpLoaderSet(...)
end
Main-->>Trainer: (train_data, val_data, DPPath)
Trainer->>DL: get_data_loader(train_data)
DL->>DS: initialize (sampler, collate_fn)
DS->>Reader: Reader.__init__(lmdb_path)
Reader->>LMDB: open read-only txn and read __metadata__
loop per-batch
DL->>Reader: Reader.__getitem__(frame_id)
Reader->>LMDB: fetch frame bytes
Reader->>Reader: decode arrays & remap keys
Reader-->>DL: frame dict (np arrays)
DL->>DL: _collate_lmdb_batch(batch)
DL-->>Trainer: batched tensors
end
sequenceDiagram
participant App as Test CLI
participant LD as LmdbTestData
participant View as _LmdbTestDataNlocView
participant TestRunner as test functions
App->>LD: LmdbTestData(lmdb_path)
LD->>LD: compute nloc_groups
alt multiple nloc groups
LD->>View: create per-nloc views
loop per nloc
App->>TestRunner: run tests(view, label)
end
else single group
App->>TestRunner: run tests(LD, label)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pt/train/training.py (1)
529-542:⚠️ Potential issue | 🔴 CriticalMulti-task LMDB path still assumes
.sampler.weightsand can crash.In LMDB mode, dataloaders are batch-sampler based and do not provide sampler weights. The current multi-task code still dereferences
.sampler.weightsin both summary printing andnum_epoch_dicttotal-batch computation.Proposed fix
training_data[model_key].print_summary( f"training in {model_key}", - to_numpy_array(self.training_dataloader[model_key].sampler.weights), + to_numpy_array(self.training_dataloader[model_key].sampler.weights) + if not isinstance(training_data[model_key], LmdbDataset) + else [1.0], ) @@ validation_data[model_key].print_summary( f"validation in {model_key}", - to_numpy_array( - self.validation_dataloader[model_key].sampler.weights - ), + to_numpy_array( + self.validation_dataloader[model_key].sampler.weights + ) + if not isinstance(validation_data[model_key], LmdbDataset) + else [1.0], ) @@ for model_key in self.model_keys: - sampler_weights = to_numpy_array( - self.training_dataloader[model_key].sampler.weights - ) - per_task_total.append( - compute_total_numb_batch( - training_data[model_key].index, - sampler_weights, - ) - ) + if isinstance(training_data[model_key], LmdbDataset): + per_task_total.append(training_data[model_key].total_batch) + else: + sampler_weights = to_numpy_array( + self.training_dataloader[model_key].sampler.weights + ) + per_task_total.append( + compute_total_numb_batch( + training_data[model_key].index, + sampler_weights, + ) + )Also applies to: 583-590
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/train/training.py` around lines 529 - 542, The code assumes sampler.weights exists for training/validation dataloaders (see training_data[model_key].print_summary calls and the num_epoch_dict total-batch computation) which breaks for LMDB batch-sampler dataloaders; fix by guarding access to sampler.weights: when printing call to_numpy_array(self.training_dataloader[model_key].sampler.weights) only if hasattr(self.training_dataloader[model_key].sampler, "weights") (or use getattr with a default None) and pass None or an empty array otherwise, and for num_epoch_dict total-batch computation replace uses of sum(sampler.weights) with a safe fallback that uses len(self.training_dataloader[model_key]) (number of batches) when weights are absent; update the same logic for validation_dataloader and any other places that reference .sampler.weights (e.g., the block around num_epoch_dict).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 242-253: The batch_size handling silently coerces unknown strings
and accepts non-positive integers; update the branch in the batch_size setter so
that only "auto" or "auto:<positive int>" strings are allowed (parse "auto:<n>"
into an int and validate n>0), otherwise raise a ValueError with a clear
message; likewise, when batch_size is numeric (the else branch that sets
self.batch_size = int(batch_size)) validate the parsed integer is >0 and raise a
ValueError if not; keep using _compute_batch_size(self._natoms, self._auto_rule)
for the "auto" case but ensure _auto_rule is a positive int after parsing.
- Line 374: The local variable unique_nlocs assigned from
sorted(self._nloc_groups.keys()) is unused and triggers a Ruff F841 lint error;
remove this assignment statement (the unused local unique_nlocs) from the method
where it appears so the code no longer creates unique_nlocs but retains any
needed logic that uses self._nloc_groups directly.
- Around line 880-882: The merge logic assumes src_nlocs has at least nframes
entries and does frame_nlocs.append(int(src_nlocs[i])) without bounds checking;
change this to guard the index: if src_nlocs is not None and i < len(src_nlocs)
then append int(src_nlocs[i]) else compute/derive the per-frame natoms (e.g., by
parsing the frame from the source data or using the fallback natoms value
already available) so that malformed/truncated metadata does not raise
IndexError during the merge; update the block around the frame_nlocs population
(references: frame_nlocs, src_nlocs, nframes) to perform this fallback behavior.
- Around line 705-719: In get_test(), handle the empty-dataset case before using
max on self._nloc_groups: if self._nloc_groups is empty (or len(self._frames) ==
0) raise a clear ValueError (e.g. "No frames in LMDB / no nloc groups
available") instead of letting max(...) raise a cryptic error; otherwise proceed
with the existing logic that sets natoms and frame_indices from the largest
group in self._nloc_groups.
- Around line 263-267: Before building natoms_vec/vec, validate that all atom
type ids in atype are within [0, self._ntypes-1]; detect any ids <0 or >=
self._ntypes and raise a descriptive ValueError (including the offending ids or
their min/max and, if available, the frame index) instead of silently truncating
via np.bincount(... )[: self._ntypes]. Place this check immediately before the
np.bincount call that computes counts (the code using atype and self._ntypes and
filling vec/natoms_vec) so you fail fast and keep natoms_vec consistent with the
frame contents.
In `@deepmd/entrypoints/test.py`:
- Around line 233-266: The append_detail flag is incorrectly derived from cc so
LMDB sub-groups for the first system overwrite detail files; instead track
whether a given sys_label has already been written and pass append_detail=True
for subsequent sub-groups. Modify the loop over data_items: introduce a small
seen map (e.g., seen_systems dict keyed by sys_label) and compute append_detail
= seen_systems.get(sys_label, False) when calling
test_ener/test_dos/test_property (replace the current append_detail=(cc != 0));
after the call set seen_systems[sys_label] = True so later sub-groups append
rather than overwrite.
In `@source/tests/consistent/test_lmdb_data.py`:
- Line 128: The loop variable j in the loop "for j in range(count):" is unused;
rename it to "_" (or "_i"/"_count" per project convention) to satisfy Ruff B007
and avoid lint failures, update any related references in that loop (none
expected), and re-run ruff check . and ruff format . to ensure formatting and
linting pass.
- Around line 405-407: The current test contains a tautological assertion
self.assertEqual(atype.shape[1], atype.shape[1]) which does not validate batch
consistency; replace it with an assertion that compares the nloc dimension
across frames in the batch (use atype as the per-batch array) — e.g. compute a
reference_nloc from the first frame (reference_nloc = atype[0].shape[1]) and
assert every frame's shape[1] equals reference_nloc (or use numpy/all to check
atype[:,1?] consistency) so the test in test_lmdb_data.py actually verifies "All
frames in batch have same nloc".
In `@source/tests/pt/test_lmdb_dataloader.py`:
- Line 480: The test contains Ruff violations: the loop uses an unused loop
variable i and the variable s_default is assigned but never used; fix by
changing the loop to use a throwaway name (e.g., replace "for i in range(10):"
with "for _ in range(10):") or otherwise use i, and remove or use the s_default
assignment (remove the dead "s_default = ..." line if it's not needed or
reference it where intended) so there are no unused variables left.
---
Outside diff comments:
In `@deepmd/pt/train/training.py`:
- Around line 529-542: The code assumes sampler.weights exists for
training/validation dataloaders (see training_data[model_key].print_summary
calls and the num_epoch_dict total-batch computation) which breaks for LMDB
batch-sampler dataloaders; fix by guarding access to sampler.weights: when
printing call
to_numpy_array(self.training_dataloader[model_key].sampler.weights) only if
hasattr(self.training_dataloader[model_key].sampler, "weights") (or use getattr
with a default None) and pass None or an empty array otherwise, and for
num_epoch_dict total-batch computation replace uses of sum(sampler.weights) with
a safe fallback that uses len(self.training_dataloader[model_key]) (number of
batches) when weights are absent; update the same logic for
validation_dataloader and any other places that reference .sampler.weights
(e.g., the block around num_epoch_dict).
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
deepmd/dpmodel/utils/__init__.pydeepmd/dpmodel/utils/lmdb_data.pydeepmd/entrypoints/test.pydeepmd/pt/entrypoints/main.pydeepmd/pt/train/training.pydeepmd/pt/utils/lmdb_dataset.pyexamples/lmdb_data/input_lmdb.jsonsource/tests/consistent/test_lmdb_data.pysource/tests/pt/test_lmdb_dataloader.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5283 +/- ##
==========================================
+ Coverage 82.38% 82.39% +0.01%
==========================================
Files 812 814 +2
Lines 83611 84574 +963
Branches 4091 4090 -1
==========================================
+ Hits 68882 69687 +805
- Misses 13508 13667 +159
+ Partials 1221 1220 -1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (5)
deepmd/dpmodel/utils/lmdb_data.py (5)
248-260:⚠️ Potential issue | 🟡 MinorValidate
batch_sizeto prevent silent misconfiguration.Unknown string specs (e.g.,
"auto_invalid") fall through toself._auto_rule = 32at line 256, hiding config mistakes. Non-positive integers are also accepted at line 260.Suggested fix
self._auto_rule: int | None = None if isinstance(batch_size, str): if batch_size == "auto": self._auto_rule = 32 elif batch_size.startswith("auto:"): - self._auto_rule = int(batch_size.split(":")[1]) + try: + self._auto_rule = int(batch_size.split(":", 1)[1]) + except ValueError as exc: + raise ValueError( + f"Invalid batch_size spec: {batch_size!r}" + ) from exc + if self._auto_rule <= 0: + raise ValueError("auto batch_size rule must be > 0") else: - self._auto_rule = 32 + raise ValueError( + f"batch_size must be 'auto', 'auto:N', or a positive int, got {batch_size!r}" + ) self.batch_size = _compute_batch_size(self._natoms, self._auto_rule) else: self.batch_size = int(batch_size) + if self.batch_size <= 0: + raise ValueError("batch_size must be > 0")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 248 - 260, The batch_size parsing in the constructor silently treats unknown string specs and non-positive integers as valid (setting self._auto_rule = 32 or accepting non-positive ints), so add explicit validation: in the batch_size handling for strings (function/class using self._auto_rule and _compute_batch_size) only accept "auto" or "auto:<positive_int>" and raise a ValueError for any other string spec; when batch_size is numeric (the else branch that sets self.batch_size = int(batch_size)) validate that the parsed int is > 0 and raise ValueError if not; ensure error messages reference the provided batch_size value and keep use of _compute_batch_size(self._natoms, self._auto_rule) unchanged when the auto rule is valid.
974-976:⚠️ Potential issue | 🟡 MinorGuard against truncated
frame_nlocsmetadata during merge.If
src_nlocshas fewer entries thannframes(malformed metadata),src_nlocs[i]raisesIndexError.Suggested fix
- if src_nlocs is not None: + if src_nlocs is not None and i < len(src_nlocs): frame_nlocs.append(int(src_nlocs[i])) else:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 974 - 976, The code indexing src_nlocs[i] can raise IndexError when src_nlocs is shorter than nframes; in the block that checks "if src_nlocs is not None:" (where frame_nlocs.append(int(src_nlocs[i])) is called), first verify i < len(src_nlocs) (or use a safe iterator) and if the entry is missing append a sensible default (e.g., 0) or otherwise handle the malformed metadata (log/warn and append 0) instead of indexing past the end; update the logic around frame_nlocs, src_nlocs and nframes to perform this bounds check before casting/indexing.
785-808:⚠️ Potential issue | 🟡 MinorHandle empty datasets in
get_test().If the LMDB has zero frames,
self._nloc_groupsis empty andmax(self._nloc_groups, ...)at line 798 raisesValueError: max() arg is an empty sequence.Suggested fix
if nloc is not None: if nloc not in self._nloc_groups: raise ValueError( f"No frames with nloc={nloc}. Available: {sorted(self._nloc_groups.keys())}" ) frame_indices = self._nloc_groups[nloc] natoms = nloc + elif not self._frames: + raise ValueError("LMDB dataset contains no frames") elif len(self._nloc_groups) == 1:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 785 - 808, The method that selects frames (inside get_test / the function using self._nloc_groups and _frames) fails when the LMDB is empty because max(...) is called on an empty self._nloc_groups; add an early check for empty dataset (e.g., if not self._nloc_groups or not self._frames) and handle it by either raising a clear ValueError ("No frames in LMDB") or returning an appropriate empty result before reaching the mixed-nloc branch; update the branches that reference natoms and frame_indices (the code that computes natoms = max(...) and frame_indices = ...) so they only run when groups exist to avoid the max() on empty sequence.
436-447:⚠️ Potential issue | 🟡 MinorRemove unused
unique_nlocsvariable (Ruff F841).
unique_nlocsis assigned but never used. This will cause CI failure per coding guidelines.Suggested fix
def print_summary(self, name: str, prob: Any) -> None: """Print basic dataset info.""" - unique_nlocs = sorted(self._nloc_groups.keys()) nloc_info = ", ".join( f"{nloc}({len(idxs)})" for nloc, idxs in sorted(self._nloc_groups.items()) )As per coding guidelines,
**/*.py: Always runruff check .andruff format .before committing changes or CI will fail.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 436 - 447, In print_summary (method print_summary) the local variable unique_nlocs is assigned from self._nloc_groups but never used; remove the unused assignment (unique_nlocs = sorted(self._nloc_groups.keys())) so the method only builds nloc_info from self._nloc_groups and logs it, eliminating the Ruff F841 unused-variable error.
265-276:⚠️ Potential issue | 🟡 MinorGuard against out-of-range atom type IDs.
np.bincount(...)[: self._ntypes]truncates counts for type IDs>= self._ntypes, makingnatoms_vecinconsistent with frame contents.Suggested fix
def _compute_natoms_vec(self, atype: np.ndarray) -> np.ndarray: nloc = len(atype) + if atype.size and (np.min(atype) < 0 or np.max(atype) >= self._ntypes): + raise ValueError( + f"atype contains out-of-range IDs; expected [0, {self._ntypes - 1}]" + ) counts = np.bincount(atype, minlength=self._ntypes)[: self._ntypes]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 265 - 276, _compute_natoms_vec currently truncates counts for atom type IDs >= self._ntypes; to fix, validate the atype array before counting: compute max_id = atype.max() (handle empty atype properly) and if max_id >= self._ntypes or any id < 0 raise a clear ValueError referencing _compute_natoms_vec and include the offending max_id and self._ntypes in the message; otherwise use np.bincount(atype, minlength=self._ntypes) to build counts and keep the existing vec[0]=nloc, vec[1]=nloc, vec[2:]=counts logic.
🧹 Nitpick comments (2)
deepmd/pt/train/training.py (1)
263-265: Consider documenting the_readeraccess pattern.
_data._readeraccesses a private attribute ofLmdbDataset. While acceptable within the same package, consider either making this a public property or adding a brief comment explaining this coupling is intentional for distributed sampler construction.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/train/training.py` around lines 263 - 265, The code is accessing a private attribute _data._reader of LmdbDataset when constructing the DistributedSameNlocBatchSampler; update the codebase to avoid relying on a private member by either exposing a public property/method on LmdbDataset (e.g., reader or get_reader) and using that, or add a concise inline comment next to the DistributedSameNlocBatchSampler construction documenting that accessing _reader is an intentional coupling for distributed sampler creation and should not be refactored without updating the sampler; reference symbols: _data._reader, DistributedSameNlocBatchSampler, LmdbDataset.deepmd/dpmodel/utils/lmdb_data.py (1)
869-869: Rename unused loop variable to_req_info.
req_infois not used within the loop body (Ruff B007). Rename to_req_infoto indicate it's intentionally unused.- for key, req_info in all_keys.items(): + for key, _req_info in all_keys.items():🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` at line 869, The loop over all_keys uses an unused variable name req_info which triggers a lint warning; update the for statement in lmdb_data.py (the loop "for key, req_info in all_keys.items():") to rename req_info to _req_info so it becomes "for key, _req_info in all_keys.items():" to indicate the variable is intentionally unused (no other changes needed if req_info is not referenced elsewhere).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pt/train/training.py`:
- Around line 554-555: The reported mismatch comes from using
LmdbDataset.total_batch (which does nframes // batch_size) to set
total_numb_batch while SameNlocBatchSampler.__len__ uses per-nloc ceiling
division; fix by computing total_numb_batch from the actual sampler instead of
total_batch: when using SameNlocBatchSampler (or when auto batch sizing / mixed
nloc are possible) set total_numb_batch = len(sampler) or recompute it by
summing ceil(group_size / batch_size_for_that_nloc) so num_steps (used later)
matches the sampler’s actual number of batches; update the block in training.py
that currently uses training_data.total_batch to use
SameNlocBatchSampler.__len__ (or equivalent ceiling-sum logic) so the two agree.
---
Duplicate comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 248-260: The batch_size parsing in the constructor silently treats
unknown string specs and non-positive integers as valid (setting self._auto_rule
= 32 or accepting non-positive ints), so add explicit validation: in the
batch_size handling for strings (function/class using self._auto_rule and
_compute_batch_size) only accept "auto" or "auto:<positive_int>" and raise a
ValueError for any other string spec; when batch_size is numeric (the else
branch that sets self.batch_size = int(batch_size)) validate that the parsed int
is > 0 and raise ValueError if not; ensure error messages reference the provided
batch_size value and keep use of _compute_batch_size(self._natoms,
self._auto_rule) unchanged when the auto rule is valid.
- Around line 974-976: The code indexing src_nlocs[i] can raise IndexError when
src_nlocs is shorter than nframes; in the block that checks "if src_nlocs is not
None:" (where frame_nlocs.append(int(src_nlocs[i])) is called), first verify i <
len(src_nlocs) (or use a safe iterator) and if the entry is missing append a
sensible default (e.g., 0) or otherwise handle the malformed metadata (log/warn
and append 0) instead of indexing past the end; update the logic around
frame_nlocs, src_nlocs and nframes to perform this bounds check before
casting/indexing.
- Around line 785-808: The method that selects frames (inside get_test / the
function using self._nloc_groups and _frames) fails when the LMDB is empty
because max(...) is called on an empty self._nloc_groups; add an early check for
empty dataset (e.g., if not self._nloc_groups or not self._frames) and handle it
by either raising a clear ValueError ("No frames in LMDB") or returning an
appropriate empty result before reaching the mixed-nloc branch; update the
branches that reference natoms and frame_indices (the code that computes natoms
= max(...) and frame_indices = ...) so they only run when groups exist to avoid
the max() on empty sequence.
- Around line 436-447: In print_summary (method print_summary) the local
variable unique_nlocs is assigned from self._nloc_groups but never used; remove
the unused assignment (unique_nlocs = sorted(self._nloc_groups.keys())) so the
method only builds nloc_info from self._nloc_groups and logs it, eliminating the
Ruff F841 unused-variable error.
- Around line 265-276: _compute_natoms_vec currently truncates counts for atom
type IDs >= self._ntypes; to fix, validate the atype array before counting:
compute max_id = atype.max() (handle empty atype properly) and if max_id >=
self._ntypes or any id < 0 raise a clear ValueError referencing
_compute_natoms_vec and include the offending max_id and self._ntypes in the
message; otherwise use np.bincount(atype, minlength=self._ntypes) to build
counts and keep the existing vec[0]=nloc, vec[1]=nloc, vec[2:]=counts logic.
---
Nitpick comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Line 869: The loop over all_keys uses an unused variable name req_info which
triggers a lint warning; update the for statement in lmdb_data.py (the loop "for
key, req_info in all_keys.items():") to rename req_info to _req_info so it
becomes "for key, _req_info in all_keys.items():" to indicate the variable is
intentionally unused (no other changes needed if req_info is not referenced
elsewhere).
In `@deepmd/pt/train/training.py`:
- Around line 263-265: The code is accessing a private attribute _data._reader
of LmdbDataset when constructing the DistributedSameNlocBatchSampler; update the
codebase to avoid relying on a private member by either exposing a public
property/method on LmdbDataset (e.g., reader or get_reader) and using that, or
add a concise inline comment next to the DistributedSameNlocBatchSampler
construction documenting that accessing _reader is an intentional coupling for
distributed sampler creation and should not be refactored without updating the
sampler; reference symbols: _data._reader, DistributedSameNlocBatchSampler,
LmdbDataset.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: e42a3a8f-33f3-46a5-b818-47bcea016afe
📒 Files selected for processing (3)
deepmd/dpmodel/utils/lmdb_data.pydeepmd/pt/train/training.pypyproject.toml
wanghan-iapcm
requested changes
Mar 11, 2026
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
deepmd/dpmodel/utils/lmdb_data.py (2)
892-896: Rename unused loop variablereq_infoto_req_info(Ruff B007).The loop variable
req_infois not used within the loop body. Rename it to_req_infoto signal intentional non-use and satisfy Ruff.🧹 Proposed fix
- for key, req_info in all_keys.items(): + for key, _req_info in all_keys.items(): has_key = any( key in f and isinstance(f.get(key), np.ndarray) for f in frames )As per coding guidelines,
**/*.py: Always runruff check .andruff format .before committing changes or CI will fail.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 892 - 896, The loop in lmdb_data.py iterates with an unused variable req_info (for key, req_info in all_keys.items()) — rename it to _req_info to indicate intentional non-use and satisfy Ruff B007; update the for-loop variable in that block (the loop that sets result[f"find_{key}"]) and run ruff check/ruff format before committing.
962-966: Consider safer directory removal approach.
shutil.rmtree(dst_path)on line 966 unconditionally removes the existing path. While this enables idempotent merges, consider adding a confirmation parameter or logging a warning when overwriting to prevent accidental data loss.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 962 - 966, The code unconditionally deletes dst_path using shutil.rmtree which can cause accidental data loss; modify the enclosing function (look for the routine that uses dst_path and shutil.rmtree) to accept an explicit boolean flag like overwrite or force_remove (default False), add logging (import logging) and call logging.warning with the dst_path before deletion, and only call shutil.rmtree(dst_path) when the flag is True (otherwise raise an error or return); this makes removal explicit and auditable while keeping the idempotent behavior when explicitly requested.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 892-896: The loop in lmdb_data.py iterates with an unused variable
req_info (for key, req_info in all_keys.items()) — rename it to _req_info to
indicate intentional non-use and satisfy Ruff B007; update the for-loop variable
in that block (the loop that sets result[f"find_{key}"]) and run ruff check/ruff
format before committing.
- Around line 962-966: The code unconditionally deletes dst_path using
shutil.rmtree which can cause accidental data loss; modify the enclosing
function (look for the routine that uses dst_path and shutil.rmtree) to accept
an explicit boolean flag like overwrite or force_remove (default False), add
logging (import logging) and call logging.warning with the dst_path before
deletion, and only call shutil.rmtree(dst_path) when the flag is True (otherwise
raise an error or return); this makes removal explicit and auditable while
keeping the idempotent behavior when explicitly requested.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 7be23553-3c34-461f-896f-f5f2b69bf477
📒 Files selected for processing (1)
deepmd/dpmodel/utils/lmdb_data.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (5)
deepmd/dpmodel/utils/lmdb_data.py (5)
367-371:⚠️ Potential issue | 🟠 MajorFail fast on invalid
atypeids before remap/counting.At Line 371 and Line 291,
atypeis used without bounds validation. Invalid ids can corrupt type remapping/counts or raise runtime errors in non-obvious paths. Validate range once before remap and beforenatoms_veccomputation.Suggested fix
if "atype" in frame and isinstance(frame["atype"], np.ndarray): frame["atype"] = frame["atype"].reshape(-1).astype(np.int64) + if frame["atype"].size and np.min(frame["atype"]) < 0: + raise ValueError("atype contains negative type ids") # Remap atom types from LMDB's type_map to model's type_map if self._type_remap is not None: + if frame["atype"].size and np.max(frame["atype"]) >= len(self._type_remap): + raise ValueError( + f"atype id out of range for LMDB type_map (max allowed {len(self._type_remap)-1})" + ) frame["atype"] = self._type_remap[frame["atype"]].astype(np.int64) + if frame["atype"].size and np.max(frame["atype"]) >= self._ntypes: + raise ValueError( + f"atype contains out-of-range ids; expected [0, {self._ntypes - 1}]" + )Also applies to: 290-292
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 367 - 371, Validate that all atom type IDs in frame["atype"] are within the valid range before any remapping or any use in natoms_vec computation: check they are non-negative integers and less than the length/shape of self._type_remap (when self._type_remap is not None) or less than expected num_types; if any id is out of range raise a clear ValueError including the offending ids and context (frame id), then proceed to reshape/astype, apply self._type_remap on valid indices, and only after successful remap compute natoms_vec; refer to symbols frame["atype"], self._type_remap and natoms_vec to locate where to insert the validation and error handling.
270-281:⚠️ Potential issue | 🟠 MajorValidate
batch_sizestrictly; don’t silently coerce invalid specs.At Line 276, unknown string specs are silently treated as
auto:32, and at Line 280 non-positive integers are accepted. This masks config mistakes and can break batching downstream.Suggested fix
self._auto_rule: int | None = None if isinstance(batch_size, str): if batch_size == "auto": self._auto_rule = 32 elif batch_size.startswith("auto:"): - self._auto_rule = int(batch_size.split(":")[1]) + try: + self._auto_rule = int(batch_size.split(":", 1)[1]) + except ValueError as exc: + raise ValueError( + "batch_size must be 'auto', 'auto:N', or a positive integer" + ) from exc + if self._auto_rule <= 0: + raise ValueError("auto batch_size rule N must be > 0") else: - self._auto_rule = 32 + raise ValueError( + "batch_size must be 'auto', 'auto:N', or a positive integer" + ) # Default batch_size uses first frame's nloc (for total_batch estimate) self.batch_size = _compute_batch_size(self._natoms, self._auto_rule) else: self.batch_size = int(batch_size) + if self.batch_size <= 0: + raise ValueError("batch_size must be > 0")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 270 - 281, The batch_size handling currently silently coerces invalid string specs and accepts non-positive integers; update the block that sets self._auto_rule and self.batch_size (around batch_size, self._auto_rule, and _compute_batch_size in lmdb_data.py) to strictly validate inputs: accept only "auto" or "auto:<positive int>" for string specs (parse and set self._auto_rule or raise ValueError on any other format), and when batch_size is numeric ensure int(batch_size) > 0 (raise ValueError otherwise), also catch and re-raise parsing errors with a clear message so invalid configs fail fast instead of being coerced.
461-461:⚠️ Potential issue | 🟠 MajorRemove unused local
unique_nlocs(F841).Line 461 assigns
unique_nlocsbut never uses it.Suggested fix
- unique_nlocs = sorted(self._nloc_groups.keys()) nloc_info = ", ".join( f"{nloc}({len(idxs)})" for nloc, idxs in sorted(self._nloc_groups.items()) )As per coding guidelines,
**/*.py: Always runruff check .andruff format .before committing changes or CI will fail.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` at line 461, The local variable unique_nlocs is assigned via sorted(self._nloc_groups.keys()) but never used, causing an F841; remove the unused assignment (delete the unique_nlocs line) or replace its usage if it was intended to be used later in the same function/method that contains self._nloc_groups; ensure you reference the variable name unique_nlocs and the expression sorted(self._nloc_groups.keys()) when making the change, then run ruff check . and ruff format . before committing.
848-855:⚠️ Potential issue | 🟡 MinorHandle empty datasets explicitly in
get_test().If the LMDB has zero frames, Line 854 calls
max(...)on an empty dict and raises a cryptic error. Add an explicit empty-dataset guard.Suggested fix
if nloc is not None: if nloc not in self._nloc_groups: raise ValueError( f"No frames with nloc={nloc}. Available: {sorted(self._nloc_groups.keys())}" ) frame_indices = self._nloc_groups[nloc] natoms = nloc + elif not self._frames: + raise ValueError("LMDB dataset contains no frames") elif len(self._nloc_groups) == 1: # Uniform nloc — use all frames natoms = next(iter(self._nloc_groups))🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 848 - 855, get_test() can call max(...) on an empty self._nloc_groups when the LMDB has zero frames; add an explicit guard at the start of get_test() (or before the branching that uses self._nloc_groups) to detect an empty dataset (e.g., if not self._frames or not self._nloc_groups) and handle it cleanly by returning an empty result or raising a clear exception; update the branch that sets natoms/frame_indices (the block using self._nloc_groups, natoms, and frame_indices) to assume non-empty only after this guard so max(...) is never called on an empty dict.
1030-1032:⚠️ Potential issue | 🟠 MajorGuard
frame_nlocsmetadata indexing in merge path.Line 1031 assumes
src_nlocshasnframesentries. Truncated/malformed metadata raisesIndexErrorand aborts merge.Suggested fix
- if src_nlocs is not None: - frame_nlocs.append(int(src_nlocs[i])) - else: + if src_nlocs is not None and i < len(src_nlocs): + frame_nlocs.append(int(src_nlocs[i])) + else: frame_raw = msgpack.unpackb(raw, raw=False) atype_raw = frame_raw.get("atom_types") if isinstance(atype_raw, dict):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 1030 - 1032, The code unconditionally indexes src_nlocs[i] when building frame_nlocs in the merge path, which can raise IndexError for truncated/malformed metadata; update the merge logic (the block that appends to frame_nlocs using src_nlocs, referencing variables frame_nlocs, src_nlocs, nframes, and loop index i) to guard access: check that src_nlocs is not None and i < len(src_nlocs) before converting/append; if missing, append a sensible default (e.g., 0) or log a warning and continue so the merge does not abort.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Line 925: The loop uses an unused variable `req_info` in the iteration `for
key, req_info in all_keys.items()` — rename `req_info` to `_req_info` or `_` to
satisfy Ruff B007 and avoid unused-variable warnings; update that loop
occurrence (and any identical loops in the same file) and run `ruff check .` /
`ruff format .` before committing.
---
Duplicate comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 367-371: Validate that all atom type IDs in frame["atype"] are
within the valid range before any remapping or any use in natoms_vec
computation: check they are non-negative integers and less than the length/shape
of self._type_remap (when self._type_remap is not None) or less than expected
num_types; if any id is out of range raise a clear ValueError including the
offending ids and context (frame id), then proceed to reshape/astype, apply
self._type_remap on valid indices, and only after successful remap compute
natoms_vec; refer to symbols frame["atype"], self._type_remap and natoms_vec to
locate where to insert the validation and error handling.
- Around line 270-281: The batch_size handling currently silently coerces
invalid string specs and accepts non-positive integers; update the block that
sets self._auto_rule and self.batch_size (around batch_size, self._auto_rule,
and _compute_batch_size in lmdb_data.py) to strictly validate inputs: accept
only "auto" or "auto:<positive int>" for string specs (parse and set
self._auto_rule or raise ValueError on any other format), and when batch_size is
numeric ensure int(batch_size) > 0 (raise ValueError otherwise), also catch and
re-raise parsing errors with a clear message so invalid configs fail fast
instead of being coerced.
- Line 461: The local variable unique_nlocs is assigned via
sorted(self._nloc_groups.keys()) but never used, causing an F841; remove the
unused assignment (delete the unique_nlocs line) or replace its usage if it was
intended to be used later in the same function/method that contains
self._nloc_groups; ensure you reference the variable name unique_nlocs and the
expression sorted(self._nloc_groups.keys()) when making the change, then run
ruff check . and ruff format . before committing.
- Around line 848-855: get_test() can call max(...) on an empty
self._nloc_groups when the LMDB has zero frames; add an explicit guard at the
start of get_test() (or before the branching that uses self._nloc_groups) to
detect an empty dataset (e.g., if not self._frames or not self._nloc_groups) and
handle it cleanly by returning an empty result or raising a clear exception;
update the branch that sets natoms/frame_indices (the block using
self._nloc_groups, natoms, and frame_indices) to assume non-empty only after
this guard so max(...) is never called on an empty dict.
- Around line 1030-1032: The code unconditionally indexes src_nlocs[i] when
building frame_nlocs in the merge path, which can raise IndexError for
truncated/malformed metadata; update the merge logic (the block that appends to
frame_nlocs using src_nlocs, referencing variables frame_nlocs, src_nlocs,
nframes, and loop index i) to guard access: check that src_nlocs is not None and
i < len(src_nlocs) before converting/append; if missing, append a sensible
default (e.g., 0) or log a warning and continue so the merge does not abort.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: ed63ac4d-0db4-446a-aaf9-6d84603c1c04
📒 Files selected for processing (1)
deepmd/dpmodel/utils/lmdb_data.py
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (3)
deepmd/dpmodel/utils/lmdb_data.py (3)
762-765: Consider deterministic seeding for test reproducibility.The shuffle uses
np.random.default_rng()without a seed, making test results non-reproducible. While randomization is often desirable for testing, consider adding an optionalseedparameter for debugging reproducibility.♻️ Proposed enhancement
def __init__( self, lmdb_path: str, type_map: list[str] | None = None, shuffle_test: bool = True, + seed: int | None = None, **kwargs: Any, ) -> None: ... # Shuffle if requested if shuffle_test: - rng = np.random.default_rng() + rng = np.random.default_rng(seed) indices = rng.permutation(len(self._frames)) self._frames = [self._frames[i] for i in indices]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 762 - 765, The test shuffle is non-deterministic because it calls np.random.default_rng() without a seed; update the method (the code that checks shuffle_test and shuffles self._frames) to accept an optional seed parameter (default None) and pass it into np.random.default_rng(seed) so permutation(self._frames) becomes reproducible when a seed is provided; ensure callers that invoke this behavior can pass the seed (or leave None for non-deterministic behavior) and keep using indices = rng.permutation(len(self._frames)) and self._frames = [self._frames[i] for i in indices] unchanged.
1007-1045: Consider using context manager for source LMDB environments.While
src_env.close()is called, an exception during the loop could leave environments open. Using a context manager would ensure cleanup.♻️ Proposed enhancement
for src_path in src_paths: - src_env = _open_lmdb(src_path) - with src_env.begin() as txn: - meta = _read_metadata(txn) - nframes, src_fmt, natoms_per_type = _parse_metadata(meta) - ... - src_env.close() + with lmdb.open(src_path, readonly=True, lock=False, readahead=False, meminit=False) as src_env: + with src_env.begin() as txn: + meta = _read_metadata(txn) + nframes, src_fmt, natoms_per_type = _parse_metadata(meta) + ...Note: Verify that
lmdb.Environmentsupports context manager protocol (it does since lmdb 0.94).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 1007 - 1045, Wrap source LMDB environments opened by _open_lmdb in a context manager to guarantee they are closed on exceptions: replace the pattern "src_env = _open_lmdb(src_path); with src_env.begin() as txn: ...; src_env.close()" with "with _open_lmdb(src_path) as src_env: ..." and keep the existing nested "with src_env.begin() as src_txn, dst_env.begin(write=True) as dst_txn:" block intact; remove the explicit src_env.close() call and ensure code still references src_nlocs, src_txn, frame_idx, and other variables unchanged.
245-247: Thread-safety assumption for persistent transaction.The comment at Line 246 states the persistent transaction is "Safe because we use num_workers=0 in DataLoader." This is a critical assumption—if a caller uses
num_workers > 0with PyTorch DataLoader, the LMDB transaction will be shared across processes (via fork), which can cause corruption or hangs.Consider adding a runtime check or explicit documentation in the class docstring warning against multi-worker usage, or lazily initializing the transaction per-worker.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 762-765: The test shuffle is non-deterministic because it calls
np.random.default_rng() without a seed; update the method (the code that checks
shuffle_test and shuffles self._frames) to accept an optional seed parameter
(default None) and pass it into np.random.default_rng(seed) so
permutation(self._frames) becomes reproducible when a seed is provided; ensure
callers that invoke this behavior can pass the seed (or leave None for
non-deterministic behavior) and keep using indices =
rng.permutation(len(self._frames)) and self._frames = [self._frames[i] for i in
indices] unchanged.
- Around line 1007-1045: Wrap source LMDB environments opened by _open_lmdb in a
context manager to guarantee they are closed on exceptions: replace the pattern
"src_env = _open_lmdb(src_path); with src_env.begin() as txn: ...;
src_env.close()" with "with _open_lmdb(src_path) as src_env: ..." and keep the
existing nested "with src_env.begin() as src_txn, dst_env.begin(write=True) as
dst_txn:" block intact; remove the explicit src_env.close() call and ensure code
still references src_nlocs, src_txn, frame_idx, and other variables unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 96e35842-9118-4449-b995-2d43315d81b6
📒 Files selected for processing (1)
deepmd/dpmodel/utils/lmdb_data.py
| del _ENV_CACHE[resolved] | ||
| try: | ||
| env.close() | ||
| except Exception: |
Check notice
Code scanning / CodeQL
Empty except Note
wanghan-iapcm
approved these changes
Apr 2, 2026
njzjz
requested changes
Apr 6, 2026
Member
njzjz
left a comment
njzjz
left a comment
There was a problem hiding this comment.
I feel that there are too many if...else blocks added in this PR and they don't have a good abstract.
Comment on lines
+637
to
+639
| from deepmd.utils.data_system import ( | ||
| prob_sys_size_ext, | ||
| ) |
Check notice
Code scanning / CodeQL
Cyclic import Note
Comment on lines
+809
to
+811
| from deepmd.dpmodel.utils.lmdb_data import ( | ||
| is_lmdb, | ||
| ) |
Check notice
Code scanning / CodeQL
Cyclic import Note
njzjz
approved these changes
Apr 6, 2026
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
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.
Summary by CodeRabbit
New Features
Tests
Documentation
Chores