From c95b4d676eb0b33b7661e150adc1405aa3b5f8bc Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 17:27:34 +0800 Subject: [PATCH 1/4] fix(pt): allow concurrent stat cache reads Add explicit read-only access while preserving update mode by default. --- deepmd/pt/entrypoints/main.py | 68 +++++++++++++++--- deepmd/utils/argcheck.py | 16 +++++ deepmd/utils/path.py | 4 ++ examples/water/dpa4/input.json | 1 + source/tests/pt/test_stat_file_mode.py | 97 ++++++++++++++++++++++++++ 5 files changed, 177 insertions(+), 9 deletions(-) create mode 100644 source/tests/pt/test_stat_file_mode.py diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 560ea5a1ba..2aa4888765 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/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", None) 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 diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index fe5f542933..bf28fc7e3d 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5296,6 +5296,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 " @@ -5331,6 +5338,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 diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index b9e48fe531..44bbe34edf 100644 --- a/deepmd/utils/path.py +++ b/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") diff --git a/examples/water/dpa4/input.json b/examples/water/dpa4/input.json index 1126c7107a..d8891bac64 100644 --- a/examples/water/dpa4/input.json +++ b/examples/water/dpa4/input.json @@ -83,6 +83,7 @@ }, "training": { "stat_file": "./dpa4.hdf5", + "stat_file_mode": "update", "training_data": { "systems": [ "../data/data_0", diff --git a/source/tests/pt/test_stat_file_mode.py b/source/tests/pt/test_stat_file_mode.py new file mode 100644 index 0000000000..8bff16ded1 --- /dev/null +++ b/source/tests/pt/test_stat_file_mode.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import h5py +import numpy as np +import pytest +from dargs.dargs import ( + ArgumentValueError, +) + +from deepmd.pt.entrypoints.main import ( + _prepare_stat_file_path, +) +from deepmd.utils.argcheck import ( + normalize, +) +from deepmd.utils.path import ( + DPH5Path, + DPPath, +) + + +def _close_stat_path(stat_path: DPPath) -> None: + """Close a test HDF5 path and reset its process-local caches.""" + assert isinstance(stat_path, DPH5Path) + stat_path.root.close() + DPH5Path._load_h5py.cache_clear() + DPH5Path._file_keys.cache_clear() + + +def _load_dpa4_example() -> dict[str, Any]: + """Load the DPA4 example used by configuration validation tests.""" + example_path = ( + Path(__file__).resolve().parents[3] + / "examples" + / "water" + / "dpa4" + / "input.json" + ) + with example_path.open(encoding="utf-8") as stream: + return json.load(stream) + + +def test_default_stat_file_mode_remains_writable(tmp_path: Path) -> None: + stat_file = tmp_path / "stat.hdf5" + stat_path = _prepare_stat_file_path(str(stat_file)) + assert isinstance(stat_path, DPH5Path) + try: + assert stat_path.mode == "a" + (stat_path / "value").save_numpy(np.array([1.0])) + assert (stat_path / "value").load_numpy().tolist() == [1.0] + finally: + _close_stat_path(stat_path) + + +def test_read_stat_file_mode_is_read_only(tmp_path: Path) -> None: + stat_file = tmp_path / "stat.hdf5" + with h5py.File(stat_file, "w") as file: + file.create_dataset("value", data=[1.0]) + + stat_path = _prepare_stat_file_path(str(stat_file), "read") + assert isinstance(stat_path, DPH5Path) + try: + assert stat_path.mode == "r" + assert stat_path.root.mode == "r" + assert (stat_path / "value").load_numpy().tolist() == [1.0] + with pytest.raises(ValueError, match="read-only"): + (stat_path / "other").save_numpy(np.array([2.0])) + finally: + _close_stat_path(stat_path) + + +def test_read_stat_file_mode_requires_existing_cache(tmp_path: Path) -> None: + stat_file = tmp_path / "missing.hdf5" + with pytest.raises(FileNotFoundError, match="does not exist in read mode"): + _prepare_stat_file_path(str(stat_file), "read") + + +def test_stat_file_mode_configuration_validation() -> None: + config = _load_dpa4_example() + config["training"].pop("stat_file_mode") + normalized = normalize(config) + assert normalized["training"]["stat_file_mode"] == "update" + + config["training"]["stat_file_mode"] = "read" + normalized = normalize(config) + assert normalized["training"]["stat_file_mode"] == "read" + + config["training"]["stat_file_mode"] = "invalid" + with pytest.raises(ArgumentValueError, match="stat_file_mode"): + normalize(config) From 68642596424857ae376965fed3dfaeae9b892cff Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 12 Jul 2026 19:44:23 +0800 Subject: [PATCH 2/4] fix(pt): validate read-only statistics caches --- deepmd/dpmodel/fitting/general_fitting.py | 5 ++ deepmd/dpmodel/utils/stat.py | 34 ++++++++++ .../model/atomic_model/sezm_atomic_model.py | 4 ++ deepmd/pt/model/task/fitting.py | 5 ++ deepmd/pt/utils/stat.py | 5 ++ deepmd/utils/argcheck.py | 14 ++++ source/tests/common/test_path.py | 23 +++++++ source/tests/pt/test_stat_file_mode.py | 67 ++++++++++++++++++- 8 files changed, 154 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 4734a6be66..2b2ef2693b 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/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() diff --git a/deepmd/dpmodel/utils/stat.py b/deepmd/dpmodel/utils/stat.py index ca00f6c064..bcb74695b2 100644 --- a/deepmd/dpmodel/utils/stat.py +++ b/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 diff --git a/deepmd/pt/model/atomic_model/sezm_atomic_model.py b/deepmd/pt/model/atomic_model/sezm_atomic_model.py index 201f1a7715..f3e0914e47 100644 --- a/deepmd/pt/model/atomic_model/sezm_atomic_model.py +++ b/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: diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 0bc322f331..71b947fedb 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/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() diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index f8c3685b78..b268ca2ba7 100644 --- a/deepmd/pt/utils/stat.py +++ b/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 diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index bf28fc7e3d..1ee9f296ed 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5525,9 +5525,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", {}) diff --git a/source/tests/common/test_path.py b/source/tests/common/test_path.py index 4eb74edc53..ec3f175c0a 100644 --- a/source/tests/common/test_path.py +++ b/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() diff --git a/source/tests/pt/test_stat_file_mode.py b/source/tests/pt/test_stat_file_mode.py index 8bff16ded1..4f01e90e81 100644 --- a/source/tests/pt/test_stat_file_mode.py +++ b/source/tests/pt/test_stat_file_mode.py @@ -17,6 +17,9 @@ from deepmd.pt.entrypoints.main import ( _prepare_stat_file_path, ) +from deepmd.pt.utils.stat import ( + compute_output_stats, +) from deepmd.utils.argcheck import ( normalize, ) @@ -59,7 +62,7 @@ def test_default_stat_file_mode_remains_writable(tmp_path: Path) -> None: _close_stat_path(stat_path) -def test_read_stat_file_mode_is_read_only(tmp_path: Path) -> None: +def test_read_stat_file_mode_reads_existing_cache(tmp_path: Path) -> None: stat_file = tmp_path / "stat.hdf5" with h5py.File(stat_file, "w") as file: file.create_dataset("value", data=[1.0]) @@ -70,8 +73,6 @@ def test_read_stat_file_mode_is_read_only(tmp_path: Path) -> None: assert stat_path.mode == "r" assert stat_path.root.mode == "r" assert (stat_path / "value").load_numpy().tolist() == [1.0] - with pytest.raises(ValueError, match="read-only"): - (stat_path / "other").save_numpy(np.array([2.0])) finally: _close_stat_path(stat_path) @@ -82,6 +83,61 @@ def test_read_stat_file_mode_requires_existing_cache(tmp_path: Path) -> None: _prepare_stat_file_path(str(stat_file), "read") +def test_read_stat_file_mode_requires_cache_path() -> None: + with pytest.raises(ValueError, match="requires `stat_file`"): + _prepare_stat_file_path(None, "read") + + +def test_read_stat_file_mode_rejects_incomplete_statistics_cache( + tmp_path: Path, +) -> None: + stat_file = tmp_path / "stat.hdf5" + with h5py.File(stat_file, "w") as file: + file.create_dataset("bias_atom_energy", data=np.zeros((1, 1))) + + stat_path = _prepare_stat_file_path(str(stat_file), "read") + try: + with pytest.raises(FileNotFoundError, match="std_atom_energy"): + compute_output_stats( + lambda: pytest.fail("read-only statistics must not sample data"), + ntypes=1, + keys=["energy"], + stat_file_path=stat_path, + ) + finally: + _close_stat_path(stat_path) + + +def test_read_stat_file_mode_loads_complete_cache_from_two_readers( + tmp_path: Path, +) -> None: + stat_file = tmp_path / "stat.hdf5" + with h5py.File(stat_file, "w") as file: + file.create_dataset("bias_atom_energy", data=np.zeros((1, 1))) + file.create_dataset("std_atom_energy", data=np.ones((1, 1))) + + reader_one = _prepare_stat_file_path(str(stat_file), "read") + DPH5Path._load_h5py.cache_clear() + reader_two = _prepare_stat_file_path(str(stat_file), "read") + assert isinstance(reader_one, DPH5Path) + assert isinstance(reader_two, DPH5Path) + try: + for reader in (reader_one, reader_two): + bias, std = compute_output_stats( + lambda: pytest.fail("complete read-only cache must not sample data"), + ntypes=1, + keys=["energy"], + stat_file_path=reader, + ) + assert bias["energy"].shape == (1, 1) + assert std["energy"].shape == (1, 1) + finally: + reader_one.root.close() + reader_two.root.close() + DPH5Path._load_h5py.cache_clear() + DPH5Path._file_keys.cache_clear() + + def test_stat_file_mode_configuration_validation() -> None: config = _load_dpa4_example() config["training"].pop("stat_file_mode") @@ -95,3 +151,8 @@ def test_stat_file_mode_configuration_validation() -> None: config["training"]["stat_file_mode"] = "invalid" with pytest.raises(ArgumentValueError, match="stat_file_mode"): normalize(config) + + config["training"]["stat_file_mode"] = "read" + config["training"].pop("stat_file") + with pytest.raises(ValueError, match="stat_file_mode"): + normalize(config) From 0246850bd72cb968e69e41dfc43c2cb54f646c1c Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 15 Jul 2026 22:27:17 +0800 Subject: [PATCH 3/4] test(stat): distinguish missing cache behavior by mode Keep update-mode fallback coverage while asserting that read-only caches reject missing observed-type statistics. --- source/tests/common/dpmodel/test_observed_type.py | 9 +++++++-- source/tests/pt/test_observed_type.py | 9 +++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/source/tests/common/dpmodel/test_observed_type.py b/source/tests/common/dpmodel/test_observed_type.py index c08b672415..2449700dbb 100644 --- a/source/tests/common/dpmodel/test_observed_type.py +++ b/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) diff --git a/source/tests/pt/test_observed_type.py b/source/tests/pt/test_observed_type.py index c834ec567e..a2c7e37a73 100644 --- a/source/tests/pt/test_observed_type.py +++ b/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) From d62b215eef43260898ad6e97ba00a8a10713ff7a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 16 Jul 2026 09:41:56 +0800 Subject: [PATCH 4/4] fix(stat): exclude derived outputs from cache requirements --- .../model/atomic_model/base_atomic_model.py | 4 +- source/tests/pt/test_stat_file_mode.py | 82 +++++++++++++++++++ 2 files changed, 84 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index e700a49b32..13c8e8ae8a 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/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, diff --git a/source/tests/pt/test_stat_file_mode.py b/source/tests/pt/test_stat_file_mode.py index 4f01e90e81..8c9bb2aa35 100644 --- a/source/tests/pt/test_stat_file_mode.py +++ b/source/tests/pt/test_stat_file_mode.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +from copy import ( + deepcopy, +) from pathlib import ( Path, ) @@ -10,6 +13,7 @@ import h5py import numpy as np import pytest +import torch from dargs.dargs import ( ArgumentValueError, ) @@ -17,6 +21,12 @@ from deepmd.pt.entrypoints.main import ( _prepare_stat_file_path, ) +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.utils.env import ( + DEVICE, +) from deepmd.pt.utils.stat import ( compute_output_stats, ) @@ -138,6 +148,78 @@ def test_read_stat_file_mode_loads_complete_cache_from_two_readers( DPH5Path._file_keys.cache_clear() +def test_energy_model_reloads_update_cache_in_read_mode(tmp_path: Path) -> None: + model_params = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [4, 4], + "rcut": 3.0, + "rcut_smth": 2.5, + "neuron": [4, 8], + "axis_neuron": 4, + "precision": "float64", + }, + "fitting_net": { + "type": "ener", + "neuron": [8], + "precision": "float64", + }, + } + sampled = [ + { + "coord": torch.tensor( + [ + [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], + ], + dtype=torch.float64, + device=DEVICE, + ), + "atype": torch.tensor( + [[0, 0], [1, 1]], + dtype=torch.int64, + device=DEVICE, + ), + "box": None, + "natoms": torch.tensor( + [[2, 2, 2, 0], [2, 2, 0, 2]], + dtype=torch.int64, + device=DEVICE, + ), + "energy": torch.tensor( + [[2.0], [4.0]], + dtype=torch.float64, + device=DEVICE, + ), + "find_energy": np.float32(1.0), + } + ] + stat_file = tmp_path / "stat.hdf5" + + update_path = _prepare_stat_file_path(str(stat_file), "update") + assert isinstance(update_path, DPH5Path) + try: + update_model = get_model(deepcopy(model_params)).to(DEVICE) + update_model.compute_or_load_stat(lambda: sampled, update_path) + stat_root = update_path / "O H" + assert (stat_root / "bias_atom_energy").is_file() + assert not (stat_root / "bias_atom_mask").is_file() + finally: + _close_stat_path(update_path) + + read_path = _prepare_stat_file_path(str(stat_file), "read") + assert isinstance(read_path, DPH5Path) + try: + read_model = get_model(deepcopy(model_params)).to(DEVICE) + read_model.compute_or_load_stat( + lambda: pytest.fail("complete read-only cache must not sample data"), + read_path, + ) + finally: + _close_stat_path(read_path) + + def test_stat_file_mode_configuration_validation() -> None: config = _load_dpa4_example() config["training"].pop("stat_file_mode")