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/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index 543b918b42..f5a97f1297 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/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. 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/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 9da161580b..98a354cba9 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") 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/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/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/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py index e29fe01ac6..6898e5be43 100644 --- a/deepmd/pt/model/descriptor/repflows.py +++ b/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: diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index ca8da0ac1b..440dcbe664 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/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: diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 6cafd43e2f..89b88be151 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/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: diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 808515702a..3f18b069d8 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/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: diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 654e2e16bc..da4d2fa500 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/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: diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index bd2a9a60cd..cc21c836cc 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/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: diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 6937bb99e8..9e3ab52d44 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/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: 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 ab14e59bb1..91dfab4824 100644 --- a/deepmd/utils/argcheck.py +++ b/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", {}) diff --git a/deepmd/utils/env_mat_stat.py b/deepmd/utils/env_mat_stat.py index 626950a3a1..193f5b8e62 100644 --- a/deepmd/utils/env_mat_stat.py +++ b/deepmd/utils/env_mat_stat.py @@ -8,6 +8,7 @@ defaultdict, ) from collections.abc import ( + Callable, Iterator, ) @@ -137,6 +138,10 @@ def iter(self, data: list[dict[str, np.ndarray]]) -> Iterator[dict[str, StatItem The statistics of the environment matrix. """ + def get_stat_keys(self) -> list[str]: + """Get the dataset names required for a complete statistics cache.""" + return [] + def save_stats(self, path: DPPath) -> None: """Save the statistics of the environment matrix. @@ -170,25 +175,52 @@ def load_stats(self, path: DPPath) -> None: ) def load_or_compute_stats( - self, data: list[dict[str, np.ndarray]], path: DPPath | None = None + self, + data: (Callable[[], list[dict[str, np.ndarray]]] | list[dict[str, np.ndarray]]), + path: DPPath | None = None, ) -> None: """Load the statistics of the environment matrix if it exists, otherwise compute and save it. Parameters ---------- + data : Callable or list[dict[str, np.ndarray]] + The environment-matrix data or a lazy callable that returns it. path : DPPath The path to load the statistics of the environment matrix. - data : list[dict[str, np.ndarray]] - The environment matrix. + + Raises + ------ + FileNotFoundError + If a read-only cache is missing its statistics group or any + required dataset. """ - if path is not None and path.is_dir(): + cache_exists = path is not None and path.is_dir() + missing = ( + [key for key in self.get_stat_keys() if not (path / key).is_file()] + if cache_exists + else [] + ) + if cache_exists and not missing: self.load_stats(path) log.info(f"Load stats from {path}.") - else: - self.compute_stats(data) - if path is not None: - self.save_stats(path) - log.info(f"Save stats to {path}.") + return + if path is not None and getattr(path, "mode", None) == "r": + if not cache_exists: + raise FileNotFoundError( + f"Read-only statistics cache {path} is missing the required " + "environment statistics group." + ) + missing_items = ", ".join(repr(item) for item in missing) + raise FileNotFoundError( + f"Read-only statistics cache {path} is missing required " + f"environment statistics item(s): {missing_items}." + ) + + sampled = data() if callable(data) else data + self.compute_stats(sampled) + if path is not None: + self.save_stats(path) + log.info(f"Save stats to {path}.") def get_avg(self, default: float = 0) -> dict[str, float]: """Get the average of the environment matrix. 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 1819a3afad..c5ea37a6cb 100644 --- a/examples/water/dpa4/input.json +++ b/examples/water/dpa4/input.json @@ -84,6 +84,7 @@ }, "training": { "stat_file": "./dpa4.hdf5", + "stat_file_mode": "update", "training_data": { "systems": [ "../data/data_0", 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/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_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) 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..58f0e09c7f --- /dev/null +++ b/source/tests/pt/test_stat_file_mode.py @@ -0,0 +1,303 @@ +# 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 +import torch +from dargs.dargs import ( + ArgumentValueError, +) + +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, +) +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 _energy_model_params() -> dict[str, Any]: + """Build a minimal PT energy-model configuration.""" + return { + "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", + }, + } + + +def _energy_stat_sample() -> list[dict[str, Any]]: + """Build statistics data for the minimal PT energy model.""" + return [ + { + "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), + } + ] + + +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_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]) + + 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] + 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_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_energy_model_reloads_update_cache_in_read_mode(tmp_path: Path) -> None: + sampled = _energy_stat_sample() + 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(_energy_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(_energy_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_read_mode_rejects_missing_descriptor_stats_before_sampling( + tmp_path: Path, +) -> None: + stat_file = tmp_path / "stat.hdf5" + with h5py.File(stat_file, "w") as file: + file.create_group("O H") + + read_path = _prepare_stat_file_path(str(stat_file), "read") + assert isinstance(read_path, DPH5Path) + try: + model = get_model(_energy_model_params()).to(DEVICE) + with pytest.raises(FileNotFoundError, match="environment statistics"): + model.compute_or_load_stat( + lambda: pytest.fail("read-only cache miss must not sample data"), + read_path, + ) + finally: + _close_stat_path(read_path) + + with h5py.File(stat_file, "r") as file: + assert list(file.keys()) == ["O H"] + assert len(file["O H"]) == 0 + + +def test_read_mode_rejects_partial_descriptor_stats_before_sampling( + tmp_path: Path, +) -> None: + stat_file = tmp_path / "stat.hdf5" + update_path = _prepare_stat_file_path(str(stat_file), "update") + assert isinstance(update_path, DPH5Path) + try: + model = get_model(_energy_model_params()).to(DEVICE) + model.compute_or_load_stat(_energy_stat_sample, update_path) + finally: + _close_stat_path(update_path) + + with h5py.File(stat_file, "r+") as file: + type_map_group = file["O H"] + descriptor_groups = [ + value for value in type_map_group.values() if isinstance(value, h5py.Group) + ] + assert len(descriptor_groups) == 1 + del descriptor_groups[0]["r_0"] + + read_path = _prepare_stat_file_path(str(stat_file), "read") + assert isinstance(read_path, DPH5Path) + try: + model = get_model(_energy_model_params()).to(DEVICE) + with pytest.raises(FileNotFoundError, match="'r_0'"): + model.compute_or_load_stat( + lambda: pytest.fail("partial 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") + 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) + + config["training"]["stat_file_mode"] = "read" + config["training"].pop("stat_file") + with pytest.raises(ValueError, match="stat_file_mode"): + normalize(config)