resolve comments

Author: iProzd

deepmd/dpmodel/utils/lmdb_data.py

@@ -146,11 +146,9 @@ def _remap_keys(frame: dict[str, Any]) -> dict[str, Any]:
     return out
 
 
-def is_lmdb(systems: Any) -> bool:
+def is_lmdb(systems: str) -> bool:
     """Check if systems points to an LMDB dataset."""
-    if not isinstance(systems, str):
-        return False
-    return systems.endswith(".lmdb") or Path(systems, "data.mdb").exists()
+    return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file()
 
 
 def _parse_metadata(meta: dict) -> tuple[int, str, list[int]]:

deepmd/entrypoints/test.py

@@ -13,7 +13,6 @@
 import numpy as np
 
 from deepmd.common import (
-    expand_sys_str,
     j_loader,
 )
 from deepmd.dpmodel.utils.lmdb_data import (
@@ -148,11 +147,8 @@ def test(
             systems = str((root / Path(systems)).resolve())
         else:
             systems = [str((root / Path(ss)).resolve()) for ss in systems]
-        if is_lmdb(systems):
-            all_sys = [systems]
-        else:
-            patterns = data_params.get("rglob_patterns", None)
-            all_sys = process_systems(systems, patterns=patterns)
+        patterns = data_params.get("rglob_patterns", None)
+        all_sys = process_systems(systems, patterns=patterns)
     elif valid_json is not None:
         jdata = j_loader(valid_json)
         jdata = update_deepmd_input(jdata)
@@ -165,19 +161,13 @@ def test(
             systems = str((root / Path(systems)).resolve())
         else:
             systems = [str((root / Path(ss)).resolve()) for ss in systems]
-        if is_lmdb(systems):
-            all_sys = [systems]
-        else:
-            patterns = data_params.get("rglob_patterns", None)
-            all_sys = process_systems(systems, patterns=patterns)
+        patterns = data_params.get("rglob_patterns", None)
+        all_sys = process_systems(systems, patterns=patterns)
     elif datafile is not None:
         with open(datafile) as datalist:
             all_sys = datalist.read().splitlines()
     elif system is not None:
-        if is_lmdb(system):
-            all_sys = [system]
-        else:
-            all_sys = expand_sys_str(system)
+        all_sys = process_systems(system)
     else:
         raise RuntimeError("No data source specified for testing")
 

deepmd/pt/entrypoints/main.py

@@ -147,6 +147,23 @@ def prepare_trainer_input_single(
                     Path(stat_file_path_single).mkdir()
             stat_file_path_single = DPPath(stat_file_path_single, "a")
 
+        rank_seed = [rank, seed % (2**32)] if seed is not None else None
+
+        def _make_dp_loader_set(
+            systems: str | list[str],
+            dataset_params: dict[str, Any],
+        ) -> DpLoaderSet:
+            """Create a DpLoaderSet from systems with pattern expansion."""
+            patterns = dataset_params.get("rglob_patterns", None)
+            systems = process_systems(systems, patterns=patterns)
+            return DpLoaderSet(
+                systems,
+                dataset_params["batch_size"],
+                model_params_single["type_map"],
+                seed=rank_seed,
+                modifier=modifier,
+            )
+
         # LMDB path: single string → LmdbDataset
         if is_lmdb(training_systems):
             auto_prob = training_dataset_params.get("auto_prob", None)
@@ -163,46 +180,21 @@ def prepare_trainer_input_single(
                     validation_dataset_params["batch_size"],
                 )
             elif validation_systems is not None:
-                val_patterns = validation_dataset_params.get("rglob_patterns", None)
-                validation_systems = process_systems(validation_systems, val_patterns)
-                rank_seed = [rank, seed % (2**32)] if seed is not None else None
-                validation_data_single = DpLoaderSet(
-                    validation_systems,
-                    validation_dataset_params["batch_size"],
-                    model_params_single["type_map"],
-                    seed=rank_seed,
-                    modifier=modifier,
+                validation_data_single = _make_dp_loader_set(
+                    validation_systems, validation_dataset_params
                 )
             else:
                 validation_data_single = None
         else:
             # Standard npy path
-            trn_patterns = training_dataset_params.get("rglob_patterns", None)
-            training_systems = process_systems(training_systems, patterns=trn_patterns)
-            if validation_systems is not None:
-                val_patterns = validation_dataset_params.get("rglob_patterns", None)
-                validation_systems = process_systems(validation_systems, val_patterns)
-
-            # avoid the same batch sequence among devices
-            rank_seed = [rank, seed % (2**32)] if seed is not None else None
+            train_data_single = _make_dp_loader_set(
+                training_systems, training_dataset_params
+            )
             validation_data_single = (
-                DpLoaderSet(
-                    validation_systems,
-                    validation_dataset_params["batch_size"],
-                    model_params_single["type_map"],
-                    seed=rank_seed,
-                    modifier=modifier,
-                )
+                _make_dp_loader_set(validation_systems, validation_dataset_params)
                 if validation_systems
                 else None
             )
-            train_data_single = DpLoaderSet(
-                training_systems,
-                training_dataset_params["batch_size"],
-                model_params_single["type_map"],
-                seed=rank_seed,
-                modifier=modifier,
-            )
         return (
             train_data_single,
             validation_data_single,

deepmd/utils/data_system.py

@@ -791,6 +791,7 @@ def process_systems(
 
     If it is a single directory, search for all the systems in the directory.
     If it is a list, each item in the list is treated as a directory to search.
+    If it is a single LMDB path, return it directly without expansion.
     Check if the systems are valid.
 
     Parameters
@@ -805,6 +806,14 @@ def process_systems(
     result_systems: list of str
         The valid systems
     """
+    from deepmd.dpmodel.utils.lmdb_data import (
+        is_lmdb,
+    )
+
+    # LMDB path: return directly without expansion
+    if isinstance(systems, str) and is_lmdb(systems):
+        return [systems]
+
     # Normalize input to a list of paths to search
     if isinstance(systems, str):
         search_paths = [systems]

source/tests/common/dpmodel/test_lmdb_data.py

@@ -309,7 +309,6 @@ def test_is_lmdb(self):
         self.assertTrue(is_lmdb(self._lmdb_path))
         self.assertTrue(is_lmdb("something.lmdb"))
         self.assertFalse(is_lmdb("/some/npy/system"))
-        self.assertFalse(is_lmdb(["list", "of", "systems"]))
 
     def test_lmdb_test_data(self):
         td = LmdbTestData(self._lmdb_path, type_map=self._type_map, shuffle_test=False)