fix(dpmodel): reject incompatible LMDB merge type maps - #5838
Conversation
Preflight all source metadata before touching the destination. Permit identical explicit maps and all-mapless legacy inputs, but reject reordered maps and explicit/missing mixtures that byte-for-byte frame copying cannot preserve safely. Cover incompatible and ambiguous sources, destination preservation on validation failure, and frame types from each segment of a valid merge. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
📝 WalkthroughWalkthroughChanges
LMDB type-map-safe merging
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant merge_lmdb
participant SourceLMDBs
participant TypeMapValidator
participant DestinationLMDB
merge_lmdb->>SourceLMDBs: Read metadata from all sources
merge_lmdb->>TypeMapValidator: Validate type_map compatibility
TypeMapValidator-->>merge_lmdb: Return shared mapping or ValueError
merge_lmdb->>DestinationLMDB: Create or replace destination after validation
merge_lmdb->>SourceLMDBs: Copy frames and release environments
merge_lmdb->>DestinationLMDB: Write validated metadata
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 |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #5838 +/- ##
==========================================
- Coverage 79.47% 79.22% -0.25%
==========================================
Files 1072 1072
Lines 125041 125056 +15
Branches 4536 4541 +5
==========================================
- Hits 99373 99075 -298
- Misses 24044 24356 +312
- Partials 1624 1625 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The fix is correct and the fail-fast-before-touching-the-destination ordering is well tested (marker-preserved / dst-not-created). One test-coverage suggestion inline. Also flagging a coordination point: this PR and the open #5797 both rewrite merge_lmdb's env teardown and change the same src_env.close() -> _close_lmdb(src_path) line, so they will conflict — worth deciding a merge order (this PR already fixes #5797's shared-reader .close() bug; #5797 additionally wraps the frame-copy loop in try/finally for #5635).
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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/lmdb_data.py (1)
1822-1887: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMerge loop and
dst_envaren't exception-safe; a mid-loop failure leaks the cached source env and leavesdst_envunclosed.Unlike the pre-load loop (which wraps
_close_lmdbintry/finally), this loop calls_close_lmdb(src_path)at Line 1873 only after thewithblock finishes normally. If anything inside the loop raises (e.g.msgpack.unpackbon a malformedrawvalue, or an unexpected metadata shape), the ref-counted cached env for that source is never released, anddst_env(opened at Line 1814) is never closed either — both leak on the error path.🔒️ Proposed fix
for src_path, meta in source_metadata: src_env = _open_lmdb(src_path) - nframes, src_fmt, natoms_per_type = _parse_metadata(meta) - fallback_natoms = sum(natoms_per_type) - ... - with src_env.begin() as src_txn, dst_env.begin(write=True) as dst_txn: - ... - ... - _close_lmdb(src_path) + try: + nframes, src_fmt, natoms_per_type = _parse_metadata(meta) + fallback_natoms = sum(natoms_per_type) + ... + with src_env.begin() as src_txn, dst_env.begin(write=True) as dst_txn: + ... + ... + finally: + _close_lmdb(src_path)Similarly wrap the body from
dst_env = lmdb.open(...)through the metadata write intry/finally: dst_env.close().🤖 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/lmdb_data.py` around lines 1822 - 1887, Make the merge operation exception-safe by wrapping the full destination workflow—from the `dst_env` open through the merged metadata write—in `try/finally` and always closing `dst_env`. Within the `for src_path, meta in source_metadata` loop, ensure each source environment opened by `_open_lmdb` is released with `_close_lmdb(src_path)` in a per-source `finally` block so failures during transaction processing or metadata parsing cannot leak either environment.
🤖 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/lmdb_data.py`:
- Around line 1822-1887: Make the merge operation exception-safe by wrapping the
full destination workflow—from the `dst_env` open through the merged metadata
write—in `try/finally` and always closing `dst_env`. Within the `for src_path,
meta in source_metadata` loop, ensure each source environment opened by
`_open_lmdb` is released with `_close_lmdb(src_path)` in a per-source `finally`
block so failures during transaction processing or metadata parsing cannot leak
either environment.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ef126776-0db4-45d5-be26-d0173ad8e5ef
📒 Files selected for processing (2)
deepmd/dpmodel/utils/lmdb_data.pysource/tests/pt/test_lmdb_dataloader.py
Resolve the LMDB merge conflict by combining fail-fast type-map validation with master’s exception-safe environment leasing. Reuse preflight metadata during copying to avoid a second metadata decode. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/dpmodel/utils/lmdb_data.py (1)
2847-2879: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConsider removing the partial destination when a merge fails.
The
finallyblock closesdst_env, but the partially written LMDB directory stays atdst_path. A later reader can open it and read an incomplete dataset without metadata, or with stale metadata from a previous run. The validation preflight already protects an existing destination; this gap only affects the newly created output.♻️ Proposed cleanup on failure
dst_env = lmdb.open(dst_path, map_size=map_size) frame_idx = 0 fmt = "012d" frame_nlocs: list[int] = [] frame_system_ids: list[int] = [] first_system_info: dict | None = None sys_id_offset = 0 + merge_succeeded = False try: for src_path, metadata in source_metadata: @@ with dst_env.begin(write=True) as transaction: transaction.put( b"__metadata__", msgpack.packb(merged_meta, use_bin_type=True), ) + merge_succeeded = True finally: dst_env.close() + if not merge_succeeded: + shutil.rmtree(dst_path, ignore_errors=True)Note:
source/tests/pt/test_lmdb_dataloader.py::test_failed_merge_releases_source_leaseasserts only the source lease refcount, so this change stays compatible with it.🤖 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/lmdb_data.py` around lines 2847 - 2879, Update the merge failure handling around the destination creation and the existing try/finally that closes dst_env so any newly created destination at dst_path is removed when the merge fails. Preserve existing destinations protected by validation preflight, and ensure cleanup occurs only after an unsuccessful merge without affecting source lease release or successful output.
🤖 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 `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 2847-2879: Update the merge failure handling around the
destination creation and the existing try/finally that closes dst_env so any
newly created destination at dst_path is removed when the merge fails. Preserve
existing destinations protected by validation preflight, and ensure cleanup
occurs only after an unsuccessful merge without affecting source lease release
or successful output.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7133280b-c11c-4a67-828f-e5642356d989
📒 Files selected for processing (1)
deepmd/dpmodel/utils/lmdb_data.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Addressed exactly as asked, including the part I cared about most.
The parametrization uses ["O", "H", "N"] against a base of ["O", "H"], which is genuinely the prefix-compatible case -- indices 0 and 1 denote the same species in both maps, so a byte-for-byte copy would in fact be safe there, and the PR rejects it anyway. That is the boundary worth pinning, and the ids=["reordered", "prefix-compatible-superset"] labels make it obvious at a glance which case is which when one of them fails. The comment records the reasoning in the right terms too: requiring identical metadata rather than proving frame-by-frame safety is a deliberate conservative choice, and now a future attempt to relax it has to argue with a red test instead of slipping through.
I ran the class against a clean master baseline rather than reading it. Three fail unpatched -- both parametrizations plus test_merge_rejects_mixed_explicit_and_missing_type_maps -- and all eight pass after. Worth noting the tests also assert the rejection happens before any output is written (marker.read_text() == "preserve me" and not dst.exists()), which is what makes this a data-safety fix rather than just a nicer error message; a version that validated after opening the destination would still pass a naive pytest.raises check but fail these.
Approving.
3 participants
Fixes #5634
Summary
This chooses fail-fast validation instead of decoding and rewriting frames. A remapping implementation would also need to keep
atom_types,atom_names,atom_numbs, andsystem_infomutually consistent; rejecting incompatible metadata is the smaller safe change.Why existing tests missed this
The merge coverage used either two mapless legacy databases or sources with the same
["O", "H"]map. Its type-map test inspected only the merged metadata and never read frames from both source segments. Because every fixture generated the same numeric 0/1 type ordering, raw frame copying preserved all tested shapes, counts, and IDs even though a reversed source map would change their species meaning.The regular LMDB remapping tests cover one database mapped to a model type map, not multiple source maps being collapsed into one merged metadata map. Reader/dataset consistency checks would also interpret the same corrupted merged metadata and therefore agree with each other.
Validation
source/tests/pt/test_lmdb_dataloader.py: 53 passedruff format .ruff check .Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit