feat(pt_expt): resume checkpointed training data progress

deepmd/dpmodel/train/__init__.py

@@ -12,6 +12,12 @@
     make_task_maps,
     print_data_summaries,
 )
+from .data_progress import (
+    DATA_PROGRESS_CHECKPOINT_KEY,
+    collect_training_data_progress,
+    preserve_training_data_progress,
+    restore_training_data_progress,
+)
 from .entrypoint import (
     AbstractTrainEntrypoint,
     TrainEntrypointOptions,
@@ -44,6 +50,7 @@
 )
 
 __all__ = [
+    "DATA_PROGRESS_CHECKPOINT_KEY",
     "DEFAULT_TASK_KEY",
     "AbstractTrainEntrypoint",
     "AbstractTrainer",
@@ -64,9 +71,12 @@
     "build_checkpoint_stores",
     "change_model_out_bias",
     "change_model_out_bias_by_task",
+    "collect_training_data_progress",
     "iter_training_task_configs",
     "make_task_maps",
+    "preserve_training_data_progress",
     "print_data_summaries",
     "resolve_keep_ckpt_count",
     "resolve_step_schedule",
+    "restore_training_data_progress",
 ]

deepmd/dpmodel/train/data_progress.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Small, backend-independent checkpoints of the next training batches.
+
+Data sources own their cursors and deterministic sampler state. This module
+adds task names and the shared data-source RNG, not decoded data or permutations.
+"""
+
+from collections.abc import (
+    Iterator,
+    Mapping,
+)
+from contextlib import (
+    contextmanager,
+)
+from typing import (
+    Any,
+)
+
+from deepmd.utils import random as dp_random
+
+DATA_PROGRESS_CHECKPOINT_KEY = "data_progress"
+DATA_PROGRESS_VERSION = 1
+
+
+def collect_training_data_progress(
+    training_data_by_task: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Collect supported data-source cursors and the shared data-source RNG."""
+    return {
+        "version": DATA_PROGRESS_VERSION,
+        "tasks": {
+            key: data.state_dict()
+            for key, data in training_data_by_task.items()
+            if callable(getattr(data, "state_dict", None))
+            and callable(getattr(data, "load_state_dict", None))
+        },
+        "rng": dp_random.get_state(),
+    }
+
+
+def restore_training_data_progress(
+    training_data_by_task: Mapping[str, Any],
+    progress: Mapping[str, Any] | None,
+) -> None:
+    """Restore cursors lazily; missing progress preserves legacy restarts.
+
+    The source implementations validate their sampling configuration and must
+    not read or decode data here. Restoring the shared RNG last also isolates
+    it from any source-specific reconstruction.
+    """
+    if progress is None:
+        return
+    if progress.get("version") != DATA_PROGRESS_VERSION:
+        raise ValueError("Unsupported data-progress checkpoint version.")
+    tasks = progress.get("tasks")
+    if not isinstance(tasks, Mapping):
+        raise ValueError("data-progress 'tasks' must be a mapping.")
+    supported = {
+        key: data
+        for key, data in training_data_by_task.items()
+        if callable(getattr(data, "load_state_dict", None))
+        and callable(getattr(data, "state_dict", None))
+    }
+    if set(tasks) != set(supported):
+        raise ValueError("data-progress tasks do not match the training data sources.")
+    if "rng" not in progress:
+        raise ValueError("data-progress checkpoint is missing the selection RNG.")
+    for key, data in supported.items():
+        data.load_state_dict(tasks[key])
+    dp_random.set_state(progress["rng"])
+
+
+@contextmanager
+def preserve_training_data_progress(
+    training_data_by_task: Mapping[str, Any],
+) -> Iterator[None]:
+    """Keep diagnostics and compilation from consuming the training stream."""
+    progress = collect_training_data_progress(training_data_by_task)
+    try:
+        yield
+    finally:
+        restore_training_data_progress(training_data_by_task, progress)