From 37b3b760c0b9ec9b71ca1e462eabb895a0c68a55 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 21 Jun 2026 00:01:08 +0800 Subject: [PATCH 01/12] feat(data): add dpdata format conversion --- deepmd/pd/entrypoints/main.py | 33 +- deepmd/pt/entrypoints/main.py | 86 +-- deepmd/pt_expt/entrypoints/main.py | 38 +- deepmd/utils/argcheck.py | 54 ++ deepmd/utils/data_system.py | 529 +++++++++++++++++- pyproject.toml | 2 +- .../common/test_data_system_conversion.py | 205 +++++++ 7 files changed, 895 insertions(+), 52 deletions(-) create mode 100644 source/tests/common/test_data_system_conversion.py diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index f397bc358b..72c9ef7677 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -23,6 +23,9 @@ expand_sys_str, j_loader, ) +from deepmd.dpmodel.utils.lmdb_data import ( + is_lmdb, +) from deepmd.loggers.loggers import ( set_log_handles, ) @@ -109,10 +112,36 @@ def prepare_trainer_input_single( ) training_systems = training_dataset_params["systems"] trn_patterns = training_dataset_params.get("rglob_patterns", None) - training_systems = process_systems(training_systems, patterns=trn_patterns) + training_systems = process_systems( + training_systems, + patterns=trn_patterns, + fmt=training_dataset_params.get("format", None), + out_fmt=training_dataset_params.get( + "out_format", training_dataset_params.get("output_format", None) + ), + ) + if len(training_systems) == 1 and is_lmdb(training_systems[0]): + raise NotImplementedError( + "Paddle backend does not support LMDB training data yet. " + "Set training_data.out_format to 'deepmd/hdf5' when using " + "training_data.format for automatic conversion." + ) if validation_systems is not None: val_patterns = validation_dataset_params.get("rglob_patterns", None) - validation_systems = process_systems(validation_systems, val_patterns) + validation_systems = process_systems( + validation_systems, + val_patterns, + fmt=validation_dataset_params.get("format", None), + out_fmt=validation_dataset_params.get( + "out_format", validation_dataset_params.get("output_format", None) + ), + ) + if len(validation_systems) == 1 and is_lmdb(validation_systems[0]): + raise NotImplementedError( + "Paddle backend does not support LMDB validation data yet. " + "Set validation_data.out_format to 'deepmd/hdf5' when using " + "validation_data.format for automatic conversion." + ) # stat files stat_file_path_single = data_dict_single.get("stat_file", None) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 560ea5a1ba..8bd147697d 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -183,10 +183,24 @@ def prepare_trainer_input_single( def _make_dp_loader_set( systems: str | list[str], dataset_params: dict[str, Any], - ) -> DpLoaderSet: - """Create a DpLoaderSet from systems with pattern expansion.""" + ) -> DpLoaderSet | LmdbDataset: + """Create a dataset from systems with pattern expansion/conversion.""" patterns = dataset_params.get("rglob_patterns", None) - systems = process_systems(systems, patterns=patterns) + systems = process_systems( + systems, + patterns=patterns, + fmt=dataset_params.get("format", None), + out_fmt=dataset_params.get( + "out_format", dataset_params.get("output_format", None) + ), + ) + if len(systems) == 1 and is_lmdb(systems[0]): + return LmdbDataset( + systems[0], + model_params_single["type_map"], + dataset_params["batch_size"], + auto_prob_style=dataset_params.get("auto_prob", None), + ) return DpLoaderSet( systems, dataset_params["batch_size"], @@ -196,7 +210,11 @@ def _make_dp_loader_set( ) # LMDB path: single string → LmdbDataset - if isinstance(training_systems, str) and is_lmdb(training_systems): + if ( + training_dataset_params.get("format", None) is None + and isinstance(training_systems, str) + and is_lmdb(training_systems) + ): auto_prob = training_dataset_params.get("auto_prob", None) train_data_single = LmdbDataset( training_systems, @@ -206,6 +224,7 @@ def _make_dp_loader_set( ) if ( validation_systems is not None + and validation_dataset_params.get("format", None) is None and isinstance(validation_systems, str) and is_lmdb(validation_systems) ): @@ -397,23 +416,40 @@ def train( "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)" ) - if not multi_task: - type_map = config["model"].get("type_map") - training_systems = config["training"]["training_data"].get("systems") + def _get_neighbor_stat_data_from_params( + dataset_params: dict[str, Any], + type_map: list[str] | None, + ) -> Any: + training_systems = dataset_params.get("systems") if ( - training_systems is not None + dataset_params.get("format", None) is None + and training_systems is not None and isinstance(training_systems, str) and is_lmdb(training_systems) ): + systems = [training_systems] + else: + systems = process_systems( + training_systems, + patterns=dataset_params.get("rglob_patterns", None), + fmt=dataset_params.get("format", None), + out_fmt=dataset_params.get( + "out_format", dataset_params.get("output_format", None) + ), + ) + if len(systems) == 1 and is_lmdb(systems[0]): from deepmd.dpmodel.utils.lmdb_data import ( make_neighbor_stat_data, ) - train_data = make_neighbor_stat_data(training_systems, type_map) - else: - train_data = get_data( - config["training"]["training_data"], 0, type_map, None - ) + return make_neighbor_stat_data(systems[0], type_map) + return get_data(dataset_params, 0, type_map, None) + + if not multi_task: + type_map = config["model"].get("type_map") + train_data = _get_neighbor_stat_data_from_params( + config["training"]["training_data"], type_map + ) config["model"], min_nbor_dist = BaseModel.update_sel( train_data, type_map, config["model"] ) @@ -421,26 +457,10 @@ def train( min_nbor_dist = {} for model_item in config["model"]["model_dict"]: type_map = config["model"]["model_dict"][model_item].get("type_map") - training_systems = config["training"]["data_dict"][model_item][ - "training_data" - ].get("systems") - if ( - training_systems is not None - and isinstance(training_systems, str) - and is_lmdb(training_systems) - ): - from deepmd.dpmodel.utils.lmdb_data import ( - make_neighbor_stat_data, - ) - - train_data = make_neighbor_stat_data(training_systems, type_map) - else: - train_data = get_data( - config["training"]["data_dict"][model_item]["training_data"], - 0, - type_map, - None, - ) + train_data = _get_neighbor_stat_data_from_params( + config["training"]["data_dict"][model_item]["training_data"], + type_map, + ) config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( BaseModel.update_sel( train_data, type_map, config["model"]["model_dict"][model_item] diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index da28229bf4..38f16095f4 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -104,13 +104,31 @@ def _get_neighbor_stat_data( ``make_neighbor_stat_data``; falls back to the legacy ``get_data`` for npy/HDF5 directories. """ - lmdb_path = _detect_lmdb_path(dataset_params.get("systems")) + lmdb_path = ( + None + if dataset_params.get("format", None) is not None + else _detect_lmdb_path(dataset_params.get("systems")) + ) if lmdb_path is not None: from deepmd.dpmodel.utils.lmdb_data import ( make_neighbor_stat_data, ) return make_neighbor_stat_data(lmdb_path, type_map) + systems = process_systems( + dataset_params["systems"], + patterns=dataset_params.get("rglob_patterns", None), + fmt=dataset_params.get("format", None), + out_fmt=dataset_params.get( + "out_format", dataset_params.get("output_format", None) + ), + ) + if len(systems) == 1 and is_lmdb(systems[0]): + from deepmd.dpmodel.utils.lmdb_data import ( + make_neighbor_stat_data, + ) + + return make_neighbor_stat_data(systems[0], type_map) return get_data(dataset_params, 0, type_map, None) @@ -126,7 +144,11 @@ def _build_data_system( :class:`DeepmdDataSystem` path with system expansion. """ systems_raw = dataset_params["systems"] - lmdb_path = _detect_lmdb_path(systems_raw) + lmdb_path = ( + None + if dataset_params.get("format", None) is not None + else _detect_lmdb_path(systems_raw) + ) if lmdb_path is not None: return LmdbDataSystem( lmdb_path=lmdb_path, @@ -138,7 +160,19 @@ def _build_data_system( systems = process_systems( systems_raw, patterns=dataset_params.get("rglob_patterns", None), + fmt=dataset_params.get("format", None), + out_fmt=dataset_params.get( + "out_format", dataset_params.get("output_format", None) + ), ) + if len(systems) == 1 and is_lmdb(systems[0]): + return LmdbDataSystem( + lmdb_path=systems[0], + type_map=type_map, + batch_size=dataset_params["batch_size"], + auto_prob_style=dataset_params.get("auto_prob"), + seed=seed, + ) return DeepmdDataSystem( systems=systems, batch_size=dataset_params["batch_size"], diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 48c711a10b..6bfab83e62 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -4745,6 +4745,19 @@ def training_data_args() -> list[ doc_patterns = ( "The customized patterns used in `rglob` to collect all training systems. " ) + doc_format = ( + "The input data format passed to dpdata for automatic conversion. " + "If this key is not set, `systems` must already point to DeePMD data. " + "If this key is set to a non-DeePMD format, each selected input path is " + "loaded by dpdata and converted before training. Use dpdata format names " + "such as `extxyz`, `ase/structure`, `ase/traj`, or `auto`." + ) + doc_out_format = ( + "The output data format passed to dpdata for automatic conversion. " + "When `format` requests conversion from a non-DeePMD format, this key " + "defaults to `lmdb`. Use a DeePMD format supported by dpdata, such as " + "`lmdb`, `deepmd/hdf5`, or `deepmd/npy`." + ) doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ - int: all {link_sys} use the same batch size.\n\n\ @@ -4783,6 +4796,20 @@ def training_data_args() -> list[ default=None, doc=doc_patterns + doc_only_pt_supported, ), + Argument( + "format", + [str, None], + optional=True, + doc=doc_format, + ), + Argument( + "out_format", + [str, None], + optional=True, + default="lmdb", + doc=doc_out_format, + alias=["output_format"], + ), Argument( "batch_size", [list[int], int, str], @@ -4842,6 +4869,19 @@ def validation_data_args() -> list[ doc_patterns = ( "The customized patterns used in `rglob` to collect all validation systems. " ) + doc_format = ( + "The input data format passed to dpdata for automatic conversion. " + "If this key is not set, `systems` must already point to DeePMD data. " + "If this key is set to a non-DeePMD format, each selected input path is " + "loaded by dpdata and converted before validation. Use dpdata format names " + "such as `extxyz`, `ase/structure`, `ase/traj`, or `auto`." + ) + doc_out_format = ( + "The output data format passed to dpdata for automatic conversion. " + "When `format` requests conversion from a non-DeePMD format, this key " + "defaults to `lmdb`. Use a DeePMD format supported by dpdata, such as " + "`lmdb`, `deepmd/hdf5`, or `deepmd/npy`." + ) doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ - int: all {link_sys} use the same batch size.\n\n\ @@ -4867,6 +4907,20 @@ def validation_data_args() -> list[ default=None, doc=doc_patterns + doc_only_pt_supported, ), + Argument( + "format", + [str, None], + optional=True, + doc=doc_format, + ), + Argument( + "out_format", + [str, None], + optional=True, + default="lmdb", + doc=doc_out_format, + alias=["output_format"], + ), Argument( "batch_size", [list[int], int, str], diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 9d13cb4699..e6a795c616 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -1,10 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import collections +import hashlib import logging +import os +import shutil +import time import warnings from functools import ( cached_property, ) +from pathlib import ( + Path, +) from typing import ( Any, ) @@ -30,6 +37,10 @@ log = logging.getLogger(__name__) +_DPDATA_CACHE_DIR = ".deepmd_dpdata_cache" +_DPDATA_DEFAULT_OUT_FORMAT = "lmdb" +_DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str], list[str]] = {} + class DeepmdDataSystem: """Class for manipulating many data systems. @@ -672,6 +683,263 @@ def _check_type_map_consistency( return ret +class LmdbDataSystem: + """A DeepmdDataSystem-compatible adapter for LMDB datasets. + + The adapter returns raw DeePMD-style numpy batches (``type``, + ``natoms_vec``, ``default_mesh``) so it can be consumed by the legacy + TensorFlow/JAX training paths. Consumers that need the dpmodel canonical + format can still call ``normalize_batch`` on its output. + """ + + def __init__( + self, + lmdb_path: str, + type_map: list[str], + batch_size: int | str = "auto", + auto_prob_style: str | None = None, + seed: int | None = None, + ) -> None: + if not type_map: + raise ValueError( + "LMDB datasets require a non-empty model/type_map because " + "LMDB stores atom type indices and the training data adapter " + "must map them to element names." + ) + + from deepmd.dpmodel.utils.lmdb_data import ( + LmdbDataReader, + SameNlocBatchSampler, + compute_block_targets, + ) + + self.lmdb_path = lmdb_path + self._reader = LmdbDataReader( + lmdb_path, type_map, batch_size, mixed_batch=False + ) + self._type_map = list(type_map) + self.mixed_type = self._detect_mixed_type() + self.nsystems = 1 + self.system_dirs = [lmdb_path] + self.natoms = [max(self._reader.frame_nlocs) if self._reader.frame_nlocs else 0] + self.batch_size = [self._reader.batch_size] + self.nbatches = [self._reader.total_batch] + self.sys_probs = [1.0] + self.data_systems = [self] + self._nloc_set_indices = { + f"{self.lmdb_path}#nloc={nloc}": indices + for nloc, indices in sorted(self._reader.nloc_groups.items()) + } + self.dirs = list(self._nloc_set_indices) + self.pbc = self._detect_pbc() + self._data_dict = { + "box": { + "ndof": 9, + "atomic": False, + "must": False, + "high_prec": False, + "type_sel": None, + "repeat": 1, + "default": 0.0, + "dtype": None, + "output_natoms_for_type_sel": False, + }, + "coord": { + "ndof": 3, + "atomic": True, + "must": True, + "high_prec": False, + "type_sel": None, + "repeat": 1, + "default": 0.0, + "dtype": None, + "output_natoms_for_type_sel": False, + }, + "numb_copy": { + "ndof": 1, + "atomic": False, + "must": False, + "high_prec": False, + "type_sel": None, + "repeat": 1, + "default": 1, + "dtype": int, + "output_natoms_for_type_sel": False, + }, + } + + block_targets = None + if auto_prob_style is not None and self._reader.frame_system_ids is not None: + block_targets = compute_block_targets( + auto_prob_style, + self._reader.nsystems, + self._reader.system_nframes, + ) + self._sampler = SameNlocBatchSampler( + self._reader, + shuffle=True, + seed=seed, + block_targets=block_targets, + ) + self._iter = iter(self._sampler) + + def _detect_mixed_type(self) -> bool: + """Return True when frames cannot be represented as fixed-type data.""" + if len(self._reader.nloc_groups) > 1: + return True + if len(self._reader) == 0: + return False + ref_type = self._reader[0]["atype"] + for idx in range(1, len(self._reader)): + if not np.array_equal(self._reader[idx]["atype"], ref_type): + return True + return False + + def _detect_pbc(self) -> bool: + """Return True when LMDB frames contain a non-zero simulation box.""" + if len(self._reader) == 0: + return False + box = self._reader[0].get("box") + return box is not None and not np.allclose(box, 0.0) + + def add_data_requirements( + self, data_requirements: list[DataRequirementItem] + ) -> None: + """Add label/auxiliary data requirements.""" + for item in data_requirements: + self._data_dict[item.key] = item.dict + self._reader.add_data_requirement(data_requirements) + + def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: + """Alias used by DataLoader-style backends.""" + self.add_data_requirements(data_requirement) + + def add( + self, + key: str, + ndof: int, + atomic: bool = False, + must: bool = False, + high_prec: bool = False, + type_sel: list[int] | None = None, + repeat: int = 1, + default: float = 0.0, + dtype: np.dtype | None = None, + output_natoms_for_type_sel: bool = False, + ) -> None: + item = DataRequirementItem( + key, + ndof, + atomic=atomic, + must=must, + high_prec=high_prec, + type_sel=type_sel, + repeat=repeat, + default=default, + dtype=dtype, + output_natoms_for_type_sel=output_natoms_for_type_sel, + ) + self.add_data_requirements([item]) + + def get_data_dict(self, ii: int = 0) -> dict[str, dict[str, Any]]: + del ii + return self._data_dict + + def _load_set(self, set_name: str) -> dict[str, Any]: + """Load one same-nloc LMDB group for legacy neighbor-stat code.""" + indices = self._nloc_set_indices[str(set_name)] + frames = [self._reader[int(idx)] for idx in indices] + return self._stack_frames(frames) + + def _next_indices(self) -> list[int]: + try: + return next(self._iter) + except StopIteration: + self._iter = iter(self._sampler) + return next(self._iter) + + def _stack_frames(self, frames: list[dict[str, Any]]) -> dict[str, Any]: + out: dict[str, Any] = {} + structural_keys = {"coord", "box"} + for key in frames[0]: + if key in {"atype", "fid", "natoms", "real_natoms_vec"}: + continue + if key.startswith("find_") and key[5:] not in self._data_dict: + continue + if ( + not key.startswith("find_") + and key not in structural_keys + and key not in self._data_dict + ): + continue + if key.startswith("find_"): + out[key] = np.asarray(frames[0][key], dtype=np.float32) + elif frames[0][key] is None: + out[key] = None + else: + arr = np.stack([frame[key] for frame in frames]) + data_info = self._data_dict.get(key) + if data_info is not None and data_info["atomic"] and arr.ndim >= 3: + arr = arr.reshape(arr.shape[0], -1) + out[key] = arr + + atype = np.stack([frame["atype"] for frame in frames]).astype(np.int32) + real_natoms_vec = np.stack([frame["natoms"] for frame in frames]).astype( + np.int32 + ) + nloc = int(real_natoms_vec[:, 0].max()) + natoms_vec = np.concatenate( + ( + np.array([nloc, nloc], dtype=np.int32), + real_natoms_vec[:, 2:].max(axis=0).astype(np.int32), + ) + ) + + out["type"] = atype + out["natoms_vec"] = natoms_vec + out["real_natoms_vec"] = real_natoms_vec + if "box" not in out or out["box"] is None: + out["box"] = np.zeros((len(frames), 9), dtype=GLOBAL_NP_FLOAT_PRECISION) + out["find_box"] = np.float32(0.0) + elif "find_box" not in out: + out["find_box"] = np.float32(0.0 if np.allclose(out["box"], 0.0) else 1.0) + if "find_coord" not in out: + out["find_coord"] = np.float32(1.0) + if "numb_copy" not in out: + out["numb_copy"] = np.ones((len(frames), 1), dtype=np.int64) + out["find_numb_copy"] = np.float32(0.0) + out["default_mesh"] = np.asarray( + make_default_mesh(bool(float(out["find_box"]) > 0.5), self.mixed_type), + dtype=np.int32, + ) + return out + + def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: + del sys_idx + indices = self._next_indices() + frames = [self._reader[int(idx)] for idx in indices] + return self._stack_frames(frames) + + def get_nsystems(self) -> int: + return self.nsystems + + def get_natoms(self) -> int: + return self.natoms[0] + + def get_ntypes(self) -> int: + return len(self._type_map) + + def get_type_map(self) -> list[str]: + return self._type_map + + def get_batch_size(self) -> list[int]: + return self.batch_size + + def print_summary(self, name: str, prob: Any | None = None) -> None: + del prob + self._reader.print_summary(name, self.sys_probs) + + def _format_name_length(name: str, width: int) -> str: if len(name) <= width: return "{: >{}}".format(name, width) @@ -815,14 +1083,212 @@ def prob_sys_size_ext(keywords: str, nsystems: int, nbatch: int) -> list[float]: return sys_probs +def _is_deepmd_data_format(fmt: str) -> bool: + return fmt in { + "deepmd", + "deepmd/raw", + "deepmd/npy", + "deepmd/comp", + "deepmd/npy/mixed", + "deepmd/hdf5", + "lmdb", + } + + +def _looks_like_extxyz(path: Path) -> bool: + if not path.is_file(): + return False + try: + with path.open() as fp: + fp.readline() + comment = fp.readline() + except OSError: + return False + return "Properties=" in comment or "Lattice=" in comment + + +def _normalize_dpdata_format(fmt: str, source: Path) -> str: + fmt = fmt.lower() + if fmt == "ase": + return "ase/structure" + if fmt != "auto": + return fmt + suffix = source.suffix.lower().lstrip(".") + if suffix == "traj": + return "ase/traj" + if suffix == "extxyz" or (suffix == "xyz" and _looks_like_extxyz(source)): + return "extxyz" + return suffix or fmt + + +def _iter_conversion_inputs(path: str, patterns: list[str] | None) -> list[str]: + if patterns is None: + return [path] + root = Path(path) + if not root.is_dir(): + return [path] + matches = [] + for pattern in patterns: + matches.extend(str(match) for match in root.rglob(pattern)) + return sorted(set(matches)) + + +def _conversion_cache_path(source: Path, fmt: str, out_fmt: str) -> Path: + source_resolved = source.resolve(strict=False) + digest = hashlib.sha1(f"{source_resolved}|{fmt}|{out_fmt}".encode()).hexdigest()[ + :16 + ] + stem = source_resolved.stem or source_resolved.name or "dataset" + safe_out_fmt = out_fmt.replace("/", "-") + suffix = ".lmdb" if out_fmt == "lmdb" else "" + return Path.cwd() / _DPDATA_CACHE_DIR / f"{stem}-{safe_out_fmt}-{digest}{suffix}" + + +def _source_mtime(source: Path, cache_file: Path) -> float: + if source.is_file(): + return source.stat().st_mtime + if not source.is_dir(): + return 0.0 + cache_dir = cache_file.parent.resolve(strict=False) + latest = source.stat().st_mtime + for item in source.rglob("*"): + try: + item_resolved = item.resolve(strict=False) + if item_resolved == cache_file or cache_dir in item_resolved.parents: + continue + latest = max(latest, item.stat().st_mtime) + except OSError: + continue + return latest + + +def _is_conversion_current(source: Path, output: Path) -> bool: + if not output.exists(): + return False + return output.stat().st_mtime >= _source_mtime(source, output) + + +def _wait_for_conversion(source: Path, output: Path, lock_path: Path) -> bool: + for _ in range(300): + if not lock_path.exists(): + return _is_conversion_current(source, output) + if _is_conversion_current(source, output): + return True + time.sleep(1.0) + return False + + +def _remove_path(path: Path) -> None: + if path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + + +def _write_dpdata_conversion( + source: Path, fmt: str, out_fmt: str, output: Path +) -> None: + try: + import dpdata + except ImportError as exc: + raise ImportError( + "dpdata is required when training_data.format or " + "validation_data.format is specified. Install dpdata to enable " + "automatic dataset conversion." + ) from exc + + tmp_output = output.with_name(f".{output.name}.{os.getpid()}.tmp") + _remove_path(tmp_output) + try: + multi_systems = dpdata.MultiSystems() + try: + multi_systems.load_systems_from_file(str(source), fmt=fmt) + except NotImplementedError: + labeled_system = dpdata.LabeledSystem(str(source), fmt=fmt) + multi_systems = dpdata.MultiSystems(labeled_system) + if len(multi_systems) == 0: + raise RuntimeError(f"No frames were loaded by dpdata from {source}") + multi_systems.to(out_fmt, str(tmp_output)) + _remove_path(output) + os.replace(tmp_output, output) + except Exception: + _remove_path(tmp_output) + raise + + +def _convert_system_by_dpdata( + source_path: str, fmt: str, out_fmt: str | None +) -> list[str]: + if out_fmt is None: + out_fmt = _DPDATA_DEFAULT_OUT_FORMAT + source = Path(source_path) + fmt = _normalize_dpdata_format(fmt, source) + out_fmt = out_fmt.lower() + cache_key = ( + str(Path.cwd().resolve(strict=False)), + str(source.resolve(strict=False)), + fmt, + out_fmt, + ) + if cache_key in _DPDATA_CONVERSION_CACHE: + return _DPDATA_CONVERSION_CACHE[cache_key] + + output = _conversion_cache_path(source, fmt, out_fmt) + output.parent.mkdir(parents=True, exist_ok=True) + lock_path = output.with_suffix(output.suffix + ".lock") + if not _is_conversion_current(source, output): + while True: + try: + lock_fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + if _wait_for_conversion(source, output, lock_path): + break + raise TimeoutError( + f"Timed out waiting for dpdata conversion lock {lock_path}" + ) from None + else: + with os.fdopen(lock_fd, "w") as fp: + fp.write(str(os.getpid())) + try: + if not _is_conversion_current(source, output): + log.info( + "Converting %s from dpdata format %s to %s at %s", + source, + fmt, + out_fmt, + output, + ) + _write_dpdata_conversion(source, fmt, out_fmt, output) + finally: + try: + lock_path.unlink() + except FileNotFoundError: + pass + break + + if out_fmt == "lmdb": + converted_systems = [str(output)] + else: + converted_systems = expand_sys_str(str(output)) + if not converted_systems: + raise RuntimeError(f"No DeePMD systems were found in converted file {output}") + _DPDATA_CONVERSION_CACHE[cache_key] = converted_systems + return converted_systems + + def process_systems( - systems: str | list[str], patterns: list[str] | None = None + systems: str | list[str], + patterns: list[str] | None = None, + fmt: str | None = None, + out_fmt: str | None = None, ) -> list[str]: """Process the user-input systems. If it is a single directory, search for all the systems in the directory. If it is a list, each item in the list is treated as a directory to search. If it is a single LMDB path, return it directly without expansion. + If fmt is specified and is not a DeePMD data format, each input path is + converted by dpdata and the converted systems are returned. Check if the systems are valid. Parameters @@ -831,20 +1297,17 @@ def process_systems( The user-input systems patterns : list of str, optional The patterns to match the systems, by default None + fmt : str, optional + The dpdata input format. If None, no conversion is performed. + out_fmt : str, optional + The dpdata output format. If None, ``lmdb`` is used when fmt triggers + conversion. Returns ------- result_systems: list of str The valid systems """ - from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, - ) - - # LMDB path: return directly without expansion - if isinstance(systems, str) and is_lmdb(systems): - return [systems] - # Normalize input to a list of paths to search if isinstance(systems, str): search_paths = [systems] @@ -856,15 +1319,31 @@ def process_systems( f"Invalid systems type: {type(systems)}. Must be str or list[str]." ) + if fmt is not None: + fmt = fmt.lower() + if _is_deepmd_data_format(fmt): + fmt = None + + from deepmd.dpmodel.utils.lmdb_data import ( + is_lmdb, + ) + # Iterate over the search_paths list and apply expansion logic to each path result_systems = [] for path in search_paths: - if patterns is None: + if fmt is not None: + for input_path in _iter_conversion_inputs(path, patterns): + result_systems.extend( + _convert_system_by_dpdata(input_path, fmt, out_fmt) + ) + elif is_lmdb(path): + result_systems.append(path) + elif patterns is None: expanded_paths = expand_sys_str(path) + result_systems.extend(expanded_paths) else: expanded_paths = rglob_sys_str(path, patterns) - - result_systems.extend(expanded_paths) + result_systems.extend(expanded_paths) return result_systems @@ -875,7 +1354,7 @@ def get_data( type_map: list[str] | None, modifier: Any | None, multi_task_mode: bool = False, -) -> DeepmdDataSystem: +) -> DeepmdDataSystem | LmdbDataSystem: """Get the data system. Parameters @@ -898,13 +1377,35 @@ def get_data( """ systems = jdata["systems"] rglob_patterns = jdata.get("rglob_patterns", None) - systems = process_systems(systems, patterns=rglob_patterns) + data_format = jdata.get("format", None) + out_format = jdata.get("out_format", jdata.get("output_format", None)) + systems = process_systems( + systems, patterns=rglob_patterns, fmt=data_format, out_fmt=out_format + ) batch_size = jdata["batch_size"] sys_probs = jdata.get("sys_probs", None) auto_prob = jdata.get("auto_prob", "prob_sys_size") optional_type_map = not multi_task_mode + from deepmd.dpmodel.utils.lmdb_data import ( + is_lmdb, + ) + + if len(systems) == 1 and is_lmdb(systems[0]): + if type_map is None: + raise ValueError( + "LMDB training data requires model/type_map to be set. " + "Set model/type_map or choose training_data.out_format=" + "'deepmd/hdf5' for automatic conversion." + ) + return LmdbDataSystem( + lmdb_path=systems[0], + type_map=type_map, + batch_size=batch_size, + auto_prob_style=auto_prob, + ) + data = DeepmdDataSystem( systems=systems, batch_size=batch_size, diff --git a/pyproject.toml b/pyproject.toml index 35fc0fdb18..37312b0a3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,7 @@ dependencies = [ 'array-api-compat', 'lmdb', 'msgpack', + 'dpdata>=1.0.1', ] requires-python = ">=3.10" keywords = ["deepmd"] @@ -81,7 +82,6 @@ repository = "https://github.com/deepmodeling/deepmd-kit" # which can be read by the build backend. [tool.deepmd_build_backend.optional-dependencies] test = [ - "dpdata>=0.2.7", # ASE issue: https://gitlab.com/ase/ase/-/merge_requests/2843 # fixed in 3.23.0 "ase>=3.23.0", diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py new file mode 100644 index 0000000000..37b8b7bddf --- /dev/null +++ b/source/tests/common/test_data_system_conversion.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os +import sys +import tempfile +import types +import unittest +from pathlib import ( + Path, +) +from typing import ( + ClassVar, +) +from unittest.mock import ( + patch, +) + +import h5py +import lmdb +import msgpack +import numpy as np + +from deepmd.dpmodel.utils.lmdb_data import ( + is_lmdb, +) +from deepmd.utils import ( + data_system, +) +from deepmd.utils.data_system import ( + LmdbDataSystem, + get_data, + process_systems, +) + + +def _write_minimal_deepmd_hdf5(file_name: str) -> None: + with h5py.File(file_name, "w") as fp: + system = fp.create_group("H") + system.create_dataset("type.raw", data=np.array([0], dtype=np.int32)) + string_dtype = h5py.string_dtype(encoding="utf-8") + system.create_dataset("type_map.raw", data=np.array(["H"], dtype=string_dtype)) + set_dir = system.create_group("set.000") + set_dir.create_dataset("coord.npy", data=np.zeros((1, 3), dtype=np.float32)) + set_dir.create_dataset( + "box.npy", data=np.eye(3, dtype=np.float32).reshape(1, 9) + ) + + +def _encode_array(arr: np.ndarray) -> dict: + return { + "type": str(arr.dtype), + "shape": list(arr.shape), + "data": arr.tobytes(), + } + + +def _write_minimal_lmdb(path: str) -> None: + env = lmdb.open(path, map_size=10 * 1024 * 1024) + frame = { + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": _encode_array(np.array([0], dtype=np.int64)), + "cells": _encode_array(np.eye(3, dtype=np.float64) * 8.0), + "coords": _encode_array(np.zeros((1, 3), dtype=np.float64)), + "energies": _encode_array(np.array([0.0], dtype=np.float64)), + "forces": _encode_array(np.zeros((1, 3), dtype=np.float64)), + } + metadata = { + "nframes": 1, + "frame_idx_fmt": "012d", + "type_map": ["H"], + "system_info": { + "formula": "H", + "natoms": [1], + "nframes": 1, + }, + } + with env.begin(write=True) as txn: + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + txn.put(b"000000000000", msgpack.packb(frame, use_bin_type=True)) + env.close() + + +class _FakeMultiSystems: + write_count = 0 + load_calls: ClassVar[list[tuple[str, str]]] = [] + + def __init__(self, *systems) -> None: + self.systems = list(systems) + self.loaded = False + + def load_systems_from_file(self, file_name: str, fmt: str): + self.load_calls.append((file_name, fmt)) + self.loaded = True + return self + + def __len__(self) -> int: + return 1 if self.loaded or self.systems else 0 + + def to(self, fmt: str, file_name: str) -> None: + type(self).write_count += 1 + if fmt == "deepmd/hdf5": + _write_minimal_deepmd_hdf5(file_name) + elif fmt == "lmdb": + _write_minimal_lmdb(file_name) + else: + raise AssertionError(fmt) + + +class _FakeLabeledSystem: + def __init__(self, file_name: str, fmt: str) -> None: + self.file_name = file_name + self.fmt = fmt + + +class TestDpdataFormatConversion(unittest.TestCase): + def setUp(self) -> None: + self.tmpdir = tempfile.TemporaryDirectory() + self.root = Path(self.tmpdir.name) + self.old_cwd = Path.cwd() + os.chdir(self.root) + self.source = self.root / "data.extxyz" + self.source.write_text("1\nProperties=species:S:1:pos:R:3\nH 0 0 0\n") + _FakeMultiSystems.write_count = 0 + _FakeMultiSystems.load_calls = [] + data_system._DPDATA_CONVERSION_CACHE.clear() + self.fake_dpdata = types.SimpleNamespace( + MultiSystems=_FakeMultiSystems, + LabeledSystem=_FakeLabeledSystem, + ) + + def tearDown(self) -> None: + os.chdir(self.old_cwd) + self.tmpdir.cleanup() + data_system._DPDATA_CONVERSION_CACHE.clear() + + def test_process_systems_defaults_to_lmdb_and_reuses_cache(self) -> None: + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems(str(self.source), fmt="extxyz") + systems_again = process_systems(str(self.source), fmt="extxyz") + + self.assertEqual(systems, systems_again) + self.assertEqual(_FakeMultiSystems.write_count, 1) + self.assertEqual(_FakeMultiSystems.load_calls, [(str(self.source), "extxyz")]) + self.assertEqual(len(systems), 1) + self.assertTrue(systems[0].endswith(".lmdb")) + self.assertTrue(is_lmdb(systems[0])) + self.assertTrue(Path(systems[0]).is_relative_to(self.root)) + self.assertEqual(Path(systems[0]).parent, self.root / ".deepmd_dpdata_cache") + + def test_process_systems_cache_is_scoped_to_cwd(self) -> None: + other_cwd = self.root / "run2" + other_cwd.mkdir() + + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems(str(self.source), fmt="extxyz") + os.chdir(other_cwd) + systems_other = process_systems(str(self.source), fmt="extxyz") + + self.assertNotEqual(systems, systems_other) + self.assertEqual(_FakeMultiSystems.write_count, 2) + self.assertEqual(Path(systems[0]).parent, self.root / ".deepmd_dpdata_cache") + self.assertEqual( + Path(systems_other[0]).parent, + other_cwd / ".deepmd_dpdata_cache", + ) + + def test_process_systems_converts_to_explicit_hdf5(self) -> None: + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems( + str(self.source), fmt="extxyz", out_fmt="deepmd/hdf5" + ) + + self.assertEqual(_FakeMultiSystems.write_count, 1) + self.assertEqual(_FakeMultiSystems.load_calls, [(str(self.source), "extxyz")]) + self.assertEqual(len(systems), 1) + self.assertTrue(systems[0].endswith("#/H")) + + def test_get_data_uses_format_conversion(self) -> None: + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + data = get_data( + { + "systems": str(self.source), + "format": "auto", + "batch_size": 1, + }, + 0.0, + ["H"], + None, + ) + + self.assertEqual(data.get_nsystems(), 1) + self.assertIsInstance(data, LmdbDataSystem) + self.assertEqual(_FakeMultiSystems.load_calls, [(str(self.source), "extxyz")]) + batch = data.get_batch() + self.assertIn("type", batch) + self.assertIn("natoms_vec", batch) + self.assertEqual(batch["coord"].shape, (1, 3)) + self.assertEqual(data.data_systems, [data]) + stat_set = data._load_set(data.dirs[0]) + self.assertEqual(stat_set["coord"].shape, (1, 3)) + self.assertEqual(stat_set["type"].shape, (1, 1)) + + +if __name__ == "__main__": + unittest.main() From 2feb07e7f72d8d522a261203c5cd0ed96e2bd7bd Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 23:23:03 +0800 Subject: [PATCH 02/12] fix(data): validate converted LMDB system lists Break the LMDB/data-system import cycle, reject ambiguous multi-LMDB results across backends, reject LMDB on Paddle, and guard empty LMDB frame batches with direct regressions. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/lmdb_data.py | 18 ++++-- deepmd/pd/entrypoints/main.py | 20 ++----- deepmd/pt/entrypoints/main.py | 11 ++-- deepmd/pt_expt/entrypoints/main.py | 15 +++-- deepmd/utils/data_system.py | 56 +++++++++++++------ .../common/test_data_system_conversion.py | 35 ++++++++++++ source/tests/pt_expt/test_lmdb_training.py | 28 ++++++++++ 7 files changed, 138 insertions(+), 45 deletions(-) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index d4e4c65b23..220ae8ccac 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -869,10 +869,6 @@ def compute_block_targets( Each element is ``(system_indices_in_block, target_frame_count)``. Returns empty list if no expansion is needed (all targets == actual). """ - from deepmd.utils.data_system import ( - prob_sys_size_ext, - ) - # Parse block definitions from the auto_prob string # Format: "prob_sys_size;stt:end:weight;stt:end:weight;..." block_str = auto_prob_style.split(";")[1:] @@ -918,8 +914,18 @@ def compute_block_targets( ) blocks = nonempty - # Compute per-system probabilities using the standard function - sys_probs = prob_sys_size_ext(auto_prob_style, nsystems, system_nframes) + # Compute the same per-system probabilities as prob_sys_size_ext locally. + # Keeping this framework-agnostic LMDB module independent of data_system + # avoids an import cycle when the legacy adapter imports the LMDB reader. + block_weights = np.asarray([weight for _, _, weight in blocks], dtype=float) + assert np.all(block_weights >= 0), "the weight of a block should be no less than 0" + block_probs = block_weights / np.sum(block_weights) + sys_probs = np.zeros(nsystems, dtype=np.float64) + for block_idx, (stt, end, _weight) in enumerate(blocks): + block_frames = np.asarray(system_nframes[stt:end], dtype=float) + sys_probs[stt:end] = ( + block_frames / np.sum(block_frames) * block_probs[block_idx] + ) # Group systems by block, compute block-level frames and prob block_info: list[tuple[list[int], int, float]] = [] # (sys_ids, frames, prob) diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index 72c9ef7677..e9b6396b1e 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -23,9 +23,6 @@ expand_sys_str, j_loader, ) -from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, -) from deepmd.loggers.loggers import ( set_log_handles, ) @@ -72,6 +69,7 @@ from deepmd.utils.data_system import ( get_data, process_systems, + validate_lmdb_systems, ) from deepmd.utils.path import ( DPPath, @@ -120,12 +118,7 @@ def prepare_trainer_input_single( "out_format", training_dataset_params.get("output_format", None) ), ) - if len(training_systems) == 1 and is_lmdb(training_systems[0]): - raise NotImplementedError( - "Paddle backend does not support LMDB training data yet. " - "Set training_data.out_format to 'deepmd/hdf5' when using " - "training_data.format for automatic conversion." - ) + validate_lmdb_systems(training_systems, backend_name="Paddle", supported=False) if validation_systems is not None: val_patterns = validation_dataset_params.get("rglob_patterns", None) validation_systems = process_systems( @@ -136,12 +129,9 @@ def prepare_trainer_input_single( "out_format", validation_dataset_params.get("output_format", None) ), ) - if len(validation_systems) == 1 and is_lmdb(validation_systems[0]): - raise NotImplementedError( - "Paddle backend does not support LMDB validation data yet. " - "Set validation_data.out_format to 'deepmd/hdf5' when using " - "validation_data.format for automatic conversion." - ) + validate_lmdb_systems( + validation_systems, backend_name="Paddle", supported=False + ) # stat files stat_file_path_single = data_dict_single.get("stat_file", None) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 8bd147697d..8cfb9edacf 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -91,6 +91,7 @@ from deepmd.utils.data_system import ( get_data, process_systems, + validate_lmdb_systems, ) from deepmd.utils.path import ( DPPath, @@ -194,9 +195,10 @@ def _make_dp_loader_set( "out_format", dataset_params.get("output_format", None) ), ) - if len(systems) == 1 and is_lmdb(systems[0]): + lmdb_path = validate_lmdb_systems(systems, backend_name="PyTorch") + if lmdb_path is not None: return LmdbDataset( - systems[0], + lmdb_path, model_params_single["type_map"], dataset_params["batch_size"], auto_prob_style=dataset_params.get("auto_prob", None), @@ -437,12 +439,13 @@ def _get_neighbor_stat_data_from_params( "out_format", dataset_params.get("output_format", None) ), ) - if len(systems) == 1 and is_lmdb(systems[0]): + lmdb_path = validate_lmdb_systems(systems, backend_name="PyTorch") + if lmdb_path is not None: from deepmd.dpmodel.utils.lmdb_data import ( make_neighbor_stat_data, ) - return make_neighbor_stat_data(systems[0], type_map) + return make_neighbor_stat_data(lmdb_path, type_map) return get_data(dataset_params, 0, type_map, None) if not multi_task: diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 38f16095f4..131ce012b1 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -33,6 +33,7 @@ DeepmdDataSystem, get_data, process_systems, + validate_lmdb_systems, ) from deepmd.utils.path import ( DPPath, @@ -123,12 +124,15 @@ def _get_neighbor_stat_data( "out_format", dataset_params.get("output_format", None) ), ) - if len(systems) == 1 and is_lmdb(systems[0]): + converted_lmdb_path = validate_lmdb_systems( + systems, backend_name="PyTorch exportable" + ) + if converted_lmdb_path is not None: from deepmd.dpmodel.utils.lmdb_data import ( make_neighbor_stat_data, ) - return make_neighbor_stat_data(systems[0], type_map) + return make_neighbor_stat_data(converted_lmdb_path, type_map) return get_data(dataset_params, 0, type_map, None) @@ -165,9 +169,12 @@ def _build_data_system( "out_format", dataset_params.get("output_format", None) ), ) - if len(systems) == 1 and is_lmdb(systems[0]): + converted_lmdb_path = validate_lmdb_systems( + systems, backend_name="PyTorch exportable" + ) + if converted_lmdb_path is not None: return LmdbDataSystem( - lmdb_path=systems[0], + lmdb_path=converted_lmdb_path, type_map=type_map, batch_size=dataset_params["batch_size"], auto_prob_style=dataset_params.get("auto_prob"), diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index e6a795c616..348dabb410 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -24,6 +24,12 @@ make_default_mesh, rglob_sys_str, ) +from deepmd.dpmodel.utils.lmdb_data import ( + LmdbDataReader, + SameNlocBatchSampler, + compute_block_targets, + is_lmdb, +) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -42,6 +48,33 @@ _DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str], list[str]] = {} +def validate_lmdb_systems( + systems: list[str], + *, + backend_name: str, + supported: bool = True, +) -> str | None: + """Validate expanded systems and return the sole resolved LMDB path. + + LMDB stores multiple logical systems inside one database, so mixing an + LMDB path with other expanded paths is ambiguous and unsupported. + """ + lmdb_paths = [path for path in systems if is_lmdb(path)] + if not lmdb_paths: + return None + if not supported: + raise NotImplementedError( + f"{backend_name} backend does not support LMDB data yet. " + "Choose out_format='deepmd/hdf5' for automatic conversion." + ) + if len(systems) != 1: + raise ValueError( + f"{backend_name} backend requires an LMDB dataset to resolve to " + "exactly one path; LMDB paths cannot be mixed with other systems." + ) + return lmdb_paths[0] + + class DeepmdDataSystem: """Class for manipulating many data systems. @@ -707,12 +740,6 @@ def __init__( "must map them to element names." ) - from deepmd.dpmodel.utils.lmdb_data import ( - LmdbDataReader, - SameNlocBatchSampler, - compute_block_targets, - ) - self.lmdb_path = lmdb_path self._reader = LmdbDataReader( lmdb_path, type_map, batch_size, mixed_batch=False @@ -859,6 +886,8 @@ def _next_indices(self) -> list[int]: return next(self._iter) def _stack_frames(self, frames: list[dict[str, Any]]) -> dict[str, Any]: + if not frames: + raise ValueError("Cannot stack an empty LMDB frame batch.") out: dict[str, Any] = {} structural_keys = {"coord", "box"} for key in frames[0]: @@ -1263,6 +1292,8 @@ def _convert_system_by_dpdata( try: lock_path.unlink() except FileNotFoundError: + # Concurrent conversion cleanup is idempotent; another + # process may already have removed the shared lock. pass break @@ -1324,10 +1355,6 @@ def process_systems( if _is_deepmd_data_format(fmt): fmt = None - from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, - ) - # Iterate over the search_paths list and apply expansion logic to each path result_systems = [] for path in search_paths: @@ -1388,11 +1415,8 @@ def get_data( auto_prob = jdata.get("auto_prob", "prob_sys_size") optional_type_map = not multi_task_mode - from deepmd.dpmodel.utils.lmdb_data import ( - is_lmdb, - ) - - if len(systems) == 1 and is_lmdb(systems[0]): + lmdb_path = validate_lmdb_systems(systems, backend_name="legacy data loader") + if lmdb_path is not None: if type_map is None: raise ValueError( "LMDB training data requires model/type_map to be set. " @@ -1400,7 +1424,7 @@ def get_data( "'deepmd/hdf5' for automatic conversion." ) return LmdbDataSystem( - lmdb_path=systems[0], + lmdb_path=lmdb_path, type_map=type_map, batch_size=batch_size, auto_prob_style=auto_prob, diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index 37b8b7bddf..2d606d2dde 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -29,6 +29,7 @@ LmdbDataSystem, get_data, process_systems, + validate_lmdb_systems, ) @@ -200,6 +201,40 @@ def test_get_data_uses_format_conversion(self) -> None: self.assertEqual(stat_set["coord"].shape, (1, 3)) self.assertEqual(stat_set["type"].shape, (1, 1)) + def test_multiple_lmdb_paths_are_rejected(self) -> None: + lmdb_a = self.root / "a.lmdb" + lmdb_b = self.root / "b.lmdb" + _write_minimal_lmdb(str(lmdb_a)) + _write_minimal_lmdb(str(lmdb_b)) + + with self.assertRaisesRegex(ValueError, "exactly one path"): + get_data( + { + "systems": [str(lmdb_a), str(lmdb_b)], + "batch_size": 1, + }, + 0.0, + ["H"], + None, + ) + + def test_backend_without_lmdb_support_rejects_any_resolved_path(self) -> None: + lmdb_path = self.root / "unsupported.lmdb" + _write_minimal_lmdb(str(lmdb_path)) + + with self.assertRaisesRegex(NotImplementedError, "Paddle backend"): + validate_lmdb_systems( + [str(lmdb_path)], backend_name="Paddle", supported=False + ) + + def test_lmdb_stack_frames_rejects_empty_batch(self) -> None: + lmdb_path = self.root / "empty-batch.lmdb" + _write_minimal_lmdb(str(lmdb_path)) + data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=1) + + with self.assertRaisesRegex(ValueError, "empty LMDB frame batch"): + data._stack_frames([]) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 94673a8761..acbeb95824 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -13,6 +13,9 @@ import shutil import tempfile import unittest +from unittest.mock import ( + patch, +) import lmdb import msgpack @@ -23,6 +26,8 @@ split_batch, ) from deepmd.pt_expt.entrypoints.main import ( + _build_data_system, + _get_neighbor_stat_data, get_trainer, ) from deepmd.pt_expt.utils.lmdb_dataset import ( @@ -46,6 +51,29 @@ def _encode_array(arr: np.ndarray) -> dict: } +class TestConvertedLmdbValidation(unittest.TestCase): + """Reject format conversion that resolves to multiple LMDB databases.""" + + def test_neighbor_stat_and_training_data_reject_multiple_lmdb(self) -> None: + params = { + "systems": "input.extxyz", + "format": "extxyz", + "batch_size": 1, + } + converted = ["first.lmdb", "second.lmdb"] + with ( + patch( + "deepmd.pt_expt.entrypoints.main.process_systems", + return_value=converted, + ), + patch("deepmd.utils.data_system.is_lmdb", return_value=True), + ): + with self.assertRaisesRegex(ValueError, "exactly one path"): + _get_neighbor_stat_data(params, ["O", "H"]) + with self.assertRaisesRegex(ValueError, "exactly one path"): + _build_data_system(params, ["O", "H"]) + + def _make_frame(natoms: int, seed: int) -> dict: """Synthetic LMDB frame matching the on-disk schema used by LmdbDataReader.""" rng = np.random.RandomState(seed) From f99f0c1ac0299aaeef906050090bdaa60485234f Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 18 Jul 2026 14:47:16 +0800 Subject: [PATCH 03/12] fix(data): avoid LMDB import cycle Load the dpmodel LMDB helpers only after the legacy data-system module has initialized, preventing backend imports from re-entering a partially initialized module. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/utils/data_system.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 63e4a6a67a..e54f964393 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -24,12 +24,6 @@ make_default_mesh, rglob_sys_str, ) -from deepmd.dpmodel.utils.lmdb_data import ( - LmdbDataReader, - SameNlocBatchSampler, - compute_block_targets, - is_lmdb, -) from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -59,6 +53,12 @@ def validate_lmdb_systems( LMDB stores multiple logical systems inside one database, so mixing an LMDB path with other expanded paths is ambiguous and unsupported. """ + # Import after data_system has initialized. Importing the dpmodel package + # at module load time re-enters this module through descriptor utilities. + from deepmd.dpmodel.utils.lmdb_data import ( + is_lmdb, + ) + lmdb_paths = [path for path in systems if is_lmdb(path)] if not lmdb_paths: return None @@ -733,6 +733,14 @@ def __init__( auto_prob_style: str | None = None, seed: int | None = None, ) -> None: + # Keep the framework-agnostic LMDB implementation lazy so importing a + # legacy backend cannot create a data_system <-> dpmodel import cycle. + from deepmd.dpmodel.utils.lmdb_data import ( + LmdbDataReader, + SameNlocBatchSampler, + compute_block_targets, + ) + if not type_map: raise ValueError( "LMDB datasets require a non-empty model/type_map because " @@ -1339,6 +1347,12 @@ def process_systems( result_systems: list of str The valid systems """ + # See validate_lmdb_systems: this must remain a local import because + # deepmd.dpmodel initializes descriptors that depend on data_system. + from deepmd.dpmodel.utils.lmdb_data import ( + is_lmdb, + ) + # Normalize input to a list of paths to search if isinstance(systems, str): search_paths = [systems] From c2e7727f628a70f18c5a3b31f8fa537b7c07a6b8 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Mon, 27 Jul 2026 20:14:32 +0800 Subject: [PATCH 04/12] test(pt_expt): patch is_lmdb at its source module deepmd.utils.data_system imports is_lmdb inside the validating function rather than at module scope, so patching the import site no longer resolves and mock raises AttributeError. --- source/tests/pt_expt/test_lmdb_training.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index e39bdf719f..5b61c2144a 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -69,7 +69,12 @@ def test_neighbor_stat_and_training_data_reject_multiple_lmdb(self) -> None: "deepmd.pt_expt.entrypoints.main.process_systems", return_value=converted, ), - patch("deepmd.utils.data_system.is_lmdb", return_value=True), + # is_lmdb is imported inside the validating function, so patch + # it at the source module rather than at an import site. + patch( + "deepmd.dpmodel.utils.lmdb_data.is_lmdb", + return_value=True, + ), ): with self.assertRaisesRegex(ValueError, "exactly one path"): _get_neighbor_stat_data(params, ["O", "H"]) From b16b18067e22c3731d4ece2b807f4d18e3f5c9cc Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 1 Aug 2026 22:45:51 +0800 Subject: [PATCH 05/12] fix(data): refresh stale conversion caches Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/utils/data_system.py | 28 ++++++++----------- .../common/test_data_system_conversion.py | 13 +++++++++ 2 files changed, 25 insertions(+), 16 deletions(-) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index e54f964393..c124d2d356 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -753,7 +753,9 @@ def __init__( lmdb_path, type_map, batch_size, mixed_batch=False ) self._type_map = list(type_map) - self.mixed_type = self._detect_mixed_type() + # LMDB is defined as mixed-type by its reader contract; determining + # this must not scan every frame during data-system initialization. + self.mixed_type = self._reader.mixed_type self.nsystems = 1 self.system_dirs = [lmdb_path] self.natoms = [max(self._reader.frame_nlocs) if self._reader.frame_nlocs else 0] @@ -818,18 +820,6 @@ def __init__( ) self._iter = iter(self._sampler) - def _detect_mixed_type(self) -> bool: - """Return True when frames cannot be represented as fixed-type data.""" - if len(self._reader.nloc_groups) > 1: - return True - if len(self._reader) == 0: - return False - ref_type = self._reader[0]["atype"] - for idx in range(1, len(self._reader)): - if not np.array_equal(self._reader[idx]["atype"], ref_type): - return True - return False - def _detect_pbc(self) -> bool: """Return True when LMDB frames contain a non-zero simulation box.""" if len(self._reader) == 0: @@ -1267,10 +1257,16 @@ def _convert_system_by_dpdata( fmt, out_fmt, ) - if cache_key in _DPDATA_CONVERSION_CACHE: - return _DPDATA_CONVERSION_CACHE[cache_key] - output = _conversion_cache_path(source, fmt, out_fmt) + cached_systems = _DPDATA_CONVERSION_CACHE.get(cache_key) + if cached_systems is not None: + if _is_conversion_current(source, output): + return cached_systems + # A long-lived training/validation process may observe source files + # rewritten in place. Drop the fast-path entry so the normal locked + # conversion flow refreshes the on-disk result before it is reused. + del _DPDATA_CONVERSION_CACHE[cache_key] + output.parent.mkdir(parents=True, exist_ok=True) lock_path = output.with_suffix(output.suffix + ".lock") if not _is_conversion_current(source, output): diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index 2d606d2dde..404d733f0f 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -165,6 +165,18 @@ def test_process_systems_cache_is_scoped_to_cwd(self) -> None: other_cwd / ".deepmd_dpdata_cache", ) + def test_process_systems_revalidates_in_memory_cache(self) -> None: + """A source rewrite in one process must refresh its converted output.""" + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems(str(self.source), fmt="extxyz") + output_mtime = Path(systems[0]).stat().st_mtime + refreshed_mtime = output_mtime + 1.0 + os.utime(self.source, (refreshed_mtime, refreshed_mtime)) + systems_again = process_systems(str(self.source), fmt="extxyz") + + self.assertEqual(systems, systems_again) + self.assertEqual(_FakeMultiSystems.write_count, 2) + def test_process_systems_converts_to_explicit_hdf5(self) -> None: with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): systems = process_systems( @@ -191,6 +203,7 @@ def test_get_data_uses_format_conversion(self) -> None: self.assertEqual(data.get_nsystems(), 1) self.assertIsInstance(data, LmdbDataSystem) + self.assertTrue(data.mixed_type) self.assertEqual(_FakeMultiSystems.load_calls, [(str(self.source), "extxyz")]) batch = data.get_batch() self.assertIn("type", batch) From 7ae5c2e32f8f857d64bb46378580598138d3b388 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 2 Aug 2026 00:13:17 +0800 Subject: [PATCH 06/12] fix(data): harden LMDB conversion validation Reject invalid auto-probability weights before normalization and prevent cache cleanup from following directory symlinks. Add focused regression coverage for both validation paths. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/lmdb_data.py | 19 ++++++++++++++----- deepmd/utils/data_system.py | 6 +++++- source/tests/common/dpmodel/test_lmdb_data.py | 16 ++++++++++++++++ .../common/test_data_system_conversion.py | 16 ++++++++++++++++ 4 files changed, 51 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 91ac3e9d32..bf3288c32e 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -936,14 +936,23 @@ def compute_block_targets( # Keeping this framework-agnostic LMDB module independent of data_system # avoids an import cycle when the legacy adapter imports the LMDB reader. block_weights = np.asarray([weight for _, _, weight in blocks], dtype=float) - assert np.all(block_weights >= 0), "the weight of a block should be no less than 0" - block_probs = block_weights / np.sum(block_weights) + if not np.all(np.isfinite(block_weights)): + raise ValueError("block weights must be finite") + if np.any(block_weights < 0): + raise ValueError("the weight of a block should be no less than 0") + total_block_weight = np.sum(block_weights) + if total_block_weight <= 0: + raise ValueError("the sum of block weights should be greater than 0") + block_probs = block_weights / total_block_weight sys_probs = np.zeros(nsystems, dtype=np.float64) for block_idx, (stt, end, _weight) in enumerate(blocks): block_frames = np.asarray(system_nframes[stt:end], dtype=float) - sys_probs[stt:end] = ( - block_frames / np.sum(block_frames) * block_probs[block_idx] - ) + total_block_frames = np.sum(block_frames) + if total_block_frames <= 0: + raise ValueError( + f"block {stt}:{end} must contain at least one retained frame" + ) + sys_probs[stt:end] = block_frames / total_block_frames * block_probs[block_idx] # Group systems by block, compute block-level frames and prob block_info: list[tuple[list[int], int, float]] = [] # (sys_ids, frames, prob) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index c124d2d356..1bfe4f09fe 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -1206,7 +1206,11 @@ def _wait_for_conversion(source: Path, output: Path, lock_path: Path) -> bool: def _remove_path(path: Path) -> None: - if path.is_dir(): + # Check links before directories: Path.is_dir follows a directory symlink, + # while cache cleanup must never recurse into a target outside the cache. + if path.is_symlink(): + path.unlink() + elif path.is_dir(): shutil.rmtree(path) elif path.exists(): path.unlink() diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 59db455a76..c5d2b49574 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -716,6 +716,22 @@ def test_compute_block_targets_logs_dropped_block(self): self.assertTrue(any("empty blocks" in msg for msg in cm.output)) self.assertEqual(result, []) + def test_compute_block_targets_rejects_invalid_weights(self): + """Weights must define a finite, nonnegative probability mass.""" + invalid_styles = ( + ("prob_sys_size;0:1:-0.1;1:2:1.1", "no less than 0"), + ("prob_sys_size;0:1:0;1:2:0", "greater than 0"), + ("prob_sys_size;0:1:nan;1:2:1", "finite"), + ) + for style, message in invalid_styles: + with self.subTest(style=style): + with self.assertRaisesRegex(ValueError, message): + compute_block_targets( + style, + nsystems=2, + system_nframes=[100, 100], + ) + def test_expand_indices_basic(self): frame_system_ids = [0] * 5 + [1] * 5 block_targets = [([0], 25), ([1], 25)] diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index 404d733f0f..856bd28df7 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -177,6 +177,22 @@ def test_process_systems_revalidates_in_memory_cache(self) -> None: self.assertEqual(systems, systems_again) self.assertEqual(_FakeMultiSystems.write_count, 2) + def test_cache_cleanup_unlinks_directory_symlink(self) -> None: + """Cleanup must not recurse through a symlink outside the cache.""" + target = self.root / "outside" + target.mkdir() + sentinel = target / "keep.txt" + sentinel.write_text("keep") + link = self.root / ".deepmd_dpdata_cache" / "stale.tmp" + link.parent.mkdir() + link.symlink_to(target, target_is_directory=True) + + data_system._remove_path(link) + + self.assertFalse(link.exists()) + self.assertTrue(target.is_dir()) + self.assertEqual(sentinel.read_text(), "keep") + def test_process_systems_converts_to_explicit_hdf5(self) -> None: with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): systems = process_systems( From e548f74ac93cb7c0bec3d0c7d95f686f5bd45826 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 30 Aug 2026 01:25:04 +0800 Subject: [PATCH 07/12] refactor(data): use dpdata transactional LMDB writer Require dpdata 1.1.0, use the canonical deepmd/lmdb format, and delegate LMDB overwrite publication to dpdata. Add mock and real conversion coverage for compatibility and refreshes. Coding-Agent: Codex Codex-Version: codex-cli 0.151.0 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/utils/argcheck.py | 12 ++--- deepmd/utils/data_system.py | 49 ++++++++++++----- pyproject.toml | 2 +- .../common/test_data_system_conversion.py | 53 +++++++++++++++++-- 4 files changed, 93 insertions(+), 23 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 86e3e55253..b2ed7377c2 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5014,8 +5014,8 @@ def training_data_args() -> list[ doc_out_format = ( "The output data format passed to dpdata for automatic conversion. " "When `format` requests conversion from a non-DeePMD format, this key " - "defaults to `lmdb`. Use a DeePMD format supported by dpdata, such as " - "`lmdb`, `deepmd/hdf5`, or `deepmd/npy`." + "defaults to `deepmd/lmdb`. Use a DeePMD format supported by dpdata, " + "such as `deepmd/lmdb`, `deepmd/hdf5`, or `deepmd/npy`." ) doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ @@ -5065,7 +5065,7 @@ def training_data_args() -> list[ "out_format", [str, None], optional=True, - default="lmdb", + default="deepmd/lmdb", doc=doc_out_format, alias=["output_format"], ), @@ -5138,8 +5138,8 @@ def validation_data_args() -> list[ doc_out_format = ( "The output data format passed to dpdata for automatic conversion. " "When `format` requests conversion from a non-DeePMD format, this key " - "defaults to `lmdb`. Use a DeePMD format supported by dpdata, such as " - "`lmdb`, `deepmd/hdf5`, or `deepmd/npy`." + "defaults to `deepmd/lmdb`. Use a DeePMD format supported by dpdata, " + "such as `deepmd/lmdb`, `deepmd/hdf5`, or `deepmd/npy`." ) doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ @@ -5176,7 +5176,7 @@ def validation_data_args() -> list[ "out_format", [str, None], optional=True, - default="lmdb", + default="deepmd/lmdb", doc=doc_out_format, alias=["output_format"], ), diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 1bfe4f09fe..827ab9d507 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -38,7 +38,7 @@ log = logging.getLogger(__name__) _DPDATA_CACHE_DIR = ".deepmd_dpdata_cache" -_DPDATA_DEFAULT_OUT_FORMAT = "lmdb" +_DPDATA_DEFAULT_OUT_FORMAT = "deepmd/lmdb" _DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str], list[str]] = {} @@ -1118,10 +1118,16 @@ def _is_deepmd_data_format(fmt: str) -> bool: "deepmd/comp", "deepmd/npy/mixed", "deepmd/hdf5", + "deepmd/lmdb", "lmdb", } +def _is_dpdata_lmdb_format(fmt: str) -> bool: + """Return whether *fmt* names dpdata's DeePMD-compatible LMDB format.""" + return fmt in {"deepmd/lmdb", "lmdb"} + + def _looks_like_extxyz(path: Path) -> bool: if not path.is_file(): return False @@ -1167,7 +1173,7 @@ def _conversion_cache_path(source: Path, fmt: str, out_fmt: str) -> Path: ] stem = source_resolved.stem or source_resolved.name or "dataset" safe_out_fmt = out_fmt.replace("/", "-") - suffix = ".lmdb" if out_fmt == "lmdb" else "" + suffix = ".lmdb" if _is_dpdata_lmdb_format(out_fmt) else "" return Path.cwd() / _DPDATA_CACHE_DIR / f"{stem}-{safe_out_fmt}-{digest}{suffix}" @@ -1219,6 +1225,14 @@ def _remove_path(path: Path) -> None: def _write_dpdata_conversion( source: Path, fmt: str, out_fmt: str, output: Path ) -> None: + """Load *source* with dpdata and publish it in a DeePMD format. + + dpdata 1.1.0 makes ``deepmd/lmdb`` writes transactional: it stages and + validates the complete database before atomically publishing it. Use that + writer directly so DeepMD-kit does not duplicate or weaken its overwrite + guarantees. Other dpdata formats do not share that contract, so they keep + the cache-level temporary output used for failure isolation. + """ try: import dpdata except ImportError as exc: @@ -1228,17 +1242,22 @@ def _write_dpdata_conversion( "automatic dataset conversion." ) from exc + multi_systems = dpdata.MultiSystems() + try: + multi_systems.load_systems_from_file(str(source), fmt=fmt) + except NotImplementedError: + labeled_system = dpdata.LabeledSystem(str(source), fmt=fmt) + multi_systems = dpdata.MultiSystems(labeled_system) + if len(multi_systems) == 0: + raise RuntimeError(f"No frames were loaded by dpdata from {source}") + + if _is_dpdata_lmdb_format(out_fmt): + multi_systems.to(out_fmt, str(output), overwrite=True) + return + tmp_output = output.with_name(f".{output.name}.{os.getpid()}.tmp") _remove_path(tmp_output) try: - multi_systems = dpdata.MultiSystems() - try: - multi_systems.load_systems_from_file(str(source), fmt=fmt) - except NotImplementedError: - labeled_system = dpdata.LabeledSystem(str(source), fmt=fmt) - multi_systems = dpdata.MultiSystems(labeled_system) - if len(multi_systems) == 0: - raise RuntimeError(f"No frames were loaded by dpdata from {source}") multi_systems.to(out_fmt, str(tmp_output)) _remove_path(output) os.replace(tmp_output, output) @@ -1255,6 +1274,10 @@ def _convert_system_by_dpdata( source = Path(source_path) fmt = _normalize_dpdata_format(fmt, source) out_fmt = out_fmt.lower() + if out_fmt == "lmdb": + # dpdata keeps ``lmdb`` as an alias, but ``deepmd/lmdb`` is the + # canonical name for the DeePMD-compatible schema since dpdata 1.1.0. + out_fmt = _DPDATA_DEFAULT_OUT_FORMAT cache_key = ( str(Path.cwd().resolve(strict=False)), str(source.resolve(strict=False)), @@ -1305,7 +1328,7 @@ def _convert_system_by_dpdata( pass break - if out_fmt == "lmdb": + if _is_dpdata_lmdb_format(out_fmt): converted_systems = [str(output)] else: converted_systems = expand_sys_str(str(output)) @@ -1339,8 +1362,8 @@ def process_systems( fmt : str, optional The dpdata input format. If None, no conversion is performed. out_fmt : str, optional - The dpdata output format. If None, ``lmdb`` is used when fmt triggers - conversion. + The dpdata output format. If None, ``deepmd/lmdb`` is used when fmt + triggers conversion. Returns ------- diff --git a/pyproject.toml b/pyproject.toml index 8787db3952..7098b188a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,7 @@ dependencies = [ 'array-api-compat', 'lmdb', 'msgpack', - 'dpdata>=1.0.1', + 'dpdata>=1.1.0', ] requires-python = ">=3.10" keywords = ["deepmd"] diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index 856bd28df7..2c95e8224b 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -84,6 +84,7 @@ def _write_minimal_lmdb(path: str) -> None: class _FakeMultiSystems: write_count = 0 load_calls: ClassVar[list[tuple[str, str]]] = [] + to_calls: ClassVar[list[tuple[str, str, dict]]] = [] def __init__(self, *systems) -> None: self.systems = list(systems) @@ -97,11 +98,12 @@ def load_systems_from_file(self, file_name: str, fmt: str): def __len__(self) -> int: return 1 if self.loaded or self.systems else 0 - def to(self, fmt: str, file_name: str) -> None: + def to(self, fmt: str, file_name: str, **kwargs) -> None: type(self).write_count += 1 + self.to_calls.append((fmt, file_name, kwargs)) if fmt == "deepmd/hdf5": _write_minimal_deepmd_hdf5(file_name) - elif fmt == "lmdb": + elif fmt == "deepmd/lmdb": _write_minimal_lmdb(file_name) else: raise AssertionError(fmt) @@ -123,6 +125,7 @@ def setUp(self) -> None: self.source.write_text("1\nProperties=species:S:1:pos:R:3\nH 0 0 0\n") _FakeMultiSystems.write_count = 0 _FakeMultiSystems.load_calls = [] + _FakeMultiSystems.to_calls = [] data_system._DPDATA_CONVERSION_CACHE.clear() self.fake_dpdata = types.SimpleNamespace( MultiSystems=_FakeMultiSystems, @@ -134,7 +137,7 @@ def tearDown(self) -> None: self.tmpdir.cleanup() data_system._DPDATA_CONVERSION_CACHE.clear() - def test_process_systems_defaults_to_lmdb_and_reuses_cache(self) -> None: + def test_process_systems_defaults_to_deepmd_lmdb_and_reuses_cache(self) -> None: with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): systems = process_systems(str(self.source), fmt="extxyz") systems_again = process_systems(str(self.source), fmt="extxyz") @@ -147,6 +150,10 @@ def test_process_systems_defaults_to_lmdb_and_reuses_cache(self) -> None: self.assertTrue(is_lmdb(systems[0])) self.assertTrue(Path(systems[0]).is_relative_to(self.root)) self.assertEqual(Path(systems[0]).parent, self.root / ".deepmd_dpdata_cache") + self.assertEqual( + _FakeMultiSystems.to_calls, + [("deepmd/lmdb", systems[0], {"overwrite": True})], + ) def test_process_systems_cache_is_scoped_to_cwd(self) -> None: other_cwd = self.root / "run2" @@ -176,6 +183,46 @@ def test_process_systems_revalidates_in_memory_cache(self) -> None: self.assertEqual(systems, systems_again) self.assertEqual(_FakeMultiSystems.write_count, 2) + self.assertTrue( + all( + call == ("deepmd/lmdb", systems[0], {"overwrite": True}) + for call in _FakeMultiSystems.to_calls + ) + ) + + def test_lmdb_alias_uses_canonical_dpdata_writer(self) -> None: + """The legacy alias must share the canonical dpdata cache entry.""" + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems(str(self.source), fmt="extxyz", out_fmt="lmdb") + + self.assertEqual( + _FakeMultiSystems.to_calls, + [("deepmd/lmdb", systems[0], {"overwrite": True})], + ) + + def test_real_dpdata_lmdb_writer_is_deepmd_compatible(self) -> None: + """Verify dpdata 1.1 writes the schema consumed by DeepMD's LMDB reader.""" + extxyz = ( + "1\n" + "Properties=species:S:1:pos:R:3:forces:R:3 energy=0.0 " + 'Lattice="8 0 0 0 8 0 0 0 8"\n' + "H 0 0 0 0 0 0\n" + ) + self.source.write_text(extxyz) + + systems = process_systems(str(self.source), fmt="extxyz") + # A changed source exercises dpdata's own transactional overwrite path; + # DeepMD-kit must not remove or rename the LMDB directory around it. + self.source.write_text(extxyz.replace("energy=0.0", "energy=1.0")) + systems_again = process_systems(str(self.source), fmt="extxyz") + + self.assertEqual(len(systems), 1) + self.assertEqual(systems_again, systems) + self.assertTrue(is_lmdb(systems[0])) + data = LmdbDataSystem(systems[0], ["H"], batch_size=1) + batch = data.get_batch() + self.assertEqual(batch["coord"].shape, (1, 3)) + self.assertEqual(batch["type"].shape, (1, 1)) def test_cache_cleanup_unlinks_directory_symlink(self) -> None: """Cleanup must not recurse through a symlink outside the cache.""" From b3a2715a923e918b0c9aa6aeef6c94e628f2b69b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sun, 30 Aug 2026 02:36:32 +0800 Subject: [PATCH 08/12] fix(data): harden dpdata LMDB conversion Fix legacy LMDB requirement registration, bounded statistics, full validation, DDP routing, sampling validation, conversion locking, publication rollback, and resource cleanup. Coding-Agent: Codex Codex-Version: codex-cli 0.151.0 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/lmdb_data.py | 141 ++-- deepmd/entrypoints/test.py | 14 +- deepmd/jax/entrypoints/train.py | 23 +- deepmd/pd/entrypoints/main.py | 11 + deepmd/pt/entrypoints/main.py | 14 +- deepmd/pt_expt/entrypoints/main.py | 11 + deepmd/tf/entrypoints/train.py | 53 +- deepmd/tf2/entrypoints/train.py | 27 +- deepmd/utils/data_system.py | 650 ++++++++++++++---- source/tests/common/dpmodel/test_lmdb_data.py | 7 + .../common/test_data_system_conversion.py | 395 ++++++++++- source/tests/pt_expt/test_lmdb_training.py | 36 + 12 files changed, 1173 insertions(+), 209 deletions(-) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 389ca09089..759ddf3bc9 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -393,7 +393,9 @@ def _raw_frame_availability( explicit_true |= bit else: bit = key_bits.get(name) - if bit is not None: + if bit is not None and ( + name != "box" or not np.allclose(_decode_value(value), 0.0) + ): present |= bit return present & (~explicit_known | explicit_true) @@ -737,6 +739,11 @@ def _frame_source_available(frame: dict[str, Any], key: str) -> bool: source_present = _is_encoded_array(value) or isinstance( value, (np.ndarray, np.generic, int, float, bool) ) + if source_present and key == "box": + # A zero cell is DeePMD's non-periodic sentinel, not an available + # periodic box. Treating it as present would mix PBC and non-PBC + # frames under one scalar find_box flag. + source_present = not np.allclose(_decode_value(value), 0.0) find_key = f"find_{key}" if find_key in frame: return source_present and bool( @@ -1970,8 +1977,17 @@ def __init__( self, lmdb_path: str, type_map: list[str], - batch_size: int | str = "auto", + batch_size: int | str | Sequence[int | str] = "auto", ) -> None: + if isinstance(batch_size, (Sequence, np.ndarray)) and not isinstance( + batch_size, str + ): + if len(batch_size) != 1: + raise ValueError( + "One LMDB path is one training dataset and therefore " + "requires exactly one batch_size value." + ) + batch_size = batch_size[0] self.lmdb_path = str(Path(lmdb_path).resolve()) self._type_map = type_map # Read before opening the frame-serving environment, which disables @@ -2235,14 +2251,8 @@ def get_batch_size_for_nloc(self, nloc: int) -> int: def __len__(self) -> int: return self.nframes - def __getitem__(self, index: int) -> dict[str, Any]: - """Read frame from LMDB, decode, remap keys, return dict of numpy arrays. - - ``index`` is a dataset-level index in ``[0, len(self))``. Under - ``filter:N`` the LMDB key space may have gaps (dropped frames), so - we translate through ``self._retained_keys`` before hitting LMDB. - """ - self._data_requirements_frozen = True + def _read_frame(self, index: int) -> dict[str, Any]: + """Decode one frame without changing requirement-registration state.""" if index < 0 or index >= self.nframes: raise IndexError(f"dataset index {index} out of range [0, {self.nframes})") original_key = int(self._retained_keys[index]) @@ -2259,6 +2269,25 @@ def __getitem__(self, index: int) -> dict[str, Any]: copy_arrays=True, ) + def peek_frame(self, index: int) -> dict[str, Any]: + """Inspect one frame without freezing later requirement registration. + + Structural probes such as periodic-boundary detection happen before a + model supplies its label requirements. They may inspect a frame, but + must not turn that inspection into the first training read. + """ + return self._read_frame(index) + + def __getitem__(self, index: int) -> dict[str, Any]: + """Read and decode one frame, freezing the registered data contract. + + ``index`` is a dataset-level index in ``[0, len(self))``. Under + ``filter:N`` the LMDB key space may have gaps, which :meth:`_read_frame` + translates through ``self._retained_keys``. + """ + self._data_requirements_frozen = True + return self._read_frame(index) + def original_keys(self, indices: Sequence[int]) -> list[int]: """Translate dataset indices to original integer LMDB keys.""" keys: list[int] = [] @@ -2854,13 +2883,17 @@ def compute_block_targets( Each element is ``(system_indices_in_block, target_frame_count)``. Returns empty list if no expansion is needed (all targets == actual). """ - # Parse block definitions from the auto_prob string - # Format: "prob_sys_size;stt:end:weight;stt:end:weight;..." - block_str = auto_prob_style.split(";")[1:] + # ``prob_uniform`` is one equal-weight block per original system. The + # extended ``prob_sys_size`` form names arbitrary ranges explicitly. blocks: list[tuple[int, int, float]] = [] - for part in block_str: - stt, end, weight = part.split(":") - blocks.append((int(stt), int(end), float(weight))) + if auto_prob_style == "prob_uniform": + blocks = [(system_id, system_id + 1, 1.0) for system_id in range(nsystems)] + elif auto_prob_style.startswith("prob_sys_size"): + for part in auto_prob_style.split(";")[1:]: + stt, end, weight = part.split(":") + blocks.append((int(stt), int(end), float(weight))) + else: + raise RuntimeError(f"Unknown auto prob style: {auto_prob_style}") # A bare ``prob_sys_size`` names no blocks: it asks for a probability # proportional to system size, which is what sampling the merged frames @@ -2901,9 +2934,6 @@ def compute_block_targets( f"0 frames, likely after filter:N): {dropped}. Remaining block " "weights will be renormalised to sum to 1.0." ) - auto_prob_style = "prob_sys_size;" + ";".join( - f"{stt}:{end}:{weight}" for stt, end, weight in nonempty - ) blocks = nonempty # Compute the same per-system probabilities as prob_sys_size_ext locally. @@ -3688,22 +3718,29 @@ def make_neighbor_stat_data( ) reader = LmdbDataReader(lmdb_path, type_map=type_map) - nframes = len(reader) - rng = np.random.RandomState(42) - if nframes > max_frames: - indices = np.sort(rng.choice(nframes, max_frames, replace=False)) - else: - indices = np.arange(nframes, dtype=np.int64) - - # Read sampled frames, group by nloc - nloc_frames: dict[int, list[tuple[np.ndarray, np.ndarray, np.ndarray | None]]] = {} - for idx in indices: - frame = reader[int(idx)] - atype = frame["atype"] - nloc = len(atype) - nloc_frames.setdefault(nloc, []).append( - (frame["coord"], atype, frame.get("box")) - ) + try: + nframes = len(reader) + rng = np.random.RandomState(42) + if nframes > max_frames: + indices = np.sort(rng.choice(nframes, max_frames, replace=False)) + else: + indices = np.arange(nframes, dtype=np.int64) + + # The copied arrays remain valid after the reader is closed, so this + # helper does not leave an LMDB transaction or mmap alive in callers. + nloc_frames: dict[ + int, list[tuple[np.ndarray, np.ndarray, np.ndarray | None]] + ] = {} + for idx in indices: + frame = reader[int(idx)] + atype = frame["atype"] + nloc = len(atype) + nloc_frames.setdefault(nloc, []).append( + (frame["coord"], atype, frame.get("box")) + ) + ntypes = len(type_map) if type_map else reader._ntypes + finally: + reader.close() # Build per-nloc data_system proxies data_systems = [] @@ -3725,7 +3762,6 @@ def make_neighbor_stat_data( data_systems.append(proxy) system_dirs.append(label) - ntypes = len(type_map) if type_map else reader._ntypes return SimpleNamespace( system_dirs=system_dirs, data_systems=data_systems, @@ -3903,19 +3939,17 @@ def _read_frames(self, frame_indices: Sequence[int]) -> list[dict[str, Any]]: ) return frames - def __del__(self) -> None: - """Release the LMDB environment ref-count on garbage collection. - - The count is released only once, and only if construction got as far - as taking it: an instance that failed earlier holds no reference, and - releasing one it never took would close the environment underneath - whichever reader does hold it. - """ + def close(self) -> None: + """Release the LMDB environment ref-count idempotently.""" if getattr(self, "_env", None) is None: return self._env = None _close_lmdb(self.lmdb_path) + def __del__(self) -> None: + """Release the LMDB environment ref-count on garbage collection.""" + self.close() + @property def nloc_groups(self) -> dict[int, np.ndarray]: """Nloc → the LMDB frame indices retained for that atom count.""" @@ -4278,14 +4312,33 @@ def __init__( lmdb_test_data: "LmdbTestData", nloc: int, frame_indices: Sequence[int] | None = None, + *, + pbc: bool | None = None, + stat_groups: dict[str, Sequence[int]] | None = None, ) -> None: self._inner = lmdb_test_data self._nloc = nloc self._frame_indices = frame_indices + self._pbc = pbc + self._stat_groups = stat_groups or {} + self.dirs = list(self._stat_groups) def __getattr__(self, name: str) -> Any: return getattr(self._inner, name) + @property + def pbc(self) -> bool: + """Whether every frame represented by this view is periodic.""" + return self._inner.pbc if self._pbc is None else self._pbc + + def get_natoms(self) -> int: + """Return the fixed atom count of this stack-compatible view.""" + return self._nloc + + def _load_set(self, set_name: str) -> dict[str, Any]: + """Load one bounded neighbor-stat chunk from this view.""" + return self._inner.get_test_by_indices(self._stat_groups[str(set_name)]) + def get_test(self) -> dict[str, Any]: if self._frame_indices is not None: return self._inner.get_test_by_indices(self._frame_indices) diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index d9972b967e..c08751cbf4 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -111,7 +111,12 @@ def test( else: systems = [str((root / Path(ss)).resolve()) for ss in systems] patterns = data_params.get("rglob_patterns", None) - all_sys = process_systems(systems, patterns=patterns) + all_sys = process_systems( + systems, + patterns=patterns, + fmt=data_params.get("format"), + out_fmt=data_params.get("out_format", data_params.get("output_format")), + ) elif valid_json is not None: jdata = j_loader(valid_json) jdata = update_deepmd_input(jdata) @@ -125,7 +130,12 @@ def test( else: systems = [str((root / Path(ss)).resolve()) for ss in systems] patterns = data_params.get("rglob_patterns", None) - all_sys = process_systems(systems, patterns=patterns) + all_sys = process_systems( + systems, + patterns=patterns, + fmt=data_params.get("format"), + out_fmt=data_params.get("out_format", data_params.get("output_format")), + ) elif datafile is not None: with open(datafile) as datalist: all_sys = datalist.read().splitlines() diff --git a/deepmd/jax/entrypoints/train.py b/deepmd/jax/entrypoints/train.py index 61d415a340..1cad2af5bc 100644 --- a/deepmd/jax/entrypoints/train.py +++ b/deepmd/jax/entrypoints/train.py @@ -36,6 +36,7 @@ ) from deepmd.utils import random as dp_random from deepmd.utils.data_system import ( + close_data_systems, get_data, ) from deepmd.utils.summary import SummaryPrinter as BaseSummaryPrinter @@ -201,11 +202,14 @@ def factory( train_data_map, valid_data_map, _ = make_task_maps(config, factory) print_data_summaries(train_data_map, valid_data_map) - start_time = time.time() - model.train(train_data_map, valid_data_map) - end_time = time.time() - log.info("finished training") - log.info(f"wall time: {(end_time - start_time):.3f} s") + try: + start_time = time.time() + model.train(train_data_map, valid_data_map) + end_time = time.time() + log.info("finished training") + log.info(f"wall time: {(end_time - start_time):.3f} s") + finally: + close_data_systems(train_data_map, valid_data_map) def train( @@ -296,9 +300,12 @@ def update_sel( type_map, None, # not used ) - updated_model, task_min_nbor_dist = BaseModel.update_sel( - train_data, type_map, dict(task_config.model_params) - ) + try: + updated_model, task_min_nbor_dist = BaseModel.update_sel( + train_data, type_map, dict(task_config.model_params) + ) + finally: + close_data_systems(train_data) if multi_task: jdata_cpy["model"]["model_dict"][task_config.key] = updated_model min_nbor_dist[task_config.key] = task_min_nbor_dist diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index 15ed2371a4..60698416a7 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -69,6 +69,7 @@ from deepmd.utils.data_system import ( get_data, process_systems, + validate_backend_data_config, validate_lmdb_systems, ) from deepmd.utils.path import ( @@ -109,6 +110,11 @@ def prepare_trainer_input_single( validation_dataset_params["systems"] if validation_dataset_params else None ) training_systems = training_dataset_params["systems"] + validate_backend_data_config( + training_dataset_params, + backend_name="Paddle", + lmdb_supported=False, + ) trn_patterns = training_dataset_params.get("rglob_patterns", None) training_systems = process_systems( training_systems, @@ -120,6 +126,11 @@ def prepare_trainer_input_single( ) validate_lmdb_systems(training_systems, backend_name="Paddle", supported=False) if validation_systems is not None: + validate_backend_data_config( + validation_dataset_params, + backend_name="Paddle", + lmdb_supported=False, + ) val_patterns = validation_dataset_params.get("rglob_patterns", None) validation_systems = process_systems( validation_systems, diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index d8e03bbe6f..08a3a183c3 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -88,8 +88,10 @@ update_deepmd_input, ) from deepmd.utils.data_system import ( + conversion_will_write_lmdb, get_data, process_systems, + validate_lmdb_sampling_options, validate_lmdb_systems, ) from deepmd.utils.stat_file import ( @@ -179,6 +181,8 @@ def _make_dp_loader_set( dataset_params: dict[str, Any], ) -> DpLoaderSet | LmdbDataset: """Create a dataset from systems with pattern expansion/conversion.""" + if conversion_will_write_lmdb(dataset_params): + validate_lmdb_sampling_options(dataset_params) patterns = dataset_params.get("rglob_patterns") systems = process_systems( systems, @@ -190,6 +194,7 @@ def _make_dp_loader_set( ) lmdb_path = validate_lmdb_systems(systems, backend_name="PyTorch") if lmdb_path is not None: + validate_lmdb_sampling_options(dataset_params) return LmdbDataset( lmdb_path, model_params_single["type_map"], @@ -210,6 +215,7 @@ def _make_dp_loader_set( and isinstance(training_systems, str) and is_lmdb(training_systems) ): + validate_lmdb_sampling_options(training_dataset_params) auto_prob = training_dataset_params.get("auto_prob", None) train_data_single = LmdbDataset( training_systems, @@ -223,6 +229,7 @@ def _make_dp_loader_set( and isinstance(validation_systems, str) and is_lmdb(validation_systems) ): + validate_lmdb_sampling_options(validation_dataset_params) validation_data_single = LmdbDataset( validation_systems, model_params_single["type_map"], @@ -416,12 +423,15 @@ def _get_neighbor_stat_data_from_params( type_map: list[str] | None, ) -> Any: training_systems = dataset_params.get("systems") - if ( + direct_lmdb = ( dataset_params.get("format") is None and training_systems is not None and isinstance(training_systems, str) and is_lmdb(training_systems) - ): + ) + if direct_lmdb or conversion_will_write_lmdb(dataset_params): + validate_lmdb_sampling_options(dataset_params) + if direct_lmdb: systems = [training_systems] else: systems = process_systems( diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 85c94937ff..1381750e0d 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -33,8 +33,10 @@ ) from deepmd.utils.data_system import ( DeepmdDataSystem, + conversion_will_write_lmdb, get_data, process_systems, + validate_lmdb_sampling_options, validate_lmdb_systems, ) from deepmd.utils.stat_file import ( @@ -123,11 +125,14 @@ def _get_neighbor_stat_data( else _detect_lmdb_path(dataset_params.get("systems")) ) if lmdb_path is not None: + validate_lmdb_sampling_options(dataset_params) from deepmd.dpmodel.utils.lmdb_data import ( make_neighbor_stat_data, ) return make_neighbor_stat_data(lmdb_path, type_map) + if conversion_will_write_lmdb(dataset_params): + validate_lmdb_sampling_options(dataset_params) systems = process_systems( dataset_params["systems"], patterns=dataset_params.get("rglob_patterns"), @@ -168,6 +173,7 @@ def _build_data_system( else _detect_lmdb_path(systems_raw) ) if lmdb_path is not None: + validate_lmdb_sampling_options(dataset_params) return LmdbDataSystem( lmdb_path=lmdb_path, type_map=type_map, @@ -177,6 +183,8 @@ def _build_data_system( rank=rank, world_size=world_size, ) + if conversion_will_write_lmdb(dataset_params): + validate_lmdb_sampling_options(dataset_params) systems = process_systems( systems_raw, patterns=dataset_params.get("rglob_patterns"), @@ -187,12 +195,15 @@ def _build_data_system( systems, backend_name="PyTorch exportable" ) if converted_lmdb_path is not None: + validate_lmdb_sampling_options(dataset_params) return LmdbDataSystem( lmdb_path=converted_lmdb_path, type_map=type_map, batch_size=dataset_params["batch_size"], auto_prob_style=dataset_params.get("auto_prob"), seed=seed, + rank=rank, + world_size=world_size, ) return DeepmdDataSystem( systems=systems, diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index d7b40ebfee..2dc1feba9f 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -52,6 +52,7 @@ replace_model_params_with_pretrained_model, ) from deepmd.utils.data_system import ( + close_data_systems, get_data, ) from deepmd.utils.path import ( @@ -308,26 +309,31 @@ def _do_work( if ( origin_type_map is not None and not origin_type_map ): # get the type_map from data if not provided - origin_type_map = get_data( - jdata["training"]["training_data"], rcut, None, modifier - ).get_type_map() - model.build( - train_data, - stop_batch, - origin_type_map=origin_type_map, - stat_file_path=stat_file_path, - ) + origin_data = get_data(jdata["training"]["training_data"], rcut, None, modifier) + try: + origin_type_map = origin_data.get_type_map() + finally: + close_data_systems(origin_data) + try: + model.build( + train_data, + stop_batch, + origin_type_map=origin_type_map, + stat_file_path=stat_file_path, + ) - if not is_compress: - # train the model with the provided systems in a cyclic way - start_time = time.time() - model.train(train_data, valid_data) - end_time = time.time() - log.info("finished training") - log.info(f"wall time: {(end_time - start_time):.3f} s") - else: - model.save_compressed() - log.info("finished compressing") + if not is_compress: + # train the model with the provided systems in a cyclic way + start_time = time.time() + model.train(train_data, valid_data) + end_time = time.time() + log.info("finished training") + log.info(f"wall time: {(end_time - start_time):.3f} s") + else: + model.save_compressed() + log.info("finished compressing") + finally: + close_data_systems(train_data, valid_data) def get_modifier(modi_data: dict | None = None) -> BaseModifier | None: @@ -355,9 +361,12 @@ def update_sel(jdata: dict) -> dict: type_map, None, # not used ) - jdata_cpy["model"], min_nbor_dist = Model.update_sel( - train_data, type_map, jdata["model"] - ) + try: + jdata_cpy["model"], min_nbor_dist = Model.update_sel( + train_data, type_map, jdata["model"] + ) + finally: + close_data_systems(train_data) if min_nbor_dist is not None: tf.constant( diff --git a/deepmd/tf2/entrypoints/train.py b/deepmd/tf2/entrypoints/train.py index a026bcc03f..941bca8984 100644 --- a/deepmd/tf2/entrypoints/train.py +++ b/deepmd/tf2/entrypoints/train.py @@ -34,6 +34,7 @@ ) from deepmd.utils import random as dp_random from deepmd.utils.data_system import ( + close_data_systems, get_data, ) from deepmd.utils.summary import SummaryPrinter as BaseSummaryPrinter @@ -192,11 +193,14 @@ def factory( shared_links=self.shared_links, min_nbor_dist=neighbor_stat, ) - start_time = time.time() - trainer.run() - end_time = time.time() - log.info("finished training") - log.info("wall time: %.3f s", end_time - start_time) + try: + start_time = time.time() + trainer.run() + end_time = time.time() + log.info("finished training") + log.info("wall time: %.3f s", end_time - start_time) + finally: + close_data_systems(train_data_map, valid_data_map) def train( @@ -249,11 +253,14 @@ def update_sel( type_map, None, ) - updated_model, task_min_nbor_dist = BaseModel.update_sel( - train_data, - type_map, - dict(task_config.model_params), - ) + try: + updated_model, task_min_nbor_dist = BaseModel.update_sel( + train_data, + type_map, + dict(task_config.model_params), + ) + finally: + close_data_systems(train_data) min_nbor_dist[task_config.key] = task_min_nbor_dist if multi_task: jdata_cpy["model"]["model_dict"][task_config.key] = updated_model diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 8720315ea5..90ec9b31da 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import collections import hashlib +import importlib.metadata +import json import logging import os import shutil +import socket +import threading import time import warnings from functools import ( @@ -14,6 +18,7 @@ ) from typing import ( Any, + Self, ) import numpy as np @@ -39,7 +44,15 @@ _DPDATA_CACHE_DIR = ".deepmd_dpdata_cache" _DPDATA_DEFAULT_OUT_FORMAT = "deepmd/lmdb" -_DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str], list[str]] = {} +_DPDATA_CONVERSION_SCHEMA_VERSION = "2" +_DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str, str], list[str]] = {} +_DPDATA_SOURCE_MTIME_CACHE: dict[tuple[str, str], tuple[float, float]] = {} +# Neighbor-stat, trainer construction, and multi-task routing commonly query +# the same directory within one startup. Reuse that O(file-count) scan while +# still letting a long-lived process notice later source rewrites. +_DPDATA_SOURCE_MTIME_CACHE_TTL = 60.0 +_CONVERSION_LOCK_HEARTBEAT_SECONDS = 5.0 +_CONVERSION_LOCK_STALE_SECONDS = 30.0 def validate_lmdb_systems( @@ -758,7 +771,7 @@ def __init__( self, lmdb_path: str, type_map: list[str], - batch_size: int | str = "auto", + batch_size: int | str | list[int | str] = "auto", auto_prob_style: str | None = None, seed: int | None = None, ) -> None: @@ -767,6 +780,7 @@ def __init__( from deepmd.dpmodel.utils.lmdb_data import ( LmdbBatchSampler, LmdbDataReader, + LmdbTestData, compute_block_targets, ) @@ -778,32 +792,16 @@ def __init__( ) self.lmdb_path = lmdb_path - self._reader = LmdbDataReader(lmdb_path, type_map, batch_size) self._type_map = list(type_map) - # LMDB is defined as mixed-type by its reader contract; determining - # this must not scan every frame during data-system initialization. - self.mixed_type = self._reader.mixed_type - self.nsystems = 1 - self.system_dirs = [lmdb_path] - self.natoms = [ - int(self._reader.frame_nlocs.max()) if len(self._reader.frame_nlocs) else 0 - ] - self.batch_size = [self._reader.batch_size] - self.sys_probs = [1.0] - self.data_systems = [self] - self.pbc = self._detect_pbc() + self._closed = False self._data_dict = { - "box": { - "ndof": 9, - "atomic": False, - "must": False, - "high_prec": False, - "type_sel": None, - "repeat": 1, - "default": 0.0, - "dtype": None, - "output_natoms_for_type_sel": False, - }, + "box": DataRequirementItem( + "box", + 9, + atomic=False, + must=False, + default=np.zeros(9, dtype=GLOBAL_NP_FLOAT_PRECISION), + ).dict, "coord": { "ndof": 3, "atomic": True, @@ -828,6 +826,34 @@ def __init__( }, } + self._reader = LmdbDataReader(lmdb_path, type_map, batch_size) + # Box availability is part of stack compatibility. Register it before + # any grouping so periodic and non-periodic frames never share the + # scalar ``find_box`` flag of one legacy batch. + box_requirement = DataRequirementItem( + "box", + 9, + atomic=False, + must=False, + default=np.zeros(9, dtype=GLOBAL_NP_FLOAT_PRECISION), + ) + self._reader.add_data_requirement([box_requirement]) + self._test_data = LmdbTestData( + lmdb_path, + type_map=type_map, + shuffle_test=False, + ) + self._test_data.add_data_requirement([box_requirement]) + # LMDB is defined as mixed-type by its reader contract; determining + # this must not scan every frame during data-system initialization. + self.mixed_type = self._reader.mixed_type + self.nsystems = 1 + self.natoms = [ + int(self._reader.frame_nlocs.max()) if len(self._reader.frame_nlocs) else 0 + ] + self.batch_size = [self._reader.batch_size] + self.sys_probs = [1.0] + block_targets = None if auto_prob_style is not None and self._reader.frame_system_ids is not None: block_targets = compute_block_targets( @@ -843,36 +869,94 @@ def __init__( ) self.nbatches = [self._sampler.total_batches] self._iter = iter(self._sampler) - self._refresh_set_indices() + self._refresh_groups() - def _refresh_set_indices(self) -> None: - """Refresh stack-compatible groups after data requirements change.""" + def _refresh_groups(self) -> None: + """Refresh bounded statistics chunks and full-validation views.""" from deepmd.dpmodel.utils.lmdb_data import ( + LmdbTestDataNlocView, collect_lmdb_sampling_groups, ) groups = collect_lmdb_sampling_groups(self._reader) - self._nloc_set_indices = { - f"{self.lmdb_path}#group={group_idx}:nloc={nloc}": indices - for group_idx, (nloc, indices) in enumerate(groups) - } + self._stat_groups = groups + self._stat_offsets = [0] * len(groups) + + # Neighbor statistics are a bounded sample, matching the dedicated + # LMDB path. Chunks additionally cap decoded atoms so one large-nloc + # group cannot create a large transient Python/NumPy allocation. + selected = np.zeros(len(self._reader), dtype=bool) + max_frames = min(len(self._reader), 2000) + if max_frames: + rng = np.random.RandomState(42) + chosen = ( + rng.choice(len(self._reader), max_frames, replace=False) + if max_frames < len(self._reader) + else np.arange(len(self._reader), dtype=np.int64) + ) + selected[np.asarray(chosen, dtype=np.int64)] = True + + self._nloc_set_indices: dict[str, np.ndarray] = {} + data_systems = [] + system_dirs: list[str] = [] + any_pbc = False + for group_idx, (nloc, indices) in enumerate(groups): + group_label = f"{self.lmdb_path}#group={group_idx}:nloc={nloc}" + original_indices = np.asarray( + self._reader.original_keys(indices), dtype=np.int64 + ) + frame = self._reader.peek_frame(int(indices[0])) + group_pbc = bool(float(frame.get("find_box", 0.0)) > 0.5) + any_pbc = any_pbc or group_pbc + + stat_groups: dict[str, np.ndarray] = {} + sampled_indices = np.asarray(indices)[selected[np.asarray(indices)]] + chunk_size = max(1, min(128, 20000 // max(int(nloc), 1))) + for chunk_idx, start in enumerate( + range(0, len(sampled_indices), chunk_size) + ): + chunk = sampled_indices[start : start + chunk_size] + set_name = f"{group_label}:chunk={chunk_idx}" + self._nloc_set_indices[set_name] = np.asarray(chunk, dtype=np.int64) + stat_groups[set_name] = np.asarray( + self._reader.original_keys(chunk), dtype=np.int64 + ) + + data_systems.append( + LmdbTestDataNlocView( + self._test_data, + int(nloc), + original_indices, + pbc=group_pbc, + stat_groups=stat_groups, + ) + ) + system_dirs.append(group_label) + + # These views do not point back to this adapter, avoiding the + # ``data_systems=[self]`` reference cycle while satisfying both the + # neighbor-stat and JAX/TF2 full-validation contracts. + self.data_systems = data_systems + self.system_dirs = system_dirs self.dirs = list(self._nloc_set_indices) + self.pbc = any_pbc def _detect_pbc(self) -> bool: """Return True when LMDB frames contain a non-zero simulation box.""" if len(self._reader) == 0: return False - box = self._reader[0].get("box") - return box is not None and not np.allclose(box, 0.0) + frame = self._reader.peek_frame(0) + return bool(float(frame.get("find_box", 0.0)) > 0.5) def add_data_requirements( self, data_requirements: list[DataRequirementItem] ) -> None: """Add label/auxiliary data requirements.""" + self._reader.add_data_requirement(data_requirements) + self._test_data.add_data_requirement(data_requirements) for item in data_requirements: self._data_dict[item.key] = item.dict - self._reader.add_data_requirement(data_requirements) - self._refresh_set_indices() + self._refresh_groups() self.nbatches = [self._sampler.total_batches] self._iter = iter(self._sampler) @@ -912,10 +996,9 @@ def get_data_dict(self, ii: int = 0) -> dict[str, dict[str, Any]]: return self._data_dict def _load_set(self, set_name: str) -> dict[str, Any]: - """Load one same-nloc LMDB group for legacy neighbor-stat code.""" + """Load one bounded same-nloc chunk for legacy neighbor statistics.""" indices = self._nloc_set_indices[str(set_name)] - frames = [self._reader[int(idx)] for idx in indices] - return self._stack_frames(frames) + return self._legacy_batch(self._reader.decode_batch(indices, ragged=False)) def _next_indices(self) -> list[int]: try: @@ -924,13 +1007,21 @@ def _next_indices(self) -> list[int]: self._iter = iter(self._sampler) return next(self._iter) - def _stack_frames(self, frames: list[dict[str, Any]]) -> dict[str, Any]: - if not frames: - raise ValueError("Cannot stack an empty LMDB frame batch.") + def _legacy_batch(self, batch: dict[str, Any]) -> dict[str, Any]: + """Translate a canonical LMDB batch to the legacy data-system shape.""" + coord = np.asarray(batch["coord"]) + nframes = coord.shape[0] out: dict[str, Any] = {} structural_keys = {"coord", "box"} - for key in frames[0]: - if key in {"atype", "fid", "natoms", "real_natoms_vec"}: + for key, value in batch.items(): + if key in { + "atype", + "natoms", + "real_natoms_vec", + "fid", + "sid", + "n_node", + }: continue if key.startswith("find_") and key[5:] not in self._data_dict: continue @@ -940,25 +1031,27 @@ def _stack_frames(self, frames: list[dict[str, Any]]) -> dict[str, Any]: and key not in self._data_dict ): continue - if key.startswith("find_"): - out[key] = np.asarray(frames[0][key], dtype=np.float32) - elif frames[0][key] is None: + if value is None: out[key] = None - else: - arr = np.stack([frame[key] for frame in frames]) - data_info = self._data_dict.get(key) - if data_info is not None and data_info["atomic"] and arr.ndim >= 3: - arr = arr.reshape(arr.shape[0], -1) - out[key] = arr - - atype = np.stack([frame["atype"] for frame in frames]).astype(np.int32) - real_natoms_vec = np.stack([frame["natoms"] for frame in frames]).astype( - np.int32 + continue + array = np.asarray(value) + data_info = self._data_dict.get(key) + if key == "coord" or ( + data_info is not None and data_info["atomic"] and array.ndim >= 3 + ): + array = array.reshape(nframes, -1) + out[key] = array + + atype = np.asarray(batch["atype"], dtype=np.int32) + real_natoms_vec = np.asarray( + batch.get("real_natoms_vec", batch["natoms"]), dtype=np.int32 ) - nloc = int(real_natoms_vec[:, 0].max()) + if real_natoms_vec.ndim == 1: + real_natoms_vec = np.tile(real_natoms_vec, (nframes, 1)) + pad_nloc = int(atype.shape[1]) natoms_vec = np.concatenate( ( - np.array([nloc, nloc], dtype=np.int32), + np.array([pad_nloc, pad_nloc], dtype=np.int32), real_natoms_vec[:, 2:].max(axis=0).astype(np.int32), ) ) @@ -967,14 +1060,13 @@ def _stack_frames(self, frames: list[dict[str, Any]]) -> dict[str, Any]: out["natoms_vec"] = natoms_vec out["real_natoms_vec"] = real_natoms_vec if "box" not in out or out["box"] is None: - out["box"] = np.zeros((len(frames), 9), dtype=GLOBAL_NP_FLOAT_PRECISION) + out["box"] = np.zeros((nframes, 9), dtype=GLOBAL_NP_FLOAT_PRECISION) out["find_box"] = np.float32(0.0) elif "find_box" not in out: out["find_box"] = np.float32(0.0 if np.allclose(out["box"], 0.0) else 1.0) - if "find_coord" not in out: - out["find_coord"] = np.float32(1.0) + out.setdefault("find_coord", np.float32(1.0)) if "numb_copy" not in out: - out["numb_copy"] = np.ones((len(frames), 1), dtype=np.int64) + out["numb_copy"] = np.ones((nframes, 1), dtype=np.int64) out["find_numb_copy"] = np.float32(0.0) out["default_mesh"] = np.asarray( make_default_mesh(bool(float(out["find_box"]) > 0.5), self.mixed_type), @@ -982,11 +1074,54 @@ def _stack_frames(self, frames: list[dict[str, Any]]) -> dict[str, Any]: ) return out + def _stack_frames(self, frames: list[dict[str, Any]]) -> dict[str, Any]: + """Collate already-decoded frames with the reader's flag semantics.""" + if not frames: + raise ValueError("Cannot stack an empty LMDB frame batch.") + from deepmd.dpmodel.utils.lmdb_data import ( + collate_lmdb_frames, + resolve_per_atom_keys, + ) + + per_atom_keys = resolve_per_atom_keys(frames[0], self._reader.decode_config) + return self._legacy_batch(collate_lmdb_frames(frames, per_atom_keys)) + def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: del sys_idx indices = self._next_indices() - frames = [self._reader[int(idx)] for idx in indices] - return self._stack_frames(frames) + return self._legacy_batch(self._reader.decode_batch(indices, ragged=False)) + + def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: + """Return one bounded batch from a homogeneous statistical group.""" + if not 0 <= sys_idx < len(self._stat_groups): + raise IndexError(f"Statistical system index {sys_idx} is out of range") + nloc, indices = self._stat_groups[sys_idx] + batch_size = self._get_stat_batch_size(nloc) + start = self._stat_offsets[sys_idx] + if start >= len(indices): + start = 0 + stop = min(start + batch_size, len(indices)) + self._stat_offsets[sys_idx] = stop + return self._legacy_batch( + self._reader.decode_batch(indices[start:stop], ragged=False) + ) + + def get_stat_nsystems(self) -> int: + """Return the number of stack-compatible statistical groups.""" + return len(self._stat_groups) + + def _get_stat_batch_size(self, nloc: int) -> int: + """Cap model-stat decoding by both frames and decoded atom rows.""" + configured = self._reader.get_batch_size_for_nloc(nloc) + return max(1, min(configured, 128, 20000 // max(int(nloc), 1))) + + def get_stat_numb_batches(self, sys_idx: int) -> int: + """Return the finite batch count of one statistical group.""" + if not 0 <= sys_idx < len(self._stat_groups): + raise IndexError(f"Statistical system index {sys_idx} is out of range") + nloc, indices = self._stat_groups[sys_idx] + batch_size = self._get_stat_batch_size(nloc) + return (len(indices) + batch_size - 1) // batch_size def get_nsystems(self) -> int: return self.nsystems @@ -1000,6 +1135,11 @@ def get_ntypes(self) -> int: def get_type_map(self) -> list[str]: return self._type_map + @property + def type_map(self) -> list[str]: + """Model-side atom names exposed by the legacy data-system API.""" + return self._type_map + def get_batch_size(self) -> list[int]: return self.batch_size @@ -1007,6 +1147,28 @@ def print_summary(self, name: str, prob: Any | None = None) -> None: del prob self._reader.print_summary(name, self.sys_probs) + def close(self) -> None: + """Release LMDB readers idempotently.""" + if getattr(self, "_closed", True): + return + self.data_systems = [] + test_data = getattr(self, "_test_data", None) + if test_data is not None: + test_data.close() + reader = getattr(self, "_reader", None) + if reader is not None: + reader.close() + self._closed = True + + def __enter__(self) -> Self: + return self + + def __exit__(self, *args: object) -> None: + self.close() + + def __del__(self) -> None: + self.close() + def _format_name_length(name: str, width: int) -> str: if len(name) <= width: @@ -1169,6 +1331,78 @@ def _is_dpdata_lmdb_format(fmt: str) -> bool: return fmt in {"deepmd/lmdb", "lmdb"} +def _canonical_dpdata_out_format(out_fmt: str | None) -> str: + """Return the canonical dpdata output format used by conversion caches.""" + if out_fmt is None: + return _DPDATA_DEFAULT_OUT_FORMAT + out_fmt = out_fmt.lower() + return _DPDATA_DEFAULT_OUT_FORMAT if out_fmt == "lmdb" else out_fmt + + +def conversion_will_write_lmdb(data_config: dict[str, Any]) -> bool: + """Whether a non-DeePMD input config will be converted to LMDB.""" + data_format = data_config.get("format") + if data_format is None or _is_deepmd_data_format(data_format.lower()): + return False + out_format = data_config.get("out_format", data_config.get("output_format")) + return _is_dpdata_lmdb_format(_canonical_dpdata_out_format(out_format)) + + +def validate_backend_data_config( + data_config: dict[str, Any], + *, + backend_name: str, + lmdb_supported: bool, +) -> None: + """Reject unsupported converted output before dpdata performs any I/O.""" + if not lmdb_supported and conversion_will_write_lmdb(data_config): + raise NotImplementedError( + f"{backend_name} backend does not support LMDB data yet. " + "Choose out_format='deepmd/hdf5' for automatic conversion." + ) + + +def validate_lmdb_sampling_options(data_config: dict[str, Any]) -> None: + """Reject sampling options that cannot be represented by one LMDB route.""" + if data_config.get("sys_probs") is not None: + raise ValueError( + "LMDB data does not support explicit sys_probs yet. Use auto_prob " + "('prob_sys_size', 'prob_uniform', or block weights) so sampling " + "can be derived from LMDB frame_system_ids." + ) + + +def close_data_systems(*values: Any) -> None: + """Close nested data-system mappings/sequences, ignoring shared objects.""" + seen: set[int] = set() + + def close_one(value: Any) -> None: + if value is None or id(value) in seen: + return + seen.add(id(value)) + if isinstance(value, dict): + for child in value.values(): + close_one(child) + return + if isinstance(value, (list, tuple)): + for child in value: + close_one(child) + return + close = getattr(value, "close", None) + if callable(close): + try: + close() + except Exception: + log.warning( + "Failed to close data system %r during cleanup", + value, + exc_info=True, + ) + + for value in values: + close_one(value) + + def _looks_like_extxyz(path: Path) -> bool: if not path.is_file(): return False @@ -1209,20 +1443,39 @@ def _iter_conversion_inputs(path: str, patterns: list[str] | None) -> list[str]: def _conversion_cache_path(source: Path, fmt: str, out_fmt: str) -> Path: source_resolved = source.resolve(strict=False) - digest = hashlib.sha1(f"{source_resolved}|{fmt}|{out_fmt}".encode()).hexdigest()[ - :16 - ] + try: + dpdata_version = importlib.metadata.version("dpdata") + except importlib.metadata.PackageNotFoundError: + dpdata_version = "unknown" + digest = hashlib.sha1( + ( + f"{source_resolved}|{fmt}|{out_fmt}|" + f"schema={_DPDATA_CONVERSION_SCHEMA_VERSION}|dpdata={dpdata_version}" + ).encode() + ).hexdigest()[:16] stem = source_resolved.stem or source_resolved.name or "dataset" safe_out_fmt = out_fmt.replace("/", "-") suffix = ".lmdb" if _is_dpdata_lmdb_format(out_fmt) else "" return Path.cwd() / _DPDATA_CACHE_DIR / f"{stem}-{safe_out_fmt}-{digest}{suffix}" -def _source_mtime(source: Path, cache_file: Path) -> float: +def _source_mtime(source: Path, cache_file: Path, *, force: bool = False) -> float: if source.is_file(): return source.stat().st_mtime if not source.is_dir(): return 0.0 + cache_key = ( + str(source.resolve(strict=False)), + str(cache_file.parent.resolve(strict=False)), + ) + now = time.monotonic() + cached = _DPDATA_SOURCE_MTIME_CACHE.get(cache_key) + if ( + not force + and cached is not None + and now - cached[0] < _DPDATA_SOURCE_MTIME_CACHE_TTL + ): + return cached[1] cache_dir = cache_file.parent.resolve(strict=False) latest = source.stat().st_mtime for item in source.rglob("*"): @@ -1233,23 +1486,143 @@ def _source_mtime(source: Path, cache_file: Path) -> float: latest = max(latest, item.stat().st_mtime) except OSError: continue + _DPDATA_SOURCE_MTIME_CACHE[cache_key] = (now, latest) return latest -def _is_conversion_current(source: Path, output: Path) -> bool: +def _is_conversion_current( + source: Path, output: Path, *, force_source_scan: bool = False +) -> bool: if not output.exists(): return False - return output.stat().st_mtime >= _source_mtime(source, output) + return output.stat().st_mtime >= _source_mtime( + source, output, force=force_source_scan + ) + + +def _process_start_time(pid: int) -> str | None: + """Return Linux's stable process start token, if available.""" + try: + stat_text = Path(f"/proc/{pid}/stat").read_text() + except OSError: + return None + fields_after_name = stat_text.rsplit(")", 1)[1].split() + return fields_after_name[19] if len(fields_after_name) > 19 else None + + +def _same_lock_file(lock_path: Path, expected: os.stat_result) -> bool: + """Whether *lock_path* still names the inode originally acquired/read.""" + try: + current = lock_path.stat(follow_symlinks=False) + except FileNotFoundError: + return False + return (current.st_dev, current.st_ino) == (expected.st_dev, expected.st_ino) + + +class _ConversionLock: + """Owned conversion lock with a heartbeat for cross-host stale recovery.""" + + def __init__(self, lock_path: Path, lock_fd: int) -> None: + self.path = lock_path + self._stat = os.fstat(lock_fd) + payload = { + "hostname": socket.gethostname(), + "pid": os.getpid(), + "process_start": _process_start_time(os.getpid()), + "created": time.time(), + } + with os.fdopen(lock_fd, "w") as fp: + json.dump(payload, fp) + fp.flush() + os.fsync(fp.fileno()) + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._heartbeat, + name="deepmd-dpdata-conversion-lock", + daemon=True, + ) + try: + self._thread.start() + except Exception: + if _same_lock_file(self.path, self._stat): + self.path.unlink(missing_ok=True) + raise + + def _heartbeat(self) -> None: + while not self._stop.wait(_CONVERSION_LOCK_HEARTBEAT_SECONDS): + if not _same_lock_file(self.path, self._stat): + return + try: + os.utime(self.path, None, follow_symlinks=False) + except FileNotFoundError: + return + + def release(self) -> None: + """Stop heartbeating and remove only the lock inode we own.""" + self._stop.set() + self._thread.join() + if _same_lock_file(self.path, self._stat): + try: + self.path.unlink() + except FileNotFoundError: + pass + + +def _lock_owner_is_alive(payload: dict[str, Any]) -> bool | None: + """Return owner liveness locally, or None for another host/invalid data.""" + if payload.get("hostname") != socket.gethostname(): + return None + try: + pid = int(payload["pid"]) + except (KeyError, TypeError, ValueError): + return None + expected_start = payload.get("process_start") + current_start = _process_start_time(pid) + if expected_start is not None and current_start is not None: + return str(expected_start) == current_start + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _recover_stale_conversion_lock(lock_path: Path) -> bool: + """Remove a dead-owner or expired-lease lock without touching a successor.""" + try: + lock_stat = lock_path.stat(follow_symlinks=False) + except FileNotFoundError: + return True + except OSError: + return False + try: + payload = json.loads(lock_path.read_text()) + except (OSError, json.JSONDecodeError): + payload = {} + + owner_alive = _lock_owner_is_alive(payload) + lease_expired = time.time() - lock_stat.st_mtime > _CONVERSION_LOCK_STALE_SECONDS + stale = owner_alive is False or (owner_alive is None and lease_expired) + if not stale or not _same_lock_file(lock_path, lock_stat): + return False + log.warning("Recovering stale dpdata conversion lock %s", lock_path) + try: + lock_path.unlink() + except FileNotFoundError: + pass + return True def _wait_for_conversion(source: Path, output: Path, lock_path: Path) -> bool: - for _ in range(300): - if not lock_path.exists(): - return _is_conversion_current(source, output) - if _is_conversion_current(source, output): - return True + """Wait without rescanning the source tree while a valid writer owns it.""" + while lock_path.exists(): + if _recover_stale_conversion_lock(lock_path): + continue time.sleep(1.0) - return False + # Freshness is checked once after publication, not once per waiter-second. + return _is_conversion_current(source, output, force_source_scan=True) def _remove_path(path: Path) -> None: @@ -1263,6 +1636,28 @@ def _remove_path(path: Path) -> None: path.unlink() +def _publish_conversion_output(tmp_output: Path, output: Path) -> None: + """Publish a non-LMDB conversion without discarding a valid old cache.""" + if not output.exists() or output.is_symlink() or not output.is_dir(): + os.replace(tmp_output, output) + backup = output.with_name(f".{output.name}.backup") + _remove_path(backup) + return + + # POSIX cannot replace a non-empty directory directly. Preserve the old + # cache under a sibling name and restore it if the second rename fails. + backup = output.with_name(f".{output.name}.backup") + _remove_path(backup) + os.replace(output, backup) + try: + os.replace(tmp_output, output) + except Exception: + os.replace(backup, output) + raise + else: + _remove_path(backup) + + def _write_dpdata_conversion( source: Path, fmt: str, out_fmt: str, output: Path ) -> None: @@ -1270,7 +1665,7 @@ def _write_dpdata_conversion( dpdata 1.1.0 makes ``deepmd/lmdb`` writes transactional: it stages and validates the complete database before atomically publishing it. Use that - writer directly so DeepMD-kit does not duplicate or weaken its overwrite + writer directly so DeePMD-kit does not duplicate or weaken its overwrite guarantees. Other dpdata formats do not share that contract, so they keep the cache-level temporary output used for failure isolation. """ @@ -1286,9 +1681,24 @@ def _write_dpdata_conversion( multi_systems = dpdata.MultiSystems() try: multi_systems.load_systems_from_file(str(source), fmt=fmt) - except NotImplementedError: - labeled_system = dpdata.LabeledSystem(str(source), fmt=fmt) - multi_systems = dpdata.MultiSystems(labeled_system) + except (NotImplementedError, ValueError) as labeled_error: + # dpdata 1.1 exposes an explicit unlabeled path. This matters for + # structure-only EXTXYZ/ASE inputs, which are valid descriptor data + # even though a supervised loss may later require labels. + unlabeled_systems = dpdata.MultiSystems() + try: + unlabeled_systems.load_systems_from_file( + str(source), fmt=fmt, labeled=False + ) + except (NotImplementedError, TypeError, ValueError): + try: + labeled_system = dpdata.LabeledSystem(str(source), fmt=fmt) + except Exception: + raise labeled_error from None + multi_systems = dpdata.MultiSystems(labeled_system) + else: + log.info("Loaded unlabeled dpdata input %s using format %s", source, fmt) + multi_systems = unlabeled_systems if len(multi_systems) == 0: raise RuntimeError(f"No frames were loaded by dpdata from {source}") @@ -1300,8 +1710,7 @@ def _write_dpdata_conversion( _remove_path(tmp_output) try: multi_systems.to(out_fmt, str(tmp_output)) - _remove_path(output) - os.replace(tmp_output, output) + _publish_conversion_output(tmp_output, output) except Exception: _remove_path(tmp_output) raise @@ -1310,22 +1719,17 @@ def _write_dpdata_conversion( def _convert_system_by_dpdata( source_path: str, fmt: str, out_fmt: str | None ) -> list[str]: - if out_fmt is None: - out_fmt = _DPDATA_DEFAULT_OUT_FORMAT source = Path(source_path) fmt = _normalize_dpdata_format(fmt, source) - out_fmt = out_fmt.lower() - if out_fmt == "lmdb": - # dpdata keeps ``lmdb`` as an alias, but ``deepmd/lmdb`` is the - # canonical name for the DeePMD-compatible schema since dpdata 1.1.0. - out_fmt = _DPDATA_DEFAULT_OUT_FORMAT + out_fmt = _canonical_dpdata_out_format(out_fmt) + output = _conversion_cache_path(source, fmt, out_fmt) cache_key = ( str(Path.cwd().resolve(strict=False)), str(source.resolve(strict=False)), fmt, out_fmt, + str(output), ) - output = _conversion_cache_path(source, fmt, out_fmt) cached_systems = _DPDATA_CONVERSION_CACHE.get(cache_key) if cached_systems is not None: if _is_conversion_current(source, output): @@ -1344,14 +1748,13 @@ def _convert_system_by_dpdata( except FileExistsError: if _wait_for_conversion(source, output, lock_path): break - raise TimeoutError( - f"Timed out waiting for dpdata conversion lock {lock_path}" - ) from None + continue else: - with os.fdopen(lock_fd, "w") as fp: - fp.write(str(os.getpid())) + conversion_lock = _ConversionLock(lock_path, lock_fd) try: - if not _is_conversion_current(source, output): + if not _is_conversion_current( + source, output, force_source_scan=True + ): log.info( "Converting %s from dpdata format %s to %s at %s", source, @@ -1361,12 +1764,7 @@ def _convert_system_by_dpdata( ) _write_dpdata_conversion(source, fmt, out_fmt, output) finally: - try: - lock_path.unlink() - except FileNotFoundError: - # Concurrent conversion cleanup is idempotent; another - # process may already have removed the shared lock. - pass + conversion_lock.release() break if _is_dpdata_lmdb_format(out_fmt): @@ -1433,22 +1831,35 @@ def process_systems( if _is_deepmd_data_format(fmt): fmt = None + conversion_inputs: list[str] = [] + if fmt is not None: + for path in search_paths: + conversion_inputs.extend(_iter_conversion_inputs(path, patterns)) + if ( + _is_dpdata_lmdb_format(_canonical_dpdata_out_format(out_fmt)) + and len(conversion_inputs) != 1 + ): + raise ValueError( + "Automatic LMDB conversion requires exactly one resolved input " + "path. Merge multiple inputs with dpdata first or choose " + "out_format='deepmd/hdf5'." + ) + # Iterate over the search_paths list and apply expansion logic to each path result_systems = [] - for path in search_paths: - if fmt is not None: - for input_path in _iter_conversion_inputs(path, patterns): - result_systems.extend( - _convert_system_by_dpdata(input_path, fmt, out_fmt) - ) - elif is_lmdb(path): - result_systems.append(path) - elif patterns is None: - expanded_paths = expand_sys_str(path) - result_systems.extend(expanded_paths) - else: - expanded_paths = rglob_sys_str(path, patterns) - result_systems.extend(expanded_paths) + if fmt is not None: + for input_path in conversion_inputs: + result_systems.extend(_convert_system_by_dpdata(input_path, fmt, out_fmt)) + else: + for path in search_paths: + if is_lmdb(path): + result_systems.append(path) + elif patterns is None: + expanded_paths = expand_sys_str(path) + result_systems.extend(expanded_paths) + else: + expanded_paths = rglob_sys_str(path, patterns) + result_systems.extend(expanded_paths) return result_systems @@ -1484,6 +1895,8 @@ def get_data( rglob_patterns = jdata.get("rglob_patterns") data_format = jdata.get("format") out_format = jdata.get("out_format", jdata.get("output_format")) + if conversion_will_write_lmdb(jdata): + validate_lmdb_sampling_options(jdata) systems = process_systems( systems, patterns=rglob_patterns, fmt=data_format, out_fmt=out_format ) @@ -1495,6 +1908,7 @@ def get_data( lmdb_path = validate_lmdb_systems(systems, backend_name="legacy data loader") if lmdb_path is not None: + validate_lmdb_sampling_options(jdata) if type_map is None: raise ValueError( "LMDB training data requires model/type_map to be set. " diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index afa2a79420..bc17b876d4 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -2021,6 +2021,13 @@ def test_compute_block_targets_unequal(self): self.assertEqual(result[0], ([0], 500)) self.assertEqual(result[1], ([1], 500)) + def test_compute_block_targets_prob_uniform(self): + """Bare prob_uniform equalizes unequal original system sizes.""" + result = compute_block_targets( + "prob_uniform", nsystems=2, system_nframes=[100, 500] + ) + self.assertEqual(result, [([0], 500), ([1], 500)]) + def test_compute_block_targets_multi_sys_block(self): result = compute_block_targets( "prob_sys_size;0:2:0.5;2:3:0.5", diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index 2c95e8224b..658a97519f 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import json import os import sys import tempfile @@ -22,13 +23,18 @@ from deepmd.dpmodel.utils.lmdb_data import ( is_lmdb, ) +from deepmd.entrypoints.test import test as run_model_test from deepmd.utils import ( data_system, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from deepmd.utils.data_system import ( LmdbDataSystem, get_data, process_systems, + validate_backend_data_config, validate_lmdb_systems, ) @@ -81,6 +87,101 @@ def _write_minimal_lmdb(path: str) -> None: env.close() +def _write_repeated_lmdb(path: str, nframes: int) -> None: + """Write a compact same-nloc LMDB suitable for bounded-I/O assertions.""" + env = lmdb.open(path, map_size=64 * 1024 * 1024) + frame = { + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": _encode_array(np.array([0], dtype=np.int64)), + "cells": _encode_array(np.eye(3, dtype=np.float64) * 8.0), + "coords": _encode_array(np.zeros((1, 3), dtype=np.float64)), + "energies": _encode_array(np.array([0.0], dtype=np.float64)), + "forces": _encode_array(np.zeros((1, 3), dtype=np.float64)), + } + packed_frame = msgpack.packb(frame, use_bin_type=True) + metadata = { + "nframes": nframes, + "frame_idx_fmt": "012d", + "frame_nlocs": [1] * nframes, + "type_map": ["H"], + "system_info": { + "formula": "H", + "natoms": [1], + "nframes": nframes, + }, + } + with env.begin(write=True) as txn: + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + for index in range(nframes): + txn.put(f"{index:012d}".encode(), packed_frame) + env.close() + + +def _write_mixed_nloc_lmdb(path: str) -> None: + """Write two frames that exercise legacy mix:N padding.""" + env = lmdb.open(path, map_size=10 * 1024 * 1024) + frames = [] + for nloc in (1, 2): + frames.append( + { + "atom_names": ["H"], + "atom_numbs": [nloc], + "atom_types": _encode_array(np.zeros(nloc, dtype=np.int64)), + "cells": _encode_array(np.eye(3, dtype=np.float64) * 8.0), + "coords": _encode_array(np.zeros((nloc, 3), dtype=np.float64)), + "energies": _encode_array(np.array([0.0], dtype=np.float64)), + "forces": _encode_array(np.zeros((nloc, 3), dtype=np.float64)), + } + ) + metadata = { + "nframes": 2, + "frame_idx_fmt": "012d", + "frame_nlocs": [1, 2], + "type_map": ["H"], + "system_info": { + "formula": "H", + "natoms": [1], + "nframes": 2, + }, + } + with env.begin(write=True) as txn: + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + for index, frame in enumerate(frames): + txn.put(f"{index:012d}".encode(), msgpack.packb(frame, use_bin_type=True)) + env.close() + + +def _write_mixed_pbc_lmdb(path: str) -> None: + """Write equal-nloc periodic and non-periodic frames.""" + env = lmdb.open(path, map_size=10 * 1024 * 1024) + frames = [] + for cell in (np.eye(3, dtype=np.float64) * 8.0, np.zeros((3, 3))): + frames.append( + { + "atom_names": ["H"], + "atom_numbs": [1], + "atom_types": _encode_array(np.array([0], dtype=np.int64)), + "cells": _encode_array(cell), + "coords": _encode_array(np.zeros((1, 3), dtype=np.float64)), + "energies": _encode_array(np.array([0.0], dtype=np.float64)), + "forces": _encode_array(np.zeros((1, 3), dtype=np.float64)), + } + ) + metadata = { + "nframes": 2, + "frame_idx_fmt": "012d", + "frame_nlocs": [1, 1], + "type_map": ["H"], + "system_info": {"formula": "H", "natoms": [1], "nframes": 2}, + } + with env.begin(write=True) as txn: + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + for index, frame in enumerate(frames): + txn.put(f"{index:012d}".encode(), msgpack.packb(frame, use_bin_type=True)) + env.close() + + class _FakeMultiSystems: write_count = 0 load_calls: ClassVar[list[tuple[str, str]]] = [] @@ -127,6 +228,7 @@ def setUp(self) -> None: _FakeMultiSystems.load_calls = [] _FakeMultiSystems.to_calls = [] data_system._DPDATA_CONVERSION_CACHE.clear() + data_system._DPDATA_SOURCE_MTIME_CACHE.clear() self.fake_dpdata = types.SimpleNamespace( MultiSystems=_FakeMultiSystems, LabeledSystem=_FakeLabeledSystem, @@ -136,6 +238,7 @@ def tearDown(self) -> None: os.chdir(self.old_cwd) self.tmpdir.cleanup() data_system._DPDATA_CONVERSION_CACHE.clear() + data_system._DPDATA_SOURCE_MTIME_CACHE.clear() def test_process_systems_defaults_to_deepmd_lmdb_and_reuses_cache(self) -> None: with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): @@ -201,7 +304,7 @@ def test_lmdb_alias_uses_canonical_dpdata_writer(self) -> None: ) def test_real_dpdata_lmdb_writer_is_deepmd_compatible(self) -> None: - """Verify dpdata 1.1 writes the schema consumed by DeepMD's LMDB reader.""" + """Verify dpdata 1.1 writes the schema consumed by DeePMD's LMDB reader.""" extxyz = ( "1\n" "Properties=species:S:1:pos:R:3:forces:R:3 energy=0.0 " @@ -212,7 +315,7 @@ def test_real_dpdata_lmdb_writer_is_deepmd_compatible(self) -> None: systems = process_systems(str(self.source), fmt="extxyz") # A changed source exercises dpdata's own transactional overwrite path; - # DeepMD-kit must not remove or rename the LMDB directory around it. + # DeePMD-kit must not remove or rename the LMDB directory around it. self.source.write_text(extxyz.replace("energy=0.0", "energy=1.0")) systems_again = process_systems(str(self.source), fmt="extxyz") @@ -272,10 +375,14 @@ def test_get_data_uses_format_conversion(self) -> None: self.assertIn("type", batch) self.assertIn("natoms_vec", batch) self.assertEqual(batch["coord"].shape, (1, 3)) - self.assertEqual(data.data_systems, [data]) + self.assertIsNot(data.data_systems[0], data) + self.assertEqual(data.type_map, ["H"]) + self.assertEqual(data.data_systems[0].get_test()["coord"].shape, (1, 3)) stat_set = data._load_set(data.dirs[0]) self.assertEqual(stat_set["coord"].shape, (1, 3)) self.assertEqual(stat_set["type"].shape, (1, 1)) + data.close() + data.close() def test_multiple_lmdb_paths_are_rejected(self) -> None: lmdb_a = self.root / "a.lmdb" @@ -311,6 +418,288 @@ def test_lmdb_stack_frames_rejects_empty_batch(self) -> None: with self.assertRaisesRegex(ValueError, "empty LMDB frame batch"): data._stack_frames([]) + def test_requirements_can_be_registered_after_adapter_construction(self) -> None: + """PBC probing must not freeze the model's label contract.""" + lmdb_path = self.root / "requirements.lmdb" + _write_minimal_lmdb(str(lmdb_path)) + data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=[1]) + + data.add_data_requirements( + [DataRequirementItem("energy", 1, atomic=False, must=True)] + ) + batch = data.get_batch() + + self.assertEqual(batch["energy"].shape, (1, 1)) + self.assertEqual(float(batch["find_energy"]), 1.0) + + def test_legacy_collation_ands_find_flags(self) -> None: + lmdb_path = self.root / "find-flags.lmdb" + _write_minimal_lmdb(str(lmdb_path)) + data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=2) + data.add_data_requirements( + [DataRequirementItem("energy", 1, atomic=False, must=False)] + ) + first = data._reader.peek_frame(0) + second = { + key: value.copy() if isinstance(value, np.ndarray) else value + for key, value in first.items() + } + second["find_energy"] = np.float32(0.0) + + batch = data._stack_frames([first, second]) + + self.assertEqual(float(batch["find_energy"]), 0.0) + + def test_neighbor_stat_reads_are_sampled_and_chunked(self) -> None: + """A large same-nloc group must never be decoded as one Python list.""" + lmdb_path = self.root / "bounded.lmdb" + _write_repeated_lmdb(str(lmdb_path), 2101) + data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=1000) + + total_sampled = sum(len(indices) for indices in data._nloc_set_indices.values()) + self.assertEqual(total_sampled, 2000) + self.assertLessEqual( + max(len(indices) for indices in data._nloc_set_indices.values()), 128 + ) + first = data._load_set(data.dirs[0]) + self.assertLessEqual(first["coord"].shape[0], 128) + self.assertEqual(data.get_stat_nsystems(), 1) + self.assertEqual(data.get_stat_numb_batches(0), 17) + self.assertEqual(data.get_stat_batch(0)["coord"].shape, (128, 3)) + + def test_mix_batch_pads_different_atom_counts(self) -> None: + lmdb_path = self.root / "mixed.lmdb" + _write_mixed_nloc_lmdb(str(lmdb_path)) + data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size="mix:4", seed=0) + + batch = data.get_batch() + + self.assertEqual(batch["coord"].shape, (2, 6)) + self.assertEqual(batch["type"].shape, (2, 2)) + self.assertIn(-1, batch["type"]) + + def test_periodic_and_nonperiodic_frames_are_separate_views(self) -> None: + lmdb_path = self.root / "mixed-pbc.lmdb" + _write_mixed_pbc_lmdb(str(lmdb_path)) + data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=2, seed=0) + + self.assertEqual(len(data.data_systems), 2) + self.assertEqual({view.pbc for view in data.data_systems}, {False, True}) + batches = [data.get_batch(), data.get_batch()] + self.assertEqual({float(batch["find_box"]) for batch in batches}, {0.0, 1.0}) + + def test_multiple_conversion_inputs_fail_before_dpdata_io(self) -> None: + second = self.root / "second.extxyz" + second.write_text(self.source.read_text()) + + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + with self.assertRaisesRegex(ValueError, "exactly one resolved input"): + process_systems( + [str(self.source), str(second)], + fmt="extxyz", + ) + + self.assertEqual(_FakeMultiSystems.load_calls, []) + self.assertEqual(_FakeMultiSystems.write_count, 0) + + def test_lmdb_sys_probs_fail_before_conversion(self) -> None: + with self.assertRaisesRegex(ValueError, "does not support explicit sys_probs"): + get_data( + { + "systems": str(self.source), + "format": "extxyz", + "batch_size": 1, + "sys_probs": [1.0], + }, + 0.0, + ["H"], + None, + ) + + self.assertEqual(_FakeMultiSystems.write_count, 0) + + def test_paddle_rejects_default_lmdb_before_conversion(self) -> None: + with self.assertRaisesRegex(NotImplementedError, "Paddle backend"): + validate_backend_data_config( + {"systems": str(self.source), "format": "extxyz"}, + backend_name="Paddle", + lmdb_supported=False, + ) + + self.assertEqual(_FakeMultiSystems.write_count, 0) + + def test_waiter_checks_source_only_after_lock_release(self) -> None: + output = self.root / "wait.lmdb" + output.mkdir() + lock_path = output.with_suffix(".lmdb.lock") + lock_path.write_text( + json.dumps( + { + "hostname": data_system.socket.gethostname(), + "pid": os.getpid(), + "process_start": data_system._process_start_time(os.getpid()), + } + ) + ) + waits = 0 + + def release_after_long_valid_conversion(_seconds: float) -> None: + nonlocal waits + waits += 1 + if waits == 301: + lock_path.unlink() + + with ( + patch.object( + data_system.time, "sleep", release_after_long_valid_conversion + ), + patch.object( + data_system, "_is_conversion_current", return_value=True + ) as is_current, + ): + self.assertTrue( + data_system._wait_for_conversion(self.source, output, lock_path) + ) + + self.assertEqual(waits, 301) + is_current.assert_called_once_with(self.source, output, force_source_scan=True) + + def test_dead_conversion_owner_is_recovered(self) -> None: + lock_path = self.root / "dead.lock" + lock_path.write_text( + json.dumps( + { + "hostname": data_system.socket.gethostname(), + "pid": 2147483647, + "process_start": "missing", + } + ) + ) + + self.assertTrue(data_system._recover_stale_conversion_lock(lock_path)) + self.assertFalse(lock_path.exists()) + + def test_directory_publication_rolls_back_on_failure(self) -> None: + output = self.root / "published" + output.mkdir() + (output / "old.txt").write_text("old") + staged = self.root / "staged" + staged.mkdir() + (staged / "new.txt").write_text("new") + real_replace = os.replace + calls = 0 + + def fail_publication(source: str | Path, target: str | Path) -> None: + nonlocal calls + calls += 1 + if calls == 2: + raise OSError("publish failed") + real_replace(source, target) + + with patch.object(data_system.os, "replace", fail_publication): + with self.assertRaisesRegex(OSError, "publish failed"): + data_system._publish_conversion_output(staged, output) + + self.assertEqual((output / "old.txt").read_text(), "old") + self.assertFalse((output / "new.txt").exists()) + + def test_unlabeled_input_uses_dpdata_unlabeled_loader(self) -> None: + class UnlabeledMultiSystems(_FakeMultiSystems): + calls: ClassVar[list[bool]] = [] + + def load_systems_from_file( + self, file_name: str, fmt: str, *, labeled: bool = True + ): + type(self).calls.append(labeled) + if labeled: + raise ValueError("missing labels") + self.loaded = True + return self + + UnlabeledMultiSystems.calls = [] + fake_dpdata = types.SimpleNamespace( + MultiSystems=UnlabeledMultiSystems, + LabeledSystem=_FakeLabeledSystem, + ) + + with patch.dict(sys.modules, {"dpdata": fake_dpdata}): + systems = process_systems(str(self.source), fmt="extxyz") + + self.assertEqual(UnlabeledMultiSystems.calls, [True, False]) + self.assertTrue(is_lmdb(systems[0])) + + def test_cache_path_changes_with_dpdata_version(self) -> None: + with patch.object( + data_system.importlib.metadata, "version", side_effect=["1.1.0", "1.2.0"] + ): + first = data_system._conversion_cache_path( + self.source, "extxyz", "deepmd/lmdb" + ) + second = data_system._conversion_cache_path( + self.source, "extxyz", "deepmd/lmdb" + ) + + self.assertNotEqual(first, second) + + def test_directory_freshness_scan_is_reused_within_one_routing_pass(self) -> None: + source_dir = self.root / "source-dir" + source_dir.mkdir() + (source_dir / "frame.xyz").write_text("frame") + output = self.root / ".deepmd_dpdata_cache" / "cached.lmdb" + output.parent.mkdir() + + first = data_system._source_mtime(source_dir, output) + with patch.object( + Path, "rglob", side_effect=AssertionError("unexpected rescan") + ): + second = data_system._source_mtime(source_dir, output) + + self.assertEqual(first, second) + + def test_dp_test_forwards_conversion_format_from_training_config(self) -> None: + config_path = self.root / "input.json" + config_path.write_text("{}") + config = { + "training": { + "training_data": { + "systems": "data.extxyz", + "format": "extxyz", + "out_format": "deepmd/lmdb", + "rglob_patterns": ["*.extxyz"], + } + } + } + with ( + patch("deepmd.entrypoints.test.j_loader", return_value=config), + patch( + "deepmd.entrypoints.test.update_deepmd_input", + side_effect=lambda value: value, + ), + patch( + "deepmd.entrypoints.test.process_systems", + side_effect=RuntimeError("stop after routing"), + ) as process, + ): + with self.assertRaisesRegex(RuntimeError, "stop after routing"): + run_model_test( + model="unused.pb", + system=None, + datafile=None, + train_json=str(config_path), + numb_test=1, + rand_seed=None, + shuffle_test=False, + detail_file="detail", + atomic=False, + ) + + process.assert_called_once_with( + str(self.source.resolve()), + patterns=["*.extxyz"], + fmt="extxyz", + out_fmt="deepmd/lmdb", + ) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index c6b79e6165..0e350daa5a 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -98,6 +98,42 @@ def test_neighbor_stat_and_training_data_reject_multiple_lmdb(self) -> None: with self.assertRaisesRegex(ValueError, "exactly one path"): _build_data_system(params, ["O", "H"]) + def test_converted_lmdb_forwards_distributed_sharding(self) -> None: + """Converted and direct LMDB routes must pass identical DDP metadata.""" + params = { + "systems": "input.extxyz", + "format": "extxyz", + "batch_size": 2, + } + with ( + patch( + "deepmd.pt_expt.entrypoints.main.process_systems", + return_value=["converted.lmdb"], + ), + patch( + "deepmd.pt_expt.entrypoints.main.validate_lmdb_systems", + return_value="converted.lmdb", + ), + patch("deepmd.pt_expt.entrypoints.main.LmdbDataSystem") as lmdb_data_system, + ): + _build_data_system( + params, + ["O", "H"], + seed=7, + rank=2, + world_size=4, + ) + + lmdb_data_system.assert_called_once_with( + lmdb_path="converted.lmdb", + type_map=["O", "H"], + batch_size=2, + auto_prob_style=None, + seed=7, + rank=2, + world_size=4, + ) + def _make_frame(natoms: int, seed: int, *, include_spin: bool = False) -> dict: """Synthetic LMDB frame matching the on-disk schema used by LmdbDataReader.""" From 75f49ce45739ccb5bfd60fe4f75f5bb06a64eaeb Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Sat, 12 Sep 2026 07:39:18 +0800 Subject: [PATCH 09/12] fix(data): address conversion and LMDB review findings Validate conversion caches with recorded source signatures, preserve seeded LMDB sampling, reject unsupported modifiers, and clean up partial training setup. Document LMDB defaults and protect remote writers during lock recovery. Coding-Agent: Codex Codex-Version: codex-cli 0.154.0 Model: gpt-6-astra Reasoning-Effort: xhigh --- deepmd/dpmodel/train/data.py | 23 +- deepmd/dpmodel/utils/lmdb_data.py | 23 +- deepmd/jax/entrypoints/train.py | 26 +- deepmd/pt/entrypoints/main.py | 14 +- deepmd/pt/utils/lmdb_dataset.py | 12 + deepmd/tf/entrypoints/train.py | 110 ++++---- deepmd/tf2/entrypoints/train.py | 50 ++-- deepmd/utils/argcheck.py | 18 +- deepmd/utils/data_system.py | 247 +++++++++++------ .../tests/common/dpmodel/test_train_data.py | 37 +++ .../common/test_data_system_conversion.py | 252 ++++++++++++++++-- source/tests/jax/test_training.py | 42 +++ source/tests/pt/test_lmdb_dataloader.py | 14 + source/tests/pt_expt/test_lmdb_training.py | 2 + source/tests/tf/test_train_cleanup.py | 51 ++++ source/tests/tf2/test_training.py | 45 ++++ 16 files changed, 751 insertions(+), 215 deletions(-) create mode 100644 source/tests/tf/test_train_cleanup.py diff --git a/deepmd/dpmodel/train/data.py b/deepmd/dpmodel/train/data.py index 3e91dcfc53..62a59a5541 100644 --- a/deepmd/dpmodel/train/data.py +++ b/deepmd/dpmodel/train/data.py @@ -14,6 +14,9 @@ Any, ) +from deepmd.utils.data_system import ( + close_data_systems, +) from deepmd.utils.stat_file import ( StatFileMode, StatFileSpec, @@ -98,15 +101,23 @@ def make_task_maps( config: Mapping[str, Any], factory: Callable[[TrainingTaskConfig], tuple[Any, Any | None, Any | None]], ) -> tuple[dict[str, Any], dict[str, Any | None], dict[str, Any | None]]: - """Build training, validation, and stat maps from normalized task configs.""" + """Build task maps, closing completed data objects if a later task fails. + + Factories own partial results until they return and must clean those up + themselves. Successful maps transfer ownership to the caller. + """ training_data: dict[str, Any] = {} validation_data: dict[str, Any | None] = {} stat_data: dict[str, Any | None] = {} - for task_config in iter_training_task_configs(config): - train_item, valid_item, stat_item = factory(task_config) - training_data[task_config.key] = train_item - validation_data[task_config.key] = valid_item - stat_data[task_config.key] = stat_item + try: + for task_config in iter_training_task_configs(config): + train_item, valid_item, stat_item = factory(task_config) + training_data[task_config.key] = train_item + validation_data[task_config.key] = valid_item + stat_data[task_config.key] = stat_item + except BaseException: + close_data_systems(training_data, validation_data) + raise return training_data, validation_data, stat_data diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 759ddf3bc9..a4ef54aac1 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -2936,9 +2936,12 @@ def compute_block_targets( ) blocks = nonempty - # Compute the same per-system probabilities as prob_sys_size_ext locally. - # Keeping this framework-agnostic LMDB module independent of data_system - # avoids an import cycle when the legacy adapter imports the LMDB reader. + # Keep this import local: data_system's legacy adapter also uses LMDB. + # Share the NPY probability calculation after filtering empty blocks. + from deepmd.utils.data_system import ( + prob_sys_size_ext, + ) + block_weights = np.asarray([weight for _, _, weight in blocks], dtype=float) if not np.all(np.isfinite(block_weights)): raise ValueError("block weights must be finite") @@ -2947,16 +2950,10 @@ def compute_block_targets( total_block_weight = np.sum(block_weights) if total_block_weight <= 0: raise ValueError("the sum of block weights should be greater than 0") - block_probs = block_weights / total_block_weight - sys_probs = np.zeros(nsystems, dtype=np.float64) - for block_idx, (stt, end, _weight) in enumerate(blocks): - block_frames = np.asarray(system_nframes[stt:end], dtype=float) - total_block_frames = np.sum(block_frames) - if total_block_frames <= 0: - raise ValueError( - f"block {stt}:{end} must contain at least one retained frame" - ) - sys_probs[stt:end] = block_frames / total_block_frames * block_probs[block_idx] + keywords = "prob_sys_size;" + ";".join( + f"{stt}:{end}:{weight}" for stt, end, weight in blocks + ) + sys_probs = prob_sys_size_ext(keywords, nsystems, system_nframes) # Group systems by block, compute block-level frames and prob block_info: list[tuple[list[int], int, float]] = [] # (sys_ids, frames, prob) diff --git a/deepmd/jax/entrypoints/train.py b/deepmd/jax/entrypoints/train.py index 1cad2af5bc..cc7671da30 100644 --- a/deepmd/jax/entrypoints/train.py +++ b/deepmd/jax/entrypoints/train.py @@ -190,19 +190,25 @@ def factory( None, ) valid_data = None - if task_config.validation_data_params is not None: - valid_data = get_data( - dict(task_config.validation_data_params), - task_model.get_rcut(), - train_data.type_map, - None, - ) - return train_data, valid_data, None + try: + if task_config.validation_data_params is not None: + valid_data = get_data( + dict(task_config.validation_data_params), + task_model.get_rcut(), + train_data.type_map, + None, + ) + return train_data, valid_data, None + except BaseException: + # The factory owns these objects until make_task_maps records + # them; validation setup or stat-file validation may fail. + close_data_systems(train_data, valid_data) + raise train_data_map, valid_data_map, _ = make_task_maps(config, factory) - print_data_summaries(train_data_map, valid_data_map) - try: + print_data_summaries(train_data_map, valid_data_map) + start_time = time.time() model.train(train_data_map, valid_data_map) end_time = time.time() diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 08a3a183c3..f5627be67a 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -182,7 +182,7 @@ def _make_dp_loader_set( ) -> DpLoaderSet | LmdbDataset: """Create a dataset from systems with pattern expansion/conversion.""" if conversion_will_write_lmdb(dataset_params): - validate_lmdb_sampling_options(dataset_params) + validate_lmdb_sampling_options(dataset_params, modifier=modifier) patterns = dataset_params.get("rglob_patterns") systems = process_systems( systems, @@ -194,12 +194,13 @@ def _make_dp_loader_set( ) lmdb_path = validate_lmdb_systems(systems, backend_name="PyTorch") if lmdb_path is not None: - validate_lmdb_sampling_options(dataset_params) + validate_lmdb_sampling_options(dataset_params, modifier=modifier) return LmdbDataset( lmdb_path, model_params_single["type_map"], dataset_params["batch_size"], auto_prob_style=dataset_params.get("auto_prob"), + seed=rank_seed, ) return DpLoaderSet( systems, @@ -215,13 +216,14 @@ def _make_dp_loader_set( and isinstance(training_systems, str) and is_lmdb(training_systems) ): - validate_lmdb_sampling_options(training_dataset_params) + validate_lmdb_sampling_options(training_dataset_params, modifier=modifier) auto_prob = training_dataset_params.get("auto_prob", None) train_data_single = LmdbDataset( training_systems, model_params_single["type_map"], training_dataset_params["batch_size"], auto_prob_style=auto_prob, + seed=rank_seed, ) if ( validation_systems is not None @@ -229,11 +231,15 @@ def _make_dp_loader_set( and isinstance(validation_systems, str) and is_lmdb(validation_systems) ): - validate_lmdb_sampling_options(validation_dataset_params) + validate_lmdb_sampling_options( + validation_dataset_params, modifier=modifier + ) validation_data_single = LmdbDataset( validation_systems, model_params_single["type_map"], validation_dataset_params["batch_size"], + auto_prob_style=validation_dataset_params.get("auto_prob"), + seed=rank_seed, ) elif validation_systems is not None: validation_data_single = _make_dp_loader_set( diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index 1648adcd04..ea5e6935af 100644 --- a/deepmd/pt/utils/lmdb_dataset.py +++ b/deepmd/pt/utils/lmdb_dataset.py @@ -10,6 +10,7 @@ Any, ) +import numpy as np import torch from torch.utils.data import ( DataLoader, @@ -207,6 +208,8 @@ class LmdbDataset(Dataset): are padded with phantom atoms. auto_prob_style : str, optional ``auto_prob`` string used to reweight the original systems. + seed : int or list[int], optional + Seed for shuffled batches, including the worker rank when distributed. """ def __init__( @@ -215,6 +218,7 @@ def __init__( type_map: list[str], batch_size: int | str = "auto", auto_prob_style: str | None = None, + seed: int | list[int] | None = None, ) -> None: self._reader = LmdbDataReader(lmdb_path, type_map, batch_size) self._collate = functools.partial( @@ -237,9 +241,17 @@ def __init__( f"nsystems={self._reader.nsystems}" ) + # The common sampler advances integer seeds by epoch. Fold the + # rank/seed pair into one reproducible integer before handing it off. + sampler_seed = ( + int(np.random.SeedSequence(seed).generate_state(1)[0]) + if isinstance(seed, list) + else seed + ) sampler = LmdbBatchSampler( self._reader, shuffle=True, + seed=sampler_seed, block_targets=self._block_targets, ) self._batch_sampler = _LmdbBatchSamplerTorch(sampler) diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index 2dc1feba9f..d28a1a1bdd 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -258,63 +258,65 @@ def _do_work( # decouple the training data from the model compress process train_data = None valid_data = None - if not is_compress: - # init data - train_data = get_data( - jdata["training"]["training_data"], rcut, ipt_type_map, modifier - ) - train_data.add_data_requirements(model.data_requirements) - train_data.print_summary("training") - if jdata["training"].get("validation_data", None) is not None: - valid_data = get_data( - jdata["training"]["validation_data"], - rcut, - train_data.type_map, - modifier, + try: + if not is_compress: + # init data + train_data = get_data( + jdata["training"]["training_data"], rcut, ipt_type_map, modifier ) - valid_data.add_data_requirements(model.data_requirements) - valid_data.print_summary("validation") - else: - if modifier is not None: - modifier.build_fv_graph() - - # get training info - training_params = jdata["training"] - stop_batch = training_params.get("numb_steps") - num_epoch = training_params.get("numb_epoch") - if stop_batch is None: - if num_epoch is None: - raise ValueError( - "Either training.numb_steps or training.num_epoch must be set." + train_data.add_data_requirements(model.data_requirements) + train_data.print_summary("training") + if jdata["training"].get("validation_data", None) is not None: + valid_data = get_data( + jdata["training"]["validation_data"], + rcut, + train_data.type_map, + modifier, + ) + valid_data.add_data_requirements(model.data_requirements) + valid_data.print_summary("validation") + else: + if modifier is not None: + modifier.build_fv_graph() + + # get training info + training_params = jdata["training"] + stop_batch = training_params.get("numb_steps") + num_epoch = training_params.get("numb_epoch") + if stop_batch is None: + if num_epoch is None: + raise ValueError( + "Either training.numb_steps or training.num_epoch must be set." + ) + if num_epoch <= 0: + raise ValueError("training.num_epoch must be positive.") + if train_data is None: + raise ValueError( + "training.num_epoch requires training data to compute total_numb_batch." + ) + total_numb_batch = compute_total_numb_batch( + train_data.nbatches, train_data.sys_probs ) - if num_epoch <= 0: - raise ValueError("training.num_epoch must be positive.") - if train_data is None: - raise ValueError( - "training.num_epoch requires training data to compute total_numb_batch." + if total_numb_batch <= 0: + raise ValueError("Total number of training batches must be positive.") + stop_batch = int(np.ceil(num_epoch * total_numb_batch)) + log.info( + "Computed numb_steps=%d from num_epoch=%s and total_numb_batch=%d.", + stop_batch, + num_epoch, + total_numb_batch, ) - total_numb_batch = compute_total_numb_batch( - train_data.nbatches, train_data.sys_probs - ) - if total_numb_batch <= 0: - raise ValueError("Total number of training batches must be positive.") - stop_batch = int(np.ceil(num_epoch * total_numb_batch)) - log.info( - "Computed numb_steps=%d from num_epoch=%s and total_numb_batch=%d.", - stop_batch, - num_epoch, - total_numb_batch, - ) - origin_type_map = jdata["model"].get("origin_type_map", None) - if ( - origin_type_map is not None and not origin_type_map - ): # get the type_map from data if not provided - origin_data = get_data(jdata["training"]["training_data"], rcut, None, modifier) - try: - origin_type_map = origin_data.get_type_map() - finally: - close_data_systems(origin_data) - try: + origin_type_map = jdata["model"].get("origin_type_map", None) + if ( + origin_type_map is not None and not origin_type_map + ): # get the type_map from data if not provided + origin_data = get_data( + jdata["training"]["training_data"], rcut, None, modifier + ) + try: + origin_type_map = origin_data.get_type_map() + finally: + close_data_systems(origin_data) model.build( train_data, stop_batch, diff --git a/deepmd/tf2/entrypoints/train.py b/deepmd/tf2/entrypoints/train.py index 941bca8984..83df7ecfc9 100644 --- a/deepmd/tf2/entrypoints/train.py +++ b/deepmd/tf2/entrypoints/train.py @@ -166,34 +166,40 @@ def factory( None, ) valid_data = None - if task_config.validation_data_params is not None: - valid_data = get_data( - dict(task_config.validation_data_params), - None, - train_data.type_map, - None, - ) - return train_data, valid_data, task_config.stat_file_spec + try: + if task_config.validation_data_params is not None: + valid_data = get_data( + dict(task_config.validation_data_params), + None, + train_data.type_map, + None, + ) + return train_data, valid_data, task_config.stat_file_spec + except BaseException: + # The factory owns these objects until make_task_maps records + # them; validation setup or stat-file validation may fail. + close_data_systems(train_data, valid_data) + raise train_data_map, valid_data_map, stat_file_spec_map = make_task_maps( config, factory, ) - print_data_summaries(train_data_map, valid_data_map) - - trainer = DPTrainer( - config, - train_data_map, - stat_file_spec=stat_file_spec_map, - validation_data=valid_data_map, - init_model=options.init_model, - restart_model=options.restart, - finetune_model=options.finetune, - finetune_links=self.finetune_links, - shared_links=self.shared_links, - min_nbor_dist=neighbor_stat, - ) try: + print_data_summaries(train_data_map, valid_data_map) + + trainer = DPTrainer( + config, + train_data_map, + stat_file_spec=stat_file_spec_map, + validation_data=valid_data_map, + init_model=options.init_model, + restart_model=options.restart, + finetune_model=options.finetune, + finetune_links=self.finetune_links, + shared_links=self.shared_links, + min_nbor_dist=neighbor_stat, + ) start_time = time.time() trainer.run() end_time = time.time() diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 8b9ac03a80..e058cb91fd 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5375,9 +5375,12 @@ def training_data_args() -> list[ ) doc_out_format = ( "The output data format passed to dpdata for automatic conversion. " - "When `format` requests conversion from a non-DeePMD format, this key " - "defaults to `deepmd/lmdb`. Use a DeePMD format supported by dpdata, " - "such as `deepmd/lmdb`, `deepmd/hdf5`, or `deepmd/npy`." + "Defaults to `deepmd/lmdb`; it is used only when `format` requests " + "conversion from a non-DeePMD format. LMDB requires exactly one input " + "path after `systems` and `rglob_patterns` expansion and an explicit " + "model/type_map. LMDB does not support Paddle, data modifiers, or " + "explicit sys_probs. Choose `deepmd/hdf5` for those capabilities or " + "for multiple input paths. Other supported outputs include `deepmd/npy`." ) doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ @@ -5500,9 +5503,12 @@ def validation_data_args() -> list[ ) doc_out_format = ( "The output data format passed to dpdata for automatic conversion. " - "When `format` requests conversion from a non-DeePMD format, this key " - "defaults to `deepmd/lmdb`. Use a DeePMD format supported by dpdata, " - "such as `deepmd/lmdb`, `deepmd/hdf5`, or `deepmd/npy`." + "Defaults to `deepmd/lmdb`; it is used only when `format` requests " + "conversion from a non-DeePMD format. LMDB requires exactly one input " + "path after `systems` and `rglob_patterns` expansion and an explicit " + "model/type_map. LMDB does not support Paddle, data modifiers, or " + "explicit sys_probs. Choose `deepmd/hdf5` for those capabilities or " + "for multiple input paths. Other supported outputs include `deepmd/npy`." ) doc_batch_size = f'This key can be \n\n\ - list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 90ec9b31da..bc3fe47582 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -7,6 +7,7 @@ import os import shutil import socket +import tempfile import threading import time import warnings @@ -17,11 +18,14 @@ Path, ) from typing import ( + IO, Any, - Self, ) import numpy as np +from typing_extensions import ( + Self, +) import deepmd.utils.random as dp_random from deepmd.common import ( @@ -44,15 +48,13 @@ _DPDATA_CACHE_DIR = ".deepmd_dpdata_cache" _DPDATA_DEFAULT_OUT_FORMAT = "deepmd/lmdb" -_DPDATA_CONVERSION_SCHEMA_VERSION = "2" +_DPDATA_CONVERSION_SCHEMA_VERSION = "3" _DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str, str], list[str]] = {} -_DPDATA_SOURCE_MTIME_CACHE: dict[tuple[str, str], tuple[float, float]] = {} -# Neighbor-stat, trainer construction, and multi-task routing commonly query -# the same directory within one startup. Reuse that O(file-count) scan while -# still letting a long-lived process notice later source rewrites. -_DPDATA_SOURCE_MTIME_CACHE_TTL = 60.0 +_DPDATA_SOURCE_HASH_LIMIT = 1024 * 1024 _CONVERSION_LOCK_HEARTBEAT_SECONDS = 5.0 -_CONVERSION_LOCK_STALE_SECONDS = 30.0 +# Allow for shared-filesystem attribute caching before diagnosing a stalled +# heartbeat. Remote owner liveness still cannot be inferred from elapsed time. +_CONVERSION_LOCK_STALE_SECONDS = 300.0 def validate_lmdb_systems( @@ -941,13 +943,6 @@ def _refresh_groups(self) -> None: self.dirs = list(self._nloc_set_indices) self.pbc = any_pbc - def _detect_pbc(self) -> bool: - """Return True when LMDB frames contain a non-zero simulation box.""" - if len(self._reader) == 0: - return False - frame = self._reader.peek_frame(0) - return bool(float(frame.get("find_box", 0.0)) > 0.5) - def add_data_requirements( self, data_requirements: list[DataRequirementItem] ) -> None: @@ -1362,8 +1357,15 @@ def validate_backend_data_config( ) -def validate_lmdb_sampling_options(data_config: dict[str, Any]) -> None: - """Reject sampling options that cannot be represented by one LMDB route.""" +def validate_lmdb_sampling_options( + data_config: dict[str, Any], *, modifier: Any | None = None +) -> None: + """Reject options that cannot be represented by one LMDB route.""" + if modifier is not None: + raise ValueError( + "LMDB data does not support data modifiers. " + "Choose out_format='deepmd/hdf5' for automatic conversion." + ) if data_config.get("sys_probs") is not None: raise ValueError( "LMDB data does not support explicit sys_probs yet. Use auto_prob " @@ -1404,18 +1406,20 @@ def close_one(value: Any) -> None: def _looks_like_extxyz(path: Path) -> bool: + """Probe the XYZ comment line without requiring a valid text input.""" if not path.is_file(): return False try: - with path.open() as fp: + with path.open(encoding="utf-8") as fp: fp.readline() comment = fp.readline() - except OSError: + except (OSError, UnicodeDecodeError): return False return "Properties=" in comment or "Lattice=" in comment def _normalize_dpdata_format(fmt: str, source: Path) -> str: + """Resolve aliases and infer common dpdata formats from the input path.""" fmt = fmt.lower() if fmt == "ase": return "ase/structure" @@ -1447,7 +1451,7 @@ def _conversion_cache_path(source: Path, fmt: str, out_fmt: str) -> Path: dpdata_version = importlib.metadata.version("dpdata") except importlib.metadata.PackageNotFoundError: dpdata_version = "unknown" - digest = hashlib.sha1( + digest = hashlib.sha256( ( f"{source_resolved}|{fmt}|{out_fmt}|" f"schema={_DPDATA_CONVERSION_SCHEMA_VERSION}|dpdata={dpdata_version}" @@ -1459,45 +1463,86 @@ def _conversion_cache_path(source: Path, fmt: str, out_fmt: str) -> Path: return Path.cwd() / _DPDATA_CACHE_DIR / f"{stem}-{safe_out_fmt}-{digest}{suffix}" -def _source_mtime(source: Path, cache_file: Path, *, force: bool = False) -> float: +def _source_signature(source: Path, output: Path) -> str: + """Fingerprint source metadata, including restored files and tree changes. + + Compare nanosecond timestamps and sizes for equality, never ordering. + ctime also detects a same-size rewrite whose mtime was preserved on Unix. + Hash small files as well, including on Windows where ctime is creation time, + without rereading large trajectory files on every routing call. Scan on + each call: a time-based memo can hide edits to directory children. + Exclude our cache subtree so publication cannot invalidate its own input. + """ + digest = hashlib.sha256() + cache_dir = output.parent.resolve() + + def add_file(path: Path, name: str) -> None: + stat = path.stat() + digest.update( + json.dumps( + [name, stat.st_mode, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns] + ).encode() + ) + if stat.st_size <= _DPDATA_SOURCE_HASH_LIMIT: + with path.open("rb") as fp: + # Bound the read even if a concurrent writer grows the input. + digest.update( + hashlib.sha256(fp.read(_DPDATA_SOURCE_HASH_LIMIT + 1)).digest() + ) + if source.is_file(): - return source.stat().st_mtime - if not source.is_dir(): - return 0.0 - cache_key = ( - str(source.resolve(strict=False)), - str(cache_file.parent.resolve(strict=False)), - ) - now = time.monotonic() - cached = _DPDATA_SOURCE_MTIME_CACHE.get(cache_key) - if ( - not force - and cached is not None - and now - cached[0] < _DPDATA_SOURCE_MTIME_CACHE_TTL - ): - return cached[1] - cache_dir = cache_file.parent.resolve(strict=False) - latest = source.stat().st_mtime - for item in source.rglob("*"): - try: - item_resolved = item.resolve(strict=False) - if item_resolved == cache_file or cache_dir in item_resolved.parents: - continue - latest = max(latest, item.stat().st_mtime) - except OSError: - continue - _DPDATA_SOURCE_MTIME_CACHE[cache_key] = (now, latest) - return latest + add_file(source, source.name) + elif source.is_dir(): + + def raise_walk_error(error: OSError) -> None: + # An incomplete scan must never validate an old conversion. + raise error + + for root, dirs, files in os.walk(source, onerror=raise_walk_error): + dirs[:] = sorted( + name for name in dirs if (Path(root) / name).resolve() != cache_dir + ) + relative_root = Path(root).relative_to(source) + digest.update(json.dumps([str(relative_root), dirs]).encode()) + for name in sorted(files): + add_file(Path(root) / name, str(relative_root / name)) + else: + raise FileNotFoundError(f"dpdata input does not exist: {source}") + return digest.hexdigest() + +def _conversion_manifest_path(output: Path) -> Path: + """Return the sidecar recording the source used for a published output.""" + return output.with_name(output.name + ".source.json") -def _is_conversion_current( - source: Path, output: Path, *, force_source_scan: bool = False -) -> bool: + +def _is_conversion_current(source: Path, output: Path) -> bool: + """Reuse only a complete conversion with exactly matching source metadata.""" if not output.exists(): return False - return output.stat().st_mtime >= _source_mtime( - source, output, force=force_source_scan - ) + try: + manifest = json.loads(_conversion_manifest_path(output).read_text()) + return manifest == {"source_signature": _source_signature(source, output)} + except (OSError, ValueError): + # Missing, damaged, or legacy manifests require a fresh conversion. + return False + + +def _write_conversion_manifest(output: Path, signature: str) -> None: + """Atomically mark a successfully published conversion as reusable.""" + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", dir=output.parent, prefix=f".{output.name}.", delete=False + ) as fp: + temporary = Path(fp.name) + json.dump({"source_signature": signature}, fp) + fp.flush() + os.fsync(fp.fileno()) + os.replace(temporary, _conversion_manifest_path(output)) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) def _process_start_time(pid: int) -> str | None: @@ -1520,30 +1565,31 @@ def _same_lock_file(lock_path: Path, expected: os.stat_result) -> bool: class _ConversionLock: - """Owned conversion lock with a heartbeat for cross-host stale recovery.""" + """Owned conversion lock with a heartbeat for detecting stalled writers.""" - def __init__(self, lock_path: Path, lock_fd: int) -> None: + def __init__(self, lock_path: Path, lock_file: IO[str]) -> None: self.path = lock_path - self._stat = os.fstat(lock_fd) - payload = { - "hostname": socket.gethostname(), - "pid": os.getpid(), - "process_start": _process_start_time(os.getpid()), - "created": time.time(), - } - with os.fdopen(lock_fd, "w") as fp: - json.dump(payload, fp) - fp.flush() - os.fsync(fp.fileno()) - self._stop = threading.Event() - self._thread = threading.Thread( - target=self._heartbeat, - name="deepmd-dpdata-conversion-lock", - daemon=True, - ) + self._stat = os.fstat(lock_file.fileno()) try: + payload = { + "hostname": socket.gethostname(), + "pid": os.getpid(), + "process_start": _process_start_time(os.getpid()), + "created": time.time(), + } + # The caller owns the file context, including initialization failures. + json.dump(payload, lock_file) + lock_file.flush() + os.fsync(lock_file.fileno()) + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._heartbeat, + name="deepmd-dpdata-conversion-lock", + daemon=True, + ) self._thread.start() - except Exception: + except BaseException: + lock_file.close() if _same_lock_file(self.path, self._stat): self.path.unlink(missing_ok=True) raise @@ -1565,6 +1611,7 @@ def release(self) -> None: try: self.path.unlink() except FileNotFoundError: + # Another waiter may already have removed this lock. pass @@ -1576,6 +1623,8 @@ def _lock_owner_is_alive(payload: dict[str, Any]) -> bool | None: pid = int(payload["pid"]) except (KeyError, TypeError, ValueError): return None + if pid <= 0: + return None expected_start = payload.get("process_start") current_start = _process_start_time(pid) if expected_start is not None and current_start is not None: @@ -1590,7 +1639,12 @@ def _lock_owner_is_alive(payload: dict[str, Any]) -> bool | None: def _recover_stale_conversion_lock(lock_path: Path) -> bool: - """Remove a dead-owner or expired-lease lock without touching a successor.""" + """Recover dead local owners and abandoned, uninitialized lock files. + + A remote host's heartbeat may look stale because of NFS attribute caching + or clock skew. Never evict an identified remote writer on that evidence; + report the stalled lock for owner verification instead of racing it. + """ try: lock_stat = lock_path.stat(follow_symlinks=False) except FileNotFoundError: @@ -1599,11 +1653,23 @@ def _recover_stale_conversion_lock(lock_path: Path) -> bool: return False try: payload = json.loads(lock_path.read_text()) - except (OSError, json.JSONDecodeError): + except (OSError, ValueError): + payload = {} + if not isinstance(payload, dict): payload = {} owner_alive = _lock_owner_is_alive(payload) lease_expired = time.time() - lock_stat.st_mtime > _CONVERSION_LOCK_STALE_SECONDS + if ( + owner_alive is None + and lease_expired + and payload.get("hostname") not in (None, socket.gethostname()) + ): + raise RuntimeError( + f"Cannot verify the remote owner of stalled dpdata conversion lock " + f"{lock_path} on {payload['hostname']}. Check that the owner has " + "stopped before removing the lock and retrying." + ) stale = owner_alive is False or (owner_alive is None and lease_expired) if not stale or not _same_lock_file(lock_path, lock_stat): return False @@ -1611,6 +1677,7 @@ def _recover_stale_conversion_lock(lock_path: Path) -> bool: try: lock_path.unlink() except FileNotFoundError: + # Recovery is complete if another waiter removed the same stale lock. pass return True @@ -1622,7 +1689,7 @@ def _wait_for_conversion(source: Path, output: Path, lock_path: Path) -> bool: continue time.sleep(1.0) # Freshness is checked once after publication, not once per waiter-second. - return _is_conversion_current(source, output, force_source_scan=True) + return _is_conversion_current(source, output) def _remove_path(path: Path) -> None: @@ -1744,17 +1811,16 @@ def _convert_system_by_dpdata( if not _is_conversion_current(source, output): while True: try: - lock_fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + lock_file = lock_path.open("x") except FileExistsError: if _wait_for_conversion(source, output, lock_path): break continue else: - conversion_lock = _ConversionLock(lock_path, lock_fd) + with lock_file: + conversion_lock = _ConversionLock(lock_path, lock_file) try: - if not _is_conversion_current( - source, output, force_source_scan=True - ): + if not _is_conversion_current(source, output): log.info( "Converting %s from dpdata format %s to %s at %s", source, @@ -1762,7 +1828,17 @@ def _convert_system_by_dpdata( out_fmt, output, ) + signature = _source_signature(source, output) + # Invalidate before replacing output: a crash must never + # pair a new dataset with an older source manifest. + _conversion_manifest_path(output).unlink(missing_ok=True) _write_dpdata_conversion(source, fmt, out_fmt, output) + if signature != _source_signature(source, output): + raise RuntimeError( + f"dpdata input changed during conversion: {source}. " + "Retry with an unchanged input dataset." + ) + _write_conversion_manifest(output, signature) finally: conversion_lock.release() break @@ -1896,7 +1972,7 @@ def get_data( data_format = jdata.get("format") out_format = jdata.get("out_format", jdata.get("output_format")) if conversion_will_write_lmdb(jdata): - validate_lmdb_sampling_options(jdata) + validate_lmdb_sampling_options(jdata, modifier=modifier) systems = process_systems( systems, patterns=rglob_patterns, fmt=data_format, out_fmt=out_format ) @@ -1908,7 +1984,7 @@ def get_data( lmdb_path = validate_lmdb_systems(systems, backend_name="legacy data loader") if lmdb_path is not None: - validate_lmdb_sampling_options(jdata) + validate_lmdb_sampling_options(jdata, modifier=modifier) if type_map is None: raise ValueError( "LMDB training data requires model/type_map to be set. " @@ -1920,6 +1996,9 @@ def get_data( type_map=type_map, batch_size=batch_size, auto_prob_style=auto_prob, + # Draw from the shared, rank-seeded data RNG so LMDB preserves + # training.seed reproducibility across all legacy entrypoints. + seed=int(dp_random.choice(2**32)), ) data = DeepmdDataSystem( diff --git a/source/tests/common/dpmodel/test_train_data.py b/source/tests/common/dpmodel/test_train_data.py index 0cae6d888a..195c84155d 100644 --- a/source/tests/common/dpmodel/test_train_data.py +++ b/source/tests/common/dpmodel/test_train_data.py @@ -1,9 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from unittest.mock import ( + Mock, +) + import pytest from deepmd.dpmodel.train.data import ( _print_summary, iter_training_task_configs, + make_task_maps, ) from deepmd.utils.stat_file import ( StatFileSpec, @@ -47,3 +52,35 @@ def test_training_task_config_preserves_stat_file_mode() -> None: task = next(iter_training_task_configs(config)) assert task.stat_file_spec == StatFileSpec("stat.hdf5", "read") + + +def test_make_task_maps_closes_completed_tasks_on_factory_failure() -> None: + """A later task failure releases already-acquired training and validation data.""" + config = { + "model": {"model_dict": {"first": {}, "second": {}}}, + "training": { + "data_dict": { + "first": {"training_data": {}}, + "second": {"training_data": {}}, + } + }, + } + training, validation = Mock(), Mock() + factory = Mock( + side_effect=[(training, validation, None), ValueError("second task failed")] + ) + with pytest.raises(ValueError, match="second task failed"): + make_task_maps(config, factory) + training.close.assert_called_once_with() + validation.close.assert_called_once_with() + + +def test_make_task_maps_transfers_successful_data_to_caller() -> None: + """Successful construction keeps data open for the caller's training loop.""" + config = {"model": {}, "training": {"training_data": {}}} + training, validation = Mock(), Mock() + maps = make_task_maps(config, Mock(return_value=(training, validation, None))) + assert list(maps[0].values()) == [training] + assert list(maps[1].values()) == [validation] + training.close.assert_not_called() + validation.close.assert_not_called() diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index 658a97519f..bf8ea26265 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -219,26 +219,24 @@ def __init__(self, file_name: str, fmt: str) -> None: class TestDpdataFormatConversion(unittest.TestCase): def setUp(self) -> None: self.tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(self.tmpdir.cleanup) self.root = Path(self.tmpdir.name) self.old_cwd = Path.cwd() os.chdir(self.root) + self.addCleanup(os.chdir, self.old_cwd) self.source = self.root / "data.extxyz" self.source.write_text("1\nProperties=species:S:1:pos:R:3\nH 0 0 0\n") _FakeMultiSystems.write_count = 0 _FakeMultiSystems.load_calls = [] _FakeMultiSystems.to_calls = [] data_system._DPDATA_CONVERSION_CACHE.clear() - data_system._DPDATA_SOURCE_MTIME_CACHE.clear() self.fake_dpdata = types.SimpleNamespace( MultiSystems=_FakeMultiSystems, LabeledSystem=_FakeLabeledSystem, ) def tearDown(self) -> None: - os.chdir(self.old_cwd) - self.tmpdir.cleanup() data_system._DPDATA_CONVERSION_CACHE.clear() - data_system._DPDATA_SOURCE_MTIME_CACHE.clear() def test_process_systems_defaults_to_deepmd_lmdb_and_reuses_cache(self) -> None: with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): @@ -316,14 +314,22 @@ def test_real_dpdata_lmdb_writer_is_deepmd_compatible(self) -> None: systems = process_systems(str(self.source), fmt="extxyz") # A changed source exercises dpdata's own transactional overwrite path; # DeePMD-kit must not remove or rename the LMDB directory around it. - self.source.write_text(extxyz.replace("energy=0.0", "energy=1.0")) + original_stat = self.source.stat() + self.source.write_text(extxyz.replace("energy=0.0", "energy=42.0")) + os.utime(self.source, (original_stat.st_atime, original_stat.st_mtime - 86400)) + data_system._DPDATA_CONVERSION_CACHE.clear() systems_again = process_systems(str(self.source), fmt="extxyz") self.assertEqual(len(systems), 1) self.assertEqual(systems_again, systems) self.assertTrue(is_lmdb(systems[0])) data = LmdbDataSystem(systems[0], ["H"], batch_size=1) + self.addCleanup(data.close) + data.add_data_requirements( + [DataRequirementItem("energy", 1, atomic=False, must=True)] + ) batch = data.get_batch() + self.assertEqual(float(batch["energy"].item()), 42.0) self.assertEqual(batch["coord"].shape, (1, 3)) self.assertEqual(batch["type"].shape, (1, 1)) @@ -366,6 +372,7 @@ def test_get_data_uses_format_conversion(self) -> None: ["H"], None, ) + self.addCleanup(data.close) self.assertEqual(data.get_nsystems(), 1) self.assertIsInstance(data, LmdbDataSystem) @@ -414,6 +421,7 @@ def test_lmdb_stack_frames_rejects_empty_batch(self) -> None: lmdb_path = self.root / "empty-batch.lmdb" _write_minimal_lmdb(str(lmdb_path)) data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=1) + self.addCleanup(data.close) with self.assertRaisesRegex(ValueError, "empty LMDB frame batch"): data._stack_frames([]) @@ -423,6 +431,7 @@ def test_requirements_can_be_registered_after_adapter_construction(self) -> None lmdb_path = self.root / "requirements.lmdb" _write_minimal_lmdb(str(lmdb_path)) data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=[1]) + self.addCleanup(data.close) data.add_data_requirements( [DataRequirementItem("energy", 1, atomic=False, must=True)] @@ -436,6 +445,7 @@ def test_legacy_collation_ands_find_flags(self) -> None: lmdb_path = self.root / "find-flags.lmdb" _write_minimal_lmdb(str(lmdb_path)) data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=2) + self.addCleanup(data.close) data.add_data_requirements( [DataRequirementItem("energy", 1, atomic=False, must=False)] ) @@ -455,6 +465,7 @@ def test_neighbor_stat_reads_are_sampled_and_chunked(self) -> None: lmdb_path = self.root / "bounded.lmdb" _write_repeated_lmdb(str(lmdb_path), 2101) data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=1000) + self.addCleanup(data.close) total_sampled = sum(len(indices) for indices in data._nloc_set_indices.values()) self.assertEqual(total_sampled, 2000) @@ -471,6 +482,7 @@ def test_mix_batch_pads_different_atom_counts(self) -> None: lmdb_path = self.root / "mixed.lmdb" _write_mixed_nloc_lmdb(str(lmdb_path)) data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size="mix:4", seed=0) + self.addCleanup(data.close) batch = data.get_batch() @@ -482,6 +494,7 @@ def test_periodic_and_nonperiodic_frames_are_separate_views(self) -> None: lmdb_path = self.root / "mixed-pbc.lmdb" _write_mixed_pbc_lmdb(str(lmdb_path)) data = LmdbDataSystem(str(lmdb_path), ["H"], batch_size=2, seed=0) + self.addCleanup(data.close) self.assertEqual(len(data.data_systems), 2) self.assertEqual({view.pbc for view in data.data_systems}, {False, True}) @@ -562,7 +575,7 @@ def release_after_long_valid_conversion(_seconds: float) -> None: ) self.assertEqual(waits, 301) - is_current.assert_called_once_with(self.source, output, force_source_scan=True) + is_current.assert_called_once_with(self.source, output) def test_dead_conversion_owner_is_recovered(self) -> None: lock_path = self.root / "dead.lock" @@ -641,20 +654,227 @@ def test_cache_path_changes_with_dpdata_version(self) -> None: self.assertNotEqual(first, second) - def test_directory_freshness_scan_is_reused_within_one_routing_pass(self) -> None: - source_dir = self.root / "source-dir" - source_dir.mkdir() - (source_dir / "frame.xyz").write_text("frame") + def test_restored_sources_invalidate_memory_and_disk_caches(self) -> None: + """Restored mtimes must invalidate file and directory conversions.""" + for directory in (False, True): + for fresh_process in (False, True): + for preserve_mtime in (False, True): + with self.subTest( + directory=directory, + fresh_process=fresh_process, + preserve_mtime=preserve_mtime, + ): + source = self.source + if directory: + source = self.root / "source-dir" + child = source / "nested" + child.mkdir(parents=True, exist_ok=True) + source_file = child / "frame.extxyz" + else: + source_file = source + source_file.write_text("original") + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems(str(source), fmt="extxyz") + before = _FakeMultiSystems.write_count + stat = source_file.stat() + source_file.write_text("restored") # Same byte count. + mtime = stat.st_mtime_ns - ( + 0 if preserve_mtime else 86400 * 10**9 + ) + os.utime(source_file, ns=(stat.st_atime_ns, mtime)) + if fresh_process: + data_system._DPDATA_CONVERSION_CACHE.clear() + self.assertEqual( + process_systems(str(source), fmt="extxyz"), systems + ) + self.assertEqual(_FakeMultiSystems.write_count, before + 1) + process_systems(str(source), fmt="extxyz") + self.assertEqual(_FakeMultiSystems.write_count, before + 1) + + def test_source_signature_tracks_tree_membership_and_excludes_cache(self) -> None: + """Publishing within the source directory cannot invalidate itself.""" output = self.root / ".deepmd_dpdata_cache" / "cached.lmdb" - output.parent.mkdir() + first = data_system._source_signature(self.root, output) + output.mkdir(parents=True) + (output / "data.mdb").write_text("cache data") + self.assertEqual(first, data_system._source_signature(self.root, output)) + child = self.root / "new.extxyz" + child.write_text("new input") + self.assertNotEqual(first, data_system._source_signature(self.root, output)) + child.unlink() + self.assertEqual(first, data_system._source_signature(self.root, output)) + + def test_source_change_during_conversion_does_not_publish_manifest(self) -> None: + """A concurrent edit cannot mark a possibly mixed snapshot fresh.""" + + def changing_writer(source: Path, fmt: str, out_fmt: str, output: Path) -> None: + _write_minimal_lmdb(str(output)) + source.write_text("changed while converting") + + with patch.object(data_system, "_write_dpdata_conversion", changing_writer): + with self.assertRaisesRegex( + RuntimeError, "input changed during conversion" + ): + process_systems(str(self.source), fmt="extxyz") + output = data_system._conversion_cache_path( + self.source, "extxyz", "deepmd/lmdb" + ) + self.assertFalse(data_system._is_conversion_current(self.source, output)) + self.assertFalse(data_system._conversion_manifest_path(output).exists()) + self.assertFalse(output.with_suffix(".lmdb.lock").exists()) - first = data_system._source_mtime(source_dir, output) - with patch.object( - Path, "rglob", side_effect=AssertionError("unexpected rescan") + def test_missing_or_invalid_manifest_rebuilds_cache(self) -> None: + """Legacy and interrupted conversions cannot validate via output mtime.""" + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems(str(self.source), fmt="extxyz") + manifest = data_system._conversion_manifest_path(Path(systems[0])) + for invalid in (None, "bad json", "null", "{}"): + if invalid is None: + manifest.unlink() + else: + manifest.write_text(invalid) + before = _FakeMultiSystems.write_count + process_systems(str(self.source), fmt="extxyz") + self.assertEqual(_FakeMultiSystems.write_count, before + 1) + + def test_auto_format_handles_xyz_headers_and_decode_errors(self) -> None: + """Probe XYZ headers only when readable, otherwise use the suffix.""" + xyz = self.root / "frame.xyz" + for header, expected in ( + (b"1\nProperties=species:S:1:pos:R:3\nH 0 0 0\n", "extxyz"), + (b'1\nLattice="8 0 0 0 8 0 0 0 8"\nH 0 0 0\n', "extxyz"), + (b"1\nplain comment\nH 0 0 0\n", "xyz"), + (b"1\n\xff\xfe\n", "xyz"), + ): + xyz.write_bytes(header) + self.assertEqual( + data_system._normalize_dpdata_format("auto", xyz), expected + ) + with patch.object(Path, "open", side_effect=OSError("unreadable")): + self.assertEqual(data_system._normalize_dpdata_format("auto", xyz), "xyz") + xyz.unlink() + self.assertEqual(data_system._normalize_dpdata_format("auto", xyz), "xyz") + for name, expected in ( + ("frame.traj", "ase/traj"), + ("frame.extxyz", "extxyz"), + ("frame", "auto"), ): - second = data_system._source_mtime(source_dir, output) + self.assertEqual( + data_system._normalize_dpdata_format("auto", Path(name)), expected + ) + self.assertEqual( + data_system._normalize_dpdata_format("ase", xyz), "ase/structure" + ) - self.assertEqual(first, second) + def test_lmdb_modifier_fails_before_conversion(self) -> None: + """Never silently drop modifiers on direct or converted LMDB data.""" + lmdb_path = self.root / "direct.lmdb" + _write_minimal_lmdb(str(lmdb_path)) + for params in ( + {"systems": str(self.source), "format": "extxyz", "batch_size": 1}, + {"systems": str(lmdb_path), "batch_size": 1}, + ): + with self.assertRaisesRegex(ValueError, "does not support data modifiers"): + get_data(params, 0.0, ["H"], object()) + self.assertEqual(_FakeMultiSystems.write_count, 0) + + def test_lmdb_batch_order_respects_shared_training_seed(self) -> None: + """Shared training seeds preserve reproducible, rank-specific shuffling.""" + lmdb_path = self.root / "seed.lmdb" + _write_repeated_lmdb(str(lmdb_path), 32) + + def order(seed: list[int]) -> list[list[int]]: + data_system.dp_random.seed(seed) + data = get_data( + {"systems": str(lmdb_path), "batch_size": 2}, 0.0, ["H"], None + ) + try: + return [list(batch) for batch in data._sampler] + finally: + data.close() + + self.assertEqual(order([0, 42]), order([0, 42])) + self.assertNotEqual(order([0, 42]), order([0, 43])) + self.assertNotEqual(order([0, 42]), order([1, 42])) + + def test_lock_owner_liveness_branches(self) -> None: + """Only same-host process evidence can establish owner liveness.""" + hostname = data_system.socket.gethostname() + self.assertIsNone(data_system._lock_owner_is_alive({"hostname": "other-host"})) + for payload in ({}, {"pid": None}, {"pid": "invalid"}, {"pid": 0}, {"pid": -1}): + self.assertIsNone( + data_system._lock_owner_is_alive({"hostname": hostname, **payload}) + ) + payload = {"hostname": hostname, "pid": os.getpid(), "process_start": "123"} + with patch.object(data_system, "_process_start_time", return_value="123"): + self.assertTrue(data_system._lock_owner_is_alive(payload)) + with patch.object(data_system, "_process_start_time", return_value="456"): + self.assertFalse(data_system._lock_owner_is_alive(payload)) + for error, expected in ( + (None, True), + (PermissionError(), True), + (ProcessLookupError(), False), + ): + with ( + patch.object(data_system, "_process_start_time", return_value=None), + patch.object(data_system.os, "kill", side_effect=error), + ): + self.assertEqual(data_system._lock_owner_is_alive(payload), expected) + + def test_stale_lock_recovery_preserves_live_owners_and_successors(self) -> None: + """An old heartbeat alone cannot evict a known live local owner.""" + path = self.root / "owner.lock" + for alive, expired, expected in ( + (True, True, False), + (None, False, False), + (None, True, True), + (False, False, True), + ): + path.write_text("{}") + stamp = data_system.time.time() - ( + data_system._CONVERSION_LOCK_STALE_SECONDS + 1 if expired else 0 + ) + os.utime(path, (stamp, stamp)) + with patch.object(data_system, "_lock_owner_is_alive", return_value=alive): + self.assertEqual( + data_system._recover_stale_conversion_lock(path), expected + ) + self.assertEqual(path.exists(), not expected) + path.write_text("{}") + with ( + patch.object(data_system, "_lock_owner_is_alive", return_value=False), + patch.object(data_system, "_same_lock_file", return_value=False), + ): + self.assertFalse(data_system._recover_stale_conversion_lock(path)) + self.assertTrue(path.exists()) + + def test_remote_lock_is_never_evicted_on_stale_heartbeat(self) -> None: + """A remote heartbeat cached by NFS is not proof that its writer died.""" + path = self.root / "remote.lock" + path.write_text(json.dumps({"hostname": "other-host", "pid": 123})) + thirty_seconds_ago = data_system.time.time() - 31 + os.utime(path, (thirty_seconds_ago, thirty_seconds_ago)) + self.assertFalse(data_system._recover_stale_conversion_lock(path)) + expired = ( + data_system.time.time() - data_system._CONVERSION_LOCK_STALE_SECONDS - 1 + ) + os.utime(path, (expired, expired)) + with self.assertRaisesRegex(RuntimeError, "Cannot verify the remote owner"): + data_system._recover_stale_conversion_lock(path) + self.assertTrue(path.exists()) + + def test_lock_initialization_failure_closes_file_and_removes_lock(self) -> None: + """A payload-write or heartbeat-start failure cannot strand a live lock.""" + for target in ("json.dump", "threading.Thread.start"): + path = self.root / "init.lock" + with patch( + f"deepmd.utils.data_system.{target}", side_effect=OSError("init failed") + ): + with path.open("x") as lock_file: + with self.assertRaisesRegex(OSError, "init failed"): + data_system._ConversionLock(path, lock_file) + self.assertTrue(lock_file.closed) + self.assertFalse(path.exists()) def test_dp_test_forwards_conversion_format_from_training_config(self) -> None: config_path = self.root / "input.json" diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index 913a851fe3..27c268c018 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -1165,3 +1165,45 @@ def test_jax_multitask_state_key_normalization_preserves_numeric_task_names() -> assert 1 not in state["models"] assert 0 in state["models"]["1"]["layers"] assert 0 in state["models"]["task"]["layers"] + + +@pytest.mark.timeout(60) +@pytest.mark.parametrize("failure", ["validation", "summary"]) +def test_training_entrypoint_closes_data_on_setup_failure(monkeypatch, failure): + """Partial task setup and summary failures release acquired data systems.""" + from importlib import ( + import_module, + ) + + module = import_module("deepmd.jax.entrypoints.train") + trainer = Mock() + task_model = Mock() + task_model.get_type_map.return_value = ["H"] + task_model.get_rcut.return_value = 6.0 + trainer.models = {"Default": task_model} + trainer_factory = Mock(return_value=trainer) + monkeypatch.setattr(module, "DPTrainer", trainer_factory) + training, validation = Mock(type_map=["H"]), Mock(type_map=["H"]) + acquired = ( + [training, ValueError("validation failed")] + if failure == "validation" + else [training, validation] + ) + monkeypatch.setattr(module, "get_data", Mock(side_effect=acquired)) + summary = Mock( + side_effect=ValueError("summary failed") if failure == "summary" else None + ) + monkeypatch.setattr(module, "print_data_summaries", summary) + if failure == "trainer": + trainer_factory.side_effect = ValueError("trainer failed") + config = { + "model": {"type_map": ["H"]}, + "training": {"training_data": {}, "validation_data": {}, "seed": 42}, + } + with pytest.raises(ValueError, match=f"{failure} failed"): + JAXTrainEntrypoint().run_training( + config, TrainEntrypointOptions(input_file="input.json"), neighbor_stat=None + ) + training.close.assert_called_once_with() + if failure != "validation": + validation.close.assert_called_once_with() diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index d8f70d0dd3..6c67045f45 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -224,6 +224,20 @@ def test_remap_keys(self): class TestLmdbDataset: """Test LmdbDataset class.""" + def test_seed_controls_batch_order(self, lmdb_dir): + """Rank-aware seeds reach the NumPy sampler used by the PT adapter.""" + + def order(seed): + data = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2, seed=seed) + try: + return [list(batch) for batch in data._batch_sampler] + finally: + data.close() + + assert order([0, 42]) == order([0, 42]) + assert order([0, 42]) != order([0, 43]) + assert order([0, 42]) != order([1, 42]) + def test_len(self, lmdb_dir): ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) assert len(ds) == 10 diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 0e350daa5a..756799486f 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -20,6 +20,7 @@ import lmdb import msgpack import numpy as np +import pytest from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module from deepmd.dpmodel.utils.batch import ( @@ -71,6 +72,7 @@ def _encode_array(arr: np.ndarray) -> dict: } +@pytest.mark.timeout(60) class TestConvertedLmdbValidation(unittest.TestCase): """Reject format conversion that resolves to multiple LMDB databases.""" diff --git a/source/tests/tf/test_train_cleanup.py b/source/tests/tf/test_train_cleanup.py new file mode 100644 index 0000000000..9979d629ba --- /dev/null +++ b/source/tests/tf/test_train_cleanup.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Training data ownership during TensorFlow setup failures.""" + +import importlib +from unittest.mock import ( + Mock, +) + +import pytest + + +@pytest.mark.timeout(60) +@pytest.mark.parametrize("failure", ["validation", "summary", "epoch", "build"]) +def test_training_setup_closes_acquired_data(monkeypatch, failure): + """Preflight errors release every data system acquired before the failure.""" + module = importlib.import_module("deepmd.tf.entrypoints.train") + trainer = Mock(data_requirements=[]) + trainer.model.get_type_map.return_value = ["H"] + trainer.model.get_rcut.return_value = 6.0 + monkeypatch.setattr(module, "DPTrainer", Mock(return_value=trainer)) + train_data = Mock(type_map=["H"]) + valid_data = Mock(type_map=["H"]) + config = { + "model": {}, + "training": { + "training_data": {}, + "validation_data": {}, + "numb_steps": 1, + }, + } + if failure == "validation": + acquired = [train_data, ValueError("validation failed")] + else: + acquired = [train_data, valid_data] + if failure == "summary": + train_data.print_summary.side_effect = ValueError("summary failed") + elif failure == "epoch": + config["training"].pop("numb_steps") + config["training"]["numb_epoch"] = 0 + elif failure == "build": + trainer.build.side_effect = ValueError("build failed") + monkeypatch.setattr(module, "get_data", Mock(side_effect=acquired)) + + with pytest.raises(ValueError): + module._do_work(config, Mock(my_rank=0, is_chief=False)) + + train_data.close.assert_called_once_with() + if failure in ("epoch", "build"): + valid_data.close.assert_called_once_with() + else: + valid_data.close.assert_not_called() diff --git a/source/tests/tf2/test_training.py b/source/tests/tf2/test_training.py index ee411b5373..a3c2fbf9be 100644 --- a/source/tests/tf2/test_training.py +++ b/source/tests/tf2/test_training.py @@ -16,6 +16,9 @@ Any, ClassVar, ) +from unittest.mock import ( + Mock, +) import numpy as np import pytest @@ -1235,3 +1238,45 @@ def run(self) -> None: ] assert trainer_calls[-1]["ran"] is True assert trainer_calls[-1]["kwargs"]["min_nbor_dist"] == 0.5 + + +@pytest.mark.timeout(60) +@pytest.mark.parametrize("failure", ["validation", "summary", "trainer"]) +def test_training_entrypoint_closes_data_on_setup_failure(monkeypatch, failure): + """Partial task setup and summary failures release acquired data systems.""" + from importlib import ( + import_module, + ) + + module = import_module("deepmd.tf2.entrypoints.train") + trainer = Mock() + task_model = Mock() + task_model.get_type_map.return_value = ["H"] + task_model.get_rcut.return_value = 6.0 + trainer.models = {"Default": task_model} + trainer_factory = Mock(return_value=trainer) + monkeypatch.setattr(module, "DPTrainer", trainer_factory) + training, validation = Mock(type_map=["H"]), Mock(type_map=["H"]) + acquired = ( + [training, ValueError("validation failed")] + if failure == "validation" + else [training, validation] + ) + monkeypatch.setattr(module, "get_data", Mock(side_effect=acquired)) + summary = Mock( + side_effect=ValueError("summary failed") if failure == "summary" else None + ) + monkeypatch.setattr(module, "print_data_summaries", summary) + if failure == "trainer": + trainer_factory.side_effect = ValueError("trainer failed") + config = { + "model": {"type_map": ["H"]}, + "training": {"training_data": {}, "validation_data": {}, "seed": 42}, + } + with pytest.raises(ValueError, match=f"{failure} failed"): + TF2TrainEntrypoint().run_training( + config, TrainEntrypointOptions(input_file="input.json"), neighbor_stat=None + ) + training.close.assert_called_once_with() + if failure != "validation": + validation.close.assert_called_once_with() From 2f3798826f858c5f6a5a772f0e6cdf91e496231d Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 15 Sep 2026 23:19:01 +0800 Subject: [PATCH 10/12] fix(data): address conversion review follow-ups Handle dangling source symlinks without losing target freshness, clean up failed lock and PyTorch Exportable setup, and share probabilities without importing data loaders. Preserve independent cache checks across lock and conversion phases. Coding-Agent: Codex Codex-Version: codex-cli 0.154.0 Model: gpt-6-astra Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/lmdb_data.py | 9 +- deepmd/pt_expt/entrypoints/main.py | 73 ++++++++------ deepmd/utils/data_system.py | 63 ++++++------ deepmd/utils/probability.py | 41 ++++++++ .../common/test_data_system_conversion.py | 62 ++++++++++-- source/tests/pt_expt/test_entrypoint.py | 95 +++++++++++++++++++ 6 files changed, 271 insertions(+), 72 deletions(-) create mode 100644 deepmd/utils/probability.py diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index a4ef54aac1..4652465561 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -54,6 +54,9 @@ DataRequirementItem, DataRequirementSourcePolicy, ) +from deepmd.utils.probability import ( + prob_sys_size_ext, +) log = logging.getLogger(__name__) @@ -2936,12 +2939,6 @@ def compute_block_targets( ) blocks = nonempty - # Keep this import local: data_system's legacy adapter also uses LMDB. - # Share the NPY probability calculation after filtering empty blocks. - from deepmd.utils.data_system import ( - prob_sys_size_ext, - ) - block_weights = np.asarray([weight for _, _, weight in blocks], dtype=float) if not np.all(np.isfinite(block_weights)): raise ValueError("block weights must be finite") diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 1381750e0d..f32f89be0c 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -33,6 +33,7 @@ ) from deepmd.utils.data_system import ( DeepmdDataSystem, + close_data_systems, conversion_will_write_lmdb, get_data, process_systems, @@ -238,6 +239,7 @@ def get_trainer( def factory( task_config: TrainingTaskConfig, ) -> tuple[DeepmdDataSystem | LmdbDataSystem, Any | None, StatFileSpec]: + """Own partial datasets until all resources for this task are ready.""" type_map = list(task_config.model_params["type_map"]) train_data = _build_data_system( dict(task_config.training_data_params), @@ -247,42 +249,51 @@ def factory( world_size=world_size, ) validation_data = None - if task_config.validation_data_params is not None: - validation_data = _build_data_system( - dict(task_config.validation_data_params), type_map, seed=data_seed + try: + if task_config.validation_data_params is not None: + validation_data = _build_data_system( + dict(task_config.validation_data_params), type_map, seed=data_seed + ) + return ( + train_data, + validation_data, + task_config.stat_file_spec, ) - return ( - train_data, - validation_data, - task_config.stat_file_spec, - ) + except BaseException: + close_data_systems(train_data, validation_data) + raise train_data_map, validation_data_map, stat_file_spec_map = make_task_maps( config, factory ) - print_data_summaries(train_data_map, validation_data_map) - if multi_task: - train_data = train_data_map - validation_data = validation_data_map - stat_file_spec = stat_file_spec_map - else: - task_key = next(iter(train_data_map)) - train_data = train_data_map[task_key] - validation_data = validation_data_map[task_key] - stat_file_spec = stat_file_spec_map[task_key] - - trainer = training.Trainer( - config, - train_data, - stat_file_spec=stat_file_spec, - validation_data=validation_data, - init_model=init_model, - restart_model=restart_model, - finetune_model=finetune_model, - finetune_links=finetune_links, - shared_links=shared_links, - ) - return trainer + try: + print_data_summaries(train_data_map, validation_data_map) + if multi_task: + train_data = train_data_map + validation_data = validation_data_map + stat_file_spec = stat_file_spec_map + else: + task_key = next(iter(train_data_map)) + train_data = train_data_map[task_key] + validation_data = validation_data_map[task_key] + stat_file_spec = stat_file_spec_map[task_key] + + return training.Trainer( + config, + train_data, + stat_file_spec=stat_file_spec, + validation_data=validation_data, + init_model=init_model, + restart_model=restart_model, + finetune_model=finetune_model, + finetune_links=finetune_links, + shared_links=shared_links, + ) + except BaseException: + # Successful construction transfers ownership to the trainer. Until + # then, summary and trainer-setup failures must release every task. + close_data_systems(train_data_map, validation_data_map) + raise class SummaryPrinter(BaseSummaryPrinter): diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index bc3fe47582..ba5da41fec 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -43,6 +43,7 @@ from deepmd.utils.out_stat import ( compute_stats_from_redu, ) +from deepmd.utils.probability import prob_sys_size_ext as prob_sys_size_ext log = logging.getLogger(__name__) @@ -1285,29 +1286,6 @@ def process_sys_probs(sys_probs: list[float], nbatch: int) -> np.ndarray: return ret_prob -def prob_sys_size_ext(keywords: str, nsystems: int, nbatch: int) -> list[float]: - block_str = keywords.split(";")[1:] - block_stt = [] - block_end = [] - block_weights = [] - for ii in block_str: - stt = int(ii.split(":")[0]) - end = int(ii.split(":")[1]) - weight = float(ii.split(":")[2]) - assert weight >= 0, "the weight of a block should be no less than 0" - block_stt.append(stt) - block_end.append(end) - block_weights.append(weight) - nblocks = len(block_str) - block_probs = np.array(block_weights) / np.sum(block_weights) - sys_probs = np.zeros([nsystems], dtype=np.float64) - for ii in range(nblocks): - nbatch_block = nbatch[block_stt[ii] : block_end[ii]] - tmp_prob = [float(i) for i in nbatch_block] / np.sum(nbatch_block) - sys_probs[block_stt[ii] : block_end[ii]] = tmp_prob * block_probs[ii] - return sys_probs - - def _is_deepmd_data_format(fmt: str) -> bool: return fmt in { "deepmd", @@ -1477,7 +1455,30 @@ def _source_signature(source: Path, output: Path) -> str: cache_dir = output.parent.resolve() def add_file(path: Path, name: str) -> None: - stat = path.stat() + try: + stat = path.stat() + except FileNotFoundError: + if not path.is_symlink(): + # A disappearing regular file is an incomplete source scan. + raise + # Broken links may be unrelated to the files dpdata reads. Track + # the link itself, but keep following valid links so target edits + # still invalidate the converted dataset. + stat = path.lstat() + digest.update( + json.dumps( + [ + "dangling_symlink", + name, + str(path.readlink()), + stat.st_mode, + stat.st_size, + stat.st_mtime_ns, + stat.st_ctime_ns, + ] + ).encode() + ) + return digest.update( json.dumps( [name, stat.st_mode, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns] @@ -1569,8 +1570,9 @@ class _ConversionLock: def __init__(self, lock_path: Path, lock_file: IO[str]) -> None: self.path = lock_path - self._stat = os.fstat(lock_file.fileno()) + self._stat = None try: + self._stat = os.fstat(lock_file.fileno()) payload = { "hostname": socket.gethostname(), "pid": os.getpid(), @@ -1590,7 +1592,9 @@ def __init__(self, lock_path: Path, lock_file: IO[str]) -> None: self._thread.start() except BaseException: lock_file.close() - if _same_lock_file(self.path, self._stat): + # Before fstat succeeds this is the fresh, exclusively created + # empty lock; its lease has not expired and no owner is published. + if self._stat is None or _same_lock_file(self.path, self._stat): self.path.unlink(missing_ok=True) raise @@ -1798,8 +1802,9 @@ def _convert_system_by_dpdata( str(output), ) cached_systems = _DPDATA_CONVERSION_CACHE.get(cache_key) + conversion_current = _is_conversion_current(source, output) if cached_systems is not None: - if _is_conversion_current(source, output): + if conversion_current: return cached_systems # A long-lived training/validation process may observe source files # rewritten in place. Drop the fast-path entry so the normal locked @@ -1808,7 +1813,9 @@ def _convert_system_by_dpdata( output.parent.mkdir(parents=True, exist_ok=True) lock_path = output.with_suffix(output.suffix + ".lock") - if not _is_conversion_current(source, output): + # Reuse only this pre-lock check. Waiting, acquiring the lock and writing + # each require fresh scans because the source can change between phases. + if not conversion_current: while True: try: lock_file = lock_path.open("x") diff --git a/deepmd/utils/probability.py b/deepmd/utils/probability.py new file mode 100644 index 0000000000..32a72843fd --- /dev/null +++ b/deepmd/utils/probability.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Sampling probabilities shared by NPY and LMDB without importing loaders.""" + +from collections.abc import ( + Sequence, +) + +import numpy as np + + +def prob_sys_size_ext( + keywords: str, nsystems: int, nbatch: Sequence[int] | np.ndarray +) -> np.ndarray: + """Distribute normalized block weights in proportion to system sizes. + + ``keywords`` has the form ``prob_sys_size;start:end:weight;...``, with + exclusive end indices. ``nbatch`` supplies one size per system: batch + counts for NPY data or frame counts for LMDB data. Callers must supply + nonempty blocks with positive total size and positive total block weight. + Systems outside the specified blocks retain zero probability. + """ + block_str = keywords.split(";")[1:] + block_stt = [] + block_end = [] + block_weights = [] + for ii in block_str: + stt = int(ii.split(":")[0]) + end = int(ii.split(":")[1]) + weight = float(ii.split(":")[2]) + assert weight >= 0, "the weight of a block should be no less than 0" + block_stt.append(stt) + block_end.append(end) + block_weights.append(weight) + nblocks = len(block_str) + block_probs = np.array(block_weights) / np.sum(block_weights) + sys_probs = np.zeros([nsystems], dtype=np.float64) + for ii in range(nblocks): + nbatch_block = nbatch[block_stt[ii] : block_end[ii]] + tmp_prob = [float(i) for i in nbatch_block] / np.sum(nbatch_block) + sys_probs[block_stt[ii] : block_end[ii]] = tmp_prob * block_probs[ii] + return sys_probs diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index bf8ea26265..ef123b458b 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -704,6 +704,57 @@ def test_source_signature_tracks_tree_membership_and_excludes_cache(self) -> Non child.unlink() self.assertEqual(first, data_system._source_signature(self.root, output)) + def test_directory_conversion_tracks_dangling_symlinks(self) -> None: + """Unrelated broken links do not block conversion or hide later changes.""" + source = self.root / "source-tree" + source.mkdir() + (source / "frame.extxyz").write_text(self.source.read_text()) + link = source / "optional.extxyz" + link.symlink_to(self.root / "missing.extxyz") + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + systems = process_systems(str(source), fmt="extxyz") + self.assertEqual(process_systems(str(source), fmt="extxyz"), systems) + self.assertEqual(_FakeMultiSystems.write_count, 1) + + # Changing a broken link's destination is a source-tree change. + link.unlink() + target = self.root / "target.extxyz" + link.symlink_to(target) + process_systems(str(source), fmt="extxyz") + self.assertEqual(_FakeMultiSystems.write_count, 2) + + # Resolving a formerly broken link must also invalidate the cache. + target.write_text(self.source.read_text()) + process_systems(str(source), fmt="extxyz") + self.assertEqual(_FakeMultiSystems.write_count, 3) + + def test_valid_symlink_target_edits_invalidate_conversion(self) -> None: + """Follow valid links even when target size and mtime are unchanged.""" + source = self.root / "source-tree" + source.mkdir() + (source / "frame.extxyz").symlink_to(self.source) + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + process_systems(str(source), fmt="extxyz") + before = self.source.stat() + self.source.write_text(self.source.read_text().replace("H 0", "H 1")) + os.utime(self.source, ns=(before.st_atime_ns, before.st_mtime_ns)) + process_systems(str(source), fmt="extxyz") + self.assertEqual(_FakeMultiSystems.write_count, 2) + + def test_source_signature_rejects_disappearing_regular_file(self) -> None: + """Only dangling links may be fingerprinted after followed stat fails.""" + output = self.root / data_system._DPDATA_CACHE_DIR / "output.lmdb" + original_stat = Path.stat + + def stat_after_removal(path: Path, **kwargs: object) -> os.stat_result: + if path == self.source: + raise FileNotFoundError("source disappeared") + return original_stat(path, **kwargs) + + with patch.object(Path, "stat", stat_after_removal): + with self.assertRaisesRegex(FileNotFoundError, "source disappeared"): + data_system._source_signature(self.root, output) + def test_source_change_during_conversion_does_not_publish_manifest(self) -> None: """A concurrent edit cannot mark a possibly mixed snapshot fresh.""" @@ -785,13 +836,10 @@ def test_lmdb_batch_order_respects_shared_training_seed(self) -> None: def order(seed: list[int]) -> list[list[int]]: data_system.dp_random.seed(seed) - data = get_data( + with get_data( {"systems": str(lmdb_path), "batch_size": 2}, 0.0, ["H"], None - ) - try: + ) as data: return [list(batch) for batch in data._sampler] - finally: - data.close() self.assertEqual(order([0, 42]), order([0, 42])) self.assertNotEqual(order([0, 42]), order([0, 43])) @@ -864,8 +912,8 @@ def test_remote_lock_is_never_evicted_on_stale_heartbeat(self) -> None: self.assertTrue(path.exists()) def test_lock_initialization_failure_closes_file_and_removes_lock(self) -> None: - """A payload-write or heartbeat-start failure cannot strand a live lock.""" - for target in ("json.dump", "threading.Thread.start"): + """Metadata, payload or heartbeat failures cannot strand a fresh lock.""" + for target in ("os.fstat", "json.dump", "threading.Thread.start"): path = self.root / "init.lock" with patch( f"deepmd.utils.data_system.{target}", side_effect=OSError("init failed") diff --git a/source/tests/pt_expt/test_entrypoint.py b/source/tests/pt_expt/test_entrypoint.py index 24bd98efaf..eec56b8409 100644 --- a/source/tests/pt_expt/test_entrypoint.py +++ b/source/tests/pt_expt/test_entrypoint.py @@ -1,4 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from unittest.mock import ( + Mock, + patch, +) + import pytest from deepmd.dpmodel.train import ( @@ -7,10 +12,100 @@ from deepmd.pt_expt.entrypoints.main import ( PTExptTrainEntrypoint, _ensure_pt_expt_model_suffix, + get_trainer, train, ) +@pytest.mark.parametrize("failure", ["validation", "stat_file", "summary", "trainer"]) +def test_get_trainer_closes_data_on_setup_failure(failure: str) -> None: + """Release acquired datasets even when the factory has not returned yet.""" + config = { + "model": {"type_map": ["H"]}, + "training": {"training_data": {}, "validation_data": {}}, + } + training_data, validation_data = Mock(), Mock() + build_results = [training_data, validation_data] + expected_error = RuntimeError + if failure == "validation": + build_results[1] = RuntimeError("validation failed") + elif failure == "stat_file": + # Exercise the real stat_file_spec property after both datasets open. + config["training"]["stat_file_mode"] = "invalid" + expected_error = ValueError + with ( + patch( + "deepmd.pt_expt.entrypoints.main._build_data_system", + side_effect=build_results, + ), + patch( + "deepmd.pt_expt.entrypoints.main.print_data_summaries", + side_effect=RuntimeError("summary failed") + if failure == "summary" + else None, + ), + patch( + "deepmd.pt_expt.entrypoints.main.training.Trainer", + side_effect=RuntimeError("trainer failed") + if failure == "trainer" + else None, + ), + ): + with pytest.raises(expected_error, match=failure): + get_trainer(config) + training_data.close.assert_called_once_with() + if failure == "validation": + validation_data.close.assert_not_called() + else: + validation_data.close.assert_called_once_with() + + +def test_get_trainer_closes_completed_and_partial_tasks() -> None: + """Factory-local cleanup complements make_task_maps cleanup of prior tasks.""" + config = { + "model": { + "model_dict": {"first": {"type_map": ["H"]}, "second": {"type_map": ["H"]}} + }, + "training": { + "data_dict": { + "first": {"training_data": {}, "validation_data": {}}, + "second": {"training_data": {}, "validation_data": {}}, + } + }, + } + datasets = [Mock(), Mock(), Mock()] + with patch( + "deepmd.pt_expt.entrypoints.main._build_data_system", + side_effect=[*datasets, RuntimeError("second validation failed")], + ): + with pytest.raises(RuntimeError, match="second validation failed"): + get_trainer(config) + for dataset in datasets: + dataset.close.assert_called_once_with() + + +def test_get_trainer_transfers_data_ownership_on_success() -> None: + """Successful setup keeps both datasets open for the training loop.""" + config = { + "model": {"type_map": ["H"]}, + "training": {"training_data": {}, "validation_data": {}}, + } + training_data, validation_data = Mock(), Mock() + with ( + patch( + "deepmd.pt_expt.entrypoints.main._build_data_system", + side_effect=[training_data, validation_data], + ), + patch("deepmd.pt_expt.entrypoints.main.print_data_summaries"), + patch("deepmd.pt_expt.entrypoints.main.training.Trainer") as trainer, + ): + assert get_trainer(config) is trainer.return_value + assert trainer.call_args.args[1] is training_data + assert trainer.call_args.kwargs["validation_data"] is validation_data + training_data.close.assert_not_called() + validation_data.close.assert_not_called() + + @pytest.mark.parametrize( ("model_path", "expected"), [ From b3dc211c3ade11c244acf1d2cfa7dc2e2a31475b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 28 Oct 2026 16:43:25 +0800 Subject: [PATCH 11/12] fix(data): reuse LMDB availability partitions and vectorize key mapping Coding-Agent: Codex Codex-Version: codex-cli 0.164.0 Model: gpt-6-astra Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/lmdb_data.py | 33 ++++++-- deepmd/utils/data_system.py | 8 +- doc/train/training-advanced.md | 30 ++++---- source/tests/common/dpmodel/test_lmdb_data.py | 76 +++++++++++++++++++ 4 files changed, 121 insertions(+), 26 deletions(-) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 2da540ac03..5b5e6b2feb 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -2624,6 +2624,20 @@ def original_keys(self, indices: Sequence[int]) -> list[int]: keys.append(int(self._retained_keys[index])) return keys + def original_keys_array(self, indices: Sequence[int]) -> np.ndarray: + """Translate a one-dimensional index group without boxing each key. + + Unlike the small-batch list API, this gather keeps whole-dataset + groups in compact NumPy storage, including gaps left by ``filter:N``. + The returned int64 array is independent of the reader's mapping. + """ + indices = np.asarray(indices, dtype=np.int64) + if indices.ndim != 1: + raise ValueError("dataset indices must be one-dimensional") + if indices.size and (indices.min() < 0 or indices.max() >= self.nframes): + raise IndexError(f"dataset indices out of range [0, {self.nframes})") + return self._retained_keys[indices] + def batch_pad_nloc(self, indices: Sequence[int]) -> int: """Return the atom count every frame of one batch is padded to. @@ -2855,11 +2869,16 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N raise RuntimeError( "LMDB data requirements must be registered before reading any frame" ) + previous_keys = _availability_requirement_keys(self._data_requirements) for item in data_requirement: self._data_requirements[item["key"]] = item self._data_requirements_revision += 1 - self._uniform_availability = None - self._availability_index = None + # Availability depends only on optional tracked field names. Reuse + # the exact PBC partition when unrelated contracts change; otherwise + # legacy per-key registration would scan the entire dataset each time. + if previous_keys != _availability_requirement_keys(self._data_requirements): + self._uniform_availability = None + self._availability_index = None self._per_atom_strides = None def availability_groups(self, indices: np.ndarray) -> list[np.ndarray]: @@ -4456,10 +4475,14 @@ def add( dtype=dtype, source_policy=source_policy, ) + previous_keys = _availability_requirement_keys(self._requirements) self._requirements[key] = requirement.dict - self._uniform_availability = None - self._availability_indices.clear() - self._full_availability_index = None + # Validation partitions have the same availability-only dependency + # as training partitions, even when decoding dtypes/defaults change. + if previous_keys != _availability_requirement_keys(self._requirements): + self._uniform_availability = None + self._availability_indices.clear() + self._full_availability_index = None def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: """Register expected keys from ``DataRequirementItem`` objects. diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 7ad8c60f6b..3bcd27af34 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -976,9 +976,7 @@ def _refresh_groups(self) -> None: any_pbc = False for group_idx, (nloc, indices) in enumerate(groups): group_label = f"{self.lmdb_path}#group={group_idx}:nloc={nloc}" - original_indices = np.asarray( - self._reader.original_keys(indices), dtype=np.int64 - ) + original_indices = self._reader.original_keys_array(indices) frame = self._reader.peek_frame(int(indices[0])) group_pbc = bool(float(frame.get("find_box", 0.0)) > 0.5) any_pbc = any_pbc or group_pbc @@ -992,9 +990,7 @@ def _refresh_groups(self) -> None: chunk = sampled_indices[start : start + chunk_size] set_name = f"{group_label}:chunk={chunk_idx}" self._nloc_set_indices[set_name] = np.asarray(chunk, dtype=np.int64) - stat_groups[set_name] = np.asarray( - self._reader.original_keys(chunk), dtype=np.int64 - ) + stat_groups[set_name] = self._reader.original_keys_array(chunk) data_systems.append( LmdbTestDataNlocView( diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 0534f97e7b..8884a4d7ce 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -155,6 +155,21 @@ This frame-share definition intentionally changes older NPY/HDF5 `prob_uniform` `validation_data` remains a dictionary. Its `systems` accepts the same recursive paths and multiple-LMDB inputs on supported backends; `numb_btch` continues to apply to the whole validation dataset. +### Automatic format conversion + +Set `format` on a training or validation data group to a dpdata format such as +`extxyz`, `abacus/scf`, or `ase/traj`. Each selected input path is converted and +cached separately; `out_format` defaults to `deepmd/lmdb`. LMDB requires an +explicit model `type_map`. Paddle, data modifiers, and explicit `sys_probs` +require another output format such as `deepmd/hdf5`. + +Automatic conversion loads each input path eagerly through dpdata. Its supported +scale is limited to inputs whose parsed arrays **and intermediate copies fit in +available RAM**; it is not a streaming converter. For large trajectories, split +inputs into separate paths or convert them offline before training. Multiple +converted paths use the same data-group sampling rules described above, including +the current frame-share definition of `prob_uniform`. + ### Mixed precision The section {ref}`mixed_precision ` specifies the mixed precision settings, which will enable the mixed precision training workflow for DeePMD-kit. The keys are explained below: @@ -299,18 +314,3 @@ dp freeze -o frozen_model_adjusted_sel.pb Two models should give the same result when the input satisfies both constraints. Note: At this time, this feature is only supported by [`se_e2_a`](../model/train-se-e2-a.md) descriptor with [`set_davg_true`](./train-input.rst) enabled, or `hybrid` composed of the above descriptors. - -### Automatic format conversion - -Set `format` on a training or validation data group to a dpdata format such as -`extxyz`, `abacus/scf`, or `ase/traj`. Each selected input path is converted and -cached separately; `out_format` defaults to `deepmd/lmdb`. LMDB requires an -explicit model `type_map`. Paddle, data modifiers, and explicit `sys_probs` -require another output format such as `deepmd/hdf5`. - -Automatic conversion loads each input path eagerly through dpdata. Its supported -scale is limited to inputs whose parsed arrays **and intermediate copies fit in -available RAM**; it is not a streaming converter. For large trajectories, split -inputs into separate paths or convert them offline before training. Multiple -converted paths use the same data-group sampling rules described above, including -the current frame-share definition of `prob_uniform`. diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index aa6013ef0b..f3e444f2e2 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -1018,6 +1018,65 @@ def recording_open(path, **kwargs): [indices.tolist() for indices in second_groups], ) + def test_availability_cache_depends_only_on_tracked_keys(self): + """Repeated contracts reuse PBC scans; new optional labels repartition.""" + path = _create_lmdb( + f"{self._tmpdir.name}/availability_cache.lmdb", nframes=4, natoms=6 + ) + with lmdb.open(path) as env, env.begin(write=True) as txn: + for index, field in ((1, "cells"), (2, "forces")): + key = f"{index:012d}".encode() + frame = msgpack.unpackb(txn.get(key), raw=False) + frame.pop(field) + txn.put(key, msgpack.packb(frame, use_bin_type=True)) + indices = np.arange(4, dtype=np.int64) + for cls in (LmdbDataReader, LmdbTestData): + with self.subTest(reader=cls.__name__): + reader = cls(path, type_map=self._type_map) + try: + box = DataRequirementItem("box", 9, atomic=False, must=False) + reader.add_data_requirement([box]) + with mock.patch.object( + lmdb_data_module, + "_scan_availability_index", + wraps=lmdb_data_module._scan_availability_index, + ) as scan: + groups = reader.availability_groups(indices) + self.assertEqual(sorted(map(len, groups)), [1, 3]) + self.assertEqual(scan.call_count, 1) + for _ in range(2): + reader.add_data_requirement( + [ + DataRequirementItem("energy", 1, must=True), + box, + DataRequirementItem( + "fparam", 1, source_policy="default" + ), + ] + ) + groups = reader.availability_groups(indices) + self.assertEqual(sorted(map(len, groups)), [1, 3]) + self.assertEqual(scan.call_count, 1) + reader.add_data_requirement( + [DataRequirementItem("force", 3, atomic=True, must=False)] + ) + groups = reader.availability_groups(indices) + self.assertEqual(sorted(map(len, groups)), [1, 1, 2]) + self.assertEqual(scan.call_count, 2) + # A policy change on an existing key must invalidate too. + reader.add_data_requirement( + [ + DataRequirementItem( + "force", 3, atomic=True, source_policy="default" + ) + ] + ) + groups = reader.availability_groups(indices) + self.assertEqual(sorted(map(len, groups)), [1, 3]) + self.assertEqual(scan.call_count, 3) + finally: + reader.close() + def test_failed_exact_scan_restores_random_read_environment(self): """A failed sequential scan leaves the reader usable for training I/O.""" path = _create_lmdb( @@ -2594,6 +2653,23 @@ def test_filter_dataset_index_is_contiguous_and_live(self): with self.assertRaises(IndexError): reader[-1] + def test_vectorized_original_keys_preserve_filtered_gaps(self): + """Bulk gathers preserve key order, duplicates and a gap in the file.""" + path = _create_mixed_sid_nloc_lmdb( + f"{self._tmpdir.name}/key_gaps.lmdb", [(2, 6), (2, 20), (2, 6)] + ) + reader = LmdbDataReader(path, self._type_map, batch_size="filter:10") + self.addCleanup(reader.close) + keys = reader.original_keys_array([3, 0, 2, 3]) + np.testing.assert_array_equal(keys, [5, 0, 4, 5]) + self.assertEqual(keys.dtype, np.int64) + keys[:] = -1 + self.assertEqual(reader.original_keys([3, 0, 2, 3]), [5, 0, 4, 5]) + self.assertEqual(reader.original_keys_array([]).size, 0) + for invalid in ([-1], [len(reader)]): + with self.assertRaises(IndexError): + reader.original_keys_array(invalid) + def test_sampler_with_filter(self): """LmdbBatchSampler only emits retained, same-nloc frames.""" reader = LmdbDataReader( From 05651917c6ae62a0134ebaca6d838a30bd9ac24e Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Wed, 28 Oct 2026 17:45:14 +0800 Subject: [PATCH 12/12] refactor(data): decouple conversion from loader discovery Coding-Agent: Codex Codex-Version: codex-cli 0.164.0 Model: gpt-6-astra Reasoning-Effort: xhigh --- deepmd/utils/data_conversion.py | 582 +++++++++++++++++ deepmd/utils/data_path.py | 40 +- deepmd/utils/data_system.py | 598 +----------------- .../common/test_data_system_conversion.py | 128 ++-- source/tests/pt_expt/test_lmdb_training.py | 2 +- 5 files changed, 689 insertions(+), 661 deletions(-) create mode 100644 deepmd/utils/data_conversion.py diff --git a/deepmd/utils/data_conversion.py b/deepmd/utils/data_conversion.py new file mode 100644 index 0000000000..134305fd0c --- /dev/null +++ b/deepmd/utils/data_conversion.py @@ -0,0 +1,582 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Cached dpdata conversion independent of dataset discovery and loaders. + +This layer owns conversion manifests and writer locks. Path routing supplies +native-format discovery after conversion, keeping the dependency graph acyclic. +""" + +import hashlib +import importlib.metadata +import json +import logging +import os +import shutil +import socket +import tempfile +import threading +import time +from collections.abc import ( + Callable, +) +from pathlib import ( + Path, +) +from typing import ( + IO, + Any, +) + +log = logging.getLogger(__name__) + +_DPDATA_CACHE_DIR = ".deepmd_dpdata_cache" +_DPDATA_DEFAULT_OUT_FORMAT = "deepmd/lmdb" +_DPDATA_CONVERSION_SCHEMA_VERSION = "3" +_DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str, str], list[str]] = {} +_DPDATA_SOURCE_HASH_LIMIT = 1024 * 1024 +_CONVERSION_LOCK_HEARTBEAT_SECONDS = 5.0 +# Allow for shared-filesystem attribute caching before diagnosing a stalled +# heartbeat. Remote owner liveness still cannot be inferred from elapsed time. +_CONVERSION_LOCK_STALE_SECONDS = 300.0 + + +def _is_deepmd_data_format(fmt: str) -> bool: + """Recognize native formats that bypass the external conversion layer.""" + return fmt in { + "deepmd", + "deepmd/raw", + "deepmd/npy", + "deepmd/comp", + "deepmd/npy/mixed", + "deepmd/hdf5", + "deepmd/lmdb", + "lmdb", + } + + +def _is_dpdata_lmdb_format(fmt: str) -> bool: + """Return whether *fmt* names dpdata's DeePMD-compatible LMDB format.""" + return fmt in {"deepmd/lmdb", "lmdb"} + + +def _canonical_dpdata_out_format(out_fmt: str | None) -> str: + """Return the canonical dpdata output format used by conversion caches.""" + if out_fmt is None: + return _DPDATA_DEFAULT_OUT_FORMAT + out_fmt = out_fmt.lower() + return _DPDATA_DEFAULT_OUT_FORMAT if out_fmt == "lmdb" else out_fmt + + +def conversion_will_write_lmdb(data_config: dict[str, Any]) -> bool: + """Whether a non-DeePMD input config will be converted to LMDB.""" + data_format = data_config.get("format") + if data_format is None or _is_deepmd_data_format(data_format.lower()): + return False + out_format = data_config.get("out_format", data_config.get("output_format")) + return _is_dpdata_lmdb_format(_canonical_dpdata_out_format(out_format)) + + +def _looks_like_extxyz(path: Path) -> bool: + """Probe the XYZ comment line without requiring a valid text input.""" + if not path.is_file(): + return False + try: + with path.open(encoding="utf-8") as fp: + fp.readline() + comment = fp.readline() + except (OSError, UnicodeDecodeError): + return False + return "Properties=" in comment or "Lattice=" in comment + + +def _normalize_dpdata_format(fmt: str, source: Path) -> str: + """Resolve aliases and infer common dpdata formats from the input path.""" + fmt = fmt.lower() + if fmt == "ase": + return "ase/structure" + if fmt != "auto": + return fmt + suffix = source.suffix.lower().lstrip(".") + if suffix == "traj": + return "ase/traj" + if suffix == "extxyz" or (suffix == "xyz" and _looks_like_extxyz(source)): + return "extxyz" + return suffix or fmt + + +def _iter_conversion_inputs(path: str, patterns: list[str] | None) -> list[str]: + if patterns is None: + return [path] + root = Path(path) + if not root.is_dir(): + return [path] + matches = [] + for pattern in patterns: + matches.extend(str(match) for match in root.rglob(pattern)) + return sorted(set(matches)) + + +def _conversion_cache_path(source: Path, fmt: str, out_fmt: str) -> Path: + source_resolved = source.resolve(strict=False) + try: + dpdata_version = importlib.metadata.version("dpdata") + except importlib.metadata.PackageNotFoundError: + dpdata_version = "unknown" + digest = hashlib.sha256( + ( + f"{source_resolved}|{fmt}|{out_fmt}|" + f"schema={_DPDATA_CONVERSION_SCHEMA_VERSION}|dpdata={dpdata_version}" + ).encode() + ).hexdigest()[:16] + stem = source_resolved.stem or source_resolved.name or "dataset" + safe_out_fmt = out_fmt.replace("/", "-") + suffix = ".lmdb" if _is_dpdata_lmdb_format(out_fmt) else "" + return Path.cwd() / _DPDATA_CACHE_DIR / f"{stem}-{safe_out_fmt}-{digest}{suffix}" + + +def _source_signature(source: Path, output: Path) -> str: + """Fingerprint source metadata, including restored files and tree changes. + + Compare nanosecond timestamps and sizes for equality, never ordering. + ctime also detects a same-size rewrite whose mtime was preserved on Unix. + Hash small files as well, including on Windows where ctime is creation time, + without rereading large trajectory files on every routing call. Scan on + each call: a time-based memo can hide edits to directory children. + Exclude our cache subtree so publication cannot invalidate its own input. + """ + digest = hashlib.sha256() + cache_dir = output.parent.resolve() + + def add_file(path: Path, name: str) -> None: + try: + stat = path.stat() + except FileNotFoundError: + if not path.is_symlink(): + # A disappearing regular file is an incomplete source scan. + raise + # Broken links may be unrelated to the files dpdata reads. Track + # the link itself, but keep following valid links so target edits + # still invalidate the converted dataset. + stat = path.lstat() + digest.update( + json.dumps( + [ + "dangling_symlink", + name, + str(path.readlink()), + stat.st_mode, + stat.st_size, + stat.st_mtime_ns, + stat.st_ctime_ns, + ] + ).encode() + ) + return + digest.update( + json.dumps( + [name, stat.st_mode, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns] + ).encode() + ) + if stat.st_size <= _DPDATA_SOURCE_HASH_LIMIT: + with path.open("rb") as fp: + # Bound the read even if a concurrent writer grows the input. + digest.update( + hashlib.sha256(fp.read(_DPDATA_SOURCE_HASH_LIMIT + 1)).digest() + ) + + if source.is_file(): + add_file(source, source.name) + elif source.is_dir(): + + def raise_walk_error(error: OSError) -> None: + # An incomplete scan must never validate an old conversion. + raise error + + visited: set[tuple[int, int]] = set() + for root, dirs, files in os.walk( + source, onerror=raise_walk_error, followlinks=True + ): + directory = Path(root) + stat = directory.stat() + identity = (stat.st_dev, stat.st_ino) + relative_root = directory.relative_to(source) + # Track both retargeting and edits below directory links. Inode + # deduplication prevents cycles and repeated scans of aliases. + digest.update( + json.dumps( + [ + str(relative_root), + identity, + str(directory.readlink()) if directory.is_symlink() else None, + ] + ).encode() + ) + if identity in visited: + dirs[:] = [] + continue + visited.add(identity) + dirs[:] = sorted( + name for name in dirs if (Path(root) / name).resolve() != cache_dir + ) + digest.update(json.dumps([str(relative_root), dirs]).encode()) + for name in sorted(files): + add_file(Path(root) / name, str(relative_root / name)) + else: + raise FileNotFoundError(f"dpdata input does not exist: {source}") + return digest.hexdigest() + + +def _conversion_manifest_path(output: Path) -> Path: + """Return the sidecar recording the source used for a published output.""" + return output.with_name(output.name + ".source.json") + + +def _is_conversion_current(source: Path, output: Path) -> bool: + """Reuse only a complete conversion with exactly matching source metadata.""" + if not output.exists(): + return False + try: + manifest = json.loads(_conversion_manifest_path(output).read_text()) + return manifest == {"source_signature": _source_signature(source, output)} + except (OSError, ValueError): + # Missing, damaged, or legacy manifests require a fresh conversion. + return False + + +def _write_conversion_manifest(output: Path, signature: str) -> None: + """Atomically mark a successfully published conversion as reusable.""" + temporary = None + try: + with tempfile.NamedTemporaryFile( + mode="w", dir=output.parent, prefix=f".{output.name}.", delete=False + ) as fp: + temporary = Path(fp.name) + json.dump({"source_signature": signature}, fp) + fp.flush() + os.fsync(fp.fileno()) + os.replace(temporary, _conversion_manifest_path(output)) + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + +def _process_start_time(pid: int) -> str | None: + """Return Linux's stable process start token, if available.""" + try: + stat_text = Path(f"/proc/{pid}/stat").read_text() + except OSError: + return None + fields_after_name = stat_text.rsplit(")", 1)[1].split() + return fields_after_name[19] if len(fields_after_name) > 19 else None + + +def _same_lock_file(lock_path: Path, expected: os.stat_result) -> bool: + """Whether *lock_path* still names the inode originally acquired/read.""" + try: + current = lock_path.stat(follow_symlinks=False) + except FileNotFoundError: + return False + return (current.st_dev, current.st_ino) == (expected.st_dev, expected.st_ino) + + +class _ConversionLock: + """Owned conversion lock with a heartbeat for detecting stalled writers.""" + + def __init__(self, lock_path: Path, lock_file: IO[str]) -> None: + self.path = lock_path + self._stat = None + try: + self._stat = os.fstat(lock_file.fileno()) + payload = { + "hostname": socket.gethostname(), + "pid": os.getpid(), + "process_start": _process_start_time(os.getpid()), + "created": time.time(), + } + # The caller owns the file context, including initialization failures. + json.dump(payload, lock_file) + lock_file.flush() + os.fsync(lock_file.fileno()) + self._stop = threading.Event() + self._thread = threading.Thread( + target=self._heartbeat, + name="deepmd-dpdata-conversion-lock", + daemon=True, + ) + self._thread.start() + except BaseException: + lock_file.close() + # Before fstat succeeds this is the fresh, exclusively created + # empty lock; its lease has not expired and no owner is published. + if self._stat is None or _same_lock_file(self.path, self._stat): + self.path.unlink(missing_ok=True) + raise + + def _heartbeat(self) -> None: + while not self._stop.wait(_CONVERSION_LOCK_HEARTBEAT_SECONDS): + if not _same_lock_file(self.path, self._stat): + return + try: + os.utime(self.path, None, follow_symlinks=False) + except FileNotFoundError: + return + + def release(self) -> None: + """Stop heartbeating and remove only the lock inode we own.""" + self._stop.set() + self._thread.join() + if _same_lock_file(self.path, self._stat): + try: + self.path.unlink() + except FileNotFoundError: + # Another waiter may already have removed this lock. + pass + + +def _lock_owner_is_alive(payload: dict[str, Any]) -> bool | None: + """Return owner liveness locally, or None for another host/invalid data.""" + if payload.get("hostname") != socket.gethostname(): + return None + try: + pid = int(payload["pid"]) + except (KeyError, TypeError, ValueError): + return None + if pid <= 0: + return None + expected_start = payload.get("process_start") + current_start = _process_start_time(pid) + if expected_start is not None and current_start is not None: + return str(expected_start) == current_start + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _recover_stale_conversion_lock(lock_path: Path) -> bool: + """Recover dead local owners and abandoned, uninitialized lock files. + + A remote host's heartbeat may look stale because of NFS attribute caching + or clock skew. Never evict an identified remote writer on that evidence; + report the stalled lock for owner verification instead of racing it. + """ + try: + lock_stat = lock_path.stat(follow_symlinks=False) + except FileNotFoundError: + return True + except OSError: + return False + try: + payload = json.loads(lock_path.read_text()) + except (OSError, ValueError): + payload = {} + if not isinstance(payload, dict): + payload = {} + + owner_alive = _lock_owner_is_alive(payload) + lease_expired = time.time() - lock_stat.st_mtime > _CONVERSION_LOCK_STALE_SECONDS + if ( + owner_alive is None + and lease_expired + and payload.get("hostname") not in (None, socket.gethostname()) + ): + raise RuntimeError( + f"Cannot verify the remote owner of stalled dpdata conversion lock " + f"{lock_path} on {payload['hostname']}. Check that the owner has " + "stopped before removing the lock and retrying." + ) + stale = owner_alive is False or (owner_alive is None and lease_expired) + if not stale or not _same_lock_file(lock_path, lock_stat): + return False + log.warning("Recovering stale dpdata conversion lock %s", lock_path) + try: + lock_path.unlink() + except FileNotFoundError: + # Recovery is complete if another waiter removed the same stale lock. + pass + return True + + +def _wait_for_conversion(source: Path, output: Path, lock_path: Path) -> bool: + """Wait without rescanning the source tree while a valid writer owns it.""" + while lock_path.exists(): + if _recover_stale_conversion_lock(lock_path): + continue + time.sleep(1.0) + # Freshness is checked once after publication, not once per waiter-second. + return _is_conversion_current(source, output) + + +def _remove_path(path: Path) -> None: + # Check links before directories: Path.is_dir follows a directory symlink, + # while cache cleanup must never recurse into a target outside the cache. + if path.is_symlink(): + path.unlink() + elif path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + + +def _publish_conversion_output(tmp_output: Path, output: Path) -> None: + """Publish a non-LMDB conversion without discarding a valid old cache.""" + if not output.exists() or output.is_symlink() or not output.is_dir(): + os.replace(tmp_output, output) + backup = output.with_name(f".{output.name}.backup") + _remove_path(backup) + return + + # POSIX cannot replace a non-empty directory directly. Preserve the old + # cache under a sibling name and restore it if the second rename fails. + backup = output.with_name(f".{output.name}.backup") + _remove_path(backup) + os.replace(output, backup) + try: + os.replace(tmp_output, output) + except Exception: + os.replace(backup, output) + raise + else: + _remove_path(backup) + + +def _write_dpdata_conversion( + source: Path, fmt: str, out_fmt: str, output: Path +) -> None: + """Load *source* with dpdata and publish it in a DeePMD format. + + dpdata 1.1.0 makes ``deepmd/lmdb`` writes transactional: it stages and + validates the complete database before atomically publishing it. Use that + writer directly so DeePMD-kit does not duplicate or weaken its overwrite + guarantees. Other dpdata formats do not share that contract, so they keep + the cache-level temporary output used for failure isolation. + """ + try: + import dpdata + except ImportError as exc: + raise ImportError( + "dpdata is required when training_data.format or " + "validation_data.format is specified. Install dpdata to enable " + "automatic dataset conversion." + ) from exc + + multi_systems = dpdata.MultiSystems() + try: + multi_systems.load_systems_from_file(str(source), fmt=fmt) + except (NotImplementedError, ValueError) as labeled_error: + # dpdata 1.1 exposes an explicit unlabeled path. This matters for + # structure-only EXTXYZ/ASE inputs, which are valid descriptor data + # even though a supervised loss may later require labels. + unlabeled_systems = dpdata.MultiSystems() + try: + unlabeled_systems.load_systems_from_file( + str(source), fmt=fmt, labeled=False + ) + except (NotImplementedError, TypeError, ValueError): + try: + labeled_system = dpdata.LabeledSystem(str(source), fmt=fmt) + except Exception: + raise labeled_error from None + multi_systems = dpdata.MultiSystems(labeled_system) + else: + log.info("Loaded unlabeled dpdata input %s using format %s", source, fmt) + multi_systems = unlabeled_systems + if len(multi_systems) == 0: + raise RuntimeError(f"No frames were loaded by dpdata from {source}") + + if _is_dpdata_lmdb_format(out_fmt): + multi_systems.to(out_fmt, str(output), overwrite=True) + return + + tmp_output = output.with_name(f".{output.name}.{os.getpid()}.tmp") + _remove_path(tmp_output) + try: + multi_systems.to(out_fmt, str(tmp_output)) + _publish_conversion_output(tmp_output, output) + except Exception: + _remove_path(tmp_output) + raise + + +def _convert_system_by_dpdata( + source_path: str, + fmt: str, + out_fmt: str | None, + *, + expand_systems: Callable[[str], list[str]], +) -> list[str]: + """Convert one source while keeping native discovery outside this module. + + The caller supplies discovery for HDF5/NPY outputs so conversion does not + import loaders or their path-routing module, including during startup. + """ + source = Path(source_path) + fmt = _normalize_dpdata_format(fmt, source) + out_fmt = _canonical_dpdata_out_format(out_fmt) + output = _conversion_cache_path(source, fmt, out_fmt) + cache_key = ( + str(Path.cwd().resolve(strict=False)), + str(source.resolve(strict=False)), + fmt, + out_fmt, + str(output), + ) + cached_systems = _DPDATA_CONVERSION_CACHE.get(cache_key) + conversion_current = _is_conversion_current(source, output) + if cached_systems is not None: + if conversion_current: + return cached_systems + # A long-lived training/validation process may observe source files + # rewritten in place. Drop the fast-path entry so the normal locked + # conversion flow refreshes the on-disk result before it is reused. + del _DPDATA_CONVERSION_CACHE[cache_key] + + output.parent.mkdir(parents=True, exist_ok=True) + lock_path = output.with_suffix(output.suffix + ".lock") + # Reuse only this pre-lock check. Waiting, acquiring the lock and writing + # each require fresh scans because the source can change between phases. + if not conversion_current: + while True: + try: + lock_file = lock_path.open("x") + except FileExistsError: + if _wait_for_conversion(source, output, lock_path): + break + continue + else: + with lock_file: + conversion_lock = _ConversionLock(lock_path, lock_file) + try: + if not _is_conversion_current(source, output): + log.info( + "Converting %s from dpdata format %s to %s at %s", + source, + fmt, + out_fmt, + output, + ) + signature = _source_signature(source, output) + # Invalidate before replacing output: a crash must never + # pair a new dataset with an older source manifest. + _conversion_manifest_path(output).unlink(missing_ok=True) + _write_dpdata_conversion(source, fmt, out_fmt, output) + if signature != _source_signature(source, output): + raise RuntimeError( + f"dpdata input changed during conversion: {source}. " + "Retry with an unchanged input dataset." + ) + _write_conversion_manifest(output, signature) + finally: + conversion_lock.release() + break + + if _is_dpdata_lmdb_format(out_fmt): + converted_systems = [str(output)] + else: + converted_systems = expand_systems(str(output)) + if not converted_systems: + raise RuntimeError(f"No DeePMD systems were found in converted file {output}") + _DPDATA_CONVERSION_CACHE[cache_key] = converted_systems + return converted_systems diff --git a/deepmd/utils/data_path.py b/deepmd/utils/data_path.py index 3dd645be6e..e1d093e890 100644 --- a/deepmd/utils/data_path.py +++ b/deepmd/utils/data_path.py @@ -8,6 +8,12 @@ Literal, ) +from deepmd.utils.data_conversion import ( + _convert_system_by_dpdata, + _is_deepmd_data_format, + _iter_conversion_inputs, + conversion_will_write_lmdb, +) from deepmd.utils.path import ( DPOSPath, DPPath, @@ -116,25 +122,23 @@ def process_systems( str((Path(root_dir) / file_name).resolve()) + separator + group for file_name, separator, group in (path.partition("#") for path in paths) ] - if fmt is not None: - # Keep conversion lazy: native discovery is also imported while the - # data-system module initializes. Validate backend support before any - # converter can create cache files. - from deepmd.utils.data_system import ( - _is_deepmd_data_format, - ) - from deepmd.utils.data_system import process_systems as convert_systems - from deepmd.utils.data_system import ( - validate_backend_data_config, - ) - - if not _is_deepmd_data_format(fmt.lower()): - validate_backend_data_config( - {"format": fmt, "out_format": out_fmt}, - backend_name="Selected", - lmdb_supported="lmdb" in formats, + if fmt is not None and not _is_deepmd_data_format(fmt.lower()): + # Reject unsupported output before conversion can create cache files. + if "lmdb" not in formats and conversion_will_write_lmdb( + {"format": fmt, "out_format": out_fmt} + ): + raise NotImplementedError( + "Selected backend does not support LMDB data yet. " + "Choose out_format='deepmd/hdf5' for automatic conversion." + ) + return [ + system + for path in paths + for input_path in _iter_conversion_inputs(path, patterns) + for system in _convert_system_by_dpdata( + input_path, fmt, out_fmt, expand_systems=expand_sys_str ) - return convert_systems(paths, patterns, fmt=fmt, out_fmt=out_fmt) + ] return [ system for path in paths diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 3bcd27af34..ee4e880468 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -1,15 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import collections -import hashlib -import importlib.metadata -import json import logging -import os -import shutil -import socket -import tempfile -import threading -import time import warnings from collections.abc import ( Callable, @@ -24,7 +15,6 @@ SimpleNamespace, ) from typing import ( - IO, Any, ) @@ -51,9 +41,11 @@ resolve_data_groups, resolve_data_sources, ) +from deepmd.utils.data_conversion import ( + conversion_will_write_lmdb as conversion_will_write_lmdb, +) from deepmd.utils.data_path import ( DatasetFormat, - expand_sys_str, is_lmdb, ) from deepmd.utils.data_path import process_systems as discover_systems @@ -72,16 +64,6 @@ log = logging.getLogger(__name__) -_DPDATA_CACHE_DIR = ".deepmd_dpdata_cache" -_DPDATA_DEFAULT_OUT_FORMAT = "deepmd/lmdb" -_DPDATA_CONVERSION_SCHEMA_VERSION = "3" -_DPDATA_CONVERSION_CACHE: dict[tuple[str, str, str, str, str], list[str]] = {} -_DPDATA_SOURCE_HASH_LIMIT = 1024 * 1024 -_CONVERSION_LOCK_HEARTBEAT_SECONDS = 5.0 -# Allow for shared-filesystem attribute caching before diagnosing a stalled -# heartbeat. Remote owner liveness still cannot be inferred from elapsed time. -_CONVERSION_LOCK_STALE_SECONDS = 300.0 - class DeepmdDataSystem: """Class for manipulating many data systems. @@ -1353,41 +1335,6 @@ def print_summary( ) -def _is_deepmd_data_format(fmt: str) -> bool: - return fmt in { - "deepmd", - "deepmd/raw", - "deepmd/npy", - "deepmd/comp", - "deepmd/npy/mixed", - "deepmd/hdf5", - "deepmd/lmdb", - "lmdb", - } - - -def _is_dpdata_lmdb_format(fmt: str) -> bool: - """Return whether *fmt* names dpdata's DeePMD-compatible LMDB format.""" - return fmt in {"deepmd/lmdb", "lmdb"} - - -def _canonical_dpdata_out_format(out_fmt: str | None) -> str: - """Return the canonical dpdata output format used by conversion caches.""" - if out_fmt is None: - return _DPDATA_DEFAULT_OUT_FORMAT - out_fmt = out_fmt.lower() - return _DPDATA_DEFAULT_OUT_FORMAT if out_fmt == "lmdb" else out_fmt - - -def conversion_will_write_lmdb(data_config: dict[str, Any]) -> bool: - """Whether a non-DeePMD input config will be converted to LMDB.""" - data_format = data_config.get("format") - if data_format is None or _is_deepmd_data_format(data_format.lower()): - return False - out_format = data_config.get("out_format", data_config.get("output_format")) - return _is_dpdata_lmdb_format(_canonical_dpdata_out_format(out_format)) - - def validate_backend_data_config( data_config: DataConfig, *, @@ -1452,504 +1399,6 @@ def close_one(value: Any) -> None: close_one(value) -def _looks_like_extxyz(path: Path) -> bool: - """Probe the XYZ comment line without requiring a valid text input.""" - if not path.is_file(): - return False - try: - with path.open(encoding="utf-8") as fp: - fp.readline() - comment = fp.readline() - except (OSError, UnicodeDecodeError): - return False - return "Properties=" in comment or "Lattice=" in comment - - -def _normalize_dpdata_format(fmt: str, source: Path) -> str: - """Resolve aliases and infer common dpdata formats from the input path.""" - fmt = fmt.lower() - if fmt == "ase": - return "ase/structure" - if fmt != "auto": - return fmt - suffix = source.suffix.lower().lstrip(".") - if suffix == "traj": - return "ase/traj" - if suffix == "extxyz" or (suffix == "xyz" and _looks_like_extxyz(source)): - return "extxyz" - return suffix or fmt - - -def _iter_conversion_inputs(path: str, patterns: list[str] | None) -> list[str]: - if patterns is None: - return [path] - root = Path(path) - if not root.is_dir(): - return [path] - matches = [] - for pattern in patterns: - matches.extend(str(match) for match in root.rglob(pattern)) - return sorted(set(matches)) - - -def _conversion_cache_path(source: Path, fmt: str, out_fmt: str) -> Path: - source_resolved = source.resolve(strict=False) - try: - dpdata_version = importlib.metadata.version("dpdata") - except importlib.metadata.PackageNotFoundError: - dpdata_version = "unknown" - digest = hashlib.sha256( - ( - f"{source_resolved}|{fmt}|{out_fmt}|" - f"schema={_DPDATA_CONVERSION_SCHEMA_VERSION}|dpdata={dpdata_version}" - ).encode() - ).hexdigest()[:16] - stem = source_resolved.stem or source_resolved.name or "dataset" - safe_out_fmt = out_fmt.replace("/", "-") - suffix = ".lmdb" if _is_dpdata_lmdb_format(out_fmt) else "" - return Path.cwd() / _DPDATA_CACHE_DIR / f"{stem}-{safe_out_fmt}-{digest}{suffix}" - - -def _source_signature(source: Path, output: Path) -> str: - """Fingerprint source metadata, including restored files and tree changes. - - Compare nanosecond timestamps and sizes for equality, never ordering. - ctime also detects a same-size rewrite whose mtime was preserved on Unix. - Hash small files as well, including on Windows where ctime is creation time, - without rereading large trajectory files on every routing call. Scan on - each call: a time-based memo can hide edits to directory children. - Exclude our cache subtree so publication cannot invalidate its own input. - """ - digest = hashlib.sha256() - cache_dir = output.parent.resolve() - - def add_file(path: Path, name: str) -> None: - try: - stat = path.stat() - except FileNotFoundError: - if not path.is_symlink(): - # A disappearing regular file is an incomplete source scan. - raise - # Broken links may be unrelated to the files dpdata reads. Track - # the link itself, but keep following valid links so target edits - # still invalidate the converted dataset. - stat = path.lstat() - digest.update( - json.dumps( - [ - "dangling_symlink", - name, - str(path.readlink()), - stat.st_mode, - stat.st_size, - stat.st_mtime_ns, - stat.st_ctime_ns, - ] - ).encode() - ) - return - digest.update( - json.dumps( - [name, stat.st_mode, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns] - ).encode() - ) - if stat.st_size <= _DPDATA_SOURCE_HASH_LIMIT: - with path.open("rb") as fp: - # Bound the read even if a concurrent writer grows the input. - digest.update( - hashlib.sha256(fp.read(_DPDATA_SOURCE_HASH_LIMIT + 1)).digest() - ) - - if source.is_file(): - add_file(source, source.name) - elif source.is_dir(): - - def raise_walk_error(error: OSError) -> None: - # An incomplete scan must never validate an old conversion. - raise error - - visited: set[tuple[int, int]] = set() - for root, dirs, files in os.walk( - source, onerror=raise_walk_error, followlinks=True - ): - directory = Path(root) - stat = directory.stat() - identity = (stat.st_dev, stat.st_ino) - relative_root = directory.relative_to(source) - # Track both retargeting and edits below directory links. Inode - # deduplication prevents cycles and repeated scans of aliases. - digest.update( - json.dumps( - [ - str(relative_root), - identity, - str(directory.readlink()) if directory.is_symlink() else None, - ] - ).encode() - ) - if identity in visited: - dirs[:] = [] - continue - visited.add(identity) - dirs[:] = sorted( - name for name in dirs if (Path(root) / name).resolve() != cache_dir - ) - digest.update(json.dumps([str(relative_root), dirs]).encode()) - for name in sorted(files): - add_file(Path(root) / name, str(relative_root / name)) - else: - raise FileNotFoundError(f"dpdata input does not exist: {source}") - return digest.hexdigest() - - -def _conversion_manifest_path(output: Path) -> Path: - """Return the sidecar recording the source used for a published output.""" - return output.with_name(output.name + ".source.json") - - -def _is_conversion_current(source: Path, output: Path) -> bool: - """Reuse only a complete conversion with exactly matching source metadata.""" - if not output.exists(): - return False - try: - manifest = json.loads(_conversion_manifest_path(output).read_text()) - return manifest == {"source_signature": _source_signature(source, output)} - except (OSError, ValueError): - # Missing, damaged, or legacy manifests require a fresh conversion. - return False - - -def _write_conversion_manifest(output: Path, signature: str) -> None: - """Atomically mark a successfully published conversion as reusable.""" - temporary = None - try: - with tempfile.NamedTemporaryFile( - mode="w", dir=output.parent, prefix=f".{output.name}.", delete=False - ) as fp: - temporary = Path(fp.name) - json.dump({"source_signature": signature}, fp) - fp.flush() - os.fsync(fp.fileno()) - os.replace(temporary, _conversion_manifest_path(output)) - finally: - if temporary is not None: - temporary.unlink(missing_ok=True) - - -def _process_start_time(pid: int) -> str | None: - """Return Linux's stable process start token, if available.""" - try: - stat_text = Path(f"/proc/{pid}/stat").read_text() - except OSError: - return None - fields_after_name = stat_text.rsplit(")", 1)[1].split() - return fields_after_name[19] if len(fields_after_name) > 19 else None - - -def _same_lock_file(lock_path: Path, expected: os.stat_result) -> bool: - """Whether *lock_path* still names the inode originally acquired/read.""" - try: - current = lock_path.stat(follow_symlinks=False) - except FileNotFoundError: - return False - return (current.st_dev, current.st_ino) == (expected.st_dev, expected.st_ino) - - -class _ConversionLock: - """Owned conversion lock with a heartbeat for detecting stalled writers.""" - - def __init__(self, lock_path: Path, lock_file: IO[str]) -> None: - self.path = lock_path - self._stat = None - try: - self._stat = os.fstat(lock_file.fileno()) - payload = { - "hostname": socket.gethostname(), - "pid": os.getpid(), - "process_start": _process_start_time(os.getpid()), - "created": time.time(), - } - # The caller owns the file context, including initialization failures. - json.dump(payload, lock_file) - lock_file.flush() - os.fsync(lock_file.fileno()) - self._stop = threading.Event() - self._thread = threading.Thread( - target=self._heartbeat, - name="deepmd-dpdata-conversion-lock", - daemon=True, - ) - self._thread.start() - except BaseException: - lock_file.close() - # Before fstat succeeds this is the fresh, exclusively created - # empty lock; its lease has not expired and no owner is published. - if self._stat is None or _same_lock_file(self.path, self._stat): - self.path.unlink(missing_ok=True) - raise - - def _heartbeat(self) -> None: - while not self._stop.wait(_CONVERSION_LOCK_HEARTBEAT_SECONDS): - if not _same_lock_file(self.path, self._stat): - return - try: - os.utime(self.path, None, follow_symlinks=False) - except FileNotFoundError: - return - - def release(self) -> None: - """Stop heartbeating and remove only the lock inode we own.""" - self._stop.set() - self._thread.join() - if _same_lock_file(self.path, self._stat): - try: - self.path.unlink() - except FileNotFoundError: - # Another waiter may already have removed this lock. - pass - - -def _lock_owner_is_alive(payload: dict[str, Any]) -> bool | None: - """Return owner liveness locally, or None for another host/invalid data.""" - if payload.get("hostname") != socket.gethostname(): - return None - try: - pid = int(payload["pid"]) - except (KeyError, TypeError, ValueError): - return None - if pid <= 0: - return None - expected_start = payload.get("process_start") - current_start = _process_start_time(pid) - if expected_start is not None and current_start is not None: - return str(expected_start) == current_start - try: - os.kill(pid, 0) - except ProcessLookupError: - return False - except PermissionError: - return True - return True - - -def _recover_stale_conversion_lock(lock_path: Path) -> bool: - """Recover dead local owners and abandoned, uninitialized lock files. - - A remote host's heartbeat may look stale because of NFS attribute caching - or clock skew. Never evict an identified remote writer on that evidence; - report the stalled lock for owner verification instead of racing it. - """ - try: - lock_stat = lock_path.stat(follow_symlinks=False) - except FileNotFoundError: - return True - except OSError: - return False - try: - payload = json.loads(lock_path.read_text()) - except (OSError, ValueError): - payload = {} - if not isinstance(payload, dict): - payload = {} - - owner_alive = _lock_owner_is_alive(payload) - lease_expired = time.time() - lock_stat.st_mtime > _CONVERSION_LOCK_STALE_SECONDS - if ( - owner_alive is None - and lease_expired - and payload.get("hostname") not in (None, socket.gethostname()) - ): - raise RuntimeError( - f"Cannot verify the remote owner of stalled dpdata conversion lock " - f"{lock_path} on {payload['hostname']}. Check that the owner has " - "stopped before removing the lock and retrying." - ) - stale = owner_alive is False or (owner_alive is None and lease_expired) - if not stale or not _same_lock_file(lock_path, lock_stat): - return False - log.warning("Recovering stale dpdata conversion lock %s", lock_path) - try: - lock_path.unlink() - except FileNotFoundError: - # Recovery is complete if another waiter removed the same stale lock. - pass - return True - - -def _wait_for_conversion(source: Path, output: Path, lock_path: Path) -> bool: - """Wait without rescanning the source tree while a valid writer owns it.""" - while lock_path.exists(): - if _recover_stale_conversion_lock(lock_path): - continue - time.sleep(1.0) - # Freshness is checked once after publication, not once per waiter-second. - return _is_conversion_current(source, output) - - -def _remove_path(path: Path) -> None: - # Check links before directories: Path.is_dir follows a directory symlink, - # while cache cleanup must never recurse into a target outside the cache. - if path.is_symlink(): - path.unlink() - elif path.is_dir(): - shutil.rmtree(path) - elif path.exists(): - path.unlink() - - -def _publish_conversion_output(tmp_output: Path, output: Path) -> None: - """Publish a non-LMDB conversion without discarding a valid old cache.""" - if not output.exists() or output.is_symlink() or not output.is_dir(): - os.replace(tmp_output, output) - backup = output.with_name(f".{output.name}.backup") - _remove_path(backup) - return - - # POSIX cannot replace a non-empty directory directly. Preserve the old - # cache under a sibling name and restore it if the second rename fails. - backup = output.with_name(f".{output.name}.backup") - _remove_path(backup) - os.replace(output, backup) - try: - os.replace(tmp_output, output) - except Exception: - os.replace(backup, output) - raise - else: - _remove_path(backup) - - -def _write_dpdata_conversion( - source: Path, fmt: str, out_fmt: str, output: Path -) -> None: - """Load *source* with dpdata and publish it in a DeePMD format. - - dpdata 1.1.0 makes ``deepmd/lmdb`` writes transactional: it stages and - validates the complete database before atomically publishing it. Use that - writer directly so DeePMD-kit does not duplicate or weaken its overwrite - guarantees. Other dpdata formats do not share that contract, so they keep - the cache-level temporary output used for failure isolation. - """ - try: - import dpdata - except ImportError as exc: - raise ImportError( - "dpdata is required when training_data.format or " - "validation_data.format is specified. Install dpdata to enable " - "automatic dataset conversion." - ) from exc - - multi_systems = dpdata.MultiSystems() - try: - multi_systems.load_systems_from_file(str(source), fmt=fmt) - except (NotImplementedError, ValueError) as labeled_error: - # dpdata 1.1 exposes an explicit unlabeled path. This matters for - # structure-only EXTXYZ/ASE inputs, which are valid descriptor data - # even though a supervised loss may later require labels. - unlabeled_systems = dpdata.MultiSystems() - try: - unlabeled_systems.load_systems_from_file( - str(source), fmt=fmt, labeled=False - ) - except (NotImplementedError, TypeError, ValueError): - try: - labeled_system = dpdata.LabeledSystem(str(source), fmt=fmt) - except Exception: - raise labeled_error from None - multi_systems = dpdata.MultiSystems(labeled_system) - else: - log.info("Loaded unlabeled dpdata input %s using format %s", source, fmt) - multi_systems = unlabeled_systems - if len(multi_systems) == 0: - raise RuntimeError(f"No frames were loaded by dpdata from {source}") - - if _is_dpdata_lmdb_format(out_fmt): - multi_systems.to(out_fmt, str(output), overwrite=True) - return - - tmp_output = output.with_name(f".{output.name}.{os.getpid()}.tmp") - _remove_path(tmp_output) - try: - multi_systems.to(out_fmt, str(tmp_output)) - _publish_conversion_output(tmp_output, output) - except Exception: - _remove_path(tmp_output) - raise - - -def _convert_system_by_dpdata( - source_path: str, fmt: str, out_fmt: str | None -) -> list[str]: - source = Path(source_path) - fmt = _normalize_dpdata_format(fmt, source) - out_fmt = _canonical_dpdata_out_format(out_fmt) - output = _conversion_cache_path(source, fmt, out_fmt) - cache_key = ( - str(Path.cwd().resolve(strict=False)), - str(source.resolve(strict=False)), - fmt, - out_fmt, - str(output), - ) - cached_systems = _DPDATA_CONVERSION_CACHE.get(cache_key) - conversion_current = _is_conversion_current(source, output) - if cached_systems is not None: - if conversion_current: - return cached_systems - # A long-lived training/validation process may observe source files - # rewritten in place. Drop the fast-path entry so the normal locked - # conversion flow refreshes the on-disk result before it is reused. - del _DPDATA_CONVERSION_CACHE[cache_key] - - output.parent.mkdir(parents=True, exist_ok=True) - lock_path = output.with_suffix(output.suffix + ".lock") - # Reuse only this pre-lock check. Waiting, acquiring the lock and writing - # each require fresh scans because the source can change between phases. - if not conversion_current: - while True: - try: - lock_file = lock_path.open("x") - except FileExistsError: - if _wait_for_conversion(source, output, lock_path): - break - continue - else: - with lock_file: - conversion_lock = _ConversionLock(lock_path, lock_file) - try: - if not _is_conversion_current(source, output): - log.info( - "Converting %s from dpdata format %s to %s at %s", - source, - fmt, - out_fmt, - output, - ) - signature = _source_signature(source, output) - # Invalidate before replacing output: a crash must never - # pair a new dataset with an older source manifest. - _conversion_manifest_path(output).unlink(missing_ok=True) - _write_dpdata_conversion(source, fmt, out_fmt, output) - if signature != _source_signature(source, output): - raise RuntimeError( - f"dpdata input changed during conversion: {source}. " - "Retry with an unchanged input dataset." - ) - _write_conversion_manifest(output, signature) - finally: - conversion_lock.release() - break - - if _is_dpdata_lmdb_format(out_fmt): - converted_systems = [str(output)] - else: - converted_systems = expand_sys_str(str(output)) - if not converted_systems: - raise RuntimeError(f"No DeePMD systems were found in converted file {output}") - _DPDATA_CONVERSION_CACHE[cache_key] = converted_systems - return converted_systems - - def process_systems( systems: str | list[str], patterns: list[str] | None = None, @@ -1985,38 +1434,15 @@ def process_systems( result_systems: list of str The valid systems """ - # Normalize input to a list of paths to search - if isinstance(systems, str): - search_paths = [systems] - elif isinstance(systems, list): - search_paths = systems - else: - # Handle unsupported input types - raise ValueError( - f"Invalid systems type: {type(systems)}. Must be str or list[str]." - ) - - if fmt is not None: - fmt = fmt.lower() - if _is_deepmd_data_format(fmt): - fmt = None - if fmt is None: - return discover_systems(systems, patterns, formats=formats, root_dir=root_dir) - if root_dir is not None: - search_paths = [str((Path(root_dir) / path).resolve()) for path in search_paths] - - conversion_inputs: list[str] = [] - if fmt is not None: - for path in search_paths: - conversion_inputs.extend(_iter_conversion_inputs(path, patterns)) - - # Iterate over the search_paths list and apply expansion logic to each path - result_systems = [] - if fmt is not None: - for input_path in conversion_inputs: - result_systems.extend(_convert_system_by_dpdata(input_path, fmt, out_fmt)) - - return result_systems + # Preserve this legacy signature while sharing routing with all backends. + return discover_systems( + systems, + patterns, + formats=formats, + root_dir=root_dir, + fmt=fmt, + out_fmt=out_fmt, + ) def get_data( diff --git a/source/tests/common/test_data_system_conversion.py b/source/tests/common/test_data_system_conversion.py index bc72546b9d..aa38be2526 100644 --- a/source/tests/common/test_data_system_conversion.py +++ b/source/tests/common/test_data_system_conversion.py @@ -26,6 +26,7 @@ ) from deepmd.entrypoints.test import test as run_model_test from deepmd.utils import ( + data_conversion, data_system, ) from deepmd.utils.data import ( @@ -229,14 +230,14 @@ def setUp(self) -> None: _FakeMultiSystems.write_count = 0 _FakeMultiSystems.load_calls = [] _FakeMultiSystems.to_calls = [] - data_system._DPDATA_CONVERSION_CACHE.clear() + data_conversion._DPDATA_CONVERSION_CACHE.clear() self.fake_dpdata = types.SimpleNamespace( MultiSystems=_FakeMultiSystems, LabeledSystem=_FakeLabeledSystem, ) def tearDown(self) -> None: - data_system._DPDATA_CONVERSION_CACHE.clear() + data_conversion._DPDATA_CONVERSION_CACHE.clear() def test_process_systems_defaults_to_deepmd_lmdb_and_reuses_cache(self) -> None: with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): @@ -317,7 +318,7 @@ def test_real_dpdata_lmdb_writer_is_deepmd_compatible(self) -> None: original_stat = self.source.stat() self.source.write_text(extxyz.replace("energy=0.0", "energy=42.0")) os.utime(self.source, (original_stat.st_atime, original_stat.st_mtime - 86400)) - data_system._DPDATA_CONVERSION_CACHE.clear() + data_conversion._DPDATA_CONVERSION_CACHE.clear() systems_again = process_systems(str(self.source), fmt="extxyz") self.assertEqual(len(systems), 1) @@ -343,7 +344,7 @@ def test_cache_cleanup_unlinks_directory_symlink(self) -> None: link.parent.mkdir() link.symlink_to(target, target_is_directory=True) - data_system._remove_path(link) + data_conversion._remove_path(link) self.assertFalse(link.exists()) self.assertTrue(target.is_dir()) @@ -629,9 +630,9 @@ def test_waiter_checks_source_only_after_lock_release(self) -> None: lock_path.write_text( json.dumps( { - "hostname": data_system.socket.gethostname(), + "hostname": data_conversion.socket.gethostname(), "pid": os.getpid(), - "process_start": data_system._process_start_time(os.getpid()), + "process_start": data_conversion._process_start_time(os.getpid()), } ) ) @@ -645,14 +646,16 @@ def release_after_long_valid_conversion(_seconds: float) -> None: with ( patch.object( - data_system.time, "sleep", release_after_long_valid_conversion + data_conversion.time, "sleep", release_after_long_valid_conversion ), patch.object( - data_system, "_is_conversion_current", return_value=True + data_conversion, + "_is_conversion_current", + return_value=True, ) as is_current, ): self.assertTrue( - data_system._wait_for_conversion(self.source, output, lock_path) + data_conversion._wait_for_conversion(self.source, output, lock_path) ) self.assertEqual(waits, 301) @@ -663,14 +666,14 @@ def test_dead_conversion_owner_is_recovered(self) -> None: lock_path.write_text( json.dumps( { - "hostname": data_system.socket.gethostname(), + "hostname": data_conversion.socket.gethostname(), "pid": 2147483647, "process_start": "missing", } ) ) - self.assertTrue(data_system._recover_stale_conversion_lock(lock_path)) + self.assertTrue(data_conversion._recover_stale_conversion_lock(lock_path)) self.assertFalse(lock_path.exists()) def test_directory_publication_rolls_back_on_failure(self) -> None: @@ -690,9 +693,9 @@ def fail_publication(source: str | Path, target: str | Path) -> None: raise OSError("publish failed") real_replace(source, target) - with patch.object(data_system.os, "replace", fail_publication): + with patch.object(data_conversion.os, "replace", fail_publication): with self.assertRaisesRegex(OSError, "publish failed"): - data_system._publish_conversion_output(staged, output) + data_conversion._publish_conversion_output(staged, output) self.assertEqual((output / "old.txt").read_text(), "old") self.assertFalse((output / "new.txt").exists()) @@ -724,12 +727,14 @@ def load_systems_from_file( def test_cache_path_changes_with_dpdata_version(self) -> None: with patch.object( - data_system.importlib.metadata, "version", side_effect=["1.1.0", "1.2.0"] + data_conversion.importlib.metadata, + "version", + side_effect=["1.1.0", "1.2.0"], ): - first = data_system._conversion_cache_path( + first = data_conversion._conversion_cache_path( self.source, "extxyz", "deepmd/lmdb" ) - second = data_system._conversion_cache_path( + second = data_conversion._conversion_cache_path( self.source, "extxyz", "deepmd/lmdb" ) @@ -764,7 +769,7 @@ def test_restored_sources_invalidate_memory_and_disk_caches(self) -> None: ) os.utime(source_file, ns=(stat.st_atime_ns, mtime)) if fresh_process: - data_system._DPDATA_CONVERSION_CACHE.clear() + data_conversion._DPDATA_CONVERSION_CACHE.clear() self.assertEqual( process_systems(str(source), fmt="extxyz"), systems ) @@ -775,15 +780,15 @@ def test_restored_sources_invalidate_memory_and_disk_caches(self) -> None: def test_source_signature_tracks_tree_membership_and_excludes_cache(self) -> None: """Publishing within the source directory cannot invalidate itself.""" output = self.root / ".deepmd_dpdata_cache" / "cached.lmdb" - first = data_system._source_signature(self.root, output) + first = data_conversion._source_signature(self.root, output) output.mkdir(parents=True) (output / "data.mdb").write_text("cache data") - self.assertEqual(first, data_system._source_signature(self.root, output)) + self.assertEqual(first, data_conversion._source_signature(self.root, output)) child = self.root / "new.extxyz" child.write_text("new input") - self.assertNotEqual(first, data_system._source_signature(self.root, output)) + self.assertNotEqual(first, data_conversion._source_signature(self.root, output)) child.unlink() - self.assertEqual(first, data_system._source_signature(self.root, output)) + self.assertEqual(first, data_conversion._source_signature(self.root, output)) def test_directory_conversion_tracks_dangling_symlinks(self) -> None: """Unrelated broken links do not block conversion or hide later changes.""" @@ -851,7 +856,7 @@ def test_directory_symlinks_track_edits_retargets_and_cycles(self) -> None: def test_source_signature_rejects_disappearing_regular_file(self) -> None: """Only dangling links may be fingerprinted after followed stat fails.""" - output = self.root / data_system._DPDATA_CACHE_DIR / "output.lmdb" + output = self.root / data_conversion._DPDATA_CACHE_DIR / "output.lmdb" original_stat = Path.stat def stat_after_removal(path: Path, **kwargs: object) -> os.stat_result: @@ -861,7 +866,7 @@ def stat_after_removal(path: Path, **kwargs: object) -> os.stat_result: with patch.object(Path, "stat", stat_after_removal): with self.assertRaisesRegex(FileNotFoundError, "source disappeared"): - data_system._source_signature(self.root, output) + data_conversion._source_signature(self.root, output) def test_source_change_during_conversion_does_not_publish_manifest(self) -> None: """A concurrent edit cannot mark a possibly mixed snapshot fresh.""" @@ -870,23 +875,23 @@ def changing_writer(source: Path, fmt: str, out_fmt: str, output: Path) -> None: _write_minimal_lmdb(str(output)) source.write_text("changed while converting") - with patch.object(data_system, "_write_dpdata_conversion", changing_writer): + with patch.object(data_conversion, "_write_dpdata_conversion", changing_writer): with self.assertRaisesRegex( RuntimeError, "input changed during conversion" ): process_systems(str(self.source), fmt="extxyz") - output = data_system._conversion_cache_path( + output = data_conversion._conversion_cache_path( self.source, "extxyz", "deepmd/lmdb" ) - self.assertFalse(data_system._is_conversion_current(self.source, output)) - self.assertFalse(data_system._conversion_manifest_path(output).exists()) + self.assertFalse(data_conversion._is_conversion_current(self.source, output)) + self.assertFalse(data_conversion._conversion_manifest_path(output).exists()) self.assertFalse(output.with_suffix(".lmdb.lock").exists()) def test_missing_or_invalid_manifest_rebuilds_cache(self) -> None: """Legacy and interrupted conversions cannot validate via output mtime.""" with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): systems = process_systems(str(self.source), fmt="extxyz") - manifest = data_system._conversion_manifest_path(Path(systems[0])) + manifest = data_conversion._conversion_manifest_path(Path(systems[0])) for invalid in (None, "bad json", "null", "{}"): if invalid is None: manifest.unlink() @@ -907,22 +912,24 @@ def test_auto_format_handles_xyz_headers_and_decode_errors(self) -> None: ): xyz.write_bytes(header) self.assertEqual( - data_system._normalize_dpdata_format("auto", xyz), expected + data_conversion._normalize_dpdata_format("auto", xyz), expected ) with patch.object(Path, "open", side_effect=OSError("unreadable")): - self.assertEqual(data_system._normalize_dpdata_format("auto", xyz), "xyz") + self.assertEqual( + data_conversion._normalize_dpdata_format("auto", xyz), "xyz" + ) xyz.unlink() - self.assertEqual(data_system._normalize_dpdata_format("auto", xyz), "xyz") + self.assertEqual(data_conversion._normalize_dpdata_format("auto", xyz), "xyz") for name, expected in ( ("frame.traj", "ase/traj"), ("frame.extxyz", "extxyz"), ("frame", "auto"), ): self.assertEqual( - data_system._normalize_dpdata_format("auto", Path(name)), expected + data_conversion._normalize_dpdata_format("auto", Path(name)), expected ) self.assertEqual( - data_system._normalize_dpdata_format("ase", xyz), "ase/structure" + data_conversion._normalize_dpdata_format("ase", xyz), "ase/structure" ) def test_lmdb_modifier_fails_before_conversion(self) -> None: @@ -955,27 +962,31 @@ def order(seed: list[int]) -> list[list[int]]: def test_lock_owner_liveness_branches(self) -> None: """Only same-host process evidence can establish owner liveness.""" - hostname = data_system.socket.gethostname() - self.assertIsNone(data_system._lock_owner_is_alive({"hostname": "other-host"})) + hostname = data_conversion.socket.gethostname() + self.assertIsNone( + data_conversion._lock_owner_is_alive({"hostname": "other-host"}) + ) for payload in ({}, {"pid": None}, {"pid": "invalid"}, {"pid": 0}, {"pid": -1}): self.assertIsNone( - data_system._lock_owner_is_alive({"hostname": hostname, **payload}) + data_conversion._lock_owner_is_alive({"hostname": hostname, **payload}) ) payload = {"hostname": hostname, "pid": os.getpid(), "process_start": "123"} - with patch.object(data_system, "_process_start_time", return_value="123"): - self.assertTrue(data_system._lock_owner_is_alive(payload)) - with patch.object(data_system, "_process_start_time", return_value="456"): - self.assertFalse(data_system._lock_owner_is_alive(payload)) + with patch.object(data_conversion, "_process_start_time", return_value="123"): + self.assertTrue(data_conversion._lock_owner_is_alive(payload)) + with patch.object(data_conversion, "_process_start_time", return_value="456"): + self.assertFalse(data_conversion._lock_owner_is_alive(payload)) for error, expected in ( (None, True), (PermissionError(), True), (ProcessLookupError(), False), ): with ( - patch.object(data_system, "_process_start_time", return_value=None), - patch.object(data_system.os, "kill", side_effect=error), + patch.object(data_conversion, "_process_start_time", return_value=None), + patch.object(data_conversion.os, "kill", side_effect=error), ): - self.assertEqual(data_system._lock_owner_is_alive(payload), expected) + self.assertEqual( + data_conversion._lock_owner_is_alive(payload), expected + ) def test_stale_lock_recovery_preserves_live_owners_and_successors(self) -> None: """An old heartbeat alone cannot evict a known live local owner.""" @@ -987,36 +998,40 @@ def test_stale_lock_recovery_preserves_live_owners_and_successors(self) -> None: (False, False, True), ): path.write_text("{}") - stamp = data_system.time.time() - ( - data_system._CONVERSION_LOCK_STALE_SECONDS + 1 if expired else 0 + stamp = data_conversion.time.time() - ( + data_conversion._CONVERSION_LOCK_STALE_SECONDS + 1 if expired else 0 ) os.utime(path, (stamp, stamp)) - with patch.object(data_system, "_lock_owner_is_alive", return_value=alive): + with patch.object( + data_conversion, "_lock_owner_is_alive", return_value=alive + ): self.assertEqual( - data_system._recover_stale_conversion_lock(path), expected + data_conversion._recover_stale_conversion_lock(path), expected ) self.assertEqual(path.exists(), not expected) path.write_text("{}") with ( - patch.object(data_system, "_lock_owner_is_alive", return_value=False), - patch.object(data_system, "_same_lock_file", return_value=False), + patch.object(data_conversion, "_lock_owner_is_alive", return_value=False), + patch.object(data_conversion, "_same_lock_file", return_value=False), ): - self.assertFalse(data_system._recover_stale_conversion_lock(path)) + self.assertFalse(data_conversion._recover_stale_conversion_lock(path)) self.assertTrue(path.exists()) def test_remote_lock_is_never_evicted_on_stale_heartbeat(self) -> None: """A remote heartbeat cached by NFS is not proof that its writer died.""" path = self.root / "remote.lock" path.write_text(json.dumps({"hostname": "other-host", "pid": 123})) - thirty_seconds_ago = data_system.time.time() - 31 + thirty_seconds_ago = data_conversion.time.time() - 31 os.utime(path, (thirty_seconds_ago, thirty_seconds_ago)) - self.assertFalse(data_system._recover_stale_conversion_lock(path)) + self.assertFalse(data_conversion._recover_stale_conversion_lock(path)) expired = ( - data_system.time.time() - data_system._CONVERSION_LOCK_STALE_SECONDS - 1 + data_conversion.time.time() + - data_conversion._CONVERSION_LOCK_STALE_SECONDS + - 1 ) os.utime(path, (expired, expired)) with self.assertRaisesRegex(RuntimeError, "Cannot verify the remote owner"): - data_system._recover_stale_conversion_lock(path) + data_conversion._recover_stale_conversion_lock(path) self.assertTrue(path.exists()) def test_lock_initialization_failure_closes_file_and_removes_lock(self) -> None: @@ -1024,11 +1039,12 @@ def test_lock_initialization_failure_closes_file_and_removes_lock(self) -> None: for target in ("os.fstat", "json.dump", "threading.Thread.start"): path = self.root / "init.lock" with patch( - f"deepmd.utils.data_system.{target}", side_effect=OSError("init failed") + f"deepmd.utils.data_conversion.{target}", + side_effect=OSError("init failed"), ): with path.open("x") as lock_file: with self.assertRaisesRegex(OSError, "init failed"): - data_system._ConversionLock(path, lock_file) + data_conversion._ConversionLock(path, lock_file) self.assertTrue(lock_file.closed) self.assertFalse(path.exists()) diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 32707b2f5f..93dd1e3d81 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -99,7 +99,7 @@ def test_neighbor_stat_and_training_data_convert_each_input(self) -> None: } for factory in (get_neighbor_stat_data, _build_data_system): with patch( - "deepmd.utils.data_system._convert_system_by_dpdata", + "deepmd.utils.data_path._convert_system_by_dpdata", side_effect=[[], RuntimeError("second source reached")], ) as convert: with self.assertRaisesRegex(RuntimeError, "second source reached"):