feat(lmdb): align backend with DeePMD-kit#1012
Conversation
Replace the legacy LMDB schema with a safe, mixed-type format that preserves system partitions and interoperates with DeePMD-kit training datasets.
dosubot
Bot
added
size:XXL
for more information, see https://pre-commit.ci
|
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: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThe LMDB backend was rewritten around a flat DeePMD-kit-compatible database schema, staged batched writes, atomic publication, validated frame loading, composition grouping, atom-order canonicalization, and custom field protocols. Documentation, exports, dependencies, and interoperability and robustness tests were updated. ChangesLMDB backend rewrite
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant LMDBFormat
participant StagedDatabase
participant PublishedDatabase
participant MultiSystems
LMDBFormat->>StagedDatabase: validate and write frame records
StagedDatabase->>PublishedDatabase: atomically publish validated data
MultiSystems->>LMDBFormat: load metadata and frames
LMDBFormat->>MultiSystems: canonicalize and group compositions
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
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 |
Merging this PR will not alter performance
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/test_lmdb_custom_dtype.py (1)
313-324: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueUse a context manager for the write env to avoid a leaked handle.
Unlike the other tests here (which use
with lmdb.open(...) as env:), this block opens the environment manually and only closes it at Line 324. Ifmsgpack.unpackb/packbraises inside thewith env.begin(write=True)block,envis never closed, which can leak the handle and, on some platforms, blockshutil.rmtreeintearDown.♻️ Wrap the environment in a context manager
- env = lmdb.open(self.lmdb_path, map_size=1 << 30) - with env.begin(write=True) as txn: - key = b"000000000001" - frame = msgpack.unpackb(txn.get(key), raw=False) - replacement = np.array([3.0, 4.0]) - frame["static_data"] = { - "type": str(replacement.dtype), - "shape": list(replacement.shape), - "data": replacement.tobytes(), - } - txn.put(key, msgpack.packb(frame, use_bin_type=True)) - env.close() + with lmdb.open(self.lmdb_path, map_size=1 << 30) as env: + with env.begin(write=True) as txn: + key = b"000000000001" + frame = msgpack.unpackb(txn.get(key), raw=False) + replacement = np.array([3.0, 4.0]) + frame["static_data"] = { + "type": str(replacement.dtype), + "shape": list(replacement.shape), + "data": replacement.tobytes(), + } + txn.put(key, msgpack.packb(frame, use_bin_type=True))🤖 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 `@tests/test_lmdb_custom_dtype.py` around lines 313 - 324, Update the LMDB environment setup in this test to use lmdb.open as a context manager around the existing write transaction and frame replacement logic. Remove the manual env.close call while preserving the current key update behavior and transaction scope.
🤖 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 `@dpdata/formats/lmdb/format.py`:
- Around line 1061-1066: Update the publish flow around _open_publish_guard and
os.replace so the guard is closed before replacing staged_data with target_data,
ensuring overwrite=True works on Windows. Preserve cleanup of the guard and temp
directory, and avoid leaving the directory handle open during os.replace.
---
Nitpick comments:
In `@tests/test_lmdb_custom_dtype.py`:
- Around line 313-324: Update the LMDB environment setup in this test to use
lmdb.open as a context manager around the existing write transaction and frame
replacement logic. Remove the manual env.close call while preserving the current
key update behavior and transaction scope.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b43de27f-e10b-4eec-9e56-4a0e30df6de9
📒 Files selected for processing (8)
docs/systems/index.rstdocs/systems/lmdb.mddpdata/format.pydpdata/formats/lmdb/__init__.pydpdata/formats/lmdb/format.pypyproject.tomltests/test_lmdb.pytests/test_lmdb_custom_dtype.py
Make protocol-metadata assertions independent of global plugin registration order and reject unsafe Windows overwrite operations explicitly.
wanghan-iapcm
requested review from
iProzd,
njzjz and
wanghan-iapcm
and removed request for
iProzd
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1012 +/- ##
==========================================
- Coverage 86.87% 86.82% -0.06%
==========================================
Files 89 89
Lines 8266 9114 +848
==========================================
+ Hits 7181 7913 +732
- Misses 1085 1201 +116 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Document the core-field wire names and require python-lmdb 2.0 so duplicate environment detection reliably protects active readers.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes for four reproducible compatibility/data-integrity issues. The LMDB-focused suite passes (439 tests), and Ruff lint/format checks pass. Full discovery ran 2,201 tests with 13 unrelated missing-optional-dependency errors (rdkit/parmed); no LMDB tests failed.
Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
Read and migrate the released schema, remap every type axis consistently, and reject cross-system field ownership or unsupported scalar data.
dosubot
Bot
added
size:L
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
dpdata/formats/lmdb/format.py (1)
1538-1548: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftAvoid whole-registry rollback across concurrent readers.
One reader can snapshot the global
DTYPES, another can register a schema, and then the first reader’s early close restores its old tuples—erasing the second reader’s committed registration. Serialize registry mutations and rollback only definitions installed by this invocation when they remain unchanged.🤖 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 `@dpdata/formats/lmdb/format.py` around lines 1538 - 1548, Update the dataset iteration flow around _register_specs to serialize global DTYPES mutations across readers and avoid restoring whole registry snapshots. Track the definitions installed by this invocation, and on early termination remove or restore only those entries when their current values still match the installed definitions, preserving concurrent registrations and committed changes.
🧹 Nitpick comments (1)
tests/test_lmdb_custom_dtype.py (1)
451-461: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a writer-side type-axis permutation case.
This test only permutes types while reading because the source and write-time
type_mapboth start as["H", "O"]. Add a case that writes with["O", "H", "C"]and reads with that same map to verify_frame_payloadplaces local values into global slots correctly.🤖 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 `@tests/test_lmdb_custom_dtype.py` around lines 451 - 461, Add a writer-side permutation test alongside test_type_axes_follow_requested_type_map that writes the system using type_map=["O", "H", "C"] and reads it with the same map, then assert atom_names and the expected type_values/static_type_values ordering to verify _frame_payload places local values in the corresponding global slots.
🤖 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 `@dpdata/formats/lmdb/format.py`:
- Around line 1654-1675: The legacy type migration block should derive active
elements from the validated `atom_types` values rather than
`data["atom_numbs"]`. Validate that every atom type index is within `old_names`,
raise the established `LMDBError` for invalid indices, then compute active names
from the unique referenced indices before checking `missing_active`; preserve
the existing remapping and recount behavior for valid data.
- Around line 368-428: Update _decode_legacy_frame and its legacy object
decoding path so malformed descriptors, including missing type/data/shape keys
and truncated dtype entries, are converted to LMDBFrameError. Extend the decode
boundary’s exception handling to cover the KeyError and IndexError raised by
_decode_legacy_object or _unpack_legacy_dtype, while preserving the existing
contextual error message and exception chaining.
- Around line 1481-1489: Update the type_map handling around _meta_get(meta,
"type_map") to distinguish missing metadata from an explicitly empty stored map.
Reject an empty type_map as malformed instead of allowing the later positional
Type_0/Type_1 fallback, while preserving existing validation and behavior for
non-empty maps and genuinely missing metadata.
---
Outside diff comments:
In `@dpdata/formats/lmdb/format.py`:
- Around line 1538-1548: Update the dataset iteration flow around
_register_specs to serialize global DTYPES mutations across readers and avoid
restoring whole registry snapshots. Track the definitions installed by this
invocation, and on early termination remove or restore only those entries when
their current values still match the installed definitions, preserving
concurrent registrations and committed changes.
---
Nitpick comments:
In `@tests/test_lmdb_custom_dtype.py`:
- Around line 451-461: Add a writer-side permutation test alongside
test_type_axes_follow_requested_type_map that writes the system using
type_map=["O", "H", "C"] and reads it with the same map, then assert atom_names
and the expected type_values/static_type_values ordering to verify
_frame_payload places local values in the corresponding global slots.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 554343ec-37be-4839-a466-e365593127f2
📒 Files selected for processing (5)
docs/systems/lmdb.mddpdata/formats/lmdb/format.pydpdata/system.pytests/test_lmdb.pytests/test_lmdb_custom_dtype.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_lmdb.py
Wrap malformed legacy descriptors, reject invalid stored type maps, and derive legacy active types from validated atom indices.
3 participants
Summary
dump_systemsAPI forprob_sys_sizeTest plan
python -m unittest test_lmdb test_lmdb_custom_dtype test_deepmd_comp test_deepmd_mixed test_multisystems(438 tests)ruff check dpdata/ tests/test_lmdb.py tests/test_lmdb_custom_dtype.pygit diff --checkLmdbDataReaderFull discovery ran 2,194 tests. The remaining 46 errors are pre-existing optional-dependency/environment failures (for example, unavailable
parmed); no LMDB tests failed.Summary by CodeRabbit
mixed_type,type_map,max_frames).dump_systemsto export systems without formula-based merging.msgpack-numpytomsgpack.