diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 1463dad928..2c05c8b4cd 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -12,6 +12,11 @@ make_task_maps, print_data_summaries, ) +from .data_progress import ( + DATA_PROGRESS_CHECKPOINT_KEY, + collect_training_data_progress, + restore_training_data_progress, +) from .entrypoint import ( AbstractTrainEntrypoint, TrainEntrypointOptions, @@ -44,6 +49,7 @@ ) __all__ = [ + "DATA_PROGRESS_CHECKPOINT_KEY", "DEFAULT_TASK_KEY", "AbstractTrainEntrypoint", "AbstractTrainer", @@ -64,9 +70,11 @@ "build_checkpoint_stores", "change_model_out_bias", "change_model_out_bias_by_task", + "collect_training_data_progress", "iter_training_task_configs", "make_task_maps", "print_data_summaries", "resolve_keep_ckpt_count", "resolve_step_schedule", + "restore_training_data_progress", ] diff --git a/deepmd/dpmodel/train/data_progress.py b/deepmd/dpmodel/train/data_progress.py new file mode 100644 index 0000000000..bf774f5e3f --- /dev/null +++ b/deepmd/dpmodel/train/data_progress.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Checkpointable training-data progress for restartable optimizers. + +A training restart restores model, optimizer, learning-rate, and EMA state, +but historically rebuilt every data source from the beginning of epoch zero. +This module defines the backend-neutral payload that records each task's +logical data cursor — the next training batch the optimizer has not yet +consumed — and the process RNG that drives multi-task and directory-system +selection. + +Prefetch and decoder read-ahead are excluded from the payload. Restoring +progress rebuilds iterators at the saved cursor so the next ``get_batch`` +matches an uninterrupted run without decoding the dataset during seek. +""" + +from __future__ import ( + annotations, +) + +from collections.abc import ( + Mapping, +) +from typing import ( + Any, +) + +from deepmd.utils import random as dp_random + +DATA_PROGRESS_CHECKPOINT_KEY = "data_progress" +DATA_PROGRESS_VERSION = 1 + +__all__ = [ + "DATA_PROGRESS_CHECKPOINT_KEY", + "DATA_PROGRESS_VERSION", + "collect_training_data_progress", + "restore_training_data_progress", +] + + +def collect_training_data_progress( + training_data_by_task: Mapping[str, Any], +) -> dict[str, Any]: + """Collect per-task data-source progress and the shared selection RNG. + + Parameters + ---------- + training_data_by_task : Mapping[str, Any] + Training data sources keyed by task name. + + Returns + ------- + dict[str, Any] + Versioned checkpoint payload. Tasks whose data source does not + implement ``state_dict`` are omitted; an empty ``tasks`` map still + carries the RNG so multi-task selection can resume. + """ + tasks: dict[str, Any] = {} + for task_key, data in training_data_by_task.items(): + state_dict = getattr(data, "state_dict", None) + if callable(state_dict): + tasks[task_key] = state_dict() + return { + "version": DATA_PROGRESS_VERSION, + "tasks": tasks, + "rng": dp_random.get_state(), + } + + +def restore_training_data_progress( + training_data_by_task: Mapping[str, Any], + progress: Mapping[str, Any] | None, +) -> None: + """Restore per-task data progress and the shared selection RNG. + + Parameters + ---------- + training_data_by_task : Mapping[str, Any] + Training data sources keyed by task name. Sources without + ``load_state_dict`` are left unchanged. + progress : Mapping[str, Any] or None + Payload previously produced by :func:`collect_training_data_progress`. + ``None`` or a missing payload retains the legacy restart behaviour + (data sources stay at their freshly constructed cursor). + + Raises + ------ + ValueError + If the payload version is unsupported, a task is missing, or a data + source rejects the restored configuration. + """ + if progress is None: + return + version = int(progress.get("version", -1)) + if version != DATA_PROGRESS_VERSION: + raise ValueError( + f"Unsupported data-progress checkpoint version {version}; " + f"expected {DATA_PROGRESS_VERSION}." + ) + tasks = progress.get("tasks", {}) + if not isinstance(tasks, Mapping): + raise ValueError("data-progress 'tasks' must be a mapping.") + for task_key, data in training_data_by_task.items(): + load_state_dict = getattr(data, "load_state_dict", None) + if not callable(load_state_dict): + continue + if task_key not in tasks: + raise ValueError(f"data-progress checkpoint is missing task {task_key!r}.") + load_state_dict(tasks[task_key]) + if "rng" in progress: + dp_random.set_state(progress["rng"]) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 897292eeb7..5012f5ab03 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -17,6 +17,7 @@ Callable, Iterable, Iterator, + Mapping, Sequence, ) from concurrent.futures import ( @@ -1617,12 +1618,19 @@ def __init__( self._reader = reader self._sampler = sampler self._epoch = 0 - self._iterator = self._iter_epoch() + self._logical_epoch = 0 + self._logical_batch_index = 0 + self._logical_epoch_length = 0 + # Length of a sampler epoch opened by prefetch before the logical + # cursor wraps. LmdbBatchSampler advances on ``__iter__``, so + # ``len(self._sampler)`` at wrap would observe the *following* epoch. + self._pending_logical_epoch_length: int | None = None self._num_workers = num_workers self._pool: _LmdbPoolEntry | None = None self._pending: _PendingBatch | None = None self._deferred_indices: list[int] | None = None self._closed = False + self._iterator = self._iter_epoch() def __iter__(self) -> "LmdbBatchIterator": return self @@ -1641,8 +1649,91 @@ def __next__(self) -> dict[str, Any]: batch = self._decode(self._next_indices()) self._schedule(self._next_indices()) + self._advance_logical_cursor() return batch + def state_dict(self) -> dict[str, int]: + """Return the logical cursor of the next batch to consume. + + Prefetched decoder work is excluded: only batches already returned to + the trainer advance this cursor. + """ + return { + "epoch": self._logical_epoch, + "batch_index": self._logical_batch_index, + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Reposition to the saved logical cursor without decoding frames. + + Parameters + ---------- + state : Mapping[str, Any] + Progress previously returned by :meth:`state_dict`. + """ + if self._closed: + raise RuntimeError( + "cannot restore progress on a closed LMDB batch iterator" + ) + epoch = int(state["epoch"]) + batch_index = int(state["batch_index"]) + if epoch < 0 or batch_index < 0: + raise ValueError( + f"data-progress cursor must be non-negative, got epoch={epoch}, " + f"batch_index={batch_index}" + ) + self._cancel_prefetch() + self._pending_logical_epoch_length = None + self._logical_epoch = epoch + self._logical_batch_index = batch_index + self._epoch = epoch + self._iterator = self._iter_epoch() + if self._logical_epoch_length == 0: + if batch_index != 0: + raise ValueError( + "data-progress batch_index must be 0 for an empty epoch, " + f"got {batch_index}" + ) + return + if batch_index >= self._logical_epoch_length: + raise ValueError( + f"data-progress batch_index {batch_index} is out of range for " + f"epoch length {self._logical_epoch_length}" + ) + for _ in range(batch_index): + self._next_indices() + + def _advance_logical_cursor(self) -> None: + """Record that one batch was returned to the trainer. + + Prefetch may already have opened the next sampler epoch. That open + stashes the next logical length in ``_pending_logical_epoch_length``; + the wrap applies it. Falling back to ``len(self._sampler)`` is only + correct when prefetch has not yet opened the successor (samplers such + as ``LmdbBatchSampler`` advance on ``__iter__``, so a post-prefetch + ``len`` would name the epoch after next). + """ + if self._logical_epoch_length <= 0: + return + self._logical_batch_index += 1 + if self._logical_batch_index < self._logical_epoch_length: + return + self._logical_epoch += 1 + self._logical_batch_index = 0 + if self._pending_logical_epoch_length is not None: + self._logical_epoch_length = self._pending_logical_epoch_length + self._pending_logical_epoch_length = None + else: + self._logical_epoch_length = len(self._sampler) + + def _cancel_prefetch(self) -> None: + """Drop prefetched work so a seek can rebuild the sampler position.""" + if self._pending is not None: + for future in self._pending.futures: + future.cancel() + self._pending = None + self._deferred_indices = None + def _decode(self, indices: list[int]) -> dict[str, Any]: """Decode one batch, in the pool when that is worthwhile and possible.""" futures = self._offer(indices) @@ -1733,6 +1824,17 @@ def _iter_epoch(self) -> Iterator[list[int]]: set_epoch = getattr(self._sampler, "set_epoch", None) if callable(set_epoch): set_epoch(self._epoch) + epoch_length = len(self._sampler) + # Prefetch may open the next sampler epoch before the trainer consumes + # the last batch of the logical epoch. Bind the active logical length + # only when reconstructing the epoch the cursor still names; otherwise + # stash the length for the upcoming wrap — ``iter(self._sampler)`` may + # already advance the sampler, so a later ``len`` would be wrong. + if self._epoch == self._logical_epoch: + self._logical_epoch_length = epoch_length + self._pending_logical_epoch_length = None + else: + self._pending_logical_epoch_length = epoch_length return iter(self._sampler) def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 9130f5753b..0e132c0bcd 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -47,6 +47,7 @@ fully_shard = None # type: ignore[assignment] from deepmd.dpmodel.train import ( + DATA_PROGRESS_CHECKPOINT_KEY, DEFAULT_TASK_KEY, AbstractTrainer, RankContext, @@ -59,7 +60,9 @@ build_checkpoint_stores, change_model_out_bias, change_model_out_bias_by_task, + collect_training_data_progress, resolve_step_schedule, + restore_training_data_progress, ) from deepmd.dpmodel.utils.batch import ( normalize_batch, @@ -2102,6 +2105,7 @@ def initialize_statistics( # cannot be copied into sharded parameters. ema_state_dict = None optimizer_state_dict = None + data_progress_state = None if resuming: log.info(f"Resuming from {resume_model}.") is_pte = resume_model.endswith((".pte", ".pt2")) @@ -2117,7 +2121,7 @@ def initialize_statistics( if "model" in state_dict: # Optimizer and EMA state describe the weights of the run # they were saved by; a finetune starts a new run and keeps - # neither. + # neither. Data-source progress is likewise resume-only. continues_run = self.restart_training and finetune_model is None optimizer_state_dict = ( state_dict["optimizer"] if continues_run else None @@ -2125,9 +2129,15 @@ def initialize_statistics( ema_state_dict = ( state_dict.get(EMA_CHECKPOINT_KEY) if continues_run else None ) + data_progress_state = ( + state_dict.get(DATA_PROGRESS_CHECKPOINT_KEY) + if continues_run + else None + ) state_dict = state_dict["model"] else: optimizer_state_dict = None + data_progress_state = None self.start_step = ( state_dict["_extra_state"]["train_infos"]["step"] if self.restart_training @@ -2371,6 +2381,14 @@ def update_finetune_bias( else None ) + # Data-source progress is restored after requirements and statistics + # hooks have run, and before training can draw a batch. Missing progress + # keeps the legacy restart behaviour of rebuilding iterators at epoch 0. + restore_training_data_progress( + self.training_data_by_task, + data_progress_state, + ) + self._configure_neighbor_graph_method( training_params.get("neighbor_graph_method", "auto") ) @@ -3066,6 +3084,12 @@ def _write_checkpoint( state["optimizer"] = optim_state if include_ema_state and self.model_ema is not None: state[EMA_CHECKPOINT_KEY] = self.model_ema.state_dict() + if include_optimizer: + # Resume-capable checkpoints record the logical data cursor so a + # restart yields the same next batch; deployment snapshots omit it. + state[DATA_PROGRESS_CHECKPOINT_KEY] = collect_training_data_progress( + self.training_data_by_task + ) ckpt_path.parent.mkdir(parents=True, exist_ok=True) torch.save(state, ckpt_path) diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index f16de39a01..775ba1a76e 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -9,6 +9,9 @@ that interface. """ +from collections.abc import ( + Mapping, +) from typing import ( Any, ) @@ -80,6 +83,10 @@ def __init__( rank: int = 0, world_size: int = 1, ) -> None: + self._seed = seed + self._batch_size_spec: int | str = batch_size + self._rank = rank + self._world_size = world_size self._reader = LmdbDataReader(lmdb_path, type_map, batch_size) block_targets = None @@ -215,6 +222,62 @@ def set_frame_transform(self, transform) -> None: # noqa: ANN001 """Install a per-frame transform on the underlying reader.""" self._reader.set_frame_transform(transform) + def state_dict(self) -> dict[str, Any]: + """Return logical training progress for checkpointing. + + Returns + ------- + dict[str, Any] + Versioned cursor plus the sampling identity needed to reject an + incompatible restart (world size, seed, retained frame count, and + batch-size rule). + """ + progress = self._batch_iterator.state_dict() + return { + "version": 1, + "kind": "lmdb", + "epoch": progress["epoch"], + "batch_index": progress["batch_index"], + "world_size": self._world_size, + "seed": self._seed, + "nframes": len(self._reader), + "batch_size": self._batch_size_spec, + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore logical training progress before the next ``get_batch``. + + Parameters + ---------- + state : Mapping[str, Any] + Progress previously returned by :meth:`state_dict`. + + Raises + ------ + ValueError + If the payload does not describe this data source configuration. + """ + if state.get("kind", "lmdb") != "lmdb": + raise ValueError( + f"Cannot restore {state.get('kind')!r} data progress into an LMDB " + "data source." + ) + expected = { + "world_size": self._world_size, + "seed": self._seed, + "nframes": len(self._reader), + "batch_size": self._batch_size_spec, + } + for key, value in expected.items(): + if key not in state: + raise ValueError(f"LMDB data-progress checkpoint is missing {key!r}.") + if state[key] != value: + raise ValueError( + f"LMDB data-progress {key} mismatch: checkpoint has " + f"{state[key]!r}, current data source has {value!r}." + ) + self._batch_iterator.load_state_dict(state) + def close(self) -> None: """Cancel prefetched work and release decoder processes.""" iterator = getattr(self, "_batch_iterator", None) diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index a232aab393..fa74e8dbc2 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -139,6 +139,7 @@ def __init__( # set counters self.set_count = 0 self.iterator = 0 + self._batch_order = np.zeros(0, dtype=np.int64) self.shuffle_test = shuffle_test # set modifier self.modifier = modifier @@ -672,12 +673,75 @@ def _load_batch_set(self, set_name: DPPath) -> None: self.batch_set = self._load_set(set_name) if self.modifier is not None: self.modifier.modify_data(self.batch_set, self) - self.batch_set, _ = self._shuffle_data(self.batch_set) + self.batch_set, self._batch_order = self._shuffle_data(self.batch_set) self.reset_get_batch() def reset_get_batch(self) -> None: self.iterator = 0 + def state_dict(self) -> dict[str, Any]: + """Return the set cursor and the shuffle order of the loaded set.""" + return { + "set_count": int(self.set_count), + "iterator": int(self.iterator), + "batch_order": self._batch_order.astype(np.int64).tolist(), + } + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore the set cursor without re-drawing the shuffle RNG. + + Parameters + ---------- + state : dict[str, Any] + Progress previously returned by :meth:`state_dict`. + """ + set_count = int(state["set_count"]) + iterator = int(state["iterator"]) + batch_order = np.asarray(state["batch_order"], dtype=np.int64) + if set_count < 0 or iterator < 0: + raise ValueError( + f"directory data-progress cursor must be non-negative, got " + f"set_count={set_count}, iterator={iterator}" + ) + if set_count == 0: + if iterator != 0 or batch_order.size: + raise ValueError( + "directory data-progress with set_count=0 must have an empty " + "batch and iterator=0" + ) + self.set_count = 0 + self.iterator = 0 + self._batch_order = np.zeros(0, dtype=np.int64) + if hasattr(self, "batch_set"): + del self.batch_set + return + set_idx = (set_count - 1) % self.get_numb_set() + raw = self._load_set(self.dirs[set_idx]) + if self.modifier is not None: + self.modifier.modify_data(raw, self) + self.batch_set = self._apply_batch_order(raw, batch_order) + self._batch_order = batch_order + self.set_count = set_count + self.iterator = iterator + + def _apply_batch_order( + self, data: dict[str, Any], idx: np.ndarray + ) -> dict[str, Any]: + """Apply a saved frame permutation produced by :meth:`_shuffle_data`.""" + ret: dict[str, Any] = {} + nframes = data["coord"].shape[0] + for key, value in data.items(): + if ( + isinstance(value, np.ndarray) + and value.ndim == 2 + and value.shape[0] == nframes + and "find_" not in key + ): + ret[key] = value[idx] + else: + ret[key] = value + return ret + def _load_test_set(self, shuffle_test: bool) -> None: test_sets = [] for ii in self.dirs: diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index cea270bd37..a41e8d4f0c 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -646,6 +646,48 @@ def get_sys(self, idx: int) -> DeepmdData: """Get a certain data system.""" return self.data_systems[idx] + def state_dict(self) -> dict: + """Return per-system training cursors for checkpointing.""" + return { + "version": 1, + "kind": "directory", + "systems": [system.state_dict() for system in self.data_systems], + "pick_idx": int(self.pick_idx), + "nsystems": int(self.nsystems), + } + + def load_state_dict(self, state: dict) -> None: + """Restore per-system training cursors. + + Parameters + ---------- + state : dict + Progress previously returned by :meth:`state_dict`. + + Raises + ------ + ValueError + If the payload does not describe this data system. + """ + if state.get("kind", "directory") != "directory": + raise ValueError( + f"Cannot restore {state.get('kind')!r} data progress into a " + "directory/HDF5 data source." + ) + if int(state.get("nsystems", -1)) != self.nsystems: + raise ValueError( + "directory data-progress nsystems mismatch: checkpoint has " + f"{state.get('nsystems')!r}, current data source has {self.nsystems}." + ) + systems = state.get("systems") + if not isinstance(systems, list) or len(systems) != self.nsystems: + raise ValueError( + "directory data-progress must list one cursor per data system." + ) + for system, system_state in zip(self.data_systems, systems, strict=True): + system.load_state_dict(system_state) + self.pick_idx = int(state.get("pick_idx", 0)) + def get_batch_size(self) -> int: """Get the batch size.""" return self.batch_size diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 8710693607..284b118ea9 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -76,4 +76,33 @@ def shuffle(x: np.ndarray) -> None: _RANDOM_GENERATOR.shuffle(x) -__all__ = ["choice", "random", "seed", "shuffle"] +def get_state() -> tuple: + """Return the internal state of the shared generator. + + Returns + ------- + tuple + A ``RandomState`` state tuple whose array component is converted to a + plain list so the payload survives ``torch.save`` with + ``weights_only=True``. + """ + kind, keys, pos, has_gauss, cached_gaussian = _RANDOM_GENERATOR.get_state() + return (kind, keys.tolist(), int(pos), int(has_gauss), float(cached_gaussian)) + + +def set_state(state: tuple) -> None: + """Restore the shared generator from :func:`get_state` output. + + Parameters + ---------- + state : tuple + State previously returned by :func:`get_state` or + ``RandomState.get_state``. + """ + kind, keys, pos, has_gauss, cached_gaussian = state + _RANDOM_GENERATOR.set_state( + (kind, np.asarray(keys, dtype=np.uint32), pos, has_gauss, cached_gaussian) + ) + + +__all__ = ["choice", "get_state", "random", "seed", "set_state", "shuffle"] diff --git a/source/tests/common/dpmodel/test_data_progress.py b/source/tests/common/dpmodel/test_data_progress.py new file mode 100644 index 0000000000..ee9b00fc0c --- /dev/null +++ b/source/tests/common/dpmodel/test_data_progress.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for checkpointable training-data progress (#5814).""" + +from __future__ import ( + annotations, +) + +import tempfile +import unittest +from pathlib import ( + Path, +) + +import lmdb +import msgpack +import numpy as np + +from deepmd.dpmodel.train.data_progress import ( + DATA_PROGRESS_VERSION, + collect_training_data_progress, + restore_training_data_progress, +) +from deepmd.dpmodel.utils.lmdb_data import ( + LmdbBatchIterator, + LmdbBatchSampler, + LmdbDataReader, +) +from deepmd.utils import random as dp_random +from deepmd.utils.data import ( + DataRequirementItem, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + + +def _encode_array(arr: np.ndarray) -> dict: + return { + "nd": None, + "type": str(arr.dtype), + "kind": "", + "shape": list(arr.shape), + "data": arr.tobytes(), + } + + +def _make_frame(natoms: int, seed: int) -> dict: + rng = np.random.RandomState(seed) + half = natoms // 2 + return { + "atom_numbs": [half, natoms - half], + "atom_names": ["O", "H"], + "atom_types": _encode_array( + np.array([0] * half + [1] * (natoms - half), dtype=np.int64) + ), + "orig": _encode_array(np.zeros(3, dtype=np.float64)), + "cells": _encode_array((np.eye(3) * 10.0).astype(np.float64)), + "coords": _encode_array((rng.rand(natoms, 3) * 10.0).astype(np.float64)), + "energies": _encode_array(np.array(rng.randn(), dtype=np.float64)), + "forces": _encode_array(rng.randn(natoms, 3).astype(np.float64)), + } + + +def _create_test_lmdb(path: str, nframes: int, natoms: int) -> None: + env = lmdb.open(path, map_size=10 * 1024 * 1024) + fmt = "012d" + metadata = { + "nframes": nframes, + "frame_idx_fmt": fmt, + "system_info": { + "formula": f"O{natoms // 2}H{natoms - natoms // 2}", + "natoms": [natoms // 2, natoms - natoms // 2], + "nframes": nframes, + }, + } + with env.begin(write=True) as txn: + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + for i in range(nframes): + key = format(i, fmt).encode() + txn.put(key, msgpack.packb(_make_frame(natoms, i), use_bin_type=True)) + env.close() + + +def _create_mixed_nloc_lmdb( + path: str, + frames_spec: list[tuple[int, int]] | None = None, +) -> None: + """Write an LMDB whose frames span several atom counts. + + ``frames_spec`` lists ``(natoms, count)`` pairs. Mixed nloc plus ``mix:N`` + makes per-epoch batch counts shuffle-dependent, which is what exposes the + single-process logical-cursor wrap bug. + """ + if frames_spec is None: + frames_spec = [(6, 8), (9, 5), (12, 3), (4, 6)] + env = lmdb.open(path, map_size=20 * 1024 * 1024) + fmt = "012d" + nframes = sum(count for _, count in frames_spec) + metadata = { + "nframes": nframes, + "frame_idx_fmt": fmt, + "system_info": { + "formula": "mixed", + "natoms": [2, 4], + "nframes": nframes, + }, + } + with env.begin(write=True) as txn: + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + idx = 0 + for natoms, count in frames_spec: + for _ in range(count): + key = format(idx, fmt).encode() + txn.put( + key, + msgpack.packb(_make_frame(natoms, idx), use_bin_type=True), + ) + idx += 1 + env.close() + + +def _create_directory_system(root: Path, *, nframes: int = 8) -> Path: + system = root / "sys0" + set_dir = system / "set.000" + set_dir.mkdir(parents=True) + (system / "type.raw").write_text("0\n1\n0\n1\n0\n1\n") + (system / "type_map.raw").write_text("O\nH\n") + natoms = 6 + rng = np.random.RandomState(0) + np.save(set_dir / "coord.npy", rng.randn(nframes, natoms * 3).astype(np.float64)) + np.save(set_dir / "box.npy", np.tile(np.eye(3).reshape(9), (nframes, 1))) + np.save(set_dir / "energy.npy", rng.randn(nframes, 1).astype(np.float64)) + return system + + +class TestLmdbDataProgress(unittest.TestCase): + def setUp(self) -> None: + self.tmp = tempfile.TemporaryDirectory() + self.lmdb_path = str(Path(self.tmp.name) / "data.lmdb") + _create_test_lmdb(self.lmdb_path, nframes=12, natoms=6) + + def tearDown(self) -> None: + self.tmp.cleanup() + + def _system(self, **kwargs): + from deepmd.pt_expt.utils.lmdb_dataset import ( + LmdbDataSystem, + ) + + params = { + "lmdb_path": self.lmdb_path, + "type_map": ["O", "H"], + "batch_size": 3, + "seed": 11, + "num_workers": 0, + } + params.update(kwargs) + return LmdbDataSystem(**params) + + def test_mid_epoch_restart_matches_uninterrupted_fids(self) -> None: + reference = self._system() + try: + consumed = [reference.get_batch()["fid"] for _ in range(2)] + progress = reference.state_dict() + expected = [reference.get_batch()["fid"] for _ in range(3)] + finally: + reference.close() + + restored = self._system() + try: + restored.load_state_dict(progress) + actual = [restored.get_batch()["fid"] for _ in range(3)] + finally: + restored.close() + + self.assertEqual(progress["epoch"], 0) + self.assertEqual(progress["batch_index"], 2) + self.assertEqual(len(consumed), 2) + self.assertEqual(actual, expected) + + def test_epoch_boundary_preserves_reshuffle(self) -> None: + reference = self._system() + try: + epoch_len = len(reference._sampler) + for _ in range(epoch_len): + reference.get_batch() + progress = reference.state_dict() + expected = [reference.get_batch()["fid"] for _ in range(2)] + finally: + reference.close() + + restored = self._system() + try: + restored.load_state_dict(progress) + actual = [restored.get_batch()["fid"] for _ in range(2)] + finally: + restored.close() + + self.assertEqual(progress["epoch"], 1) + self.assertEqual(progress["batch_index"], 0) + self.assertEqual(actual, expected) + + def test_prefetch_does_not_advance_logical_cursor(self) -> None: + reader = LmdbDataReader(self.lmdb_path, ["O", "H"], batch_size=3) + sampler = LmdbBatchSampler(reader, shuffle=True, seed=5) + iterator = LmdbBatchIterator(reader, sampler, num_workers=0) + try: + batch = next(iterator) + progress = iterator.state_dict() + self.assertIsNotNone(iterator._deferred_indices) + self.assertEqual(progress["batch_index"], 1) + self.assertEqual(len(batch["fid"]), 3) + finally: + iterator.close() + reader.close() + + def test_ragged_mix_epoch_boundary_cursor_matches_uninterrupted(self) -> None: + """Resume across a ragged ``mix:N`` wrap must keep the same next batch. + + Prefetch opens the next ``LmdbBatchSampler`` epoch before the logical + cursor wraps. That sampler advances on ``__iter__``, so a naive + ``len(sampler)`` at wrap observes epoch+2. Under ragged ``mix:N`` the + lengths differ across epochs; using the wrong length wraps early and a + restored run skips the remaining batches of the true epoch. + """ + mixed_path = str(Path(self.tmp.name) / "mixed.lmdb") + _create_mixed_nloc_lmdb(mixed_path) + seed = 0 + budget = "mix:18" + + def _epoch_lengths(count: int = 4) -> list[int]: + reader = LmdbDataReader(mixed_path, ["O", "H"], batch_size=budget) + reader.add_data_requirement( + [ + DataRequirementItem("force", 3, atomic=True, must=False), + DataRequirementItem("energy", 1, atomic=False, must=False), + ] + ) + reader.use_ragged_batches(True) + sampler = LmdbBatchSampler(reader, shuffle=True, seed=seed) + lengths = [] + for epoch in range(count): + sampler.set_epoch(epoch) + lengths.append(len(sampler)) + reader.close() + return lengths + + lengths = _epoch_lengths() + self.assertNotEqual( + lengths[1], + lengths[2], + "fixture must vary mix epoch length so the wrap bug is observable", + ) + + def _open_iterator() -> tuple[LmdbDataReader, LmdbBatchIterator]: + reader = LmdbDataReader(mixed_path, ["O", "H"], batch_size=budget) + reader.add_data_requirement( + [ + DataRequirementItem("force", 3, atomic=True, must=False), + DataRequirementItem("energy", 1, atomic=False, must=False), + ] + ) + reader.use_ragged_batches(True) + sampler = LmdbBatchSampler(reader, shuffle=True, seed=seed) + return reader, LmdbBatchIterator(reader, sampler, num_workers=0) + + reference_reader, reference = _open_iterator() + try: + for _ in range(lengths[0]): + next(reference) + self.assertEqual(reference.state_dict(), {"epoch": 1, "batch_index": 0}) + self.assertEqual(reference._logical_epoch_length, lengths[1]) + # Mid-epoch-1 checkpoint under the true length, then continue past + # where a wrong stored length would have wrapped early. + mid = lengths[1] // 2 + for _ in range(mid): + next(reference) + progress = reference.state_dict() + expected = [next(reference)["fid"] for _ in range(lengths[1] - mid + 1)] + finally: + reference.close() + reference_reader.close() + + self.assertEqual(progress["epoch"], 1) + self.assertEqual(progress["batch_index"], mid) + + restored_reader, restored = _open_iterator() + try: + restored.load_state_dict(progress) + actual = [next(restored)["fid"] for _ in range(lengths[1] - mid + 1)] + finally: + restored.close() + restored_reader.close() + + self.assertEqual(actual, expected) + + def test_distributed_ranks_restore_disjoint_shards(self) -> None: + systems = [self._system(rank=rank, world_size=2, seed=19) for rank in range(2)] + try: + for system in systems: + system.get_batch() + progresses = [system.state_dict() for system in systems] + expected = [ + [system.get_batch()["fid"] for _ in range(2)] for system in systems + ] + finally: + for system in systems: + system.close() + + self.assertEqual(progresses[0]["epoch"], progresses[1]["epoch"]) + self.assertEqual(progresses[0]["batch_index"], progresses[1]["batch_index"]) + + restored = [self._system(rank=rank, world_size=2, seed=19) for rank in range(2)] + try: + for system, progress in zip(restored, progresses, strict=True): + system.load_state_dict(progress) + actual = [ + [system.get_batch()["fid"] for _ in range(2)] for system in restored + ] + finally: + for system in restored: + system.close() + + self.assertEqual(actual, expected) + self.assertFalse(set(actual[0][0]) & set(actual[1][0])) + + def test_world_size_mismatch_raises(self) -> None: + system = self._system(world_size=1) + try: + progress = system.state_dict() + progress["world_size"] = 2 + with self.assertRaises(ValueError): + system.load_state_dict(progress) + finally: + system.close() + + def test_multitask_independent_cursors(self) -> None: + first = self._system(seed=3) + second = self._system(seed=5) + try: + first.get_batch() + second.get_batch() + second.get_batch() + progress = collect_training_data_progress({"water": first, "ice": second}) + expected = { + "water": first.get_batch()["fid"], + "ice": second.get_batch()["fid"], + } + finally: + first.close() + second.close() + + restored_water = self._system(seed=3) + restored_ice = self._system(seed=5) + try: + restore_training_data_progress( + {"water": restored_water, "ice": restored_ice}, + progress, + ) + actual = { + "water": restored_water.get_batch()["fid"], + "ice": restored_ice.get_batch()["fid"], + } + finally: + restored_water.close() + restored_ice.close() + + self.assertEqual(actual, expected) + self.assertEqual(progress["tasks"]["water"]["batch_index"], 1) + self.assertEqual(progress["tasks"]["ice"]["batch_index"], 2) + + +class TestDirectoryDataProgress(unittest.TestCase): + def setUp(self) -> None: + self.tmp = tempfile.TemporaryDirectory() + self.system_path = _create_directory_system(Path(self.tmp.name)) + + def tearDown(self) -> None: + self.tmp.cleanup() + + def test_directory_mid_set_restart_matches_coords(self) -> None: + dp_random.seed(123) + reference = DeepmdDataSystem( + systems=[str(self.system_path)], + batch_size=2, + test_size=1, + type_map=["O", "H"], + trn_all_set=True, + ) + consumed = [reference.get_batch()["coord"].copy() for _ in range(2)] + progress = collect_training_data_progress({"Default": reference}) + expected = [reference.get_batch()["coord"].copy() for _ in range(2)] + + dp_random.seed(999) + restored = DeepmdDataSystem( + systems=[str(self.system_path)], + batch_size=2, + test_size=1, + type_map=["O", "H"], + trn_all_set=True, + ) + restore_training_data_progress({"Default": restored}, progress) + actual = [restored.get_batch()["coord"].copy() for _ in range(2)] + + self.assertEqual(len(consumed), 2) + for left, right in zip(actual, expected, strict=True): + np.testing.assert_array_equal(left, right) + + def test_missing_progress_keeps_legacy_start(self) -> None: + dp_random.seed(1) + data = DeepmdDataSystem( + systems=[str(self.system_path)], + batch_size=2, + test_size=1, + type_map=["O", "H"], + trn_all_set=True, + ) + restore_training_data_progress({"Default": data}, None) + self.assertEqual(data.data_systems[0].set_count, 0) + self.assertEqual(data.data_systems[0].iterator, 0) + + +class TestTaskSelectionRngProgress(unittest.TestCase): + def test_rng_roundtrip_in_progress_payload(self) -> None: + dp_random.seed(7) + _ = dp_random.choice(np.arange(4), p=[0.1, 0.2, 0.3, 0.4]) + progress = collect_training_data_progress({}) + self.assertEqual(progress["version"], DATA_PROGRESS_VERSION) + expected = [int(dp_random.choice(np.arange(4))) for _ in range(5)] + + dp_random.seed(0) + restore_training_data_progress({}, progress) + actual = [int(dp_random.choice(np.arange(4))) for _ in range(5)] + self.assertEqual(actual, expected) + + +if __name__ == "__main__": + unittest.main()