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
|
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 (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 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 extensive interoperability and robustness tests were updated. ChangesLMDB backend rewrite
Estimated code review effort: 5 (Critical) | ~120 minutes 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
| @@ -0,0 +1,100 @@ | |||
| # LMDB Format | |||
|
|
|||
| The format `lmdb` stores the frames of one or more systems in a single [LMDB](http://www.lmdb.tech/doc/) database, and can be loaded or dumped through {class}`dpdata.System`, {class}`dpdata.LabeledSystem`, and {class}`dpdata.MultiSystems`. The on-disk layout coincides with the LMDB datasets read by the DeePMD-kit data loader. Core fields and registered additional fields are therefore available to DeePMD-kit under their `deepmd_name`. | |||
There was a problem hiding this comment.
The claim that core fields are available to DeePMD-kit "under their deepmd_name" is incorrect for core fields. The write path stores core fields under their plural dpdata names (coords, cells, energies, forces, virials) via _CORE_DISK_NAMES, not their singular deepmd_name (coord, box, energy, force, virial). Only additional registered fields fall back to deepmd_name (_disk_name = _CORE_DISK_NAMES.get(dtype.name, dtype.deepmd_name)).
This also contradicts the module docstring in format.py, which correctly states core keys use the plural names and only additional fields use DataType.deepmd_name. Suggest rewording so that core fields are described as using their plural on-disk names, and only additional registered fields as using their deepmd_name.
| return lmdb.open( | ||
| str(path), | ||
| readonly=True, | ||
| lock=False, |
There was a problem hiding this comment.
_open_publish_guard (and _open_read_env above) open with lock=False, which disables LMDB's reader-lock table — the only mechanism that would report a second in-process open. With lock=False, opening a path that another environment already has open succeeds, so this except lmdb.Error branch can never fire and the advertised "Close DeePMD-kit and other readers first" protection is effectively a no-op on POSIX.
Consequences: (a) the safety guarantee is not delivered, and (b) three tests fail locally — test_existing_external_reader_has_actionable_error, test_external_lmdb_reader_blocks_overwrite, and test_publish_guard_closes_external_reader_race. Note this PR has no unit-test CI job (only CodeRabbit / readthedocs / pre-commit), so those failures are not caught automatically. Recommend reworking the detection (e.g. rely on the reader-lock table instead of lock=False) or revisiting the tests/claims.
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:
|
2 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
dump_systemsto export without formula-based merging.type_mapoverrides, batching, overwrite control, andmax_frameslimits.lmdbsystems documentation page and linked it in the systems navigation.msgpack-numpytomsgpack.