Revert "perf: use multithread to accelarate stat computing and loading"

OutisLi · OutisLi · commit 0b3dfc8f9cfb · 2025-11-10T15:55:14.000+08:00
diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
@@ -1,12 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
-import os
 from collections import (
     defaultdict,
 )
-from concurrent.futures import (
-    ThreadPoolExecutor,
-)
 from typing import (
     Any,
     Callable,
@@ -43,7 +39,7 @@
 def make_stat_input(
     datasets: list[Any], dataloaders: list[Any], nbatches: int
 ) -> dict[str, Any]:
-    """Pack data for statistics in parallel.
+    """Pack data for statistics.
 
     Args:
     - dataset: A list of dataset to analyze.
@@ -53,83 +49,49 @@ def make_stat_input(
     -------
     - a list of dicts, each of which contains data from a system
     """
-    log.info(f"Packing data for statistics from {len(datasets)} systems")
-    dataloader_lens = [len(dl) for dl in dataloaders]
-    args_list = [
-        (dataloaders[i], nbatches, dataloader_lens[i]) for i in range(len(datasets))
-    ]
-
     lst = []
-    # I/O intensive, set a larger number of workers
-    with ThreadPoolExecutor(min(128, (os.cpu_count() or 1) * 6)) as executor:
-        lst = list(executor.map(_process_one_dataset, args_list))
-    log.info("Finished packing data.")
-    return lst
-
-
-def _process_one_dataset(args: tuple[Any, int, int]) -> dict[str, Any]:
-    """
-    Helper function to process a single dataset's dataloader for statistics.
-    Designed to be called in parallel by a ThreadPoolExecutor.
-
-    Parameters
-    ----------
-    args : tuple(Any, int, int)
-        A tuple containing (dataloader, nbatches, dataloader_len)
-
-    Returns
-    -------
-    dict[str, Any]
-        The processed sys_stat dictionary for one dataset.
-    """
-    dataloader, nbatches, dataloader_len = args
-    sys_stat = {}
-
-    with torch.device("cpu"):
-        iterator = iter(dataloader)
-        numb_batches = min(nbatches, dataloader_len)
-
-        for _ in range(numb_batches):
-            try:
-                stat_data = next(iterator)
-            except StopIteration:
-                iterator = iter(dataloader)
-                stat_data = next(iterator)
-
-            if (
-                "find_fparam" in stat_data
-                and "fparam" in stat_data
-                and stat_data["find_fparam"] == 0.0
-            ):
-                # for model using default fparam
-                stat_data.pop("fparam")
-                stat_data.pop("find_fparam")
-
-            for dd in stat_data:
-                if stat_data[dd] is None:
-                    sys_stat[dd] = None
-                elif isinstance(stat_data[dd], torch.Tensor):
-                    if dd not in sys_stat:
-                        sys_stat[dd] = []
-                    sys_stat[dd].append(stat_data[dd])
-                elif isinstance(stat_data[dd], np.float32):
-                    sys_stat[dd] = stat_data[dd]
-                else:
-                    pass
-
-    for key in sys_stat:
-        if isinstance(sys_stat[key], np.float32):
-            pass
-        elif isinstance(sys_stat[key], list):
-            if len(sys_stat[key]) == 0 or sys_stat[key][0] is None:
+    log.info(f"Packing data for statistics from {len(datasets)} systems")
+    for i in range(len(datasets)):
+        sys_stat = {}
+        with torch.device("cpu"):
+            iterator = iter(dataloaders[i])
+            numb_batches = min(nbatches, len(dataloaders[i]))
+            for _ in range(numb_batches):
+                try:
+                    stat_data = next(iterator)
+                except StopIteration:
+                    iterator = iter(dataloaders[i])
+                    stat_data = next(iterator)
+                if (
+                    "find_fparam" in stat_data
+                    and "fparam" in stat_data
+                    and stat_data["find_fparam"] == 0.0
+                ):
+                    # for model using default fparam
+                    stat_data.pop("fparam")
+                    stat_data.pop("find_fparam")
+                for dd in stat_data:
+                    if stat_data[dd] is None:
+                        sys_stat[dd] = None
+                    elif isinstance(stat_data[dd], torch.Tensor):
+                        if dd not in sys_stat:
+                            sys_stat[dd] = []
+                        sys_stat[dd].append(stat_data[dd])
+                    elif isinstance(stat_data[dd], np.float32):
+                        sys_stat[dd] = stat_data[dd]
+                    else:
+                        pass
+
+        for key in sys_stat:
+            if isinstance(sys_stat[key], np.float32):
+                pass
+            elif sys_stat[key] is None or sys_stat[key][0] is None:
                 sys_stat[key] = None
-            else:
+            elif isinstance(stat_data[dd], torch.Tensor):
                 sys_stat[key] = torch.cat(sys_stat[key], dim=0)
-        elif sys_stat[key] is None:
-            pass
-
-    dict_to_device(sys_stat)
-    return sys_stat
+        dict_to_device(sys_stat)
+        lst.append(sys_stat)
+    return lst
 
 
 def _restore_from_file(
diff --git a/deepmd/utils/env_mat_stat.py b/deepmd/utils/env_mat_stat.py
@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import logging
-import os
 from abc import (
     ABC,
     abstractmethod,
@@ -11,9 +10,6 @@
 from collections.abc import (
     Iterator,
 )
-from concurrent.futures import (
-    ThreadPoolExecutor,
-)
 from typing import (
     Optional,
 )
@@ -146,7 +142,7 @@ def save_stats(self, path: DPPath) -> None:
             (path / kk).save_numpy(np.array([vv.number, vv.sum, vv.squared_sum]))
 
     def load_stats(self, path: DPPath) -> None:
-        """Load the statistics of the environment matrix in parallel.
+        """Load the statistics of the environment matrix.
 
         Parameters
         ----------
@@ -155,18 +151,13 @@ def load_stats(self, path: DPPath) -> None:
         """
         if len(self.stats) > 0:
             raise ValueError("The statistics has already been computed.")
-
-        files_to_load = list(path.glob("*"))
-
-        if not files_to_load:
-            raise ValueError(f"No statistics files found in {path}.")
-
-        with ThreadPoolExecutor(min(64, (os.cpu_count() or 1) * 4)) as executor:
-            results = executor.map(self._load_stat_file, files_to_load)
-
-        for name, stat_item in results:
-            if stat_item is not None:
-                self.stats[name] = stat_item
+        for kk in path.glob("*"):
+            arr = kk.load_numpy()
+            self.stats[kk.name] = StatItem(
+                number=arr[0],
+                sum=arr[1],
+                squared_sum=arr[2],
+            )
 
     def load_or_compute_stats(
         self, data: list[dict[str, np.ndarray]], path: Optional[DPPath] = None
@@ -225,19 +216,3 @@ def get_std(
             kk: vv.compute_std(default=default, protection=protection)
             for kk, vv in self.stats.items()
         }
-
-    @staticmethod
-    def _load_stat_file(file_path: DPPath) -> tuple[str, StatItem]:
-        """Helper function for parallel loading of stat files."""
-        try:
-            arr = file_path.load_numpy()
-            if arr.shape == (3,):
-                return file_path.name, StatItem(
-                    number=arr[0], sum=arr[1], squared_sum=arr[2]
-                )
-            else:
-                log.warning(f"Skipping malformed stat file: {file_path.name}")
-                return file_path.name, None
-        except Exception as e:
-            log.warning(f"Failed to load stat file {file_path.name}: {e}")
-            return file_path.name, None
