From 3e38522b0a7b22c4abb7d681d54f76e15ff149ae Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 02:31:05 +0800 Subject: [PATCH 1/2] fix(dpmodel): reject incompatible LMDB merge type maps 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 --- deepmd/dpmodel/utils/lmdb_data.py | 71 +++++++++++++++++++++---- source/tests/pt/test_lmdb_dataloader.py | 44 +++++++++++++++ 2 files changed, 105 insertions(+), 10 deletions(-) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index d4e4c65b23..ac43c29b31 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -1724,6 +1724,44 @@ def get_test(self) -> dict[str, Any]: return self._inner.get_test(nloc=self._nloc) +def _validate_merge_type_maps( + source_metadata: list[tuple[str, dict[str, Any]]], +) -> list[str] | None: + """Return the shared type map required for byte-for-byte frame merging. + + ``merge_lmdb`` does not decode and rewrite atom-type arrays, so every + source must use exactly the same index-to-species mapping. All-missing + legacy metadata remains supported, but mixing explicit and missing maps is + rejected because compatibility cannot be established. + """ + source_type_maps = [(path, meta.get("type_map")) for path, meta in source_metadata] + explicit_type_maps = [ + (path, list(type_map)) + for path, type_map in source_type_maps + if type_map is not None + ] + if not explicit_type_maps: + return None + + formatted_maps = ", ".join( + f"{path}: {list(type_map)!r}" if type_map is not None else f"{path}: missing" + for path, type_map in source_type_maps + ) + if len(explicit_type_maps) != len(source_type_maps): + raise ValueError( + "Cannot merge LMDB datasets with mixed type_map metadata because " + f"raw atom-type indices cannot be validated ({formatted_maps})" + ) + + canonical_type_map = explicit_type_maps[0][1] + if any(type_map != canonical_type_map for _, type_map in explicit_type_maps[1:]): + raise ValueError( + "Cannot merge LMDB datasets with incompatible type_map values " + f"because frames are copied without remapping ({formatted_maps})" + ) + return canonical_type_map + + def merge_lmdb( src_paths: list[str], dst_path: str, @@ -1748,33 +1786,46 @@ def merge_lmdb( ------- str Path to the created LMDB. + + Raises + ------ + ValueError + If sources use different explicit type maps, or mix explicit type-map + metadata with legacy metadata where the mapping is missing. """ import os import shutil + # Validate every source before replacing or creating the destination. A + # type-map validation failure must not destroy an existing dataset or + # leave a partial output. + source_metadata: list[tuple[str, dict[str, Any]]] = [] + for src_path in src_paths: + src_env = _open_lmdb(src_path) + try: + with src_env.begin() as txn: + source_metadata.append((src_path, _read_metadata(txn))) + finally: + _close_lmdb(src_path) + merged_type_map = _validate_merge_type_maps(source_metadata) + if os.path.exists(dst_path): shutil.rmtree(dst_path) - 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 - first_type_map: list[str] | None = None sys_id_offset = 0 - for src_path in src_paths: + for src_path, meta in source_metadata: 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) fallback_natoms = sum(natoms_per_type) if first_system_info is None: first_system_info = meta.get("system_info", {}) - if first_type_map is None: - first_type_map = meta.get("type_map") # Check for pre-computed frame_nlocs in source src_nlocs = meta.get("frame_nlocs") @@ -1819,7 +1870,7 @@ def merge_lmdb( else: sys_id_offset += 1 - src_env.close() + _close_lmdb(src_path) # Write merged metadata with frame_nlocs for fast init merged_meta = { @@ -1829,8 +1880,8 @@ def merge_lmdb( "frame_nlocs": frame_nlocs, "frame_system_ids": frame_system_ids, } - if first_type_map is not None: - merged_meta["type_map"] = first_type_map + if merged_type_map is not None: + merged_meta["type_map"] = merged_type_map with dst_env.begin(write=True) as txn: txn.put(b"__metadata__", msgpack.packb(merged_meta, use_bin_type=True)) dst_env.close() diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index a51b75fff5..00ece32dfd 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -751,6 +751,50 @@ def test_merge_preserves_type_map(self, tmp_path): env.close() assert meta.get("type_map") == ["O", "H"] + reader = LmdbDataReader(dst, ["O", "H"]) + expected_atype = np.array([0, 0, 0, 1, 1, 1]) + np.testing.assert_array_equal(reader[0]["atype"], expected_atype) + np.testing.assert_array_equal(reader[5]["atype"], expected_atype) + + def test_merge_rejects_incompatible_type_maps_before_creating_output( + self, tmp_path + ): + """Raw frames cannot be shared under two different type index maps.""" + src1, src2 = str(tmp_path / "tm1.lmdb"), str(tmp_path / "tm2.lmdb") + _create_lmdb_with_system_ids( + src1, system_frames=[1], natoms=6, type_map=["O", "H"] + ) + _create_lmdb_with_system_ids( + src2, system_frames=[1], natoms=6, type_map=["H", "O"] + ) + dst = tmp_path / "incompatible.lmdb" + dst.mkdir() + marker = dst / "existing-data" + marker.write_text("preserve me") + + with pytest.raises(ValueError, match="incompatible type_map values") as exc: + merge_lmdb([src1, src2], str(dst)) + + assert src1 in str(exc.value) + assert src2 in str(exc.value) + assert marker.read_text() == "preserve me" + + def test_merge_rejects_mixed_explicit_and_missing_type_maps(self, tmp_path): + """A legacy source without a map cannot be proven index-compatible.""" + src_without_map = str(tmp_path / "legacy.lmdb") + src_with_map = str(tmp_path / "typed.lmdb") + _create_test_lmdb(src_without_map, nframes=1, natoms=6) + _create_lmdb_with_system_ids( + src_with_map, system_frames=[1], natoms=6, type_map=["O", "H"] + ) + dst = tmp_path / "mixed_metadata.lmdb" + + with pytest.raises(ValueError, match="mixed type_map metadata") as exc: + merge_lmdb([src_without_map, src_with_map], str(dst)) + + assert "missing" in str(exc.value) + assert not dst.exists() + # ============================================================ # Multitask LMDB training From 61005956d198feee1a306c0de5a322038bc0806c Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 20:03:34 +0800 Subject: [PATCH 2/2] test(dpmodel): cover LMDB type-map superset Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/pt/test_lmdb_dataloader.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index 00ece32dfd..353a28199c 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -756,16 +756,23 @@ def test_merge_preserves_type_map(self, tmp_path): np.testing.assert_array_equal(reader[0]["atype"], expected_atype) np.testing.assert_array_equal(reader[5]["atype"], expected_atype) + @pytest.mark.parametrize( + "second_type_map", + [["H", "O"], ["O", "H", "N"]], + ids=["reordered", "prefix-compatible-superset"], + ) def test_merge_rejects_incompatible_type_maps_before_creating_output( - self, tmp_path + self, tmp_path, second_type_map ): """Raw frames cannot be shared under two different type index maps.""" src1, src2 = str(tmp_path / "tm1.lmdb"), str(tmp_path / "tm2.lmdb") _create_lmdb_with_system_ids( src1, system_frames=[1], natoms=6, type_map=["O", "H"] ) + # Even a prefix-compatible superset is rejected: merge_lmdb deliberately + # requires identical metadata instead of proving frame-by-frame safety. _create_lmdb_with_system_ids( - src2, system_frames=[1], natoms=6, type_map=["H", "O"] + src2, system_frames=[1], natoms=6, type_map=second_type_map ) dst = tmp_path / "incompatible.lmdb" dst.mkdir()