fix(pt): allow concurrent stat cache reads

deepmd/dpmodel/fitting/general_fitting.py

@@ -35,6 +35,9 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
+from deepmd.dpmodel.utils.stat import (
+    _require_stat_file_items,
+)
 from deepmd.env import (
     GLOBAL_NP_FLOAT_PRECISION,
 )
@@ -261,6 +264,7 @@ def compute_input_stats(
             return
         # stat fparam
         if self.numb_fparam > 0:
+            _require_stat_file_items(stat_file_path, ["fparam"])
             if (
                 stat_file_path is not None
                 and stat_file_path.is_dir()
@@ -319,6 +323,7 @@ def compute_input_stats(
             )
         # stat aparam
         if self.numb_aparam > 0:
+            _require_stat_file_items(stat_file_path, ["aparam"])
             if (
                 stat_file_path is not None
                 and stat_file_path.is_dir()

deepmd/dpmodel/utils/env_mat_stat.py

@@ -370,6 +370,15 @@ def iter(
                     env_mats[f"a_{type_i}"] = dd[:, 1:]
                 yield self.compute_stat(env_mats)
 
+    def get_stat_keys(self) -> list[str]:
+        """Get the dataset names required for a complete statistics cache."""
+        components = ("r", "a") if self.last_dim == 4 else ("r",)
+        return [
+            f"{component}_{type_i}"
+            for type_i in range(self.descriptor.get_ntypes())
+            for component in components
+        ]
+
     def get_hash(self) -> str:
         """Get the hash of the environment matrix.
 

deepmd/dpmodel/utils/stat.py

@@ -29,6 +29,35 @@
 log = logging.getLogger(__name__)
 
 
+def _require_stat_file_items(
+    stat_file_path: DPPath | None,
+    items: list[str],
+) -> None:
+    """Require named statistics items when a cache is opened read-only.
+
+    Parameters
+    ----------
+    stat_file_path : DPPath | None
+        Statistics cache path.
+    items : list[str]
+        Relative item names required by the current statistics consumer.
+
+    Raises
+    ------
+    FileNotFoundError
+        If a read-only cache does not contain one or more required items.
+    """
+    if stat_file_path is None or getattr(stat_file_path, "mode", None) != "r":
+        return
+    missing = [item for item in items if not (stat_file_path / item).is_file()]
+    if missing:
+        missing_items = ", ".join(repr(item) for item in missing)
+        raise FileNotFoundError(
+            f"Read-only statistics cache {stat_file_path} is missing "
+            f"required item(s): {missing_items}."
+        )
+
+
 def collect_observed_types(sampled: list[dict], type_map: list[str]) -> list[str]:
     """Collect observed element types from sampled training data.
 
@@ -65,6 +94,7 @@ def _restore_observed_type_from_file(
     """Try to load observed_type from stat file."""
     if stat_file_path is None:
         return None
+    _require_stat_file_items(stat_file_path, ["observed_type"])
     fp = stat_file_path / "observed_type"
     if fp.is_file():
         arr = fp.load_numpy()
@@ -92,6 +122,10 @@ def _restore_from_file(
     """Restore bias and std from stat file."""
     if stat_file_path is None:
         return None, None
+    _require_stat_file_items(
+        stat_file_path,
+        [item for key in keys for item in (f"bias_atom_{key}", f"std_atom_{key}")],
+    )
     stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys]
     if all(not (ii.is_file()) for ii in stat_files):
         return None, None

deepmd/pt/entrypoints/main.py

@@ -100,6 +100,60 @@
 log = logging.getLogger(__name__)
 
 
+def _prepare_stat_file_path(
+    stat_file: str | None,
+    stat_file_mode: str = "update",
+) -> DPPath | None:
+    """Prepare a statistics cache with the requested access mode.
+
+    Parameters
+    ----------
+    stat_file
+        Path to an HDF5 statistics file or a directory-based statistics cache.
+    stat_file_mode
+        ``"update"`` creates the cache when needed and permits missing
+        statistics to be written. ``"read"`` requires an existing cache and
+        prevents all writes.
+
+    Returns
+    -------
+    DPPath or None
+        The prepared statistics path, or ``None`` when no cache is configured.
+
+    Raises
+    ------
+    FileNotFoundError
+        If ``stat_file_mode`` is ``"read"`` and the cache does not exist.
+    ValueError
+        If the access mode is invalid or read mode has no cache path.
+    """
+    if stat_file_mode not in {"read", "update"}:
+        raise ValueError(
+            "`stat_file_mode` must be either 'read' or 'update', "
+            f"but received {stat_file_mode!r}."
+        )
+    if stat_file is None:
+        if stat_file_mode == "read":
+            raise ValueError("`stat_file_mode='read'` requires `stat_file`.")
+        return None
+
+    path = Path(stat_file)
+    if stat_file_mode == "read":
+        if not path.exists():
+            raise FileNotFoundError(
+                f"Statistics cache {stat_file!r} does not exist in read mode."
+            )
+        return DPPath(stat_file, "r")
+
+    if not path.exists():
+        if stat_file.endswith((".h5", ".hdf5")):
+            with h5py.File(stat_file, "w"):
+                pass
+        else:
+            path.mkdir()
+    return DPPath(stat_file, "a")
+
+
 def _update_changed_model_tensors(
     target_state_dict: dict[str, Any],
     source_state_dict: dict[str, Any],
@@ -166,17 +220,13 @@ def prepare_trainer_input_single(
         training_systems = training_dataset_params["systems"]
 
         # stat files
-        stat_file_path_single = data_dict_single.get("stat_file")
         if rank != 0:
             stat_file_path_single = None
-        elif stat_file_path_single is not None:
-            if not Path(stat_file_path_single).exists():
-                if stat_file_path_single.endswith((".h5", ".hdf5")):
-                    with h5py.File(stat_file_path_single, "w") as f:
-                        pass
-                else:
-                    Path(stat_file_path_single).mkdir()
-            stat_file_path_single = DPPath(stat_file_path_single, "a")
+        else:
+            stat_file_path_single = _prepare_stat_file_path(
+                data_dict_single.get("stat_file"),
+                data_dict_single.get("stat_file_mode", "update"),
+            )
 
         rank_seed = [rank, seed % (2**32)] if seed is not None else None
 

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -620,7 +620,7 @@ def change_out_bias(
             delta_bias, out_std = compute_output_stats(
                 sample_merged,
                 self.get_ntypes(),
-                keys=list(self.atomic_output_def().keys()),
+                keys=self.bias_keys,
                 stat_file_path=stat_file_path,
                 model_forward=self._get_forward_wrapper_func(),
                 rcond=self.rcond,
@@ -633,7 +633,7 @@ def change_out_bias(
             bias_out, std_out = compute_output_stats(
                 sample_merged,
                 self.get_ntypes(),
-                keys=list(self.atomic_output_def().keys()),
+                keys=self.bias_keys,
                 stat_file_path=stat_file_path,
                 rcond=self.rcond,
                 preset_bias=self.preset_out_bias,

deepmd/pt/model/atomic_model/sezm_atomic_model.py

@@ -15,6 +15,9 @@
 import numpy as np
 import torch
 
+from deepmd.dpmodel.utils.stat import (
+    _require_stat_file_items,
+)
 from deepmd.pt.model.atomic_model.dp_atomic_model import (
     DPAtomicModel,
 )
@@ -188,6 +191,7 @@ def _compute_or_load_dens_force_stat(
         force_stat_path = (
             None if stat_file_path is None else stat_file_path / "rmsd_dforce"
         )
+        _require_stat_file_items(stat_file_path, ["rmsd_dforce"])
         if force_stat_path is not None and force_stat_path.is_file():
             force_rmsd = float(np.asarray(force_stat_path.load_numpy()).reshape(-1)[0])
         else:

deepmd/pt/model/descriptor/repflows.py

@@ -707,15 +707,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/repformers.py

@@ -571,15 +571,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_a.py

@@ -677,15 +677,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_atten.py

@@ -424,15 +424,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_r.py

@@ -318,15 +318,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_t.py

@@ -721,15 +721,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_t_tebd.py

@@ -844,15 +844,7 @@ def compute_input_stats(
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
-        if path is None or not path.is_dir():
-            if callable(merged):
-                # only get data for once
-                sampled = merged()
-            else:
-                sampled = merged
-        else:
-            sampled = []
-        env_mat_stat.load_or_compute_stats(sampled, path)
+        env_mat_stat.load_or_compute_stats(merged, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/task/fitting.py

@@ -20,6 +20,9 @@
 from deepmd.dpmodel.utils.seed import (
     child_seed,
 )
+from deepmd.dpmodel.utils.stat import (
+    _require_stat_file_items,
+)
 from deepmd.pt.model.network.mlp import (
     FittingNet,
     NetworkCollection,
@@ -269,6 +272,7 @@ def compute_input_stats(
 
         # stat fparam
         if self.numb_fparam > 0:
+            _require_stat_file_items(stat_file_path, ["fparam"])
             if (
                 stat_file_path is not None
                 and stat_file_path.is_dir()
@@ -307,6 +311,7 @@ def compute_input_stats(
 
         # stat aparam
         if self.numb_aparam > 0:
+            _require_stat_file_items(stat_file_path, ["aparam"])
             if (
                 stat_file_path is not None
                 and stat_file_path.is_dir()

deepmd/pt/utils/stat.py

@@ -37,6 +37,7 @@
 
 # Re-export from dpmodel (backend-agnostic implementations)
 from deepmd.dpmodel.utils.stat import (
+    _require_stat_file_items,
     _restore_observed_type_from_file,
     _save_observed_type_to_file,
     collect_observed_types,
@@ -113,6 +114,10 @@ def _restore_from_file(
 ) -> dict | None:
     if stat_file_path is None:
         return None, None
+    _require_stat_file_items(
+        stat_file_path,
+        [item for key in keys for item in (f"bias_atom_{key}", f"std_atom_{key}")],
+    )
     stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys]
     if all(not (ii.is_file()) for ii in stat_files):
         return None, None