diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 943f9ffb05..e6124d8ce9 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -11,6 +11,10 @@ AbstractTrainEntrypoint, TrainEntrypointOptions, ) +from .schedule import ( + StepSchedule, + resolve_step_schedule, +) from .trainer import ( DEFAULT_TASK_KEY, AbstractTrainer, @@ -30,6 +34,7 @@ "AbstractTrainer", "LearningCurveWriter", "RankContext", + "StepSchedule", "TrainEntrypointOptions", "TrainStepResult", "TrainerConfig", @@ -41,4 +46,5 @@ "iter_training_task_configs", "make_task_maps", "print_data_summaries", + "resolve_step_schedule", ] diff --git a/deepmd/dpmodel/train/schedule.py b/deepmd/dpmodel/train/schedule.py new file mode 100644 index 0000000000..0e04c27466 --- /dev/null +++ b/deepmd/dpmodel/train/schedule.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-independent resolution of the training-step schedule. + +A run length is expressed either directly in optimizer steps +(``training.numb_steps``) or as a number of passes over the training data +(``training.numb_epoch`` for single-task runs, ``training.num_epoch_dict`` +for multi-task runs). Converting epochs into steps requires exactly one +backend-specific quantity, the epoch length of a task, which callers supply +through the ``epoch_length`` callback. Validation of the mutually exclusive +options, the multi-task step split and the resulting task sampling weights +are shared by every backend. +""" + +from __future__ import ( + annotations, +) + +import logging +from dataclasses import ( + dataclass, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +import numpy as np + +from deepmd.dpmodel.utils.training_utils import ( + resolve_model_prob, + resolve_model_prob_from_epochs, +) + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + Mapping, + Sequence, + ) + +log = logging.getLogger(__name__) + +__all__ = ["StepSchedule", "resolve_step_schedule"] + + +@dataclass(frozen=True) +class StepSchedule: + """Resolved run length and task sampling weights. + + Attributes + ---------- + num_steps : int + Total number of optimizer steps of the run. + model_prob : np.ndarray | None + Probability of selecting each task at a training step, with shape + ``(ntasks,)`` and ordered like the ``model_keys`` passed to + :func:`resolve_step_schedule`. ``None`` for single-task runs. + """ + + num_steps: int + model_prob: np.ndarray | None = None + + +def resolve_step_schedule( + training_params: Mapping[str, Any], + *, + multi_task: bool, + model_keys: Sequence[str], + training_data: Mapping[str, Any], + epoch_length: Callable[[str], int], + broadcast: Callable[[list[int]], Sequence[int]] | None = None, + rank: int = 0, +) -> StepSchedule: + """Resolve the run length and the task sampling weights. + + Parameters + ---------- + training_params : Mapping[str, Any] + The normalized ``training`` section. ``numb_steps``, ``numb_epoch``, + ``num_epoch_dict`` and ``model_prob`` are read from it. + multi_task : bool + Whether the run trains several task branches. Single-task runs accept + ``numb_epoch``, multi-task runs accept ``num_epoch_dict``. + model_keys : Sequence[str] + Task keys in the order used by the trainer. Single-task runs pass the + one key under which their data is registered. + training_data : Mapping[str, Any] + Training data of every task, keyed like ``model_keys``. Only consulted + when multi-task sampling weights fall back to the per-task data size. + epoch_length : Callable[[str], int] + Maps a task key to the number of steps *this rank* performs during one + epoch of that task. Backends whose data pipeline is sharded across + ranks report their per-rank batch count directly; backends that + replicate the dataset on every rank divide the dataset-wide batch count + by the world size, so that an epoch always denotes one pass over the + whole dataset. + broadcast : Callable[[list[int]], Sequence[int]], optional + Replaces the epoch lengths by rank 0's values. Epoch lengths follow + from floating-point sampling weights and may otherwise differ by one + unit across ranks, which would desynchronize the run length and + deadlock later collective calls. + rank : int, optional + Process rank, used to restrict informational logging to the chief. + + Returns + ------- + StepSchedule + The resolved run length and, for multi-task runs, the task sampling + weights. + + Raises + ------ + ValueError + If the run length is unspecified, over-specified, or non-positive. + """ + keys = list(model_keys) + num_steps = training_params.get("numb_steps") + num_epoch = training_params.get("numb_epoch") + num_epoch_dict = training_params.get("num_epoch_dict") or {} + + if not multi_task: + if num_steps is not None and num_epoch is not None: + raise ValueError( + "training.numb_steps and training.num_epoch are mutually exclusive." + ) + if num_steps is not None: + return StepSchedule(num_steps=int(num_steps)) + if num_epoch is None: + raise ValueError( + "Either training.numb_steps or training.num_epoch must be set." + ) + num_epoch = float(num_epoch) + if num_epoch <= 0.0: + raise ValueError("training.num_epoch must be positive.") + (total_numb_batch,) = _epoch_lengths(keys, epoch_length, broadcast) + steps = int(np.ceil(num_epoch * total_numb_batch)) + if rank == 0: + log.info( + "Computed num_steps=%d from num_epoch=%s and total_numb_batch=%d.", + steps, + num_epoch, + total_numb_batch, + ) + return StepSchedule(num_steps=steps) + + if num_epoch_dict: + if num_steps is not None: + raise ValueError( + "training.numb_steps and training.num_epoch_dict " + "are mutually exclusive." + ) + per_task_total = _epoch_lengths(keys, epoch_length, broadcast) + model_prob, steps, per_task_steps = resolve_model_prob_from_epochs( + keys, + num_epoch_dict, + np.asarray(per_task_total, dtype=np.float64), + ) + if rank == 0: + log.info( + "Computed model_prob=%s and num_steps=%d from num_epoch_dict=%s " + "with per-task target steps: %s.", + model_prob, + steps, + dict(num_epoch_dict), + {k: int(np.ceil(v)) for k, v in per_task_steps.items()}, + ) + return StepSchedule(num_steps=steps, model_prob=model_prob) + + if num_steps is None: + raise ValueError( + "Either training.numb_steps (multi-task only) or " + "training.num_epoch_dict must be set." + ) + model_prob = resolve_model_prob( + keys, + training_params.get("model_prob"), + training_data, + rank=rank, + ) + return StepSchedule(num_steps=int(num_steps), model_prob=model_prob) + + +def _epoch_lengths( + model_keys: Sequence[str], + epoch_length: Callable[[str], int], + broadcast: Callable[[list[int]], Sequence[int]] | None, +) -> list[int]: + """Collect the per-task epoch lengths agreed upon by every rank.""" + lengths = [int(epoch_length(model_key)) for model_key in model_keys] + if broadcast is not None: + lengths = [int(value) for value in broadcast(lengths)] + for model_key, length in zip(model_keys, lengths, strict=True): + if length <= 0: + raise ValueError( + f"Number of training batches per epoch must be positive for " + f"task '{model_key}', got {length}." + ) + return lengths diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 96173f6bf6..a85f92be87 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -7,8 +7,23 @@ import logging import math +import multiprocessing +import signal +import threading from collections.abc import ( Iterator, + Sequence, +) +from concurrent.futures import ( + Future, + ProcessPoolExecutor, +) +from concurrent.futures import wait as futures_wait +from concurrent.futures.process import ( + BrokenProcessPool, +) +from dataclasses import ( + dataclass, ) from pathlib import ( Path, @@ -114,7 +129,7 @@ def _read_metadata(txn: lmdb.Transaction) -> dict: return msgpack.unpackb(raw, raw=False) -def _decode_array(obj: dict) -> np.ndarray: +def _decode_array(obj: dict, *, copy: bool = True) -> np.ndarray: """Reconstruct ndarray from msgpack-encoded dict with {type, shape, data}. Handles both string keys ("type", "data") and byte keys (b"type", b"data"). @@ -128,7 +143,8 @@ def _decode_array(obj: dict) -> np.ndarray: shape = tuple(obj[shape_key]) else: shape = (len(data) // dtype.itemsize,) - return np.frombuffer(data, dtype=dtype).reshape(shape).copy() + array = np.frombuffer(data, dtype=dtype).reshape(shape) + return array.copy() if copy else array def _is_encoded_array(val: Any) -> bool: @@ -138,21 +154,25 @@ def _is_encoded_array(val: Any) -> bool: return ("data" in val and "type" in val) or (b"data" in val and b"type" in val) -def _decode_value(val: Any) -> Any: +def _decode_value(val: Any, *, copy_arrays: bool = True) -> Any: """Decode a value: encoded array -> ndarray, list of encoded -> list of ndarray, else pass through.""" if _is_encoded_array(val): - return _decode_array(val) + return _decode_array(val, copy=copy_arrays) elif isinstance(val, list) and len(val) > 0 and _is_encoded_array(val[0]): - return [_decode_array(item) for item in val] + return [_decode_array(item, copy=copy_arrays) for item in val] return val -def _decode_frame(raw_bytes: bytes) -> dict[str, Any]: +def _decode_frame( + raw_bytes: bytes, + *, + copy_arrays: bool = True, +) -> dict[str, Any]: """Decode a msgpack-serialized frame into a dict of numpy arrays / scalars.""" frame = msgpack.unpackb(raw_bytes, raw=False) result = {} for key, val in frame.items(): - result[key] = _decode_value(val) + result[key] = _decode_value(val, copy_arrays=copy_arrays) return result @@ -195,6 +215,747 @@ def _remap_atom_types(atype: np.ndarray, type_remap: np.ndarray) -> np.ndarray: return remapped_atype +@dataclass +class LmdbDecodeConfig: + """Serializable state required to decode one LMDB frame. + + The configuration deliberately excludes the LMDB environment and the + dataset-wide index tables. It can therefore be sent to worker processes + without duplicating the potentially very large metadata owned by + :class:`LmdbDataReader`. + + Parameters + ---------- + ntypes + Number of model atom types. + natoms + Fallback atom count for records without ``atom_types``. + type_remap + Optional LMDB-type to model-type lookup table. + data_requirements + Registered data requirements keyed by field name. + """ + + ntypes: int + natoms: int + type_remap: np.ndarray | None + data_requirements: dict[str, Any] + + +def _requirement_dtype(requirement: Any) -> np.dtype: + """Resolve the NumPy dtype associated with a data requirement.""" + if isinstance(requirement, dict): + dtype = requirement.get("dtype") + high_precision = requirement.get("high_prec", False) + else: + dtype = getattr(requirement, "dtype", None) + high_precision = getattr(requirement, "high_prec", False) + if dtype is not None: + return np.dtype(dtype) + return np.dtype( + GLOBAL_ENER_FLOAT_PRECISION if high_precision else GLOBAL_NP_FLOAT_PRECISION + ) + + +def _resolve_frame_dtype(config: LmdbDecodeConfig, key: str) -> np.dtype: + """Resolve one decoded field's output dtype.""" + requirement = config.data_requirements.get(key) + if requirement is not None: + return _requirement_dtype(requirement) + if key in _HIGH_PREC_KEYS: + return np.dtype(GLOBAL_ENER_FLOAT_PRECISION) + return np.dtype(GLOBAL_NP_FLOAT_PRECISION) + + +def _compute_frame_natoms(atype: np.ndarray, ntypes: int) -> np.ndarray: + """Build ``[nloc, nloc, count(type_0), ...]`` for one frame. + + Negative virtual types are excluded from the per-type counts, matching + mixed-type NPY data handling, as are positive indices outside the + configured type map. The leading ``nloc`` entries still count every atom + slot. + """ + nloc = len(atype) + real_atype = atype[(atype >= 0) & (atype < ntypes)] + counts = np.bincount(real_atype, minlength=ntypes) + natoms = np.empty(ntypes + 2, dtype=np.int64) + natoms[0] = nloc + natoms[1] = nloc + natoms[2:] = counts + return natoms + + +def decode_lmdb_frame( + raw: bytes, + original_key: int, + config: LmdbDecodeConfig, + *, + copy_arrays: bool, +) -> dict[str, Any]: + """Decode and normalize one LMDB record. + + Parameters + ---------- + raw + Msgpack-encoded frame payload. + original_key + Integer LMDB frame key. + config + Decoder state independent of the LMDB environment. + copy_arrays + Whether encoded arrays are copied while unpacking. Batch decoding sets + this to ``False`` because every value is copied exactly once into its + preallocated batch destination. + + Returns + ------- + dict[str, Any] + One normalized frame in DeePMD data-system convention. + """ + frame = _remap_keys(_decode_frame(raw, copy_arrays=copy_arrays)) + + for metadata_key in ("atom_numbs", "atom_names", "orig"): + frame.pop(metadata_key, None) + + if "coord" in frame and isinstance(frame["coord"], np.ndarray): + frame["coord"] = ( + frame["coord"] + .reshape(-1, 3) + .astype(_resolve_frame_dtype(config, "coord"), copy=False) + ) + if "box" in frame and isinstance(frame["box"], np.ndarray): + frame["box"] = ( + frame["box"] + .reshape(9) + .astype(_resolve_frame_dtype(config, "box"), copy=False) + ) + if "energy" in frame: + value = frame["energy"] + if isinstance(value, np.ndarray): + frame["energy"] = value.reshape(1).astype( + _resolve_frame_dtype(config, "energy"), copy=False + ) + else: + frame["energy"] = np.array( + [float(value)], dtype=_resolve_frame_dtype(config, "energy") + ) + if "force" in frame and isinstance(frame["force"], np.ndarray): + frame["force"] = ( + frame["force"] + .reshape(-1, 3) + .astype(_resolve_frame_dtype(config, "force"), copy=False) + ) + if "atype" in frame and isinstance(frame["atype"], np.ndarray): + frame["atype"] = frame["atype"].reshape(-1).astype(np.int64, copy=False) + if config.type_remap is not None: + frame["atype"] = _remap_atom_types(frame["atype"], config.type_remap) + if "virial" in frame and isinstance(frame["virial"], np.ndarray): + frame["virial"] = ( + frame["virial"] + .reshape(9) + .astype(_resolve_frame_dtype(config, "virial"), copy=False) + ) + + atype = frame.get("atype") + if atype is not None: + frame_natoms = len(atype) + natoms = _compute_frame_natoms(atype, config.ntypes) + else: + frame_natoms = config.natoms + natoms = np.array( + [config.natoms, config.natoms] + [0] * config.ntypes, + dtype=np.int64, + ) + frame["natoms"] = natoms + frame["real_natoms_vec"] = natoms + + requirements = config.data_requirements + coord = frame.get("coord") + if ( + "min_pair_dist" in requirements + and "min_pair_dist" not in frame + and isinstance(coord, np.ndarray) + and isinstance(atype, np.ndarray) + ): + box = frame.get("box") + if box is not None and np.allclose(box, 0.0): + box = None + requirement = requirements["min_pair_dist"] + default = ( + requirement.get("default", 0.0) + if isinstance(requirement, dict) + else getattr(requirement, "default", 0.0) + ) + frame["find_min_pair_dist"] = np.float32(1.0) + frame["min_pair_dist"] = np.array( + [ + compute_min_pair_dist_single( + coord, + box, + atype, + stop_below=float(default), + ) + ], + dtype=_resolve_frame_dtype(config, "min_pair_dist"), + ) + + structural_keys = frozenset( + { + "coord", + "box", + "atype", + "natoms", + "real_natoms_vec", + "fid", + } + ) + for key in list(frame): + if key.startswith("find_") or key in structural_keys or key in requirements: + continue + frame.setdefault(f"find_{key}", np.float32(1.0)) + + for key, requirement in requirements.items(): + if isinstance(requirement, dict): + ndof = requirement["ndof"] + default = requirement["default"] + atomic = requirement["atomic"] + repeat = requirement.get("repeat", 1) + else: + ndof = requirement.ndof + default = requirement.default + atomic = requirement.atomic + repeat = getattr(requirement, "repeat", 1) + dtype = _requirement_dtype(requirement) + + if key not in frame: + frame[f"find_{key}"] = np.float32(0.0) + shape = (frame_natoms, ndof) if atomic else (ndof,) + data = np.full(shape, default, dtype=dtype) + if repeat != 1: + data = np.repeat(data, repeat).reshape(-1) + frame[key] = data + else: + frame.setdefault(f"find_{key}", np.float32(1.0)) + if repeat != 1 and isinstance(frame[key], np.ndarray): + frame[key] = ( + np.repeat(frame[key], repeat).reshape(-1).astype(dtype, copy=False) + ) + + for key in ("fparam", "aparam", "spin", "charge_spin"): + frame.setdefault( + f"find_{key}", + np.float32(1.0 if key in frame else 0.0), + ) + + frame["fid"] = original_key + return frame + + +def _allocate_lmdb_batch( + frame: dict[str, Any], + batch_size: int, +) -> dict[str, Any]: + """Allocate a contiguous NumPy batch from the first decoded frame.""" + batch: dict[str, Any] = {} + for key, value in frame.items(): + if key.startswith("find_"): + batch[key] = value + elif key == "fid": + batch[key] = [None] * batch_size + batch[key][0] = value + elif key == "type": + continue + elif value is None: + batch[key] = None + else: + array = np.asarray(value) + destination = np.empty((batch_size, *array.shape), dtype=array.dtype) + destination[0] = array + batch[key] = destination + return batch + + +def decode_lmdb_batch( + transaction: lmdb.Transaction, + original_keys: Sequence[int], + frame_format: str, + config: LmdbDecodeConfig, +) -> dict[str, Any]: + """Decode LMDB records directly into preallocated contiguous arrays. + + The function keeps at most one temporary frame alive. It avoids the + decode-copy, dtype-copy, Python frame-list, and final ``numpy.stack`` + sequence used by generic collation. + """ + if not original_keys: + raise ValueError("decode_lmdb_batch requires at least one frame key") + + batch: dict[str, Any] | None = None + batch_size = len(original_keys) + expected_fields: frozenset[str] | None = None + for row, original_key in enumerate(original_keys): + key = format(int(original_key), frame_format).encode() + raw = transaction.get(key) + if raw is None: + raise IndexError(f"Frame {original_key} not found in LMDB") + frame = decode_lmdb_frame( + raw, + int(original_key), + config, + copy_arrays=False, + ) + if batch is None: + batch = _allocate_lmdb_batch(frame, batch_size) + expected_fields = frozenset(frame) + continue + + frame_fields = frozenset(frame) + if frame_fields != expected_fields: + raise ValueError( + "LMDB frames in one same-nloc batch expose inconsistent fields: " + f"frame {original_keys[0]} has {sorted(expected_fields)}, while " + f"frame {original_key} has {sorted(frame_fields)}" + ) + for field, value in frame.items(): + if field.startswith("find_"): + if not np.array_equal(batch[field], value): + raise ValueError( + f"LMDB field availability changes within one batch: " + f"{field!r} differs at frame {original_key}" + ) + continue + if field == "type" or value is None: + continue + if field == "fid": + batch[field][row] = value + else: + destination = batch[field] + array = np.asarray(value) + if destination.shape[1:] != array.shape: + raise ValueError( + f"LMDB field {field!r} changes shape within one batch: " + f"expected {destination.shape[1:]}, got {array.shape} " + f"for frame {original_key}" + ) + result_dtype = np.result_type(destination.dtype, array.dtype) + if result_dtype != destination.dtype: + promoted = np.empty(destination.shape, dtype=result_dtype) + promoted[:row] = destination[:row] + batch[field] = destination = promoted + destination[row] = array + + assert batch is not None + batch["sid"] = np.asarray([0], dtype=np.int64) + return batch + + +_WORKER_LMDB_READERS: dict[ + str, + tuple[lmdb.Environment, lmdb.Transaction], +] = {} + + +def _decode_lmdb_worker_chunk( + lmdb_path: str, + frame_format: str, + config: LmdbDecodeConfig, + original_keys: list[int], +) -> dict[str, Any]: + """Decode one chunk using process-local LMDB state.""" + reader = _WORKER_LMDB_READERS.get(lmdb_path) + if reader is None: + environment = lmdb.open( + lmdb_path, + readonly=True, + lock=False, + readahead=False, + meminit=False, + ) + reader = (environment, environment.begin()) + _WORKER_LMDB_READERS[lmdb_path] = reader + return decode_lmdb_batch( + reader[1], + original_keys, + frame_format, + config, + ) + + +def _merge_lmdb_chunks(chunks: list[dict[str, Any]]) -> dict[str, Any]: + """Merge ordered worker chunks into one contiguous batch.""" + if not chunks: + raise ValueError("cannot merge an empty LMDB chunk list") + if len(chunks) == 1: + return chunks[0] + + first = chunks[0] + expected_fields = frozenset(first) + for chunk_index, chunk in enumerate(chunks[1:], start=1): + chunk_fields = frozenset(chunk) + if chunk_fields != expected_fields: + raise ValueError( + "LMDB worker chunks expose inconsistent fields: " + f"chunk 0 has {sorted(expected_fields)}, while chunk " + f"{chunk_index} has {sorted(chunk_fields)}" + ) + + merged: dict[str, Any] = {} + for key, value in first.items(): + if key.startswith("find_"): + for chunk_index, chunk in enumerate(chunks[1:], start=1): + if not np.array_equal(value, chunk[key]): + raise ValueError( + "LMDB field availability changes across worker chunks: " + f"{key!r} differs in chunk {chunk_index}" + ) + merged[key] = value + elif key == "sid" or value is None: + merged[key] = value + elif key == "fid": + merged[key] = [frame_id for chunk in chunks for frame_id in chunk[key]] + else: + merged[key] = np.concatenate([chunk[key] for chunk in chunks], axis=0) + return merged + + +@dataclass +class _LmdbPoolEntry: + """Reference-counted process pool shared by data tasks in one rank. + + Attributes + ---------- + executor : ProcessPoolExecutor + The pool itself. + users : int + Number of iterators holding the pool, which is retired by its last one. + healthy : bool + Whether the pool still decodes. Losing a decoder disables the pool for + every iterator sharing it, and marks it as one that must not be waited + on: a pool stuck reading the partial result of a decoder killed + mid-write never finishes shutting down. + """ + + executor: ProcessPoolExecutor + users: int + healthy: bool = True + + +_LMDB_POOL_LOCK = threading.Lock() +_LMDB_POOLS: dict[int, _LmdbPoolEntry] = {} + + +def _detach_decoder_from_session() -> None: + """Shield a decoder from the hangup that ends its launching session. + + A decoder is a background helper of the training process and owns no + terminal, so the ``SIGHUP`` delivered when the session a run was launched + from goes away carries no meaning for it, while the default disposition + makes it fatal. The signals by which a run is actually stopped, ``SIGINT`` + and ``SIGTERM``, keep their disposition. + + This protects the decoders alone. It is effective because the pool is + built on the ``spawn`` start method, whose workers are direct children of + the training process with no intermediary of their own to lose. + """ + if hasattr(signal, "SIGHUP"): + signal.signal(signal.SIGHUP, signal.SIG_IGN) + + +def _create_lmdb_executor(num_workers: int) -> ProcessPoolExecutor: + """Create a CUDA-safe LMDB decoder process pool. + + The ``spawn`` start method is chosen over ``forkserver`` for the sake of + the shielding above: a fork server is an unshielded intermediary whose own + death is reported as the death of every decoder it started, which breaks + the pool however well the decoders themselves are protected. + """ + return ProcessPoolExecutor( + max_workers=num_workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=_detach_decoder_from_session, + ) + + +def _acquire_lmdb_executor(num_workers: int) -> _LmdbPoolEntry: + """Acquire the process-wide pool for one worker count.""" + with _LMDB_POOL_LOCK: + entry = _LMDB_POOLS.get(num_workers) + if entry is None: + entry = _LmdbPoolEntry( + executor=_create_lmdb_executor(num_workers), + users=0, + ) + _LMDB_POOLS[num_workers] = entry + entry.users += 1 + return entry + + +def _release_lmdb_executor(num_workers: int) -> None: + """Release one pool user and stop the pool after its final user.""" + retired: _LmdbPoolEntry | None = None + with _LMDB_POOL_LOCK: + entry = _LMDB_POOLS.get(num_workers) + if entry is None: + return + entry.users -= 1 + if entry.users == 0: + retired = entry + del _LMDB_POOLS[num_workers] + if retired is not None: + if not retired.healthy: + _dismantle_lmdb_executor(retired.executor) + retired.executor.shutdown(wait=retired.healthy, cancel_futures=True) + + +def _dismantle_lmdb_executor(executor: ProcessPoolExecutor) -> None: + """Force a pool that stopped delivering to finish shutting down. + + A pool reading the partial result of a decoder killed mid-write waits for + bytes that never arrive, because the training process itself holds the + last write end of that queue. The interpreter joins every pool manager + thread before it exits, so a run would complete its training and then + never terminate. Closing the write end delivers the awaited end of file; + the surviving decoders are stopped first so that none of them writes into + a queue about to close. The pool offers no public way to release a manager + thread already committed to a read. + """ + for process in list(getattr(executor, "_processes", {}).values()): + if process.exitcode is None: + process.kill() + writer = getattr(getattr(executor, "_result_queue", None), "_writer", None) + if writer is not None: + writer.close() + + +#: Seconds between two liveness checks while waiting on the decoder pool. The +#: wait ends as soon as the chunks arrive, so this bounds only how long a pool +#: that has stopped delivering goes unnoticed. +_DECODER_LIVENESS_INTERVAL = 5.0 + + +@dataclass +class _PendingBatch: + """A batch handed to the decoder pool, and the indices that produced it. + + The indices are retained so that the batch can still be decoded in this + process should the pool fail to deliver it. + """ + + indices: list[int] + futures: list[Future[dict[str, Any]]] + + +class LmdbBatchIterator: + """Deterministic same-nloc batches with parallel decode and one prefetch. + + The sampler remains in the parent process. Worker tasks receive only the + selected integer frame keys and compact decoder state, so the dataset-wide + metadata is never duplicated. Data tasks in the same rank share one + process pool while retaining independent samplers and pending batches. + Before a new pass is prefetched, samplers exposing ``set_epoch`` are + advanced to the next deterministic shuffle state. + + Parameters + ---------- + reader + LMDB reader that owns metadata and the synchronous transaction. + sampler + Finite iterator yielding same-nloc dataset-index batches. + num_workers + Decoder process count. Zero or one selects synchronous decoding. + """ + + def __init__( + self, + reader: "LmdbDataReader", + sampler: Any, + num_workers: int, + ) -> None: + if num_workers < 0: + raise ValueError(f"num_workers must be non-negative, got {num_workers}") + self._reader = reader + self._sampler = sampler + self._epoch = 0 + self._iterator = self._iter_epoch() + 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 + + def __iter__(self) -> "LmdbBatchIterator": + return self + + def __next__(self) -> dict[str, Any]: + if self._closed: + raise RuntimeError("cannot read from a closed LMDB batch iterator") + + if self._pending is not None: + pending, self._pending = self._pending, None + batch = self._collect(pending) + elif self._deferred_indices is not None: + indices, self._deferred_indices = self._deferred_indices, None + batch = self._reader.decode_batch(indices) + else: + batch = self._decode(self._next_indices()) + + self._schedule(self._next_indices()) + return batch + + 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) + if futures is None: + return self._reader.decode_batch(indices) + return self._collect(_PendingBatch(indices, futures)) + + def _offer(self, indices: list[int]) -> list[Future[dict[str, Any]]] | None: + """Hand a batch to the pool, or ``None`` if it will not take it.""" + if not self._worth_decoding_in_parallel(indices): + return None + if self._pool is None: + self._pool = _acquire_lmdb_executor(self._num_workers) + if not self._pool.healthy: + return None + try: + return self._submit(indices) + except BrokenProcessPool: + self._lose_the_pool() + return None + + def _collect(self, pending: _PendingBatch) -> dict[str, Any]: + """Return a batch from the pool, decoding it here if the pool cannot. + + A decoder killed from outside -- by the kernel under memory pressure, + or by a signal aimed at the session the run was launched from -- ends + the batch one of two ways. Usually the pool notices and fails every + future it holds. Should the decoder die midway through writing a + result, however, the pool reads that partial result forever instead of + reporting itself broken, and the run stops with no diagnosis and no + error; punctuating the wait with a liveness check covers that case. + """ + try: + if self._wait_for_decoder(pending.futures): + return _merge_lmdb_chunks( + [future.result() for future in pending.futures] + ) + except BrokenProcessPool: + pass + self._lose_the_pool(pending.futures) + return self._reader.decode_batch(pending.indices) + + def _wait_for_decoder( + self, + futures: "Sequence[Future[dict[str, Any]]]", + ) -> bool: + """Wait for submitted work while checking that its decoder survives.""" + remaining = set(futures) + while remaining: + _, remaining = futures_wait(remaining, timeout=_DECODER_LIVENESS_INTERVAL) + if remaining and self._decoder_exited(): + return False + return not any( + isinstance(future.exception(), BrokenProcessPool) for future in futures + ) + + def _decoder_exited(self) -> bool: + """Whether any decoder has exited, which the pool reports nowhere else.""" + processes = getattr(self._pool.executor, "_processes", None) or {} + return any(process.exitcode is not None for process in processes.values()) + + def _lose_the_pool(self, futures: "Sequence[Future[dict[str, Any]]]" = ()) -> None: + """Disable the pool for every iterator sharing it, reporting it once.""" + for future in futures: + future.cancel() + if self._pool is None or not self._pool.healthy: + return + self._pool.healthy = False + log.warning( + "An LMDB decoder process exited unexpectedly; decoding continues " + "in the training process. Throughput may drop. Set " + "DP_LMDB_NUM_WORKERS=1 to select in-process decoding from the " + "start, and launch the run under nohup or setsid so that the " + "decoders outlive the session that started it." + ) + + def _next_indices(self) -> list[int]: + """Return the next sampler batch and restart after exhaustion.""" + try: + return next(self._iterator) + except StopIteration: + self._epoch += 1 + self._iterator = self._iter_epoch() + return next(self._iterator) + + def _iter_epoch(self) -> Iterator[list[int]]: + """Create a sampler iterator for the current epoch.""" + set_epoch = getattr(self._sampler, "set_epoch", None) + if callable(set_epoch): + set_epoch(self._epoch) + return iter(self._sampler) + + def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: + """Submit one batch as balanced contiguous chunks.""" + original_keys = self._reader.original_keys(indices) + workers = min(self._num_workers, len(original_keys)) + base_size, remainder = divmod(len(original_keys), workers) + chunks: list[list[int]] = [] + start = 0 + for worker_index in range(workers): + chunk_size = base_size + int(worker_index < remainder) + stop = start + chunk_size + chunks.append(original_keys[start:stop]) + start = stop + decode_config = self._reader.worker_decode_config() + return [ + self._pool.executor.submit( + _decode_lmdb_worker_chunk, + self._reader.lmdb_path, + self._reader.frame_format, + decode_config, + chunk, + ) + for chunk in chunks + ] + + def _worth_decoding_in_parallel(self, indices: list[int]) -> bool: + """Whether process decoding amortizes its scheduling and IPC cost.""" + return self._num_workers > 1 and len(indices) >= self._num_workers + + def _schedule(self, indices: list[int]) -> None: + """Prefetch the next batch in the pool, or leave it to the caller.""" + futures = self._offer(indices) + self._pending = _PendingBatch(indices, futures) if futures else None + self._deferred_indices = None if futures else indices + + @property + def started(self) -> bool: + """Whether this iterator has acquired the shared process pool.""" + return self._pool is not None + + @property + def closed(self) -> bool: + """Whether this iterator has released its resources.""" + return self._closed + + def close(self) -> None: + """Finish or cancel prefetched work and release the decoder pool.""" + if self._closed: + return + self._closed = True + if self._pending is not None: + running = [ + future for future in self._pending.futures if not future.cancel() + ] + if running and not self._wait_for_decoder(running): + self._lose_the_pool(running) + self._pending = None + self._deferred_indices = None + if self._pool is not None: + self._pool = None + _release_lmdb_executor(self._num_workers) + + def is_lmdb(systems: str) -> bool: """Check if systems points to an LMDB dataset.""" return systems.endswith(".lmdb") or Path(systems, "data.mdb").is_file() @@ -332,7 +1093,7 @@ def __init__( batch_size: int | str = "auto", mixed_batch: bool = False, ) -> None: - self.lmdb_path = str(lmdb_path) + self.lmdb_path = str(Path(lmdb_path).resolve()) self._type_map = type_map self._env = _open_lmdb(self.lmdb_path) self.mixed_batch = mixed_batch @@ -364,9 +1125,10 @@ def __init__( f"remap={remap}" ) - # Persistent read-only transaction for __getitem__ (avoids per-read overhead). - # Safe because we use num_workers=0 in DataLoader. + # The parent transaction serves synchronous reads. Decoder workers open + # independent process-local environments and transactions. self._txn = self._env.begin() + self._closed = False # Scan per-frame nloc only when needed for same-nloc batching. # For mixed_batch=True, skip the scan entirely (future: padding handles it). @@ -508,62 +1270,49 @@ def __init__( # Data requirements tracking self._data_requirements: dict[str, DataRequirementItem] = {} + self._data_requirements_frozen = False + self._decode_config = LmdbDecodeConfig( + ntypes=self._ntypes, + natoms=self._natoms, + type_remap=self._type_remap, + data_requirements=self._data_requirements, + ) # Availability signatures are decoded lazily and reused by every # sampler epoch. Registering new requirements invalidates the cache. self._find_signature_cache: dict[int, tuple[tuple[str, bool], ...]] = {} - def _compute_natoms_vec(self, atype: np.ndarray) -> np.ndarray: - """Compute natoms_vec from a frame's atype array. - - Negative virtual types are excluded from the per-type counts, matching - mixed-type NPY data handling. This function also excludes positive - indices outside the configured type map. The leading nloc entries still - include every atom slot. - - Returns [nloc, nloc, count_type0, count_type1, ...] with length ntypes+2. - """ - nloc = len(atype) - real_atype = atype[(atype >= 0) & (atype < self._ntypes)] - counts = np.bincount(real_atype, minlength=self._ntypes) - vec = np.empty(self._ntypes + 2, dtype=np.int64) - vec[0] = nloc - vec[1] = nloc - vec[2:] = counts - return vec - def _resolve_dtype(self, key: str) -> np.dtype: """Resolve the target numpy dtype for a given key. Priority: DataRequirementItem.dtype > DataRequirementItem.high_prec > built-in defaults (energy=high, others=normal). """ - if key in self._data_requirements: - req = self._data_requirements[key] - # Support both DataRequirementItem objects and plain dicts - if isinstance(req, dict): - dtype = req.get("dtype") - if dtype is not None: - return dtype - if req.get("high_prec", False): - return GLOBAL_ENER_FLOAT_PRECISION - return GLOBAL_NP_FLOAT_PRECISION - else: - # DataRequirementItem object - if hasattr(req, "dtype") and req.dtype is not None: - return req.dtype - if hasattr(req, "high_prec") and req.high_prec: - return GLOBAL_ENER_FLOAT_PRECISION - return GLOBAL_NP_FLOAT_PRECISION - # Fall back to built-in defaults - if key in _HIGH_PREC_KEYS: - return GLOBAL_ENER_FLOAT_PRECISION - return GLOBAL_NP_FLOAT_PRECISION + return _resolve_frame_dtype(self._decode_config, key) def __del__(self) -> None: - """Release the LMDB environment ref-count on garbage collection.""" - path = getattr(self, "lmdb_path", None) - if path is not None: - _close_lmdb(path) + """Release the parent LMDB transaction and environment.""" + self.close() + + def close(self) -> None: + """Release parent-process LMDB resources idempotently.""" + if getattr(self, "_closed", False): + return + transaction = getattr(self, "_txn", None) + if transaction is not None: + transaction.abort() + self._txn = None + environment = getattr(self, "_env", None) + if environment is not None: + self._env = None + _close_lmdb(self.lmdb_path) + self._closed = True + + def _transaction(self) -> lmdb.Transaction: + """Return the active parent transaction or fail after closure.""" + transaction = self._txn + if transaction is None: + raise RuntimeError("cannot read from a closed LMDB reader") + return transaction def get_batch_size_for_nloc(self, nloc: int) -> int: """Return the per-nloc batch size for the configured rule. @@ -595,164 +1344,68 @@ def __getitem__(self, index: int) -> dict[str, Any]: ``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 if index < 0 or index >= self.nframes: raise IndexError(f"dataset index {index} out of range [0, {self.nframes})") original_key = self._retained_keys[index] key = format(original_key, self._frame_fmt).encode() - raw = self._txn.get(key) + raw = self._transaction().get(key) if raw is None: raise IndexError( f"Frame {original_key} not found in LMDB (dataset index {index})" ) - frame = _decode_frame(raw) - frame = _remap_keys(frame) - - # Remove LMDB-specific metadata keys not needed by trainer - for meta_key in ("atom_numbs", "atom_names", "orig"): - frame.pop(meta_key, None) + return decode_lmdb_frame( + raw, + original_key, + self._decode_config, + copy_arrays=True, + ) - # Flatten arrays to match DeePMD convention - if "coord" in frame and isinstance(frame["coord"], np.ndarray): - frame["coord"] = ( - frame["coord"].reshape(-1, 3).astype(self._resolve_dtype("coord")) - ) - if "box" in frame and isinstance(frame["box"], np.ndarray): - frame["box"] = frame["box"].reshape(9).astype(self._resolve_dtype("box")) - if "energy" in frame: - val = frame["energy"] - if isinstance(val, np.ndarray): - frame["energy"] = val.reshape(1).astype(self._resolve_dtype("energy")) - else: - frame["energy"] = np.array( - [float(val)], dtype=self._resolve_dtype("energy") + def original_keys(self, indices: Sequence[int]) -> list[int]: + """Translate dataset indices to original integer LMDB keys.""" + keys: list[int] = [] + for index in indices: + index = int(index) + if index < 0 or index >= self.nframes: + raise IndexError( + f"dataset index {index} out of range [0, {self.nframes})" ) - if "force" in frame and isinstance(frame["force"], np.ndarray): - frame["force"] = ( - frame["force"].reshape(-1, 3).astype(self._resolve_dtype("force")) - ) - if "atype" in frame and isinstance(frame["atype"], np.ndarray): - frame["atype"] = frame["atype"].reshape(-1).astype(np.int64) - # Remap atom types from LMDB's type_map to model's type_map - if self._type_remap is not None: - frame["atype"] = _remap_atom_types(frame["atype"], self._type_remap) - if "virial" in frame and isinstance(frame["virial"], np.ndarray): - frame["virial"] = ( - frame["virial"].reshape(9).astype(self._resolve_dtype("virial")) - ) - - # Per-frame natoms_vec from atype - atype = frame.get("atype") - if atype is not None: - frame_natoms = len(atype) - natoms_vec = self._compute_natoms_vec(atype) - frame["natoms"] = natoms_vec - frame["real_natoms_vec"] = natoms_vec - else: - frame_natoms = self._natoms - fallback = np.array( - [self._natoms, self._natoms] + [0] * self._ntypes, dtype=np.int64 - ) - frame["natoms"] = fallback - frame["real_natoms_vec"] = fallback - - if "min_pair_dist" in self._data_requirements and "min_pair_dist" not in frame: - box = frame.get("box") - if box is not None and np.allclose(box, 0.0): - box = None - req = self._data_requirements["min_pair_dist"] - min_pair_dist = float( - req.get("default", 0.0) - if isinstance(req, dict) - else getattr(req, "default", 0.0) - ) - frame["find_min_pair_dist"] = np.float32(1.0) - frame["min_pair_dist"] = np.array( - [ - compute_min_pair_dist_single( - frame["coord"], - box, - frame["atype"], - stop_below=min_pair_dist, - ) - ], - dtype=self._resolve_dtype("min_pair_dist"), - ) - - # Add find_* flags for all data keys present in the frame. - # Core structural keys and metadata are excluded — only label-like - # and auxiliary data keys get find_* flags. - for fk in list(frame.keys()): - if fk.startswith("find_") or fk in _STRUCTURAL_KEYS: - continue - # Skip keys handled by data_requirements (processed below) - if fk in self._data_requirements: - continue - if f"find_{fk}" not in frame: - frame[f"find_{fk}"] = np.float32(1.0) - - # Handle registered data requirements: fill defaults for missing keys, - # apply repeat, and cast dtype. - for req_key, req_item in self._data_requirements.items(): - # Extract requirement fields (support both dict and object) - if isinstance(req_item, dict): - ndof = req_item["ndof"] - default = req_item["default"] - atomic = req_item["atomic"] - repeat = req_item.get("repeat", 1) - req_dtype = req_item.get("dtype") - if req_dtype is None: - req_dtype = ( - GLOBAL_ENER_FLOAT_PRECISION - if req_item.get("high_prec", False) - else GLOBAL_NP_FLOAT_PRECISION - ) - else: - ndof = req_item.ndof - default = req_item.default - atomic = req_item.atomic - repeat = getattr(req_item, "repeat", 1) - req_dtype = req_item.dtype - if req_dtype is None: - req_dtype = ( - GLOBAL_ENER_FLOAT_PRECISION - if req_item.high_prec - else GLOBAL_NP_FLOAT_PRECISION - ) - - if req_key not in frame: - frame[f"find_{req_key}"] = np.float32(0.0) - if atomic: - shape = (frame_natoms, ndof) - else: - shape = (ndof,) - data = np.full(shape, default, dtype=req_dtype) - if repeat != 1: - data = np.repeat(data, repeat).reshape(-1) - frame[req_key] = data - else: - if f"find_{req_key}" not in frame: - frame[f"find_{req_key}"] = np.float32(1.0) - # Apply repeat to existing data (e.g. atom_pref repeat=3) - if repeat != 1 and isinstance(frame[req_key], np.ndarray): - frame[req_key] = ( - np.repeat(frame[req_key], repeat).reshape(-1).astype(req_dtype) - ) + keys.append(self._retained_keys[index]) + return keys + + def decode_batch(self, indices: Sequence[int]) -> dict[str, Any]: + """Decode a same-nloc batch directly into contiguous NumPy arrays.""" + self._data_requirements_frozen = True + return decode_lmdb_batch( + self._transaction(), + self.original_keys(indices), + self._frame_fmt, + self._decode_config, + ) - # Add find_* for fparam/aparam/spin/charge_spin if not already set - for extra_key in _OPTIONAL_MODEL_INPUT_KEYS: - if f"find_{extra_key}" not in frame: - frame[f"find_{extra_key}"] = ( - np.float32(1.0) if extra_key in frame else np.float32(0.0) - ) + @property + def frame_format(self) -> str: + """Format specification used for integer LMDB frame keys.""" + return self._frame_fmt - frame["fid"] = original_key + def worker_decode_config(self) -> LmdbDecodeConfig: + """Freeze and return decoder state for worker serialization.""" + self._data_requirements_frozen = True + return self._decode_config - return frame + @property + def closed(self) -> bool: + """Whether parent-process LMDB resources have been released.""" + return self._closed # --- Data requirement interface --- def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None: """Register expected keys; missing keys get default fill + find_key=0.0.""" + if self._data_requirements_frozen: + raise RuntimeError( + "LMDB data requirements must be registered before reading any frame" + ) for item in data_requirement: self._data_requirements[item["key"]] = item self._find_signature_cache.clear() @@ -856,12 +1509,9 @@ def total_batch(self) -> int: if self.mixed_batch: return math.ceil(self.nframes / self.batch_size) if self.nframes else 0 total = 0 - for nloc, indices in self._nloc_groups.items(): + for nloc, indices in collect_lmdb_sampling_groups(self): bs = self.get_batch_size_for_nloc(nloc) - signature_groups = self.group_indices_by_find_signature(indices) - total += sum( - (len(group) + bs - 1) // bs for group in signature_groups.values() - ) + total += (len(indices) + bs - 1) // bs return total @property @@ -920,6 +1570,10 @@ def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]: ``fid`` is collected as a list; ``type`` is dropped (callers should already use ``atype``); other arrays are stacked along axis 0. A ``sid`` placeholder is appended. + + The batch keeps the key order of its frames, which is the order + :func:`decode_lmdb_batch` also produces, so a batch is the same mapping + whichever of the two decode paths built it. """ import array_api_compat @@ -928,7 +1582,10 @@ def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]: xp = array_api_compat.array_namespace(frames[0]["coord"]) dev = array_api_compat.device(frames[0]["coord"]) - out: dict[str, Any] = {} + + # Availability must agree across the batch before the flags can collapse + # to one scalar per key. Frames are checked ahead of collation so a mixed + # batch is reported rather than silently reduced to its first frame. find_keys = sorted( {key for frame in frames for key in frame if key.startswith("find_")} ) @@ -943,11 +1600,11 @@ def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]: f"LMDB batch mixes {key!r} values {values}; " "SameNlocBatchSampler must group frames by label availability" ) - out[key] = frames[0][key] + out: dict[str, Any] = {} for key in frames[0]: if key.startswith("find_"): - continue + out[key] = frames[0][key] elif key == "fid": out[key] = [f[key] for f in frames] elif key == "type": @@ -1195,10 +1852,21 @@ def _expand_indices_by_blocks( return [] -def _collect_sampling_groups( +def collect_lmdb_sampling_groups( reader: "LmdbDataReader", ) -> list[tuple[int, list[int]]]: - """Collect batch groups in the stable order shared by iteration and len.""" + """Collect homogeneous LMDB groups shared by training and statistics. + + Parameters + ---------- + reader : LmdbDataReader + Reader whose frames are grouped by atom count and label availability. + + Returns + ------- + list[tuple[int, list[int]]] + Stable ``(nloc, frame indices)`` groups compatible with collation. + """ groups: list[tuple[int, list[int]]] = [] for nloc in sorted(reader.nloc_groups): signature_groups = reader.group_indices_by_find_signature( @@ -1299,7 +1967,7 @@ def _build_all_batches( list[list[int]] Each inner list has one nloc and one scalar ``find_*`` signature. """ - groups = _collect_sampling_groups(reader) + groups = collect_lmdb_sampling_groups(reader) # Build per-group batches group_batches: list[list[list[int]]] = [] @@ -1370,8 +2038,9 @@ class SameNlocBatchSampler: When auto batch_size is used, batch_size is computed per-nloc-group. - The sampler is deterministic: given the same seed, repeated calls to - ``__iter__`` produce the same batch sequence. + The sampler is deterministic for a fixed seed and epoch. Use + :meth:`set_epoch` to select a different reproducible sequence for each + training pass. Parameters ---------- @@ -1395,18 +2064,30 @@ def __init__( self._reader = reader self._shuffle = shuffle self._seed = seed + self._epoch = 0 self._block_targets = block_targets + def set_epoch(self, epoch: int) -> None: + """Set the epoch used to derive the deterministic shuffle state. + + Parameters + ---------- + epoch : int + Zero-based training epoch. + """ + self._epoch = epoch + def __iter__(self) -> Iterator[list[int]]: - """Yield batches with one nloc and one scalar find signature.""" - rng = np.random.default_rng(self._seed) + """Yield batches of frame indices, all with the same nloc.""" + seed = None if self._seed is None else self._seed + self._epoch + rng = np.random.default_rng(seed) yield from _build_all_batches( self._reader, self._shuffle, rng, self._block_targets ) def __len__(self) -> int: """Total batches across nloc and label-availability groups.""" - groups = _collect_sampling_groups(self._reader) + groups = collect_lmdb_sampling_groups(self._reader) group_block_targets = None assigned_system_ids: set[int] = set() if self._block_targets and self._reader.frame_system_ids is not None: @@ -1442,8 +2123,10 @@ class DistributedSameNlocBatchSampler: """Distributed wrapper for same-nloc batch sampling. All ranks build the same deterministic global batch list (using - ``seed + epoch``), then each rank takes a disjoint subset via - :meth:`_partition_batches`. + ``seed + epoch``). The list is padded deterministically when its length is + not divisible by the number of ranks, then each rank takes a strided + subset via :meth:`_partition_batches`. This keeps every rank on the same + sampler epoch while duplicating at most ``world_size - 1`` batches. Override :meth:`_partition_batches` for custom load-balancing strategies. The default uses strided partitioning which gives good nloc diversity per @@ -1482,6 +2165,10 @@ def __init__( self._seed = seed if seed is not None else 0 self._epoch = 0 self._block_targets = block_targets + self.refresh_batch_count() + + def refresh_batch_count(self) -> None: + """Refresh the cached global count after sampling groups change.""" self._total_batches = len( SameNlocBatchSampler( self._reader, @@ -1511,24 +2198,38 @@ def __iter__(self) -> Iterator[list[int]]: def _partition_batches(self, all_batches: list[list[int]]) -> list[list[int]]: """Partition global batches to this rank. - Default: strided partition ``all_batches[rank::world_size]``. - This gives good nloc diversity per rank since batches are - interleaved across nloc groups before shuffling. + The default pads the global list to a multiple of ``world_size`` and + then takes ``all_batches[rank::world_size]``. This gives good nloc + diversity per rank since batches are interleaved across nloc groups + before shuffling, while ensuring that every rank yields the same + number of batches. Override this method for custom load-balancing. For example, a greedy algorithm could assign batches to ranks based on estimated compute cost (``reader.frame_nlocs[batch[0]]`` gives the nloc of each batch). """ + if not all_batches: + return [] + batches_per_rank = (len(all_batches) + self._world_size - 1) // self._world_size + total_size = batches_per_rank * self._world_size + padding_size = total_size - len(all_batches) + if padding_size: + repetitions = (padding_size + len(all_batches) - 1) // len(all_batches) + all_batches = [ + *all_batches, + *(all_batches * repetitions)[:padding_size], + ] return all_batches[self._rank :: self._world_size] def __len__(self) -> int: """Number of batches for this rank.""" - return max( - 0, - (self._total_batches + self._world_size - 1 - self._rank) - // self._world_size, - ) + return (self._total_batches + self._world_size - 1) // self._world_size + + @property + def total_batches(self) -> int: + """Return the global batch count before distributed padding.""" + return self._total_batches @property def rank(self) -> int: @@ -1984,6 +2685,69 @@ def get_test(self) -> dict[str, Any]: return self._inner.get_test(nloc=self._nloc) +def _copy_lmdb_source( + src_path: str, + dst_env: lmdb.Environment, + dst_format: str, + frame_idx: int, + frame_nlocs: list[int], + frame_system_ids: list[int], + system_id_offset: int, +) -> tuple[int, dict, list[str] | None, int]: + """Copy one source under a ref-counted environment lease.""" + src_env = _open_lmdb(src_path) + try: + with src_env.begin() as transaction: + metadata = _read_metadata(transaction) + nframes, src_format, natoms_per_type = _parse_metadata(metadata) + fallback_natoms = sum(natoms_per_type) + source_nlocs = metadata.get("frame_nlocs") + source_system_ids = metadata.get("frame_system_ids") + + with src_env.begin() as src_txn, dst_env.begin(write=True) as dst_txn: + for source_index in range(nframes): + source_key = format(source_index, src_format).encode() + raw = src_txn.get(source_key) + if raw is None: + continue + destination_key = format(frame_idx, dst_format).encode() + dst_txn.put(destination_key, raw) + + if source_nlocs is not None: + frame_nlocs.append(int(source_nlocs[source_index])) + else: + frame_raw = msgpack.unpackb(raw, raw=False) + atype_raw = frame_raw.get("atom_types") + if isinstance(atype_raw, dict): + shape = atype_raw.get("shape") or atype_raw.get(b"shape") + frame_nlocs.append(int(shape[0]) if shape else fallback_natoms) + else: + frame_nlocs.append(fallback_natoms) + + if source_system_ids is not None and source_index < len( + source_system_ids + ): + frame_system_ids.append( + int(source_system_ids[source_index]) + system_id_offset + ) + else: + frame_system_ids.append(system_id_offset) + frame_idx += 1 + + if source_system_ids is not None and len(source_system_ids) > 0: + system_id_offset += max(int(value) for value in source_system_ids) + 1 + else: + system_id_offset += 1 + return ( + frame_idx, + metadata.get("system_info", {}), + metadata.get("type_map"), + system_id_offset, + ) + finally: + _close_lmdb(src_path) + + def merge_lmdb( src_paths: list[str], dst_path: str, @@ -2023,77 +2787,43 @@ def merge_lmdb( first_system_info: dict | None = None first_type_map: list[str] | None = None sys_id_offset = 0 - - for src_path in src_paths: - src_env = _open_lmdb(src_path) - with src_env.begin() as txn: - meta = _read_metadata(txn) - nframes, src_fmt, natoms_per_type = _parse_metadata(meta) - fallback_natoms = sum(natoms_per_type) - - if first_system_info is None: - first_system_info = meta.get("system_info", {}) - if first_type_map is None: - first_type_map = meta.get("type_map") - - # Check for pre-computed frame_nlocs in source - src_nlocs = meta.get("frame_nlocs") - # Check for frame_system_ids in source - src_sys_ids = meta.get("frame_system_ids") - - with src_env.begin() as src_txn, dst_env.begin(write=True) as dst_txn: - for i in range(nframes): - src_key = format(i, src_fmt).encode() - raw = src_txn.get(src_key) - if raw is None: - continue - dst_key = format(frame_idx, fmt).encode() - dst_txn.put(dst_key, raw) - - # Get nloc for this frame - if src_nlocs is not None: - frame_nlocs.append(int(src_nlocs[i])) - else: - frame_raw = msgpack.unpackb(raw, raw=False) - atype_raw = frame_raw.get("atom_types") - if isinstance(atype_raw, dict): - shape = atype_raw.get("shape") or atype_raw.get(b"shape") - if shape: - frame_nlocs.append(int(shape[0])) - else: - frame_nlocs.append(fallback_natoms) - else: - frame_nlocs.append(fallback_natoms) - - # Propagate system IDs with offset - if src_sys_ids is not None and i < len(src_sys_ids): - frame_system_ids.append(int(src_sys_ids[i]) + sys_id_offset) - else: - frame_system_ids.append(sys_id_offset) - - frame_idx += 1 - - # Update sys_id_offset for next source - if src_sys_ids is not None and len(src_sys_ids) > 0: - sys_id_offset += max(int(s) for s in src_sys_ids) + 1 - else: - sys_id_offset += 1 - - src_env.close() - - # Write merged metadata with frame_nlocs for fast init - merged_meta = { - "nframes": frame_idx, - "frame_idx_fmt": fmt, - "system_info": first_system_info or {}, - "frame_nlocs": frame_nlocs, - "frame_system_ids": frame_system_ids, - } - if first_type_map is not None: - merged_meta["type_map"] = first_type_map - with dst_env.begin(write=True) as txn: - txn.put(b"__metadata__", msgpack.packb(merged_meta, use_bin_type=True)) - dst_env.close() + try: + for src_path in src_paths: + ( + frame_idx, + source_system_info, + source_type_map, + sys_id_offset, + ) = _copy_lmdb_source( + src_path, + dst_env, + fmt, + frame_idx, + frame_nlocs, + frame_system_ids, + sys_id_offset, + ) + if first_system_info is None: + first_system_info = source_system_info + if first_type_map is None: + first_type_map = source_type_map + + merged_meta = { + "nframes": frame_idx, + "frame_idx_fmt": fmt, + "system_info": first_system_info or {}, + "frame_nlocs": frame_nlocs, + "frame_system_ids": frame_system_ids, + } + if first_type_map is not None: + merged_meta["type_map"] = first_type_map + with dst_env.begin(write=True) as transaction: + transaction.put( + b"__metadata__", + msgpack.packb(merged_meta, use_bin_type=True), + ) + finally: + dst_env.close() nloc_counts: dict[int, int] = {} for n in frame_nlocs: diff --git a/deepmd/env.py b/deepmd/env.py index c9d0fb241f..e68ac55429 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -20,6 +20,7 @@ "LRU_CACHE_SIZE", "SHARED_LIB_DIR", "SHARED_LIB_MODULE", + "get_lmdb_num_workers", "global_float_prec", ] @@ -143,6 +144,58 @@ def get_default_nthreads() -> tuple[int, int]: ) +def get_lmdb_num_workers() -> int: + """Return the per-rank LMDB decoder process count. + + ``DP_LMDB_NUM_WORKERS`` provides an explicit override. The automatic policy + limits one rank to 32 workers and approximately 64 workers per node, then + respects the process CPU affinity. Values zero and one select synchronous + decoding. + + Returns + ------- + int + Number of LMDB decoder processes. + + Raises + ------ + ValueError + If ``DP_LMDB_NUM_WORKERS`` is not a non-negative integer. + """ + configured = os.environ.get("DP_LMDB_NUM_WORKERS") + if configured is not None: + try: + workers = int(configured) + except ValueError: + raise ValueError( + "DP_LMDB_NUM_WORKERS must be a non-negative integer, " + f"got {configured!r}" + ) from None + if workers < 0: + raise ValueError( + f"DP_LMDB_NUM_WORKERS must be non-negative, got {configured!r}" + ) + return workers + + try: + available_cpus = len(os.sched_getaffinity(0)) + except AttributeError: + available_cpus = os.cpu_count() or 1 + try: + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", "1")) + except ValueError: + raise ValueError( + "LOCAL_WORLD_SIZE must be a positive integer, " + f"got {os.environ['LOCAL_WORLD_SIZE']!r}" + ) from None + if local_world_size <= 0: + raise ValueError(f"LOCAL_WORLD_SIZE must be positive, got {local_world_size}") + + cpu_share = max(1, available_cpus // local_world_size) + node_share = max(1, 64 // local_world_size) + return min(32, cpu_share, node_share) + + def _get_package_constants( config_file: Path = CONFIG_FILE, ) -> dict[str, str]: diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index ec92301b9d..40fddc51f1 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -31,11 +31,10 @@ ) from deepmd.dpmodel.train import ( change_model_out_bias, + resolve_step_schedule, ) from deepmd.dpmodel.utils import ( compute_total_numb_batch, - resolve_model_prob, - resolve_model_prob_from_epochs, ) from deepmd.loggers.training import ( format_training_message, @@ -110,9 +109,8 @@ BaseLR, ) from deepmd.pt.utils.lmdb_dataset import ( + LmdbBatchDataLoader, LmdbDataset, - _collate_lmdb_batch, - _SameNlocBatchSamplerTorch, ) from deepmd.pt.utils.stat import ( make_stat_input, @@ -217,16 +215,12 @@ def __init__( self.rank = dist.get_rank() if self.is_distributed else 0 self.world_size = dist.get_world_size() if self.is_distributed else 1 self.num_model = len(self.model_keys) - self.model_prob = None self.stat_file_specs = stat_file_specs_by_task( stat_file_spec, self.model_keys, ) # Iteration config - self.num_steps = training_params.get("numb_steps") - self.num_epoch = training_params.get("numb_epoch") - self.num_epoch_dict = training_params.get("num_epoch_dict") self.disp_file = training_params.get("disp_file", "lcurve.out") self.disp_freq = training_params.get("disp_freq", 1000) self.disp_avg = training_params.get("disp_avg", False) @@ -299,15 +293,15 @@ def get_data_loader( _validation_data: DpLoaderSet | LmdbDataset | None, _training_params: dict[str, Any], ) -> tuple[ - DataLoader, + DataLoader | LmdbBatchDataLoader, Generator[Any, None, None], - DataLoader | None, + DataLoader | LmdbBatchDataLoader | None, Generator[Any, None, None] | None, int, ]: def get_dataloader_and_iter_lmdb( _data: LmdbDataset, - ) -> tuple[DataLoader, Generator[Any, None, None]]: + ) -> tuple[LmdbBatchDataLoader, Generator[Any, None, None]]: if _data.mixed_batch: # TODO [mixed_batch=True]: Replace SameNlocBatchSampler with # RandomSampler(replacement=False) + padding collate_fn. @@ -350,13 +344,10 @@ def get_dataloader_and_iter_lmdb( block_targets=_block_targets, ) - _batch_sampler = _SameNlocBatchSamplerTorch(_inner_sampler) - _dataloader = DataLoader( + _dataloader = LmdbBatchDataLoader( _data, - batch_sampler=_batch_sampler, - num_workers=0, - collate_fn=_collate_lmdb_batch, - pin_memory=(DEVICE != "cpu"), + _inner_sampler, + pin_memory=DEVICE.type != "cpu", ) _data_iter = cycle_iterator(_dataloader) return _dataloader, _data_iter @@ -671,91 +662,36 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: ) # Resolve training steps - per_task_total = [] - if not self.multi_task: - if self.num_steps is None: - if self.num_epoch is None: - raise ValueError( - "Either training.numb_steps or training.num_epoch must be set." - ) - if self.num_epoch <= 0: - raise ValueError("training.num_epoch must be positive.") - if isinstance(training_data, LmdbDataset): - total_numb_batch = len(self.training_dataloader) - else: - sampler_weights = to_numpy_array( - self.training_dataloader.sampler.weights - ) - total_numb_batch = compute_total_numb_batch( - training_data.index, - sampler_weights, - ) - # Sampler weights carry tiny per-rank floating-point noise, so - # the rounded batch count can differ by one unit across ranks. - # Pin it to rank 0 before deriving num_steps so every rank - # shares the same training and full-validation schedule. - total_numb_batch = self._broadcast_value_from_rank0(total_numb_batch) - if total_numb_batch <= 0: - raise ValueError( - "Total number of training batches must be positive." - ) - self.num_steps = int(np.ceil(self.num_epoch * total_numb_batch)) - log.info( - "Computed num_steps=%d from num_epoch=%s and total_numb_batch=%d.", - self.num_steps, - self.num_epoch, - total_numb_batch, - ) - else: - if self.num_epoch_dict: - if self.num_steps is not None: - raise ValueError( - "training.numb_steps and training.num_epoch_dict " - "are mutually exclusive." - ) - for model_key in self.model_keys: - if isinstance(training_data[model_key], LmdbDataset): - per_task_total.append(len(self.training_dataloader[model_key])) - else: - sampler_weights = to_numpy_array( - self.training_dataloader[model_key].sampler.weights - ) - per_task_total.append( - compute_total_numb_batch( - training_data[model_key].index, - sampler_weights, - ) - ) - per_task_total = self._broadcast_value_from_rank0(per_task_total) - ( - self.model_prob, - self.num_steps, - per_task_steps, - ) = resolve_model_prob_from_epochs( - self.model_keys, - self.num_epoch_dict, - np.asarray(per_task_total, dtype=np.float64), - ) - log.info( - "Computed model_prob=%s and num_steps=%d from num_epoch_dict=%s " - "with per-task target steps: %s.", - self.model_prob, - self.num_steps, - self.num_epoch_dict, - {k: int(np.ceil(v)) for k, v in per_task_steps.items()}, - ) - else: - if self.num_steps is None: - raise ValueError( - "Either training.numb_steps (multi-task only) or " - "training.num_epoch_dict must be set." - ) - self.model_prob = resolve_model_prob( - self.model_keys, - training_params.get("model_prob"), - training_data, - rank=self.rank, - ) + def epoch_length(model_key: str) -> int: + """Return the batches this rank consumes in one epoch of a task.""" + _data = training_data[model_key] if self.multi_task else training_data + _dataloader = ( + self.training_dataloader[model_key] + if self.multi_task + else self.training_dataloader + ) + if isinstance(_data, LmdbDataset): + return len(_dataloader) + return compute_total_numb_batch( + _data.index, + to_numpy_array(_dataloader.sampler.weights), + ) + + schedule = resolve_step_schedule( + training_params, + multi_task=self.multi_task, + model_keys=self.model_keys, + training_data=( + training_data + if self.multi_task + else {self.model_keys[0]: training_data} + ), + epoch_length=epoch_length, + broadcast=self._broadcast_value_from_rank0, + rank=self.rank, + ) + self.num_steps = schedule.num_steps + self.model_prob = schedule.model_prob # === Derive checkpoint retention from ckpt_keep_ratio === # num_steps is final here (including when derived from num_epoch), so the @@ -1462,6 +1398,14 @@ def _load_optimizer_state( self.optimizer.load_state_dict(optimizer_state_dict) def run(self) -> None: + """Run training and release asynchronous data pipelines.""" + try: + self._run() + finally: + self._close_lmdb_loaders() + + def _run(self) -> None: + """Execute the PyTorch optimization loop.""" fout = ( open( self.disp_file, @@ -2091,6 +2035,24 @@ def log_loss_valid(_task_key: str = "Default") -> dict: f"The profiling trace has been saved to: {self.profiling_file}" ) + def _close_lmdb_loaders(self) -> None: + """Release LMDB pipelines owned by training and validation loaders.""" + closed: set[int] = set() + datasets: dict[int, LmdbDataset] = {} + for loaders in (self.training_dataloader, self.validation_dataloader): + values = loaders.values() if isinstance(loaders, dict) else (loaders,) + for loader in values: + if loader is None or id(loader) in closed: + continue + closed.add(id(loader)) + close = getattr(loader, "close", None) + if close is not None: + close() + if isinstance(loader, LmdbBatchDataLoader): + datasets[id(loader.dataset)] = loader.dataset + for dataset in datasets.values(): + dataset.close() + def _collect_checkpoint_states( self, *, diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py index a67a0a50e3..3fed8282dd 100644 --- a/deepmd/pt/utils/lmdb_dataset.py +++ b/deepmd/pt/utils/lmdb_dataset.py @@ -17,6 +17,7 @@ ) from deepmd.dpmodel.utils.lmdb_data import ( + LmdbBatchIterator, LmdbDataReader, LmdbTestData, SameNlocBatchSampler, @@ -24,6 +25,9 @@ compute_block_targets, is_lmdb, ) +from deepmd.env import ( + get_lmdb_num_workers, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -32,6 +36,7 @@ # Re-export for backward compatibility __all__ = [ + "LmdbBatchDataLoader", "LmdbDataset", "LmdbTestData", "_collate_lmdb_batch", @@ -76,6 +81,25 @@ def _collate_lmdb_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: return collate_lmdb_frames(torch_frames) +def _lmdb_batch_to_torch( + batch: dict[str, Any], + *, + pin_memory: bool, +) -> dict[str, Any]: + """Convert a contiguous NumPy LMDB batch to CPU tensors.""" + converted: dict[str, Any] = {} + with torch.device("cpu"): + for key, value in batch.items(): + if key.startswith("find_") or key == "fid" or key == "type": + converted[key] = value + elif value is None: + converted[key] = None + else: + tensor = torch.as_tensor(value) + converted[key] = tensor.pin_memory() if pin_memory else tensor + return converted + + class _SameNlocBatchSamplerTorch(Sampler): """Torch Sampler adapter around the framework-agnostic SameNlocBatchSampler. @@ -99,6 +123,53 @@ def set_epoch(self, epoch: int) -> None: self._inner.set_epoch(epoch) +class LmdbBatchDataLoader: + """DataLoader-compatible iterable backed by :class:`LmdbBatchIterator`. + + The parent sampler determines batch order. The shared LMDB process pool + decodes one batch and prefetches its successor, then this adapter converts + the contiguous NumPy result to pinned CPU tensors. + """ + + def __init__( + self, + dataset: "LmdbDataset", + sampler: Any, + *, + pin_memory: bool, + num_workers: int | None = None, + ) -> None: + self.dataset = dataset + self.batch_sampler = _SameNlocBatchSamplerTorch(sampler) + self.sampler = sampler + self._pin_memory = pin_memory + self._batch_iterator = LmdbBatchIterator( + dataset._reader, + sampler, + get_lmdb_num_workers() if num_workers is None else num_workers, + ) + + def __iter__(self) -> Iterator[dict[str, Any]]: + for _ in range(len(self)): + yield _lmdb_batch_to_torch( + next(self._batch_iterator), + pin_memory=self._pin_memory, + ) + + def __len__(self) -> int: + return len(self.batch_sampler) + + def close(self) -> None: + """Release this loader's prefetched batch and shared-pool reference.""" + self._batch_iterator.close() + + def __del__(self) -> None: + """Release the shared-pool reference during interpreter teardown.""" + iterator = getattr(self, "_batch_iterator", None) + if iterator is not None: + iterator.close() + + class LmdbDataset(Dataset): """PyTorch Dataset backed by LMDB via LmdbDataReader. @@ -248,6 +319,16 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N self._reader.add_data_requirement(data_requirement) self._rebuild_nloc_dataloaders() + def close(self) -> None: + """Release parent-process LMDB resources.""" + self._reader.close() + + def __del__(self) -> None: + """Release parent-process LMDB resources during teardown.""" + reader = getattr(self, "_reader", None) + if reader is not None: + reader.close() + def preload_and_modify_all_data_torch(self) -> None: """No-op: LMDB reads on demand.""" diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 06fd41af0c..66ddd1de79 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -130,12 +130,16 @@ def _build_data_system( dataset_params: dict[str, Any], type_map: list[str], seed: int | None = None, + rank: int = 0, + world_size: int = 1, ) -> DeepmdDataSystem | LmdbDataSystem: """Build a data system from dataset config, routing LMDB paths to LmdbDataSystem. A scalar ``systems`` value pointing at an LMDB directory triggers the LMDB adapter; otherwise we fall through to the legacy - :class:`DeepmdDataSystem` path with system expansion. + :class:`DeepmdDataSystem` path with system expansion. ``rank`` and + ``world_size`` shard LMDB training batches without changing legacy data + systems. """ systems_raw = dataset_params["systems"] lmdb_path = _detect_lmdb_path(systems_raw) @@ -146,6 +150,8 @@ def _build_data_system( batch_size=dataset_params["batch_size"], auto_prob_style=dataset_params.get("auto_prob"), seed=seed, + rank=rank, + world_size=world_size, ) systems = process_systems( systems_raw, @@ -171,17 +177,26 @@ def get_trainer( shared_links: dict | None = None, ) -> training.Trainer: """Build a :class:`training.Trainer` from a normalised config.""" + import torch.distributed as dist + training_params = config["training"] multi_task = "model_dict" in config["model"] data_seed = training_params.get("seed", None) + is_distributed = dist.is_available() and dist.is_initialized() + rank = dist.get_rank() if is_distributed else 0 + world_size = dist.get_world_size() if is_distributed else 1 def factory( task_config: TrainingTaskConfig, ) -> tuple[DeepmdDataSystem | LmdbDataSystem, Any | None, StatFileSpec]: type_map = list(task_config.model_params["type_map"]) train_data = _build_data_system( - dict(task_config.training_data_params), type_map, seed=data_seed + dict(task_config.training_data_params), + type_map, + seed=data_seed, + rank=rank, + world_size=world_size, ) validation_data = None if task_config.validation_data_params is not None: diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 1129696ed1..195f1f3e7e 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -38,6 +38,7 @@ TrainStepResult, change_model_out_bias, change_model_out_bias_by_task, + resolve_step_schedule, ) from deepmd.dpmodel.utils.batch import ( normalize_batch, @@ -46,6 +47,9 @@ from deepmd.dpmodel.utils.learning_rate import ( make_learning_rate_schedule, ) +from deepmd.dpmodel.utils.training_utils import ( + compute_total_numb_batch, +) from deepmd.pt.train.utils import ( resolve_best_checkpoint_dir, ) @@ -1396,7 +1400,6 @@ def __init__( self.world_size = dist.get_world_size() if self.is_distributed else 1 # Iteration config - self.num_steps = training_params["numb_steps"] self.disp_file = training_params.get("disp_file", "lcurve.out") self.disp_freq = training_params.get("disp_freq", 1000) self.save_ckpt = training_params.get("save_ckpt", "model.ckpt") @@ -1520,19 +1523,18 @@ def initialize_statistics( operation=f"statistics initialization for task {model_key!r}", ) - # Model probability (multi-task) -------------------------------------- - if self.multi_task: - from deepmd.dpmodel.utils.training_utils import ( - resolve_model_prob, - ) - - self.model_prob = resolve_model_prob( - self.model_keys, - training_params.get("model_prob"), - self.training_data_by_task, - ) - else: - self.model_prob = None + # Training schedule --------------------------------------------------- + schedule = resolve_step_schedule( + training_params, + multi_task=self.multi_task, + model_keys=self.model_keys, + training_data=self.training_data_by_task, + epoch_length=self._epoch_length, + broadcast=self._broadcast_value_from_rank0, + rank=self.rank, + ) + self.num_steps = schedule.num_steps + self.model_prob = schedule.model_prob # Learning rate ------------------------------------------------------- self.lr_schedule = make_learning_rate_schedule( @@ -2125,6 +2127,33 @@ def get_data( return input_dict, label_dict + def _epoch_length(self, model_key: str) -> int: + """Return the steps this rank takes during one epoch of a task. + + Parameters + ---------- + model_key : str + Key of the task whose training data is measured. + + Returns + ------- + int + Number of steps covering one pass over the task's training data. + + Notes + ----- + A data system reports ``nbatches[i]``, the batch count of system ``i``, + and ``sys_probs[i]``, the probability of drawing from that system, from + which ``compute_total_numb_batch`` derives the dataset-wide epoch + length ``ceil(max_i(nbatches[i] / sys_probs[i]))``. LMDB data reports + that global count while its sampler shards batches evenly across ranks; + legacy data systems remain replicated. In both cases one rank takes + ``ceil(total / world_size)`` steps per epoch. + """ + data = self.training_data_by_task[model_key] + total = compute_total_numb_batch(data.nbatches, data.sys_probs) + return int(np.ceil(total / self.world_size)) + # ------------------------------------------------------------------ # DDP helpers # ------------------------------------------------------------------ @@ -2168,6 +2197,23 @@ def _broadcast_model_stat(model: torch.nn.Module) -> None: for b in model.buffers(): dist.broadcast(b, src=0) + def _broadcast_value_from_rank0(self, value: Any) -> Any: + """Return rank 0's copy of ``value`` on every rank. + + Epoch lengths round a quotient of sampling probabilities that is often + an exact integer in real arithmetic, so a last-bit difference between + ranks -- as reduction kernels dispatched for different CPU features + produce -- flips the rounded result and hence ``num_steps``. Ranks that + disagree on ``num_steps`` also disagree on the full-validation start + step and deadlock on mismatched collective calls, so the whole world + adopts rank 0's value. + """ + if not self.is_distributed: + return value + holder = [value] + dist.broadcast_object_list(holder, src=0, device=DEVICE) + return holder[0] + # ------------------------------------------------------------------ # Checkpointing # ------------------------------------------------------------------ @@ -2266,13 +2312,31 @@ def run(self) -> None: """Run pt_expt training through the backend-independent trainer loop.""" log.info("Start to train %d steps.", self.num_steps) wall_start = time.time() - super().run(self.training_tasks) - if self.change_bias_after_training and self.num_steps > self.start_step: - self._change_bias_after_training() - if self.rank_context.is_chief: - self.save_checkpoint(self.num_steps) + try: + super().run(self.training_tasks) + if self.change_bias_after_training and self.num_steps > self.start_step: + self._change_bias_after_training() + if self.rank_context.is_chief: + self.save_checkpoint(self.num_steps) + finally: + self._close_data_systems() log.info("Training finished. Total wall time: %.2fs", time.time() - wall_start) + def _close_data_systems(self) -> None: + """Release asynchronous data pipelines owned by this trainer.""" + closed: set[int] = set() + for data_by_task in ( + self.training_data_by_task, + self.validation_data_by_task, + ): + for data_system in data_by_task.values(): + if data_system is None or id(data_system) in closed: + continue + closed.add(id(data_system)) + close = getattr(data_system, "close", None) + if close is not None: + close() + def _change_bias_after_training(self) -> None: if self.rank == 0: change_model_out_bias_by_task( diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index c4ea43168d..7608d9eb61 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -14,11 +14,16 @@ ) from deepmd.dpmodel.utils.lmdb_data import ( + DistributedSameNlocBatchSampler, + LmdbBatchIterator, LmdbDataReader, SameNlocBatchSampler, - collate_lmdb_frames, + collect_lmdb_sampling_groups, compute_block_targets, ) +from deepmd.env import ( + get_lmdb_num_workers, +) from deepmd.utils.data import ( DataRequirementItem, ) @@ -30,12 +35,15 @@ class LmdbDataSystem: """LMDB-backed data system for pt_expt. Exposes the small surface that pt_expt's trainer touches: - ``get_batch(sys_idx=None)``, ``add_data_requirements(list)``, and - ``get_nsystems()``. Internally uses :class:`LmdbDataReader` for I/O and - :class:`SameNlocBatchSampler` to draw same-nloc batches. Statistics use a - separate logical-system view in which every ``nloc`` group is sampled - independently, matching the PyTorch DataLoader adapter without changing - the identity of the LMDB as one training dataset. + ``get_batch(sys_idx=None)``, ``add_data_requirements(list)``, + ``get_nsystems()``, and the ``nbatches``/``sys_probs`` pair from which the + trainer derives an epoch length. The whole LMDB counts as one logical + system. Internally uses :class:`LmdbDataReader` for I/O and + :class:`SameNlocBatchSampler`, or its distributed wrapper, to draw + same-nloc batches. Statistics use a separate logical-system view in which + every ``(nloc, label-availability)`` group is sampled independently, + matching the training sampler without changing the identity of the LMDB as + one training dataset. Parameters ---------- @@ -50,6 +58,14 @@ class LmdbDataSystem: per-system reweighting via :func:`compute_block_targets`. seed Optional seed for the shuffle in :class:`SameNlocBatchSampler`. + num_workers + Number of LMDB decoder worker processes. ``None`` selects the + hardware-aware default; zero or one disables multiprocessing. + rank + Rank of this process in distributed training. + world_size + Number of distributed training processes. Values greater than one + select :class:`DistributedSameNlocBatchSampler`. """ def __init__( @@ -59,6 +75,9 @@ def __init__( batch_size: int | str = "auto", auto_prob_style: str | None = None, seed: int | None = None, + num_workers: int | None = None, + rank: int = 0, + world_size: int = 1, ) -> None: self._reader = LmdbDataReader( lmdb_path, type_map, batch_size, mixed_batch=False @@ -72,15 +91,38 @@ def __init__( self._reader.system_nframes, ) - self._sampler = SameNlocBatchSampler( + if world_size > 1: + distributed_sampler = DistributedSameNlocBatchSampler( + self._reader, + rank=rank, + world_size=world_size, + shuffle=True, + seed=seed, + block_targets=block_targets, + ) + self._sampler = distributed_sampler + else: + sampler = SameNlocBatchSampler( + self._reader, + shuffle=True, + seed=seed, + block_targets=block_targets, + ) + self._sampler = sampler + self._refresh_stat_groups() + num_workers = ( + get_lmdb_num_workers() if num_workers is None else int(num_workers) + ) + self._batch_iterator = LmdbBatchIterator( self._reader, - shuffle=True, - seed=seed, - block_targets=block_targets, + self._sampler, + num_workers, ) - self._iter = iter(self._sampler) - self._stat_nlocs = tuple(sorted(self._reader.nloc_groups)) - self._stat_offsets = [0] * len(self._stat_nlocs) + + def _refresh_stat_groups(self) -> None: + """Rebuild statistical systems from the training sampler's groups.""" + self._stat_groups = collect_lmdb_sampling_groups(self._reader) + self._stat_offsets = [0] * len(self._stat_groups) # ------------------------------------------------------------------ # pt_expt trainer surface @@ -93,20 +135,15 @@ def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: sampling is baked into ``block_targets`` at sampler construction. """ del sys_idx - try: - indices = next(self._iter) - except StopIteration: - self._iter = iter(self._sampler) - indices = next(self._iter) - return self._collate_indices(indices) + return next(self._batch_iterator) def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: - """Return one batch from a fixed-``nloc`` statistical system. + """Return one batch from a homogeneous statistical system. Parameters ---------- sys_idx : int - Index into the sorted ``nloc`` groups. + Index into the ``(nloc, label-availability)`` groups. Returns ------- @@ -116,57 +153,79 @@ def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: Raises ------ IndexError - If ``sys_idx`` does not identify an available ``nloc`` group. + If ``sys_idx`` does not identify an available statistical group. """ - if not 0 <= sys_idx < len(self._stat_nlocs): + if not 0 <= sys_idx < len(self._stat_groups): raise IndexError( f"Statistical system index {sys_idx} is out of range for " - f"{len(self._stat_nlocs)} nloc groups." + f"{len(self._stat_groups)} homogeneous groups." ) - nloc = self._stat_nlocs[sys_idx] - group_indices = self._reader.nloc_groups[nloc] + nloc, group_indices = self._stat_groups[sys_idx] batch_size = self._reader.get_batch_size_for_nloc(nloc) start = self._stat_offsets[sys_idx] if start >= len(group_indices): start = 0 stop = min(start + batch_size, len(group_indices)) self._stat_offsets[sys_idx] = stop - return self._collate_indices(group_indices[start:stop]) + return self._reader.decode_batch(group_indices[start:stop]) def get_stat_nsystems(self) -> int: - """Return the number of fixed-``nloc`` statistical systems.""" - return len(self._stat_nlocs) + """Return the number of homogeneous statistical systems.""" + return len(self._stat_groups) def get_stat_numb_batches(self, sys_idx: int) -> int: """Return the available batch count for one statistical system.""" - if not 0 <= sys_idx < len(self._stat_nlocs): + if not 0 <= sys_idx < len(self._stat_groups): raise IndexError( f"Statistical system index {sys_idx} is out of range for " - f"{len(self._stat_nlocs)} nloc groups." + f"{len(self._stat_groups)} homogeneous groups." ) - nloc = self._stat_nlocs[sys_idx] - nframes = len(self._reader.nloc_groups[nloc]) + nloc, group_indices = self._stat_groups[sys_idx] + nframes = len(group_indices) batch_size = self._reader.get_batch_size_for_nloc(nloc) return (nframes + batch_size - 1) // batch_size - def _collate_indices(self, indices: list[int]) -> dict[str, Any]: - """Load and collate the requested dataset indices.""" - frames = [self._reader[int(i)] for i in indices] - return collate_lmdb_frames(frames) - def add_data_requirements( self, data_requirement: list[DataRequirementItem] ) -> None: + # Batches are partitioned by label availability. The sampler derives + # the partition on its first draw; only the distributed batch count is + # cached, so it is refreshed after the requirements change. self._reader.add_data_requirement(data_requirement) - # Discard any iterator created under the previous availability - # signature so the next batch uses the newly registered labels. - self._iter = iter(self._sampler) + self._refresh_stat_groups() + if isinstance(self._sampler, DistributedSameNlocBatchSampler): + self._sampler.refresh_batch_count() + + def close(self) -> None: + """Cancel prefetched work and release decoder processes.""" + iterator = getattr(self, "_batch_iterator", None) + if iterator is not None: + iterator.close() + reader = getattr(self, "_reader", None) + if reader is not None: + reader.close() + + def __del__(self) -> None: + """Release worker processes during interpreter teardown.""" + self.close() def get_nsystems(self) -> int: """Return one logical LMDB training dataset.""" return 1 + @property + def nbatches(self) -> list[int]: + """Return the global batch count of one full pass.""" + if isinstance(self._sampler, DistributedSameNlocBatchSampler): + return [self._sampler.total_batches] + return [len(self._sampler)] + + @property + def sys_probs(self) -> list[float]: + """Return the sampling probability of each logical system.""" + return [1.0] + # ------------------------------------------------------------------ # Misc forwarders # ------------------------------------------------------------------ diff --git a/doc/env.md b/doc/env.md index 35c7177ebc..5b4c759e48 100644 --- a/doc/env.md +++ b/doc/env.md @@ -77,6 +77,28 @@ Default backend. See [PyTorch documentation](https://pytorch.org/docs/stable/data.html) for details. ::: +:::{envvar} DP_LMDB_NUM_WORKERS + +**Type**: non-negative integer + +**Default**: automatically selected from the process CPU affinity and the +number of local training ranks, with limits of 32 workers per rank and +approximately 64 workers per node + +Number of worker processes used to read, decode, and assemble one LMDB batch +in the PyTorch and PyTorch Exportable backends. Each process owns an +independent read-only LMDB transaction. The next batch is prefetched while the +current batch is consumed, and at most one batch is prefetched. Batches with +fewer frames than workers are decoded synchronously because process startup +and IPC cost more than their small decode workload. + +Set this variable to `0` or `1` to use synchronous decoding. For multi-GPU +training, this value applies to each rank. Independent jobs do not share their +worker pools, so reduce it when the aggregate reader count across concurrent +jobs would overload the storage service. The LMDB dataset must remain immutable +while any job is reading it because readers intentionally disable LMDB locking. +::: + ## C++ interface only These environment variables also apply to third-party programs using the C++ interface, such as [LAMMPS](./third-party/lammps-command.md). diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 366a5d8813..6b56ac82af 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -4,21 +4,44 @@ Pure dpmodel (NumPy/lmdb) tests — no PyTorch dependency. """ +import os +import signal +import subprocess +import sys import tempfile +import textwrap import unittest +from concurrent.futures import ( + Future, +) +from pathlib import ( + Path, +) +from types import ( + SimpleNamespace, +) +from unittest import ( + mock, +) import lmdb import msgpack import numpy as np +from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module from deepmd.dpmodel.utils.lmdb_data import ( + LmdbBatchIterator, LmdbDataReader, + LmdbDecodeConfig, LmdbTestData, LmdbTestDataNlocView, SameNlocBatchSampler, _expand_indices_by_blocks, + _merge_lmdb_chunks, _remap_atom_types, compute_block_targets, + decode_lmdb_batch, + decode_lmdb_frame, is_lmdb, make_neighbor_stat_data, ) @@ -348,6 +371,230 @@ def test_uniform_nloc_single_group(self): self.assertIn(6, reader.nloc_groups) self.assertEqual(len(reader.nloc_groups[6]), 10) + def test_batch_iterator_advances_epoch_before_prefetch(self): + """The prefetched successor uses the next epoch's sampler state.""" + reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2) + sampler = SameNlocBatchSampler(reader, shuffle=True, seed=7) + iterator = LmdbBatchIterator(reader, sampler, num_workers=2) + + expected_sampler = SameNlocBatchSampler(reader, shuffle=True, seed=7) + expected_sampler.set_epoch(1) + expected_second = [ + key for indices in expected_sampler for key in reader.original_keys(indices) + ] + + try: + first = [key for _ in range(len(sampler)) for key in next(iterator)["fid"]] + second = [key for _ in range(len(sampler)) for key in next(iterator)["fid"]] + finally: + iterator.close() + reader.close() + + self.assertNotEqual(first, second) + self.assertEqual(second, expected_second) + + def test_requirements_are_rejected_after_the_first_decode(self): + """Late requirements are refused so no partition can predate them. + + Batches are grouped by label availability, and both the sampler + partition and the worker decoders capture the requirements in force + when they start, so a later registration would silently disagree with + the batches already produced. + """ + requirement = [DataRequirementItem("custom", ndof=1, default=0.0)] + + for label, consume in ( + ("frame", lambda reader: reader[0]), + ("batch", lambda reader: reader.decode_batch([0, 1])), + ("worker config", lambda reader: reader.worker_decode_config()), + ): + with self.subTest(consumed=label): + reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2) + reader.add_data_requirement(requirement) + consume(reader) + with self.assertRaisesRegex(RuntimeError, "before reading any frame"): + reader.add_data_requirement(requirement) + reader.close() + + def test_close_preserves_other_reader(self): + """Closing one shared-path reader leaves the other transaction valid.""" + first = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2) + second = LmdbDataReader( + f"{self._lmdb_path}/.", + self._type_map, + batch_size=2, + ) + first.close() + self.assertTrue(first.closed) + self.assertEqual(second[0]["coord"].shape, (6, 3)) + second.close() + self.assertTrue(second.closed) + with self.assertRaisesRegex(RuntimeError, "closed LMDB reader"): + _ = second[0] + + def test_batch_dtype_and_field_order_are_chunk_independent(self): + """Batch promotion and schema matching do not depend on chunking.""" + path = _create_lmdb( + f"{self._tmpdir.name}/mixed_dtype.lmdb", + nframes=2, + natoms=6, + ) + frame0 = _make_frame(natoms=6, seed=0) + frame1 = _make_frame(natoms=6, seed=1) + frame0["custom"] = { + "type": "float32", + "shape": [1], + "data": np.array([1.25], dtype=np.float32).tobytes(), + } + frame1["custom"] = { + "type": "float64", + "shape": [1], + "data": np.array([2.5], dtype=np.float64).tobytes(), + } + frame1 = dict(reversed(tuple(frame1.items()))) + environment = lmdb.open(path, readonly=False, lock=False) + with environment.begin(write=True) as transaction: + transaction.put( + b"000000000000", + msgpack.packb(frame0, use_bin_type=True), + ) + transaction.put( + b"000000000001", + msgpack.packb(frame1, use_bin_type=True), + ) + environment.close() + + config = LmdbDecodeConfig( + ntypes=2, + natoms=6, + type_remap=None, + data_requirements={}, + ) + environment = lmdb.open(path, readonly=True, lock=False) + with environment.begin() as transaction: + serial = decode_lmdb_batch( + transaction, + [0, 1], + "012d", + config, + ) + chunked = _merge_lmdb_chunks( + [ + decode_lmdb_batch(transaction, [0], "012d", config), + decode_lmdb_batch(transaction, [1], "012d", config), + ] + ) + environment.close() + + self.assertEqual(serial["custom"].dtype, np.float64) + np.testing.assert_array_equal(serial["custom"], chunked["custom"]) + + def test_parallel_batch_consistency_guards(self): + """Parallel batch decoding validates schemas, shapes, and availability.""" + config = LmdbDecodeConfig( + ntypes=2, + natoms=6, + type_remap=None, + data_requirements={}, + ) + transaction = mock.Mock() + transaction.get.return_value = b"frame" + first_frame = { + "coord": np.zeros((6, 3)), + "find_energy": np.float32(1.0), + "fid": 0, + } + frame_cases = ( + ( + "fields", + {**first_frame, "energy": np.zeros(1), "fid": 1}, + "inconsistent fields", + ), + ( + "shape", + {**first_frame, "coord": np.zeros((7, 3)), "fid": 1}, + "changes shape within one batch", + ), + ) + for guard, second_frame, expected_error in frame_cases: + with ( + self.subTest(guard=f"frame {guard}"), + mock.patch.object( + lmdb_data_module, + "decode_lmdb_frame", + side_effect=(first_frame, second_frame), + ), + self.assertRaisesRegex(ValueError, expected_error), + ): + decode_lmdb_batch(transaction, [0, 1], "012d", config) + + first_chunk = { + "find_energy": np.float32(1.0), + "coord": np.zeros((1, 6, 3)), + "fid": [0], + "sid": np.array([0], dtype=np.int64), + } + chunk_cases = ( + ( + "fields", + {key: value for key, value in first_chunk.items() if key != "coord"}, + "inconsistent fields", + ), + ( + "availability", + {**first_chunk, "find_energy": np.float32(0.0), "fid": [1]}, + "availability changes across worker chunks", + ), + ) + for guard, second_chunk, expected_error in chunk_cases: + with ( + self.subTest(guard=f"chunk {guard}"), + self.assertRaisesRegex(ValueError, expected_error), + ): + _merge_lmdb_chunks([first_chunk, second_chunk]) + + def test_batch_rejects_mixed_label_availability(self): + """A scalar find flag cannot represent mixed availability in one batch.""" + path = _create_lmdb( + f"{self._tmpdir.name}/mixed_availability.lmdb", + nframes=2, + natoms=6, + ) + frame = _make_frame(natoms=6, seed=0) + frame["custom"] = { + "type": "float64", + "shape": [1], + "data": np.array([1.0], dtype=np.float64).tobytes(), + } + environment = lmdb.open(path, readonly=False, lock=False) + with environment.begin(write=True) as transaction: + transaction.put( + b"000000000000", + msgpack.packb(frame, use_bin_type=True), + ) + environment.close() + + requirement = DataRequirementItem("custom", ndof=1, default=0.0) + config = LmdbDecodeConfig( + ntypes=2, + natoms=6, + type_remap=None, + data_requirements={"custom": requirement}, + ) + environment = lmdb.open(path, readonly=True, lock=False) + with environment.begin() as transaction: + with self.assertRaisesRegex( + ValueError, + "availability changes within one batch", + ): + decode_lmdb_batch( + transaction, + [0, 1], + "012d", + config, + ) + environment.close() + def test_is_lmdb(self): self.assertTrue(is_lmdb(self._lmdb_path)) self.assertTrue(is_lmdb("something.lmdb")) @@ -386,6 +633,31 @@ def test_min_pair_dist_requirement_computed(self): self.assertEqual(frame["find_min_pair_dist"], np.float32(1.0)) np.testing.assert_allclose(frame["min_pair_dist"], np.array([1.0])) + def test_min_pair_dist_requirement_defaults_without_atype(self): + raw_frame = _make_frame(natoms=6, seed=0) + raw_frame.pop("atom_types") + requirement = DataRequirementItem( + "min_pair_dist", + ndof=1, + default=0.25, + ) + config = LmdbDecodeConfig( + ntypes=2, + natoms=6, + type_remap=None, + data_requirements={"min_pair_dist": requirement}, + ) + + frame = decode_lmdb_frame( + msgpack.packb(raw_frame, use_bin_type=True), + 0, + config, + copy_arrays=True, + ) + + self.assertEqual(frame["find_min_pair_dist"], np.float32(0.0)) + np.testing.assert_allclose(frame["min_pair_dist"], np.array([0.25])) + # ============================================================ # Mixed nloc tests @@ -1362,5 +1634,237 @@ def test_testdata_missing_key_not_found(self): tmpdir.cleanup() +class _StalledPool: + """A pool whose decoder exited, leaving its submissions unfinished. + + This is what a decoder killed mid-result looks like from the parent: the + futures never resolve and the pool never reports itself broken. + """ + + def __init__(self) -> None: + self._processes = {1: SimpleNamespace(exitcode=-1)} + self.submissions = 0 + + def submit(self, *args: object, **kwargs: object) -> Future: + self.submissions += 1 + return Future() + + +class TestDecoderPoolFailure(unittest.TestCase): + """A dead decoder must not strand the run waiting for it.""" + + def setUp(self) -> None: + self._tmpdir = tempfile.TemporaryDirectory() + self.addCleanup(self._tmpdir.cleanup) + self._path = _create_lmdb( + f"{self._tmpdir.name}/pool.lmdb", nframes=12, natoms=6 + ) + self._reader = LmdbDataReader(self._path, ["O", "H"], batch_size=4) + self.addCleanup(self._reader.close) + # Keep the liveness check from pacing the test. + patcher = mock.patch.object( + lmdb_data_module, "_DECODER_LIVENESS_INTERVAL", 0.01 + ) + patcher.start() + self.addCleanup(patcher.stop) + self._entries: dict[int, object] = {} + + def _iterator(self, pool: object) -> LmdbBatchIterator: + """Return an iterator over two four-frame batches served by ``pool``. + + Every iterator built on one stand-in pool receives the same entry, as + the iterators of a rank share the pool registered for their worker + count. + """ + entry = self._entries.setdefault( + id(pool), lmdb_data_module._LmdbPoolEntry(executor=pool, users=0) + ) + entry.users += 1 + patcher = mock.patch.object( + lmdb_data_module, "_acquire_lmdb_executor", return_value=entry + ) + patcher.start() + self.addCleanup(patcher.stop) + # The stand-in pool was never registered, so releasing it is a no-op. + return LmdbBatchIterator( + self._reader, [[0, 1, 2, 3], [4, 5, 6, 7]], num_workers=2 + ) + + def _isolated_iterator(self) -> LmdbBatchIterator: + """Return an iterator over a real pool that no other test shares. + + The decoder pool is process-wide, so a test that kills or signals its + decoders is given one of its own rather than leaving the damage behind + for its neighbours. + """ + registry: dict = {} + patcher = mock.patch.object(lmdb_data_module, "_LMDB_POOLS", registry) + patcher.start() + self.addCleanup(patcher.stop) + self.addCleanup( + lambda: [ + entry.executor.shutdown(wait=False, cancel_futures=True) + for entry in registry.values() + ] + ) + iterator = LmdbBatchIterator( + self._reader, [[0, 1, 2, 3], [4, 5, 6, 7]], num_workers=2 + ) + self.addCleanup(iterator.close) + return iterator + + def _assert_same_batch(self, batch: dict, expected: dict) -> None: + self.assertEqual(sorted(batch), sorted(expected)) + for key, value in expected.items(): + if isinstance(value, np.ndarray): + np.testing.assert_array_equal(batch[key], value) + + def test_a_stalled_decoder_falls_back_and_is_not_retried(self) -> None: + pool = _StalledPool() + iterator = self._iterator(pool) + + with self.assertLogs(lmdb_data_module.log, level="WARNING") as captured: + first = next(iterator) + submissions = pool.submissions + second = next(iterator) + + # Each batch is the one the pool was asked for, decoded here instead, + # and the pool is not offered any more work. + self._assert_same_batch(first, self._reader.decode_batch([0, 1, 2, 3])) + self._assert_same_batch(second, self._reader.decode_batch([4, 5, 6, 7])) + self.assertIn("decoder process exited", "\n".join(captured.output)) + self.assertEqual(pool.submissions, submissions) + + def test_close_detects_a_stalled_prefetch(self) -> None: + """Closing an in-flight prefetch marks a lost decoder pool unhealthy.""" + pool = _StalledPool() + iterator = self._iterator(pool) + entry = self._entries[id(pool)] + iterator._pool = entry + running = Future() + self.assertTrue(running.set_running_or_notify_cancel()) + iterator._pending = lmdb_data_module._PendingBatch([0, 1], [running]) + + with self.assertLogs(lmdb_data_module.log, level="WARNING"): + iterator.close() + + self.assertFalse(entry.healthy) + + def test_a_second_iterator_does_not_retry_a_lost_pool(self) -> None: + """Pool health is shared, so the loss is discovered once for all.""" + pool = _StalledPool() + first = self._iterator(pool) + next(first) + submissions = pool.submissions + + second = self._iterator(pool) + batch = next(second) + + self._assert_same_batch(batch, self._reader.decode_batch([0, 1, 2, 3])) + self.assertEqual(pool.submissions, submissions) + + @unittest.skipUnless( + os.name == "posix" and sys.implementation.name == "cpython", + "requires CPython's POSIX process-pool pipe", + ) + def test_a_run_that_lost_its_pool_still_exits(self) -> None: + """Losing a decoder must not leave the interpreter unable to exit. + + A pool reading the partial result of a decoder killed mid-write is + joined by the interpreter on the way out, so a run could complete its + training and then hang forever instead of terminating. + """ + # Spawned decoders re-import the main module, so the scenario has to + # live in a file rather than be passed on the command line. + script = Path(self._tmpdir.name) / "lose_the_pool.py" + script.write_text( + textwrap.dedent(""" + import os + import struct + + from deepmd.dpmodel.utils.lmdb_data import ( + _acquire_lmdb_executor, + _release_lmdb_executor, + ) + + + def idle(): + return os.getpid() + + + if __name__ == "__main__": + entry = _acquire_lmdb_executor(2) + entry.executor.submit(idle).result() + # A frame header promising more bytes than ever arrive, which + # is what a decoder killed mid-write leaves behind. + os.write( + entry.executor._result_queue._writer.fileno(), + struct.pack("!i", 4096) + b"partial", + ) + entry.healthy = False + _release_lmdb_executor(2) + """) + ) + completed = subprocess.run( + [sys.executable, str(script)], + capture_output=True, + text=True, + timeout=90, + ) + self.assertEqual(completed.returncode, 0, completed.stderr[-2000:]) + + @unittest.skipUnless( + hasattr(signal, "SIGHUP"), + "SIGHUP is not available on this platform", + ) + def test_decoders_survive_the_hangup_of_their_launching_session(self) -> None: + """A decoder outlives the session that started the run. + + The hangup delivered when that session goes away reaches every + background helper of the run. A decoder that died of it would break + the pool of an otherwise healthy run. + """ + iterator = self._isolated_iterator() + first = next(iterator) + processes = list(iterator._pool.executor._processes.values()) + self.assertTrue(processes) + + for process in processes: + os.kill(process.pid, signal.SIGHUP) + + # Delivering the next batch is what proves the decoders lived through + # it: a pool that lost one degrades instead. That is both a stronger + # statement than reading their liveness and free of any timing. + self._assert_same_batch(first, self._reader.decode_batch([0, 1, 2, 3])) + self._assert_same_batch(next(iterator), self._reader.decode_batch([4, 5, 6, 7])) + self.assertTrue(iterator._pool.healthy) + self.assertTrue(all(process.is_alive() for process in processes)) + + @unittest.skipUnless( + hasattr(signal, "SIGKILL"), + "SIGKILL is not available on this platform", + ) + def test_killing_a_real_decoder_does_not_stop_the_run(self) -> None: + """The pool reports itself broken, and the batch still arrives. + + A decoder that dies cleanly fails every future the pool holds, which + is the other way the loss of a decoder reaches the iterator. Only a + signal it cannot ignore gets it there. + """ + iterator = self._isolated_iterator() + next(iterator) + processes = list(iterator._pool.executor._processes.values()) + self.assertTrue(processes) + process = processes[0] + self.assertTrue(process.is_alive()) + os.kill(process.pid, signal.SIGKILL) + process.join(timeout=5) + + batch = next(iterator) + + self._assert_same_batch(batch, self._reader.decode_batch([4, 5, 6, 7])) + self.assertFalse(iterator._pool.healthy) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/common/dpmodel/test_train_schedule.py b/source/tests/common/dpmodel/test_train_schedule.py new file mode 100644 index 0000000000..dd85991f2c --- /dev/null +++ b/source/tests/common/dpmodel/test_train_schedule.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the backend-independent training-step schedule.""" + +import numpy as np +import pytest + +from deepmd.dpmodel.train import ( + resolve_step_schedule, +) + + +class FakeData: + """Training-data stand-in reporting a system count.""" + + def __init__(self, nsystems: int) -> None: + self._nsystems = nsystems + + def get_nsystems(self) -> int: + return self._nsystems + + +def _unreachable_epoch_length(model_key: str) -> int: + raise AssertionError(f"epoch length of '{model_key}' must not be needed") + + +def test_explicit_steps_bypass_the_data() -> None: + schedule = resolve_step_schedule( + {"numb_steps": 7}, + multi_task=False, + model_keys=["Default"], + training_data={"Default": FakeData(1)}, + epoch_length=_unreachable_epoch_length, + ) + + assert schedule.num_steps == 7 + assert schedule.model_prob is None + + +def test_single_task_rejects_steps_and_epochs() -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + resolve_step_schedule( + {"numb_steps": 7, "numb_epoch": 1.0}, + multi_task=False, + model_keys=["Default"], + training_data={"Default": FakeData(1)}, + epoch_length=_unreachable_epoch_length, + ) + + +@pytest.mark.parametrize( + ("num_epoch", "expected"), + [(1.0, 40), (2.5, 100), (0.25, 10), (1.0 / 3.0, 14)], +) +def test_single_task_epochs_round_up(num_epoch: float, expected: int) -> None: + schedule = resolve_step_schedule( + {"numb_epoch": num_epoch}, + multi_task=False, + model_keys=["Default"], + training_data={"Default": FakeData(1)}, + epoch_length=lambda _: 40, + ) + + assert schedule.num_steps == expected + assert schedule.model_prob is None + + +def test_single_task_requires_a_run_length() -> None: + with pytest.raises(ValueError, match=r"numb_steps or training\.num_epoch"): + resolve_step_schedule( + {}, + multi_task=False, + model_keys=["Default"], + training_data={"Default": FakeData(1)}, + epoch_length=_unreachable_epoch_length, + ) + + +def test_single_task_rejects_non_positive_epochs() -> None: + with pytest.raises(ValueError, match="num_epoch must be positive"): + resolve_step_schedule( + {"numb_epoch": 0.0}, + multi_task=False, + model_keys=["Default"], + training_data={"Default": FakeData(1)}, + epoch_length=_unreachable_epoch_length, + ) + + +def test_empty_training_data_is_rejected() -> None: + with pytest.raises(ValueError, match="positive for task 'Default'"): + resolve_step_schedule( + {"numb_epoch": 1.0}, + multi_task=False, + model_keys=["Default"], + training_data={"Default": FakeData(1)}, + epoch_length=lambda _: 0, + ) + + +def test_epoch_lengths_are_pinned_to_rank_zero() -> None: + """A rank derives its run length from the broadcast value, not its own.""" + schedule = resolve_step_schedule( + {"numb_epoch": 2.0}, + multi_task=False, + model_keys=["Default"], + training_data={"Default": FakeData(1)}, + epoch_length=lambda _: 41, + broadcast=lambda lengths: [40] * len(lengths), + ) + + assert schedule.num_steps == 80 + + +def test_multi_task_epoch_lengths_are_pinned_to_rank_zero() -> None: + """Every task's epoch length comes from the broadcast, not from this rank.""" + schedule = resolve_step_schedule( + {"num_epoch_dict": {"model_1": 1.0, "model_2": 4.0}}, + multi_task=True, + model_keys=["model_1", "model_2"], + training_data={"model_1": FakeData(1), "model_2": FakeData(1)}, + epoch_length=lambda _: 11, + broadcast=lambda lengths: [40, 10][: len(lengths)], + ) + + assert schedule.num_steps == 80 + np.testing.assert_allclose(schedule.model_prob, [0.5, 0.5]) + + +def test_multi_task_epoch_dict_splits_steps_by_epoch_target() -> None: + """Each task receives the steps its epoch target asks for.""" + epoch_lengths = {"model_1": 40, "model_2": 10} + num_epoch_dict = {"model_1": 1.0, "model_2": 4.0} + + schedule = resolve_step_schedule( + {"num_epoch_dict": num_epoch_dict}, + multi_task=True, + model_keys=["model_1", "model_2"], + training_data={key: FakeData(1) for key in epoch_lengths}, + epoch_length=epoch_lengths.__getitem__, + ) + + assert schedule.num_steps == 80 + np.testing.assert_allclose(schedule.model_prob, [0.5, 0.5]) + for index, model_key in enumerate(["model_1", "model_2"]): + expected_epochs = num_epoch_dict[model_key] + drawn_steps = schedule.num_steps * schedule.model_prob[index] + assert drawn_steps / epoch_lengths[model_key] == pytest.approx(expected_epochs) + + +def test_multi_task_epoch_dict_rejects_explicit_steps() -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + resolve_step_schedule( + {"numb_steps": 10, "num_epoch_dict": {"model_1": 1.0}}, + multi_task=True, + model_keys=["model_1"], + training_data={"model_1": FakeData(1)}, + epoch_length=lambda _: 40, + ) + + +def test_multi_task_steps_keep_configured_model_prob() -> None: + schedule = resolve_step_schedule( + {"numb_steps": 12, "model_prob": {"model_1": 3.0, "model_2": 1.0}}, + multi_task=True, + model_keys=["model_1", "model_2"], + training_data={"model_1": FakeData(1), "model_2": FakeData(1)}, + epoch_length=_unreachable_epoch_length, + ) + + assert schedule.num_steps == 12 + np.testing.assert_allclose(schedule.model_prob, [0.75, 0.25]) + + +def test_multi_task_requires_a_run_length() -> None: + with pytest.raises(ValueError, match="num_epoch_dict must be set"): + resolve_step_schedule( + {}, + multi_task=True, + model_keys=["model_1"], + training_data={"model_1": FakeData(1)}, + epoch_length=_unreachable_epoch_length, + ) diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py index ae26566525..d945ccc008 100644 --- a/source/tests/pt/test_lmdb_dataloader.py +++ b/source/tests/pt/test_lmdb_dataloader.py @@ -16,6 +16,7 @@ lmdb_data, ) from deepmd.dpmodel.utils.lmdb_data import ( + _ENV_CACHE, DistributedSameNlocBatchSampler, LmdbDataReader, SameNlocBatchSampler, @@ -28,6 +29,7 @@ EnergyStdLoss, ) from deepmd.pt.utils.lmdb_dataset import ( + LmdbBatchDataLoader, LmdbDataset, _collate_lmdb_batch, ) @@ -364,6 +366,101 @@ def test_inner_dataloader(self, lmdb_dir): batch = next(iter(ds.dataloaders[0])) assert batch["coord"].shape[0] == 2 + def test_parallel_batch_loader_has_finite_epoch(self, lmdb_dir): + ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) + sampler = SameNlocBatchSampler(ds._reader, shuffle=False) + loader = LmdbBatchDataLoader( + ds, + sampler, + pin_memory=False, + num_workers=2, + ) + try: + assert sum(batch["coord"].shape[0] for batch in loader) == 10 + finally: + loader.close() + + def test_parallel_loaders_share_pool_for_same_dataset(self, lmdb_dir): + first_data = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) + second_data = LmdbDataset( + f"{lmdb_dir}/.", + type_map=["O", "H"], + batch_size=2, + ) + first = LmdbBatchDataLoader( + first_data, + SameNlocBatchSampler(first_data._reader, shuffle=True, seed=1), + pin_memory=False, + num_workers=2, + ) + second = LmdbBatchDataLoader( + second_data, + SameNlocBatchSampler(second_data._reader, shuffle=True, seed=2), + pin_memory=False, + num_workers=2, + ) + first_iterator = iter(first) + second_iterator = iter(second) + try: + next(first_iterator) + next(second_iterator) + assert first._batch_iterator._pool is second._batch_iterator._pool + first.close() + assert next(second_iterator)["coord"].shape == (2, 6, 3) + finally: + first.close() + second.close() + + def test_small_batch_stays_synchronous(self, lmdb_dir): + ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) + loader = LmdbBatchDataLoader( + ds, + SameNlocBatchSampler(ds._reader, shuffle=False), + pin_memory=False, + num_workers=4, + ) + try: + assert next(iter(loader))["coord"].shape == (2, 6, 3) + assert not loader._batch_iterator.started + assert loader._batch_iterator._pending is None + finally: + loader.close() + + def test_partial_successor_is_deferred(self, lmdb_dir): + ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=4) + loader = LmdbBatchDataLoader( + ds, + SameNlocBatchSampler(ds._reader, shuffle=False), + pin_memory=False, + num_workers=4, + ) + iterator = iter(loader) + try: + assert next(iterator)["coord"].shape[0] == 4 + assert next(iterator)["coord"].shape[0] == 4 + assert loader._batch_iterator._pending is None + deferred = loader._batch_iterator._deferred_indices + assert deferred is not None + assert len(deferred) == 2 + assert next(iterator)["coord"].shape[0] == 2 + finally: + loader.close() + + def test_requirements_freeze_after_batch_read(self, lmdb_dir): + ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2) + loader = LmdbBatchDataLoader( + ds, + SameNlocBatchSampler(ds._reader, shuffle=False), + pin_memory=False, + num_workers=0, + ) + try: + next(iter(loader)) + with pytest.raises(RuntimeError, match="must be registered before reading"): + ds.add_data_requirement([DataRequirementItem("late_label", 1)]) + finally: + loader.close() + def test_full_epoch(self, lmdb_dir): ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=3) from torch.utils.data import ( @@ -650,7 +747,7 @@ class TestDistributedSameNlocBatchSampler: """Test DistributedSameNlocBatchSampler (pure logic, no torch.distributed).""" def test_disjoint_batches(self, multi_nloc_lmdb): - reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) + reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=1) s0 = DistributedSameNlocBatchSampler( reader, rank=0, world_size=2, shuffle=True, seed=42 ) @@ -679,10 +776,14 @@ def test_len(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) total = len(SameNlocBatchSampler(reader, shuffle=False)) - dist_s = DistributedSameNlocBatchSampler( - reader, rank=0, world_size=2, shuffle=False, seed=0 - ) - assert len(dist_s) == math.ceil(total / 2) + samplers = [ + DistributedSameNlocBatchSampler( + reader, rank=rank, world_size=4, shuffle=False, seed=0 + ) + for rank in range(4) + ] + assert {len(sampler) for sampler in samplers} == {math.ceil(total / 4)} + assert all(len(list(sampler)) == len(sampler) for sampler in samplers) def test_deterministic(self, multi_nloc_lmdb): reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2) @@ -850,7 +951,7 @@ def test_distributed_len_includes_auto_prob_expansion(self, auto_prob_lmdb): block_targets=ds._block_targets, ) assert len(dist_sampler_rank0) == math.ceil(global_batches / 2) - assert len(dist_sampler_rank1) == global_batches // 2 + assert len(dist_sampler_rank1) == math.ceil(global_batches / 2) assert len(dist_sampler_rank0) == len(list(dist_sampler_rank0)) assert len(dist_sampler_rank1) == len(list(dist_sampler_rank1)) @@ -884,7 +985,7 @@ def __init__(self, *args, **kwargs): ) assert calls == 1 - expected_len = len(ds._batch_sampler) // 2 + expected_len = (len(ds._batch_sampler) + 1) // 2 assert len(dist_sampler) == expected_len assert calls == 1 @@ -892,6 +993,35 @@ def __init__(self, *args, **kwargs): class TestMergeLmdbSystemIds: """Test merge_lmdb propagates frame_system_ids.""" + def test_merge_does_not_close_active_source_reader(self, tmp_path): + src1, src2 = str(tmp_path / "live1.lmdb"), str(tmp_path / "live2.lmdb") + _create_test_lmdb(src1, nframes=3, natoms=6) + _create_test_lmdb(src2, nframes=2, natoms=6) + active = LmdbDataReader(src1, ["O", "H"]) + try: + merge_lmdb([src1, src2], str(tmp_path / "live_merged.lmdb")) + assert active[0]["coord"].shape == (6, 3) + finally: + active.close() + + def test_failed_merge_releases_source_lease(self, tmp_path): + src = str(tmp_path / "overflow_source.lmdb") + _create_test_lmdb(src, nframes=3, natoms=6) + active = LmdbDataReader(src, ["O", "H"]) + resolved = active.lmdb_path + initial_refcount = _ENV_CACHE[resolved][1] + try: + with pytest.raises(lmdb.MapFullError): + merge_lmdb( + [src], + str(tmp_path / "overflow_destination.lmdb"), + map_size=4096, + ) + assert _ENV_CACHE[resolved][1] == initial_refcount + finally: + active.close() + assert resolved not in _ENV_CACHE + def test_merge_propagates_system_ids(self, tmp_path): src1, src2 = str(tmp_path / "src1.lmdb"), str(tmp_path / "src2.lmdb") _create_lmdb_with_system_ids( @@ -944,6 +1074,32 @@ def test_merge_preserves_type_map(self, tmp_path): # ============================================================ +def test_trainer_releases_lmdb_loader_after_failure() -> None: + """The PT trainer closes asynchronous loaders on exceptional exit.""" + from deepmd.pt.train.training import ( + Trainer, + ) + + class Loader: + closed = False + + def close(self) -> None: + self.closed = True + + loader = Loader() + trainer = object.__new__(Trainer) + trainer.training_dataloader = loader + trainer.validation_dataloader = None + + def fail() -> None: + raise RuntimeError("training failure") + + trainer._run = fail + with pytest.raises(RuntimeError, match="training failure"): + trainer.run() + assert loader.closed + + @pytest.fixture def multitask_lmdb_setup(tmp_path): """Create two LMDB datasets and a multitask training config.""" @@ -1075,6 +1231,7 @@ def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch): config, tmp_path = multitask_lmdb_setup monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "2") config = update_deepmd_input(deepcopy(config), warning=True) config["model"], shared_links = preprocess_shared_params(config["model"]) config = normalize(config, multi_task=True) @@ -1099,10 +1256,18 @@ def test_multitask_lmdb_end_to_end(self, multitask_lmdb_setup, monkeypatch): ) assert "coord" in input_dict assert "sid" in log_dict + assert ( + trainer.training_dataloader["model_1"]._batch_iterator._pool + is trainer.training_dataloader["model_2"]._batch_iterator._pool + ) # -- training run assertions -- trainer.run() assert len(list(tmp_path.glob("model.ckpt*.pt"))) > 0 + assert trainer.training_dataloader["model_1"]._batch_iterator.closed + assert trainer.training_dataloader["model_2"]._batch_iterator.closed + assert trainer.training_dataloader["model_1"].dataset._reader.closed + assert trainer.training_dataloader["model_2"].dataset._reader.closed # Explicit cleanup to free memory on CI import gc diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 651717f240..641804b1cd 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -13,6 +13,9 @@ import shutil import tempfile import unittest +from unittest.mock import ( + patch, +) import lmdb import msgpack @@ -22,6 +25,9 @@ normalize_batch, split_batch, ) +from deepmd.dpmodel.utils.lmdb_data import ( + collate_lmdb_frames, +) from deepmd.pt_expt.entrypoints.main import ( get_trainer, ) @@ -37,6 +43,9 @@ from deepmd.utils.compat import ( update_deepmd_input, ) +from deepmd.utils.data import ( + DataRequirementItem, +) def _encode_array(arr: np.ndarray) -> dict: @@ -183,6 +192,76 @@ def test_get_batch_shape_and_normalize(self) -> None: self.assertIn("force", labels) self.assertIn("natoms", labels) + def test_streaming_batch_matches_frame_collation(self) -> None: + """Preallocated decoding preserves the legacy per-frame contract.""" + ds = LmdbDataSystem( + lmdb_path=self.lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=0, + num_workers=0, + ) + indices = [1, 6] + expected = collate_lmdb_frames([ds._reader[index] for index in indices]) + actual = ds._reader.decode_batch(indices) + + self.assertEqual(tuple(actual), tuple(expected)) + for key, expected_value in expected.items(): + actual_value = actual[key] + if isinstance(expected_value, np.ndarray): + np.testing.assert_array_equal(actual_value, expected_value) + else: + self.assertEqual(actual_value, expected_value) + + def test_parallel_prefetch_matches_serial_order(self) -> None: + """Worker processes preserve sampler order and numerical values.""" + serial = LmdbDataSystem( + lmdb_path=self.lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=7, + num_workers=0, + ) + parallel = LmdbDataSystem( + lmdb_path=self.lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=7, + num_workers=2, + ) + try: + for _ in range(6): + expected = serial.get_batch() + actual = parallel.get_batch() + self.assertEqual(actual["fid"], expected["fid"]) + for key, expected_value in expected.items(): + actual_value = actual[key] + if isinstance(expected_value, np.ndarray): + np.testing.assert_array_equal(actual_value, expected_value) + else: + self.assertEqual(actual_value, expected_value) + pending = parallel._batch_iterator._pending + self.assertIsNotNone(pending) + self.assertLessEqual(len(pending.futures), 2) + finally: + parallel.close() + + def test_data_requirements_freeze_after_first_read(self) -> None: + """Batch schemas cannot change after a prefetched read.""" + ds = LmdbDataSystem( + lmdb_path=self.lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=0, + num_workers=2, + ) + ds.get_batch() + with self.assertRaisesRegex( + RuntimeError, + "must be registered before reading", + ): + ds.add_data_requirements([DataRequirementItem("late_label", ndof=1)]) + def test_get_batch_iterates_past_end(self) -> None: """get_batch reseeds the sampler at the end of an epoch.""" ds = LmdbDataSystem( @@ -196,11 +275,41 @@ def test_get_batch_iterates_past_end(self) -> None: batch = ds.get_batch() self.assertEqual(batch["coord"].shape, (2, 6, 3)) - def test_add_data_requirements_passthrough(self) -> None: - from deepmd.utils.data import ( - DataRequirementItem, - ) + def test_distributed_batches_are_sharded_with_equal_epoch_lengths(self) -> None: + """Ranks cover the global pass without advancing epochs at different steps.""" + systems = [ + LmdbDataSystem( + lmdb_path=self.lmdb_path, + type_map=["O", "H"], + batch_size=3, + seed=7, + num_workers=0, + rank=rank, + world_size=2, + ) + for rank in range(2) + ] + try: + self.assertEqual([system.nbatches for system in systems], [[3], [3]]) + self.assertEqual([len(system._sampler) for system in systems], [2, 2]) + batches = [ + [system.get_batch()["fid"] for _ in range(len(system._sampler))] + for system in systems + ] + finally: + for system in systems: + system.close() + + self.assertFalse(set(batches[0][0]) & set(batches[1][0])) + observed = { + frame_id + for rank_batches in batches + for batch in rank_batches + for frame_id in batch + } + self.assertEqual(observed, set(range(8))) + def test_add_data_requirements_passthrough(self) -> None: ds = LmdbDataSystem( lmdb_path=self.lmdb_path, type_map=["O", "H"], @@ -231,10 +340,6 @@ def test_partial_labels_are_batched_by_availability(self) -> None: batch_size=2, seed=0, ) - old_iter = ds._iter - # Realize the old iterator before requirements change so this checks - # replacement rather than two independently-created lazy generators. - next(old_iter) ds.add_data_requirements( [ DataRequirementItem( @@ -245,15 +350,27 @@ def test_partial_labels_are_batched_by_availability(self) -> None: ), ] ) - self.assertIsNot(ds._iter, old_iter) - batches = [ds.get_batch(), ds.get_batch()] - observed = { - (float(batch["find_energy"]), float(batch["find_force"])) - for batch in batches - } - self.assertEqual(observed, {(1.0, 0.0), (0.0, 1.0)}) - self.assertTrue(all(batch["coord"].shape[0] == 2 for batch in batches)) + try: + batches = [ds.get_batch(), ds.get_batch()] + observed = { + (float(batch["find_energy"]), float(batch["find_force"])) + for batch in batches + } + self.assertEqual(observed, {(1.0, 0.0), (0.0, 1.0)}) + self.assertTrue(all(batch["coord"].shape[0] == 2 for batch in batches)) + + stat_samples = make_stat_input(ds, nbatches=10) + stat_availability = { + (float(sample["find_energy"]), float(sample["find_force"])) + for sample in stat_samples + } + self.assertEqual(stat_availability, observed) + self.assertTrue( + all(sample["coord"].shape[0] == 2 for sample in stat_samples) + ) + finally: + ds.close() def test_stat_input_partitions_mixed_nloc_batches(self) -> None: """Statistics expose each atom-count group as one logical system.""" @@ -394,6 +511,43 @@ def test_get_trainer_routes_lmdb(self) -> None: finally: os.chdir(cwd) + def test_numb_epoch_counts_passes_over_the_lmdb(self) -> None: + """One epoch is one pass over the frames of the LMDB.""" + config = self._make_lmdb_config() + del config["training"]["numb_steps"] + config["training"]["numb_epoch"] = 2.0 + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + trainer = get_trainer(config) + finally: + os.chdir(cwd) + + # train.lmdb holds eight frames, read one frame per batch. + self.assertEqual(trainer.num_steps, 2 * 8) + + def test_training_closes_parallel_lmdb_pipeline(self) -> None: + """Trainer shutdown releases spawned decoder processes.""" + config = self._make_lmdb_config(numb_steps=2) + config["training"]["training_data"]["batch_size"] = 2 + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + cwd = os.getcwd() + os.chdir(self.tmpdir) + try: + with patch.dict(os.environ, {"DP_LMDB_NUM_WORKERS": "2"}): + trainer = get_trainer(config) + trainer.run() + self.assertTrue(trainer.training_data._batch_iterator.closed) + self.assertIsNone(trainer.training_data._batch_iterator._pool) + self.assertTrue(trainer.training_data._reader.closed) + finally: + os.chdir(cwd) + def test_mixed_nloc_statistics_and_training(self) -> None: """Trainer computes statistics and trains across fixed-nloc batches.""" config = self._make_lmdb_config(numb_steps=2) diff --git a/source/tests/pt_expt/test_multitask.py b/source/tests/pt_expt/test_multitask.py index 92ca70936f..2abda5bb02 100644 --- a/source/tests/pt_expt/test_multitask.py +++ b/source/tests/pt_expt/test_multitask.py @@ -1249,6 +1249,42 @@ def tearDown(self) -> None: shutil.rmtree("stat_files") +class TestMultiTaskEpochSchedule(unittest.TestCase): + """Test the run length and task weights derived from num_epoch_dict.""" + + @classmethod + def setUpClass(cls) -> None: + _skip_if_no_data() + # Both tasks read the same system one frame per batch, so their epoch + # lengths are equal to its frame count. + cls.epoch_length = np.load( + os.path.join(_PT_DATA, "set.000", "coord.npy") + ).shape[0] + + def setUp(self) -> None: + self.tmpdir = tempfile.mkdtemp(prefix="pt_expt_mt_epoch_") + self._old_cwd = os.getcwd() + os.chdir(self.tmpdir) + + def tearDown(self) -> None: + os.chdir(self._old_cwd) + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def test_steps_and_prob_follow_epoch_targets(self) -> None: + config = _make_multitask_config(_descriptor_se_e2_a) + del config["training"]["numb_steps"] + del config["training"]["model_prob"] + config["training"]["num_epoch_dict"] = {"model_1": 1.0, "model_2": 3.0} + config["model"], shared_links = preprocess_shared_params(config["model"]) + config = update_deepmd_input(config, warning=False) + config = normalize(config, multi_task=True) + + trainer = get_trainer(deepcopy(config), shared_links=shared_links) + + self.assertEqual(trainer.num_steps, 4 * self.epoch_length) + np.testing.assert_allclose(trainer.model_prob, [0.25, 0.75]) + + class TestMultiTaskSeA(unittest.TestCase, MultiTaskTrainTest): """Multi-task training with se_e2_a descriptor.""" diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index a0dee795e9..0db0df4977 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -10,6 +10,7 @@ import copy import datetime +import math import os import shutil import tempfile @@ -22,6 +23,7 @@ patch, ) +import numpy as np import pytest import torch @@ -733,6 +735,42 @@ def test_compiled_gradients_match_uncompiled(self) -> None: shutil.rmtree(tmpdir, ignore_errors=True) +class TestEpochSchedule(unittest.TestCase): + """Test the run length derived from training.numb_epoch.""" + + @classmethod + def setUpClass(cls) -> None: + data_dir = os.path.join(EXAMPLE_DIR, "data") + if not os.path.isdir(data_dir): + raise unittest.SkipTest(f"Example data not found: {data_dir}") + cls.data_dir = data_dir + # data_0 holds a single system read one frame per batch, so one epoch + # takes exactly one step per frame. + cls.nframes = np.load( + os.path.join(data_dir, "data_0", "set.000", "coord.npy") + ).shape[0] + + def _num_steps_for(self, num_epoch: float) -> int: + config = _make_config(self.data_dir) + del config["training"]["numb_steps"] + config["training"]["numb_epoch"] = num_epoch + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + tmpdir = tempfile.mkdtemp(prefix="pt_expt_epoch_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + return get_trainer(config).num_steps + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_num_steps_covers_requested_epochs(self) -> None: + self.assertEqual(self._num_steps_for(1.0), self.nframes) + self.assertEqual(self._num_steps_for(2.5), math.ceil(2.5 * self.nframes)) + + class TestGetData(unittest.TestCase): """Test the batch data conversion in Trainer.get_data.""" diff --git a/source/tests/pt_expt/test_training_ddp.py b/source/tests/pt_expt/test_training_ddp.py index 0d71e66870..b866db4f38 100644 --- a/source/tests/pt_expt/test_training_ddp.py +++ b/source/tests/pt_expt/test_training_ddp.py @@ -26,6 +26,9 @@ from pathlib import ( Path, ) +from unittest.mock import ( + patch, +) import numpy as np import torch @@ -35,6 +38,9 @@ from deepmd.pt_expt.entrypoints.main import ( get_trainer, ) +from deepmd.pt_expt.train.training import ( + Trainer, +) from deepmd.pt_expt.utils.finetune import ( get_finetune_rules, ) @@ -63,6 +69,9 @@ # Auto-detect DDP backend based on device availability. _DDP_BACKEND = "nccl" if torch.cuda.is_available() else "gloo" +# Epoch length reported by rank 0 when the ranks are made to disagree. +_DRIFTED_EPOCH_LENGTH = 40 + # NCCL requires at least 2 GPUs for multi-rank tests. if _DDP_BACKEND == "nccl" and torch.cuda.device_count() < 2: raise unittest.SkipTest("NCCL DDP tests require at least 2 GPUs") @@ -607,11 +616,91 @@ def _worker_finetune( dist.destroy_process_group() +def _worker_epoch_schedule(rank, world_size, port, data_dir, drifted, result_dict): + """Worker: build a trainer whose run length comes from numb_epoch. + + When *drifted* is set, each rank reports a different local epoch length, + reproducing the last-bit disagreement that floating-point sampling + probabilities can produce between ranks. + """ + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group(backend=_DDP_BACKEND, rank=rank, world_size=world_size) + try: + tmpdir = tempfile.mkdtemp(prefix=f"ddp_epoch_rank{rank}_") + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + config = _make_config(data_dir) + del config["training"]["numb_steps"] + config["training"]["numb_epoch"] = 1.0 + config = update_deepmd_input(config, warning=False) + config = normalize(config) + if drifted: + with patch.object( + Trainer, + "_epoch_length", + lambda self, model_key: _DRIFTED_EPOCH_LENGTH + rank, + ): + num_steps = get_trainer(config).num_steps + else: + num_steps = get_trainer(config).num_steps + result_dict[rank] = {"num_steps": num_steps} + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + finally: + dist.destroy_process_group() + + # --------------------------------------------------------------------------- # Test classes # --------------------------------------------------------------------------- +class TestDDPEpochSchedule(unittest.TestCase): + """An epoch spans the dataset once across the whole world, not per rank.""" + + @classmethod + def setUpClass(cls) -> None: + data_dir = os.path.join(EXAMPLE_DIR, "data") + if not os.path.isdir(data_dir): + raise unittest.SkipTest(f"Example data not found: {data_dir}") + cls.data_dir = os.path.join(data_dir, "data_0") + + def _run(self, drifted: bool) -> dict: + port = _find_free_port() + result_dict = mp.Manager().dict() + mp.spawn( + _worker_epoch_schedule, + args=(2, port, self.data_dir, drifted, result_dict), + nprocs=2, + join=True, + ) + return dict(result_dict) + + def test_ranks_share_an_epoch(self) -> None: + results = self._run(drifted=False) + + # DistributedSampler gives each rank ceil(nframes / 2) samples, and + # the system is read one frame per batch. + nframes = np.load(os.path.join(self.data_dir, "set.000", "coord.npy")).shape[0] + self.assertEqual(results[0]["num_steps"], int(np.ceil(nframes / 2))) + self.assertEqual(results[1]["num_steps"], results[0]["num_steps"]) + + def test_drifting_epoch_lengths_are_pinned_to_rank_zero(self) -> None: + """Ranks that round differently still agree on the run length. + + A run length that drifts by one step across ranks desynchronizes the + full-validation start step and deadlocks the mismatched collectives, + so rank 0's value must win everywhere. + """ + results = self._run(drifted=True) + + self.assertEqual(results[0]["num_steps"], _DRIFTED_EPOCH_LENGTH) + self.assertEqual(results[1]["num_steps"], _DRIFTED_EPOCH_LENGTH) + + class TestDDPSingleTaskTrain(unittest.TestCase): """Smoke test: single-task DDP training with 2 ranks.""" diff --git a/source/tests/pt_expt/utils/test_env.py b/source/tests/pt_expt/utils/test_env.py index a589c80ae1..c065e52530 100644 --- a/source/tests/pt_expt/utils/test_env.py +++ b/source/tests/pt_expt/utils/test_env.py @@ -2,11 +2,69 @@ import importlib import logging +import pytest import torch import deepmd.env as common_env +def test_lmdb_num_workers_override(monkeypatch) -> None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "12") + assert common_env.get_lmdb_num_workers() == 12 + + +@pytest.mark.parametrize( + ("configured", "expected_error"), + [ + ("many", "must be a non-negative integer"), + ("-1", "must be non-negative"), + ], +) +def test_lmdb_num_workers_rejects_invalid_override( + monkeypatch, + configured: str, + expected_error: str, +) -> None: + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", configured) + with pytest.raises(ValueError, match=expected_error): + common_env.get_lmdb_num_workers() + + +@pytest.mark.parametrize( + ("configured", "expected_error"), + [("many", "positive integer"), ("0", "must be positive")], +) +def test_lmdb_num_workers_rejects_invalid_local_world_size( + monkeypatch, + configured: str, + expected_error: str, +) -> None: + monkeypatch.delenv("DP_LMDB_NUM_WORKERS", raising=False) + monkeypatch.setenv("LOCAL_WORLD_SIZE", configured) + with pytest.raises(ValueError, match=expected_error): + common_env.get_lmdb_num_workers() + + +def test_lmdb_num_workers_falls_back_without_affinity(monkeypatch) -> None: + monkeypatch.delenv("DP_LMDB_NUM_WORKERS", raising=False) + monkeypatch.delenv("LOCAL_WORLD_SIZE", raising=False) + monkeypatch.delattr(common_env.os, "sched_getaffinity", raising=False) + monkeypatch.setattr(common_env.os, "cpu_count", lambda: 12) + assert common_env.get_lmdb_num_workers() == 12 + + +def test_lmdb_num_workers_partitions_node_budget(monkeypatch) -> None: + monkeypatch.delenv("DP_LMDB_NUM_WORKERS", raising=False) + monkeypatch.setenv("LOCAL_WORLD_SIZE", "4") + monkeypatch.setattr( + common_env.os, + "sched_getaffinity", + lambda _pid: set(range(80)), + raising=False, + ) + assert common_env.get_lmdb_num_workers() == 16 + + def test_env_threads_guard_handles_runtimeerror(monkeypatch) -> None: def raise_err(*_args, **_kwargs) -> None: raise RuntimeError("boom")