diff --git a/deepmd/dpmodel/train/data.py b/deepmd/dpmodel/train/data.py index 86fde438f2..b18d8c197a 100644 --- a/deepmd/dpmodel/train/data.py +++ b/deepmd/dpmodel/train/data.py @@ -22,6 +22,9 @@ DataConfig, data_group_configs, ) +from deepmd.utils.data_system import ( + close_data_systems, +) from deepmd.utils.stat_file import ( StatFileMode, StatFileSpec, @@ -106,15 +109,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 f5e39a98ac..5b5e6b2feb 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -66,6 +66,9 @@ DataRequirementSourcePolicy, ) from deepmd.utils.data_path import is_lmdb as is_lmdb +from deepmd.utils.data_sampling import ( + prob_sys_size_ext, +) log = logging.getLogger(__name__) @@ -400,7 +403,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) @@ -770,6 +775,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( @@ -2305,8 +2315,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 @@ -2556,14 +2575,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]) @@ -2580,6 +2593,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] = [] @@ -2592,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. @@ -2823,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]: @@ -2864,6 +2915,10 @@ def availability_groups(self, indices: np.ndarray) -> list[np.ndarray]: if not availability_keys: self._uniform_availability = True return [indices] + # A scalar box flag controls periodic interactions for the entire + # batch, so a bounded label probe is not a safe structural partition. + if "box" in availability_keys: + self._uniform_availability = False if self._uniform_availability is None: self._uniform_availability = _probe_uniform_availability( self._transaction(), @@ -3206,10 +3261,6 @@ def compute_block_targets( Each element is ``(system_indices_in_block, target_frame_count)``. Returns an empty list if no resampling is needed (all targets == actual). """ - from deepmd.utils.data_sampling import ( - prob_sys_size_ext, - ) - if auto_prob_style == "prob_uniform": retained = [index for index, count in enumerate(system_nframes) if count > 0] if not retained: @@ -3219,6 +3270,9 @@ def compute_block_targets( return [] return [([index], target) for index in retained] + if not auto_prob_style.startswith("prob_sys_size"): + raise RuntimeError(f"Unknown auto prob style: {auto_prob_style}") + # 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:] @@ -3271,6 +3325,14 @@ def compute_block_targets( ) blocks = nonempty + 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") + if np.any(block_weights < 0): + raise ValueError("the weight of a block should be no less than 0") + if np.sum(block_weights) <= 0: + raise ValueError("the sum of block weights should be greater than 0") + # Compute per-system probabilities using the standard function sys_probs = prob_sys_size_ext(auto_prob_style, nsystems, system_nframes) @@ -4052,7 +4114,7 @@ def make_neighbor_stat_data( ) -> Any: """Create a duck-typed DeepmdDataSystem-like object for neighbor stat from LMDB. - Samples up to *max_frames* frames, groups them by nloc, and returns an + Samples up to *max_frames* frames, groups them by nloc and periodicity, and returns an object whose attributes satisfy the interface expected by ``NeighborStat.iterator()`` and ``UpdateSel.get_nbor_stat()``. """ @@ -4061,33 +4123,44 @@ 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[ + tuple[int, bool], list[tuple[np.ndarray, np.ndarray, np.ndarray | None]] + ] = {} + for idx in indices: + frame = reader[int(idx)] + atype = frame["atype"] + nloc = len(atype) + box = frame.get("box") + # Neighbor kernels also use one PBC flag per source. Partition the + # actual sampled frames, including zero-cell nonperiodic sentinels. + pbc = box is not None and not np.allclose(box, 0.0) + pbc = pbc and bool(frame.get("find_box", 1.0)) + nloc_frames.setdefault((nloc, pbc), []).append( + (frame["coord"], atype, box if pbc else None) + ) + ntypes = len(type_map) if type_map else reader._ntypes + finally: + reader.close() # Build per-nloc data_system proxies data_systems = [] system_dirs: list[str] = [] - for nloc, frames in nloc_frames.items(): + for (nloc, has_box), frames in nloc_frames.items(): coords = np.stack([c.reshape(nloc * 3) for c, _, _ in frames]) types = np.stack([a.reshape(nloc) for _, a, _ in frames]) - has_box = frames[0][2] is not None boxes = np.stack([b.reshape(9) for _, _, b in frames]) if has_box else None set_data = {"coord": coords, "type": types, "box": boxes} - label = f"lmdb:{nloc}atoms" + label = f"lmdb:{nloc}atoms:pbc={has_box}" proxy = SimpleNamespace( dirs=[label], pbc=has_box, @@ -4098,7 +4171,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, @@ -4262,19 +4334,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.""" @@ -4306,6 +4376,10 @@ def availability_groups(self, frame_indices: np.ndarray) -> list[np.ndarray]: if not availability_keys: self._uniform_availability = True return [frame_indices] + # A scalar box flag controls periodic interactions for the entire + # batch, so a bounded label probe is not a safe structural partition. + if "box" in availability_keys: + self._uniform_availability = False if self._uniform_availability is None: groups = list(self._nloc_groups.values()) per_group = max(1, _AVAILABILITY_PROBE_FRAMES // max(len(groups), 1)) @@ -4401,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. @@ -4637,14 +4715,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 a6cdb3bfa9..dc6dfba0cd 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -118,6 +118,8 @@ def test( group["systems"], patterns=group.get("rglob_patterns"), root_dir=Path(input_json).parent, + fmt=group.get("format"), + out_fmt=group.get("out_format", group.get("output_format")), ) ) elif datafile is not None: diff --git a/deepmd/jax/entrypoints/train.py b/deepmd/jax/entrypoints/train.py index e86e5d4c8c..60db758f02 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 @@ -189,23 +190,32 @@ 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() - log.info("finished training") - log.info(f"wall time: {(end_time - start_time):.3f} s") + 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 +306,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 7495c1cc06..e21a3fd358 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -75,6 +75,7 @@ ) from deepmd.utils.data_system import ( get_data, + validate_backend_data_config, ) from deepmd.utils.model_preset import ( expand_model_preset, @@ -294,6 +295,18 @@ def train( config = update_deepmd_input(config, warning=True, dump="input_v2_compat.json") config = normalize(config, multi_task=multi_task) + # Validate every task and validation source before neighbor-stat loading + # can run the converter. Paddle has no LMDB reader. + task_data = ( + config["training"]["data_dict"].values() if multi_task else [config["training"]] + ) + for data_config in task_data: + for key in ("training_data", "validation_data"): + if data_config.get(key) is not None: + validate_backend_data_config( + data_config[key], backend_name="Paddle", lmdb_supported=False + ) + # do neighbor stat min_nbor_dist = None if not skip_neighbor_stat: @@ -304,7 +317,11 @@ def train( if not multi_task: type_map = config["model"].get("type_map") train_data = get_data( - config["training"]["training_data"], 0, type_map, None + config["training"]["training_data"], + 0, + type_map, + None, + lmdb_supported=False, ) config["model"], min_nbor_dist = BaseModel.update_sel( train_data, type_map, config["model"] @@ -318,6 +335,7 @@ def train( 0, type_map, None, + lmdb_supported=False, ) config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( BaseModel.update_sel( diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index 297895ff24..56cfcba811 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -72,6 +72,7 @@ system_probabilities, ) from deepmd.utils.data_system import ( + close_data_systems, print_summary, ) @@ -361,49 +362,70 @@ def create_dataset( modifier: BaseModifier | None = None, ) -> DpLoaderSet | LmdbDataset | GroupedDataset: """Construct native readers from the shared recursive group configuration.""" + from deepmd.utils.data_system import ( + conversion_will_write_lmdb, + validate_lmdb_sampling_options, + ) + + for params in data_group_configs(config): + if conversion_will_write_lmdb(params): + validate_lmdb_sampling_options(params, modifier=modifier) + if not type_map: + raise ValueError( + "LMDB training data requires model/type_map to be set." + ) groups = resolve_data_groups(config) configs = data_group_configs(config) datasets = [] - for group, params in zip(groups, configs, strict=True): - sources = resolve_data_sources(group) - has_lmdb = any(source.format == "lmdb" for source in sources) - children = [] - child_configs = [] - for source in sources: - if source.format == "deepmd": - child = DpLoaderSet( - list(source.systems), - source.batch_size, - type_map, - seed=seed, - modifier=modifier, + acquired = [] + try: + for group, params in zip(groups, configs, strict=True): + sources = resolve_data_sources(group) + has_lmdb = any(source.format == "lmdb" for source in sources) + if has_lmdb: + validate_lmdb_sampling_options(params, modifier=modifier) + children = [] + child_configs = [] + for source in sources: + if source.format == "deepmd": + child = DpLoaderSet( + list(source.systems), + source.batch_size, + type_map, + seed=seed, + modifier=modifier, + ) + else: + child = LmdbDataset( + source.systems[0], + type_map, + source.batch_size, + auto_prob_style=group.auto_prob, + seed=seed, + ) + children.append(child) + acquired.append(child) + child_configs.append( + {**params, "sys_probs": group.sys_probs if not has_lmdb else None} ) - else: - child = LmdbDataset( - source.systems[0], - type_map, - source.batch_size, - auto_prob_style=group.auto_prob, + datasets.append( + children[0] + if len(children) == 1 + else GroupedDataset( + children, + child_configs, + [None] * len(children), + uniform=group.auto_prob == "prob_uniform", ) - children.append(child) - child_configs.append( - {**params, "sys_probs": group.sys_probs if not has_lmdb else None} - ) - datasets.append( - children[0] - if len(children) == 1 - else GroupedDataset( - children, - child_configs, - [None] * len(children), - uniform=group.auto_prob == "prob_uniform", ) + return ( + datasets[0] + if len(datasets) == 1 + else GroupedDataset(datasets, configs, [group.prob for group in groups]) ) - return ( - datasets[0] - if len(datasets) == 1 - else GroupedDataset(datasets, configs, [group.prob for group in groups]) - ) + except BaseException: + close_data_systems(acquired) + raise def create_dataloader( diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index 70dd25b864..c654a92d3d 100644 --- a/deepmd/pt/utils/lmdb_dataset.py +++ b/deepmd/pt/utils/lmdb_dataset.py @@ -11,6 +11,7 @@ Any, ) +import numpy as np import torch from torch.utils.data import ( DataLoader, @@ -206,6 +207,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__( @@ -214,6 +217,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( @@ -236,9 +240,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/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 7bb3d6afc2..07450a2bf2 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -46,6 +46,7 @@ ) from deepmd.utils.data_system import ( DeepmdDataSystem, + close_data_systems, get_data, get_neighbor_stat_data, ) @@ -151,6 +152,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( task_config.training_data_params, @@ -160,42 +162,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/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index d7b40ebfee..d28a1a1bdd 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 ( @@ -257,77 +258,84 @@ 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.", + 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, - num_epoch, - total_numb_batch, + origin_type_map=origin_type_map, + stat_file_path=stat_file_path, ) - 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_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, - ) - 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 +363,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 3350e98c1c..fdf2bf4252 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 @@ -165,38 +166,47 @@ 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) + 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() - log.info("finished training") - log.info("wall time: %.3f s", end_time - start_time) + 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() + 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 +259,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/argcheck.py b/deepmd/utils/argcheck.py index 751b6c69b6..b4aee63cbd 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5791,6 +5791,24 @@ def training_data_args() -> Argument: 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. " + "Defaults to `deepmd/lmdb`; it is used only when `format` requests " + "conversion from a non-DeePMD format. LMDB requires an explicit " + "model/type_map and supports multiple input paths through data groups. " + "LMDB does not support Paddle, data modifiers, or explicit sys_probs. " + "Choose `deepmd/hdf5` for those capabilities. Other supported outputs " + "include `deepmd/npy`. Each input path is loaded eagerly by dpdata: " + "the parsed data and conversion copies must fit in memory. Split large " + "inputs into separate paths or preconvert them before training." + ) 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\ @@ -5840,6 +5858,20 @@ def training_data_args() -> Argument: doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") + doc_patterns, ), + Argument( + "format", + [str, None], + optional=True, + doc=doc_format, + ), + Argument( + "out_format", + [str, None], + optional=True, + default="deepmd/lmdb", + doc=doc_out_format, + alias=["output_format"], + ), Argument( "batch_size", [list[int], int, str], @@ -5920,6 +5952,24 @@ 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. " + "Defaults to `deepmd/lmdb`; it is used only when `format` requests " + "conversion from a non-DeePMD format. LMDB requires an explicit " + "model/type_map and supports multiple input paths through data groups. " + "LMDB does not support Paddle, data modifiers, or explicit sys_probs. " + "Choose `deepmd/hdf5` for those capabilities. Other supported outputs " + "include `deepmd/npy`. Each input path is loaded eagerly by dpdata: " + "the parsed data and conversion copies must fit in memory. Split large " + "inputs into separate paths or preconvert them before training." + ) 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\ @@ -5951,6 +6001,20 @@ def validation_data_args() -> list[ doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") + doc_patterns, ), + Argument( + "format", + [str, None], + optional=True, + doc=doc_format, + ), + Argument( + "out_format", + [str, None], + optional=True, + default="deepmd/lmdb", + doc=doc_out_format, + alias=["output_format"], + ), Argument( "batch_size", [list[int], int, str], diff --git a/deepmd/utils/data_config.py b/deepmd/utils/data_config.py index 917e51849e..e7c9c1af00 100644 --- a/deepmd/utils/data_config.py +++ b/deepmd/utils/data_config.py @@ -155,7 +155,11 @@ def resolve_data_groups( if "lmdb" not in formats and any(is_lmdb(path) for path in roots): raise ValueError("LMDB training data requires the pt or pt_expt backend.") systems = process_systems( - group["systems"], group.get("rglob_patterns"), formats=formats + group["systems"], + group.get("rglob_patterns"), + formats=formats, + fmt=group.get("format"), + out_fmt=group.get("out_format", group.get("output_format")), ) if not systems: raise ValueError(f"Data group {index} contains no supported data systems.") 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 4559121a3f..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, @@ -95,6 +101,8 @@ def process_systems( *, formats: tuple[DatasetFormat, ...] = ("deepmd", "lmdb"), root_dir: str | Path | None = None, + fmt: str | None = None, + out_fmt: str | None = None, ) -> list[str]: """Expand every input path through the same format-aware discovery. @@ -114,6 +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 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 [ system for path in paths diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 65379bca5c..ee4e880468 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -8,6 +8,9 @@ from functools import ( cached_property, ) +from pathlib import ( + Path, +) from types import ( SimpleNamespace, ) @@ -16,6 +19,9 @@ ) import numpy as np +from typing_extensions import ( + Self, +) import deepmd.utils.random as dp_random from deepmd.common import ( @@ -31,13 +37,18 @@ from deepmd.utils.data_config import ( DataConfig, check_type_map_consistency, + data_group_configs, 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, is_lmdb, ) -from deepmd.utils.data_path import process_systems as process_systems +from deepmd.utils.data_path import process_systems as discover_systems from deepmd.utils.data_sampling import ( data_child_seeds, ) @@ -802,6 +813,427 @@ def _make_auto_ts(self, percent: float) -> list[int]: return ts +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 | list[int | str] = "auto", + 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 ( + LmdbBatchIterator, + LmdbBatchSampler, + LmdbDataReader, + LmdbTestData, + compute_block_targets, + ) + + 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." + ) + + self.lmdb_path = lmdb_path + self._type_map = list(type_map) + self._closed = False + self._data_dict = { + "box": DataRequirementItem( + "box", + 9, + atomic=False, + must=False, + default=np.zeros(9, dtype=GLOBAL_NP_FLOAT_PRECISION), + ).dict, + "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, + }, + } + + 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( + auto_prob_style, + self._reader.nsystems, + self._reader.system_nframes, + ) + self._sampler = LmdbBatchSampler( + self._reader, + shuffle=True, + seed=seed, + block_targets=block_targets, + ) + self.nbatches = [self._sampler.total_batches] + self._batch_iterator = LmdbBatchIterator(self._reader, self._sampler, 0) + self._refresh_groups() + + 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._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 = 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 + + 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] = self._reader.original_keys_array(chunk) + + 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 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._refresh_groups() + self.nbatches = [self._sampler.total_batches] + self._batch_iterator.close() + self._batch_iterator = type(self._batch_iterator)( + self._reader, self._sampler, 0 + ) + + 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 bounded same-nloc chunk for legacy neighbor statistics.""" + indices = self._nloc_set_indices[str(set_name)] + return self._legacy_batch(self._reader.decode_batch(indices, ragged=False)) + + 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, 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 + if ( + not key.startswith("find_") + and key not in structural_keys + and key not in self._data_dict + ): + continue + if value is None: + out[key] = None + 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 + ) + if real_natoms_vec.ndim == 1: + real_natoms_vec = np.tile(real_natoms_vec, (nframes, 1)) + pad_nloc = int(atype.shape[1]) + # Legacy mixed-type consumers require the per-type slots to sum to + # nloc. Composition is frame-specific and lives in real_natoms_vec; + # independent type maxima can exceed the size of every actual frame. + natoms_vec = np.zeros(2 + len(self._type_map), dtype=np.int32) + natoms_vec[:3] = pad_nloc + + 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((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) + out.setdefault("find_coord", np.float32(1.0)) + if "numb_copy" not in out: + 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), + dtype=np.int32, + ) + 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 + return self._legacy_batch(next(self._batch_iterator)) + + @property + def uniform_sampling_mass(self) -> float: + """Weight the queued legacy batch using the shared frame-share rule.""" + return self._batch_iterator.uniform_sampling_mass + + @property + def uniform_nominal_mass(self) -> float: + """Return the source's mean systems-per-frame weight for a pass.""" + return self._batch_iterator.uniform_nominal_mass + + def get_batch_pass_length(self) -> int: + """Expose the shared iterator's pass length to grouped data sources.""" + return self._batch_iterator.batch_pass_length + + def state_dict(self) -> dict[str, Any]: + """Save training progress independently of statistical reads.""" + return self._batch_iterator.state_dict() + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore the next training batch without eager decoding.""" + self._batch_iterator.load_state_dict(state) + + 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 + + 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 + + @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 + + 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 = [] + iterator = getattr(self, "_batch_iterator", None) + if iterator is not None: + iterator.close() + 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: return "{: >{}}".format(name, width) @@ -903,6 +1335,116 @@ def print_summary( ) +def validate_backend_data_config( + data_config: DataConfig, + *, + backend_name: str, + lmdb_supported: bool, +) -> None: + """Reject unsupported converted output before dpdata performs any I/O.""" + if not lmdb_supported and any( + conversion_will_write_lmdb(group) for group in data_group_configs(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], *, 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 " + "('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 process_systems( + systems: str | list[str], + patterns: list[str] | None = None, + fmt: str | None = None, + out_fmt: str | None = None, + *, + formats: tuple[DatasetFormat, ...] = ("deepmd", "lmdb"), + root_dir: str | Path | 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 + ---------- + systems : str or list of str + 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, ``deepmd/lmdb`` is used when fmt + triggers conversion. + + Returns + ------- + result_systems: list of str + The valid 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( jdata: DataConfig, rcut: float, @@ -912,6 +1454,7 @@ def get_data( *, lmdb_factory: Callable[[str, int | str, str | None], Any] | None = None, seed: int | None = None, + lmdb_supported: bool = True, ) -> Any: """Build native data readers and compose independently configured groups. @@ -929,8 +1472,8 @@ def get_data( If in multi task mode lmdb_factory Optional backend adapter accepting an LMDB path, batch size, and - internal auto-probability rule. Without it, only NPY/HDF5 sources - are discovered. + internal auto-probability rule. Defaults to the legacy LMDB adapter + used by TensorFlow/JAX; NPY/HDF5 sources retain their native reader. seed Seed for group selection. If omitted, draws from the shared training random generator seed the group-selection streams. @@ -945,6 +1488,31 @@ def get_data( GroupedDataSystem, ) + configs = data_group_configs(jdata) + for config in configs: + validate_backend_data_config( + config, backend_name="Selected", lmdb_supported=lmdb_supported + ) + if conversion_will_write_lmdb(config): + validate_lmdb_sampling_options(config, modifier=modifier) + if not type_map: + raise ValueError( + "LMDB training data requires model/type_map to be set." + ) + + if lmdb_factory is None and lmdb_supported: + + def lmdb_factory( + path: str, batch_size: int | str, auto_prob: str | None + ) -> LmdbDataSystem: + return LmdbDataSystem( + path, + type_map, + batch_size, + auto_prob_style=auto_prob, + seed=seed if seed is not None else int(dp_random.choice(2**32)), + ) + groups = resolve_data_groups( jdata, formats=("deepmd", "lmdb") if lmdb_factory is not None else ("deepmd",) ) @@ -956,65 +1524,76 @@ def get_data( iter(data_child_seeds(seed, len(groups) + 1)) if needs_composition else iter(()) ) group_sources = [] - for group in groups: - sources = resolve_data_sources(group) - has_lmdb = any(source.format == "lmdb" for source in sources) - children = [] - for source in sources: - if source.format == "deepmd": - child = DeepmdDataSystem( - systems=list(source.systems), - batch_size=source.batch_size, - test_size=1, - shuffle_test=True, - rcut=rcut, - type_map=type_map, - optional_type_map=not multi_task_mode, - modifier=modifier, - trn_all_set=True, - sys_probs=group.sys_probs if not has_lmdb else None, - auto_prob_style=group.auto_prob or "prob_sys_size", + try: + for group in groups: + sources = resolve_data_sources(group) + has_lmdb = any(source.format == "lmdb" for source in sources) + if has_lmdb: + validate_lmdb_sampling_options( + {"sys_probs": group.sys_probs}, modifier=modifier ) - else: - assert lmdb_factory is not None - child = lmdb_factory( - source.systems[0], source.batch_size, group.auto_prob + children = [] + group_sources.append(children) + for source in sources: + if source.format == "deepmd": + child = DeepmdDataSystem( + systems=list(source.systems), + batch_size=source.batch_size, + test_size=1, + shuffle_test=True, + rcut=rcut, + type_map=type_map, + optional_type_map=not multi_task_mode, + modifier=modifier, + trn_all_set=True, + sys_probs=group.sys_probs if not has_lmdb else None, + auto_prob_style=group.auto_prob or "prob_sys_size", + ) + else: + assert lmdb_factory is not None + child = lmdb_factory( + source.systems[0], source.batch_size, group.auto_prob + ) + children.append(child) + + sources = [source for children in group_sources for source in children] + numpy_sources = [ + source for source in sources if isinstance(source, DeepmdDataSystem) + ] + if len({source.mixed_type for source in numpy_sources}) > 1: + raise ValueError("All NumPy data systems must agree on mixed_type.") + ntypes = max(source.get_ntypes() for source in sources) + for source in numpy_sources: + source.sys_ntypes = ntypes + source.natoms_vec = [ + system.get_natoms_vec(ntypes).astype(int) + for system in source.data_systems + ] + datasets = [ + ( + children[0] + if len(children) == 1 + else GroupedDataSystem( + children, + [None] * len(children), + seed=next(seeds), + uniform=group.auto_prob == "prob_uniform", ) - children.append(child) - group_sources.append(children) - - sources = [source for children in group_sources for source in children] - numpy_sources = [ - source for source in sources if isinstance(source, DeepmdDataSystem) - ] - if len({source.mixed_type for source in numpy_sources}) > 1: - raise ValueError("All NumPy data systems must agree on mixed_type.") - ntypes = max(source.get_ntypes() for source in sources) - for source in numpy_sources: - source.sys_ntypes = ntypes - source.natoms_vec = [ - system.get_natoms_vec(ntypes).astype(int) for system in source.data_systems + ) + for group, children in zip(groups, group_sources, strict=True) ] - datasets = [ - ( - children[0] - if len(children) == 1 + return ( + datasets[0] + if len(datasets) == 1 else GroupedDataSystem( - children, - [None] * len(children), - seed=next(seeds), - uniform=group.auto_prob == "prob_uniform", + datasets, [group.prob for group in groups], seed=next(seeds) ) ) - for group, children in zip(groups, group_sources, strict=True) - ] - return ( - datasets[0] - if len(datasets) == 1 - else GroupedDataSystem( - datasets, [group.prob for group in groups], seed=next(seeds) - ) - ) + except BaseException: + # Ownership transfers only after the complete hierarchy is usable. + # A later source or metadata mismatch must release earlier LMDB maps. + close_data_systems(group_sources) + raise def get_neighbor_stat_data(jdata: DataConfig, type_map: list[str] | None) -> Any: diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 353eda9bb5..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: diff --git a/pyproject.toml b/pyproject.toml index 774cd08fa5..83ed5b95fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ dependencies = [ 'array-api-compat', 'lmdb', 'msgpack', + 'dpdata>=1.1.0', ] requires-python = ">=3.10" keywords = ["deepmd"] @@ -80,7 +81,6 @@ repository = "https://github.com/deepmodeling/deepmd-kit" [tool.deepmd_build_backend.optional-dependencies] test = [ "dependency_groups", - "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/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 27b4c90cdd..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( @@ -2195,6 +2254,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], 300), ([1], 300)]) + def test_compute_block_targets_multi_sys_block(self): result = compute_block_targets( "prob_sys_size;0:2:0.5;2:3:0.5", @@ -2229,6 +2295,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)] @@ -2571,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( 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 new file mode 100644 index 0000000000..aa38be2526 --- /dev/null +++ b/source/tests/common/test_data_system_conversion.py @@ -0,0 +1,1098 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +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, + make_neighbor_stat_data, +) +from deepmd.entrypoints.test import test as run_model_test +from deepmd.utils import ( + data_conversion, + data_system, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.data_system import ( + LmdbDataSystem, + get_data, + process_systems, + validate_backend_data_config, +) + + +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() + + +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]]] = [] + to_calls: ClassVar[list[tuple[str, str, dict]]] = [] + + 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, **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 == "deepmd/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.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_conversion._DPDATA_CONVERSION_CACHE.clear() + self.fake_dpdata = types.SimpleNamespace( + MultiSystems=_FakeMultiSystems, + LabeledSystem=_FakeLabeledSystem, + ) + + def tearDown(self) -> None: + 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}): + 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") + 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" + 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_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) + 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. + 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_conversion._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)) + + 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_conversion._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( + 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.addCleanup(data.close) + + 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) + self.assertIn("natoms_vec", batch) + self.assertEqual(batch["coord"].shape, (1, 3)) + 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_use_native_group_composition(self) -> None: + """Converted support preserves upstream's multiple-source data API.""" + lmdb_a = self.root / "a.lmdb" + lmdb_b = self.root / "b.lmdb" + _write_minimal_lmdb(str(lmdb_a)) + _write_minimal_lmdb(str(lmdb_b)) + data = get_data( + {"systems": [str(lmdb_a), str(lmdb_b)], "batch_size": 1}, + 0.0, + ["H"], + None, + ) + self.addCleanup(data_system.close_data_systems, data) + self.assertEqual(len(data.data_systems), 2) + self.assertEqual(data.get_batch()["coord"].shape, (1, 3)) + + 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([]) + + 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]) + self.addCleanup(data.close) + + 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) + self.addCleanup(data.close) + 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) + self.addCleanup(data.close) + + 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) + self.addCleanup(data.close) + + 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.addCleanup(data.close) + + 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_use_grouped_readers(self) -> None: + """Each input has its own cache and composes through the native API.""" + second = self.root / "second.extxyz" + second.write_text(self.source.read_text()) + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + data = get_data( + { + "systems": [str(self.source), str(second)], + "format": "extxyz", + "batch_size": 1, + "auto_prob": "prob_uniform", + }, + 0, + ["H"], + None, + ) + self.addCleanup(data.close) + self.assertEqual(_FakeMultiSystems.write_count, 2) + self.assertEqual(len(data.children), 2) + self.assertEqual(data.get_batch()["coord"].shape, (1, 3)) + # A grouped checkpoint must restore the next batch, not its prefetch. + state = data.state_dict() + expected = data.get_batch() + data.load_state_dict(state) + np.testing.assert_array_equal(data.get_batch()["coord"], expected["coord"]) + + def test_group_construction_failure_closes_completed_readers(self) -> None: + """Factories retain ownership until every group is constructed.""" + paths = [self.root / "valid.lmdb", self.root / "later.lmdb"] + for path in paths: + _write_minimal_lmdb(str(path)) + acquired = [] + + def factory(path, batch_size, auto_prob): + if acquired: + raise ValueError("later source failed") + reader = LmdbDataSystem(path, ["H"], batch_size, auto_prob) + acquired.append(reader) + return reader + + with self.assertRaisesRegex(ValueError, "later source failed"): + get_data( + {"systems": list(map(str, paths)), "batch_size": 1}, + 0, + ["H"], + None, + lmdb_factory=factory, + ) + self.assertTrue(acquired[0]._closed) + + def test_rare_nonperiodic_frame_outside_probe_is_partitioned(self) -> None: + """A single missed box must never disable PBC on its periodic neighbors.""" + path = self.root / "rare-pbc.lmdb" + _write_repeated_lmdb(str(path), 300) + probe = {300 * i // 256 for i in range(256)} + rare = next(i for i in range(300) if i not in probe) + with lmdb.open(str(path)) as env, env.begin(write=True) as txn: + key = f"{rare:012d}".encode() + frame = msgpack.unpackb(txn.get(key), raw=False) + frame["cells"] = _encode_array(np.zeros((3, 3))) + txn.put(key, msgpack.packb(frame, use_bin_type=True)) + data = LmdbDataSystem(str(path), ["H"], batch_size=32, seed=0) + self.addCleanup(data.close) + self.assertEqual({view.pbc for view in data.data_systems}, {False, True}) + seen = {False: 0, True: 0} + for _ in range(data.get_batch_pass_length()): + batch = data.get_batch() + periodic = bool(batch["find_box"]) + cells = batch["box"].reshape(-1, 9) + self.assertTrue(np.all(np.any(cells != 0, axis=1) == periodic)) + self.assertEqual(len(batch["default_mesh"]), 7 if periodic else 1) + seen[periodic] += len(cells) + self.assertEqual(seen, {False: 1, True: 299}) + # Full-validation views use their own availability index. + groups = data._test_data.availability_groups(np.arange(300)) + self.assertEqual(sorted(map(len, groups)), [1, 299]) + neighbor_data = make_neighbor_stat_data(str(path), ["H"]) + self.assertEqual( + {view.pbc for view in neighbor_data.data_systems}, {False, True} + ) + for view in neighbor_data.data_systems: + box = view._load_set(view.dirs[0])["box"] + self.assertEqual(box is not None, view.pbc) + + def test_legacy_composition_counts_obey_mixed_type_invariant(self) -> None: + """Independent per-type maxima must not overrun the padded atom axis.""" + path = self.root / "compositions.lmdb" + _write_repeated_lmdb(str(path), 2) + with lmdb.open(str(path)) as env, env.begin(write=True) as txn: + frame = msgpack.unpackb(txn.get(b"000000000001"), raw=False) + metadata = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) + metadata["type_map"] = ["H", "O"] + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + frame["atom_names"] = ["H", "O"] + frame["atom_numbs"] = [0, 1] + frame["atom_types"] = _encode_array(np.array([1], dtype=np.int64)) + txn.put(b"000000000001", msgpack.packb(frame, use_bin_type=True)) + data = LmdbDataSystem(str(path), ["H", "O"], batch_size=2, seed=0) + self.addCleanup(data.close) + batch = data.get_batch() + np.testing.assert_array_equal(batch["natoms_vec"], [1, 1, 1, 0]) + self.assertEqual( + set(map(tuple, batch["real_natoms_vec"])), {(1, 1, 1, 0), (1, 1, 0, 1)} + ) + + 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_conversion.socket.gethostname(), + "pid": os.getpid(), + "process_start": data_conversion._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_conversion.time, "sleep", release_after_long_valid_conversion + ), + patch.object( + data_conversion, + "_is_conversion_current", + return_value=True, + ) as is_current, + ): + self.assertTrue( + data_conversion._wait_for_conversion(self.source, output, lock_path) + ) + + self.assertEqual(waits, 301) + is_current.assert_called_once_with(self.source, output) + + def test_dead_conversion_owner_is_recovered(self) -> None: + lock_path = self.root / "dead.lock" + lock_path.write_text( + json.dumps( + { + "hostname": data_conversion.socket.gethostname(), + "pid": 2147483647, + "process_start": "missing", + } + ) + ) + + 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: + 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_conversion.os, "replace", fail_publication): + with self.assertRaisesRegex(OSError, "publish failed"): + data_conversion._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_conversion.importlib.metadata, + "version", + side_effect=["1.1.0", "1.2.0"], + ): + first = data_conversion._conversion_cache_path( + self.source, "extxyz", "deepmd/lmdb" + ) + second = data_conversion._conversion_cache_path( + self.source, "extxyz", "deepmd/lmdb" + ) + + self.assertNotEqual(first, second) + + 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_conversion._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" + first = data_conversion._source_signature(self.root, output) + output.mkdir(parents=True) + (output / "data.mdb").write_text("cache data") + self.assertEqual(first, data_conversion._source_signature(self.root, output)) + child = self.root / "new.extxyz" + child.write_text("new input") + self.assertNotEqual(first, data_conversion._source_signature(self.root, output)) + child.unlink() + 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.""" + 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_directory_symlinks_track_edits_retargets_and_cycles(self) -> None: + """ABACUS-style linked output trees participate in cache freshness.""" + source = self.root / "run" + source.mkdir() + targets = [self.root / "out-a", self.root / "out-b"] + for target in targets: + target.mkdir() + (target / "running_scf.log").write_text("energy 0") + link = source / "OUT.ABACUS" + link.symlink_to(targets[0], target_is_directory=True) + # Follow a back edge without recursing forever. + (targets[0] / "parent").symlink_to(source, target_is_directory=True) + with patch.dict(sys.modules, {"dpdata": self.fake_dpdata}): + process_systems(str(source), fmt="abacus/scf") + process_systems(str(source), fmt="abacus/scf") + self.assertEqual(_FakeMultiSystems.write_count, 1) + child = targets[0] / "running_scf.log" + before = child.stat() + child.write_text("energy 42") + os.utime(child, ns=(before.st_atime_ns, before.st_mtime_ns)) + process_systems(str(source), fmt="abacus/scf") + self.assertEqual(_FakeMultiSystems.write_count, 2) + link.unlink() + link.symlink_to(targets[1], target_is_directory=True) + process_systems(str(source), fmt="abacus/scf") + self.assertEqual(_FakeMultiSystems.write_count, 3) + + def test_source_signature_rejects_disappearing_regular_file(self) -> None: + """Only dangling links may be fingerprinted after followed stat fails.""" + 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: + 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_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.""" + + 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_conversion, "_write_dpdata_conversion", changing_writer): + with self.assertRaisesRegex( + RuntimeError, "input changed during conversion" + ): + process_systems(str(self.source), fmt="extxyz") + output = data_conversion._conversion_cache_path( + self.source, "extxyz", "deepmd/lmdb" + ) + 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_conversion._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_conversion._normalize_dpdata_format("auto", xyz), expected + ) + with patch.object(Path, "open", side_effect=OSError("unreadable")): + self.assertEqual( + data_conversion._normalize_dpdata_format("auto", xyz), "xyz" + ) + xyz.unlink() + 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_conversion._normalize_dpdata_format("auto", Path(name)), expected + ) + self.assertEqual( + data_conversion._normalize_dpdata_format("ase", xyz), "ase/structure" + ) + + 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) + with get_data( + {"systems": str(lmdb_path), "batch_size": 2}, 0.0, ["H"], None + ) as data: + return [list(batch) for batch in data._sampler] + + 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_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_conversion._lock_owner_is_alive({"hostname": hostname, **payload}) + ) + payload = {"hostname": hostname, "pid": os.getpid(), "process_start": "123"} + 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_conversion, "_process_start_time", return_value=None), + patch.object(data_conversion.os, "kill", side_effect=error), + ): + 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.""" + 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_conversion.time.time() - ( + data_conversion._CONVERSION_LOCK_STALE_SECONDS + 1 if expired else 0 + ) + os.utime(path, (stamp, stamp)) + with patch.object( + data_conversion, "_lock_owner_is_alive", return_value=alive + ): + self.assertEqual( + data_conversion._recover_stale_conversion_lock(path), expected + ) + self.assertEqual(path.exists(), not expected) + path.write_text("{}") + with ( + patch.object(data_conversion, "_lock_owner_is_alive", return_value=False), + patch.object(data_conversion, "_same_lock_file", return_value=False), + ): + 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_conversion.time.time() - 31 + os.utime(path, (thirty_seconds_ago, thirty_seconds_ago)) + self.assertFalse(data_conversion._recover_stale_conversion_lock(path)) + expired = ( + 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_conversion._recover_stale_conversion_lock(path) + self.assertTrue(path.exists()) + + def test_lock_initialization_failure_closes_file_and_removes_lock(self) -> None: + """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_conversion.{target}", + side_effect=OSError("init failed"), + ): + with path.open("x") as lock_file: + with self.assertRaisesRegex(OSError, "init failed"): + data_conversion._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" + 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( + "data.extxyz", + patterns=["*.extxyz"], + root_dir=config_path.parent, + fmt="extxyz", + out_fmt="deepmd/lmdb", + ) + + +if __name__ == "__main__": + unittest.main() 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/pd/test_conversion_validation.py b/source/tests/pd/test_conversion_validation.py new file mode 100644 index 0000000000..55bda27e0b --- /dev/null +++ b/source/tests/pd/test_conversion_validation.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unsupported conversion must fail before Paddle reads training data.""" + +from copy import ( + deepcopy, +) + +import pytest + +from deepmd.pd.entrypoints import main as entrypoint + + +@pytest.mark.parametrize("multi_task", [False, True]) +@pytest.mark.parametrize("converted_key", ["training_data", "validation_data"]) +def test_lmdb_conversion_fails_before_neighbor_stat( + monkeypatch, multi_task, converted_key +): + """Validate every task, including validation-only LMDB, before conversion.""" + native = {"systems": "native", "batch_size": 1} + task = {"training_data": native, "validation_data": native} + task[converted_key] = [{"systems": "input.extxyz", "format": "extxyz"}] + model = {"type_map": ["H"]} + config = {"model": model, "training": task} + if multi_task: + config = { + "model": {"model_dict": {"a": model, "b": model}}, + "training": { + "data_dict": { + "a": {"training_data": native}, + "b": task, + } + }, + } + monkeypatch.setattr(entrypoint, "SummaryPrinter", lambda: lambda: None) + monkeypatch.setattr(entrypoint, "j_loader", lambda _: deepcopy(config)) + monkeypatch.setattr(entrypoint, "expand_model_preset", lambda value: value) + monkeypatch.setattr( + entrypoint, "preprocess_shared_params", lambda value: (value, None) + ) + monkeypatch.setattr(entrypoint, "update_deepmd_input", lambda value, **_: value) + monkeypatch.setattr(entrypoint, "normalize", lambda value, **_: value) + + def unexpected_data_access(*args, **kwargs): + raise AssertionError("Data access started before LMDB rejection") + + monkeypatch.setattr(entrypoint, "get_data", unexpected_data_access) + monkeypatch.setattr(entrypoint, "get_trainer", unexpected_data_access) + with pytest.raises( + NotImplementedError, match="Paddle backend does not support LMDB" + ): + entrypoint.train("input.json", None, None, None, None, "RANDOM") diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index eb85481f7d..1cfaf567e5 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -256,6 +256,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_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"), [ diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 0a23567a22..93dd1e3d81 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -21,6 +21,7 @@ import lmdb import msgpack import numpy as np +import pytest import deepmd.utils.random as dp_random from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module @@ -85,6 +86,61 @@ def _encode_array(arr: np.ndarray) -> dict: } +@pytest.mark.timeout(60) +class TestConvertedLmdbValidation(unittest.TestCase): + """Converted inputs retain the native LMDB routing and DDP metadata.""" + + def test_neighbor_stat_and_training_data_convert_each_input(self) -> None: + """Both entrypoints feed every selected file to the shared converter.""" + params = { + "systems": ["first.extxyz", "second.extxyz"], + "format": "extxyz", + "batch_size": 1, + } + for factory in (get_neighbor_stat_data, _build_data_system): + with patch( + "deepmd.utils.data_path._convert_system_by_dpdata", + side_effect=[[], RuntimeError("second source reached")], + ) as convert: + with self.assertRaisesRegex(RuntimeError, "second source reached"): + factory(params, ["O", "H"]) + self.assertEqual( + [call.args[0] for call in convert.call_args_list], params["systems"] + ) + + 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.utils.data_config.process_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.""" rng = np.random.RandomState(seed) 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()