diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index a85f92be87..7014552963 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -2685,20 +2685,61 @@ 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 _copy_lmdb_source( src_path: str, + metadata: dict[str, Any], dst_env: lmdb.Environment, dst_format: str, frame_idx: int, frame_nlocs: list[int], frame_system_ids: list[int], system_id_offset: int, -) -> tuple[int, dict, list[str] | None, int]: - """Copy one source under a ref-counted environment lease.""" +) -> tuple[int, dict, int]: + """Copy one validated source under a ref-counted environment lease. + + The caller supplies metadata collected during the validation preflight so + this copy pass does not read and decode it a second time. + """ src_env = _open_lmdb(src_path) try: - with src_env.begin() as transaction: - metadata = _read_metadata(transaction) nframes, src_format, natoms_per_type = _parse_metadata(metadata) fallback_natoms = sum(natoms_per_type) source_nlocs = metadata.get("frame_nlocs") @@ -2741,7 +2782,6 @@ def _copy_lmdb_source( return ( frame_idx, metadata.get("system_info", {}), - metadata.get("type_map"), system_id_offset, ) finally: @@ -2772,30 +2812,47 @@ 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 try: - for src_path in src_paths: + for src_path, metadata in source_metadata: ( frame_idx, source_system_info, - source_type_map, sys_id_offset, ) = _copy_lmdb_source( src_path, + metadata, dst_env, fmt, frame_idx, @@ -2805,8 +2862,6 @@ def merge_lmdb( ) if first_system_info is None: first_system_info = source_system_info - if first_type_map is None: - first_type_map = source_type_map merged_meta = { "nframes": frame_idx, @@ -2815,8 +2870,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 transaction: transaction.put( b"__metadata__", diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index d945ccc008..4b253aa1ae 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -1068,6 +1068,57 @@ 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) + + @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, 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=second_type_map + ) + 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