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/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/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/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

deepmd/utils/argcheck.py

@@ -5304,6 +5304,13 @@ def training_args(
         "If the file extension is .h5 or .hdf5, an HDF5 file is used to store the statistics; "
         "otherwise, a directory containing NumPy binary files are used."
     )
+    doc_stat_file_mode = (
+        doc_only_pt_supported + "The access mode for `stat_file`. "
+        "`update` creates the cache when needed and writes any missing statistics; "
+        "this is the behavior used when the option is omitted. "
+        "`read` requires a complete existing cache and opens it read-only, allowing "
+        "multiple training processes to share an HDF5 statistics file safely."
+    )
     doc_model_prob = (
         "The visiting probability of each model for each training step in the "
         "multi-task mode. Only used when num_epoch_dict is not set. If not set "
@@ -5339,6 +5346,15 @@ def training_args(
         arg_training_data,
         arg_validation_data,
         Argument("stat_file", str, optional=True, doc=doc_stat_file),
+        Argument(
+            "stat_file_mode",
+            str,
+            optional=True,
+            default="update",
+            extra_check=lambda x: x in {"read", "update"},
+            extra_check_errmsg="must be either 'read' or 'update'",
+            doc=doc_stat_file_mode,
+        ),
     ]
     args = (
         data_args
@@ -5517,9 +5533,23 @@ def training_args(
         ),
     ]
 
+    def _validate_stat_file_mode(data: dict[str, Any], scope: str) -> None:
+        if data.get("stat_file_mode") == "read" and not data.get("stat_file"):
+            raise ValueError(
+                f"{scope}.stat_file_mode='read' requires {scope}.stat_file."
+            )
+
     def training_extra_check(data: dict | None) -> bool:
         if data is None:
             return True
+        if multi_task:
+            for model_key, data_dict in data.get("data_dict", {}).items():
+                _validate_stat_file_mode(
+                    data_dict,
+                    f"training.data_dict[{model_key!r}]",
+                )
+        else:
+            _validate_stat_file_mode(data, "training")
         num_steps = data.get("numb_steps")
         num_epoch = data.get("numb_epoch")
         num_epoch_dict = data.get("num_epoch_dict", {})

deepmd/utils/path.py

@@ -364,6 +364,8 @@ def save_numpy(self, arr: np.ndarray) -> None:
         arr : np.ndarray
             NumPy array
         """
+        if self.mode == "r":
+            raise ValueError("Cannot save to read-only path")
         if self._name in self._keys:
             del self.root[self._name]
         self.root.create_dataset(self._name, data=arr)
@@ -480,6 +482,8 @@ def mkdir(self, parents: bool = False, exist_ok: bool = False) -> None:
         exist_ok : bool, optional
             If true, no error will be raised if the target directory already exists.
         """
+        if self.mode == "r":
+            raise ValueError("Cannot mkdir to read-only path")
         if self._name in self._keys:
             if not exist_ok:
                 raise FileExistsError(f"{self} already exists")

examples/water/dpa4/input.json

@@ -84,6 +84,7 @@
   },
   "training": {
     "stat_file": "./dpa4.hdf5",
+    "stat_file_mode": "update",
     "training_data": {
       "systems": [
         "../data/data_0",

source/tests/common/dpmodel/test_observed_type.py

@@ -82,11 +82,16 @@ def test_save_and_restore(self) -> None:
         restored = _restore_observed_type_from_file(DPPath(self.tmpdir))
         self.assertEqual(restored, observed)
 
-    def test_restore_missing_file(self) -> None:
-        stat_path = DPPath(self.tmpdir, mode="r")
+    def test_restore_missing_file_in_update_mode(self) -> None:
+        stat_path = DPPath(self.tmpdir, mode="a")
         result = _restore_observed_type_from_file(stat_path)
         self.assertIsNone(result)
 
+    def test_restore_missing_file_in_read_mode(self) -> None:
+        stat_path = DPPath(self.tmpdir, mode="r")
+        with self.assertRaisesRegex(FileNotFoundError, "observed_type"):
+            _restore_observed_type_from_file(stat_path)
+
     def test_restore_none_path(self) -> None:
         result = _restore_observed_type_from_file(None)
         self.assertIsNone(result)

source/tests/common/test_path.py

@@ -9,6 +9,7 @@
 import numpy as np
 
 from deepmd.utils.path import (
+    DPH5Path,
     DPPath,
 )
 
@@ -51,3 +52,25 @@ def setUp(self) -> None:
 
     def tearDown(self) -> None:
         self.tempdir.cleanup()
+
+
+class TestH5PathReadOnly(unittest.TestCase):
+    def setUp(self) -> None:
+        self.tempdir = tempfile.TemporaryDirectory()
+        h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve())
+        with h5py.File(h5file, "w"):
+            pass
+        self.path = DPPath(h5file, "r")
+
+    def tearDown(self) -> None:
+        assert isinstance(self.path, DPH5Path)
+        self.path.root.close()
+        DPH5Path._load_h5py.cache_clear()
+        DPH5Path._file_keys.cache_clear()
+        self.tempdir.cleanup()
+
+    def test_write_operations_are_rejected(self) -> None:
+        with self.assertRaisesRegex(ValueError, "read-only"):
+            (self.path / "value").save_numpy(np.ones(1))
+        with self.assertRaisesRegex(ValueError, "read-only"):
+            (self.path / "group").mkdir()

source/tests/pt/test_observed_type.py

@@ -88,11 +88,16 @@ def test_save_and_restore(self) -> None:
         restored = _restore_observed_type_from_file(DPPath(self.tmpdir))
         self.assertEqual(restored, observed)
 
-    def test_restore_missing_file(self) -> None:
-        stat_path = DPPath(self.tmpdir, mode="r")
+    def test_restore_missing_file_in_update_mode(self) -> None:
+        stat_path = DPPath(self.tmpdir, mode="a")
         result = _restore_observed_type_from_file(stat_path)
         self.assertIsNone(result)
 
+    def test_restore_missing_file_in_read_mode(self) -> None:
+        stat_path = DPPath(self.tmpdir, mode="r")
+        with self.assertRaisesRegex(FileNotFoundError, "observed_type"):
+            _restore_observed_type_from_file(stat_path)
+
     def test_restore_none_path(self) -> None:
         result = _restore_observed_type_from_file(None)
         self.assertIsNone(result)