This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 869a6d014d [python] Support CWD checkpoints with StatefulDataLoader
(#10628)
869a6d014d is described below
commit 869a6d014d404d1d92d694a04b1741e4ea98f436
Author: Yann Byron <[email protected]>
AuthorDate: Tue Oct 20 12:10:52 2026 +0800
[python] Support CWD checkpoints with StatefulDataLoader (#10628)
---
docs/docs/pypaimon/pytorch.md | 116 ++++++++
.../pypaimon/multimodal/window_dataset.py | 112 +++++++-
.../tests/contiguous_window_resume_test.py | 292 +++++++++++++++++++++
paimon-python/pyproject.toml | 1 +
4 files changed, 520 insertions(+), 1 deletion(-)
diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md
index 5deedfe8da..cc3b94e08b 100644
--- a/docs/docs/pypaimon/pytorch.md
+++ b/docs/docs/pypaimon/pytorch.md
@@ -320,6 +320,122 @@ physical file/split boundaries inside those episodes. If
context has already
been divided among ranks, use a local sampler instead of sharding it again
with a world-size `DistributedSampler`; coordinate rank step counts separately.
+### Resuming a CWD training run
+
+Install `torchdata>=0.11,<0.12` alongside your compatible PyTorch installation.
+Use `torchdata.stateful_dataloader.StatefulDataLoader` in place of
`DataLoader`.
+With `shuffle=True` it uses TorchData's stateful random sampler. It records the
+batch position delivered to training, including the worker states needed to
+reconstruct prefetched batches. Do not checkpoint a sampler's current index as
+the training position: prefetch can put it ahead of completed optimizer steps.
+
+CWD implements `state_dict()` and `load_state_dict()` for this loader. Its
state
+checks the table, branch, resolved snapshot, actual context/anchor index, and
+window settings. Rebuild the same dataset before loading; restoration does not
+silently switch snapshots. Retain the snapshot files with a training tag for
+as long as checkpoints may be resumed. A changed snapshot, context, anchors,
+offsets or padding raises `ValueError`, even when the sample count is
unchanged.
+The index identity is cached after its first checkpoint; do not mutate dataset
+configuration or index arrays during the run.
+
+CWD also saves Python, NumPy and Torch CPU global RNG state in each process.
+Transforms and adapters with private RNGs or other mutable state must implement
+both `state_dict()` and `load_state_dict(state)`. Their state is copied before
+it is returned. Keep callback code, normalization and preprocessing settings
+unchanged; CWD cannot infer these from an arbitrary callable. Custom collators
+must be stateless apart from those global RNGs. Reader caches and video decoder
+handles are rebuilt rather than checkpointed.
+
+The training loop must save model, optimizer and loader state together, after
+an optimizer update and before fetching another batch. The following pattern
+uses `make_dataset(snapshot_id)` to rebuild your CWD with the same window and
+preprocessing arguments, via `frames.scan(snapshot_id=snapshot_id)`:
+
+```python
+import os
+import random
+
+import numpy as np
+import torch
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+loader_config = dict(batch_size=32, shuffle=True, drop_last=False,
+ num_workers=4, prefetch_factor=2,
+ persistent_workers=True, in_order=True)
+# Include the application's preprocessing/normalization version in this
contract.
+run_config = dict(loader=loader_config, preprocessing="train-v1", world_size=1,
+ sampler="torchdata.RandomSampler", seed=42)
+generator = torch.Generator().manual_seed(42)
+# Load only trusted checkpoint files, and choose their snapshot before
planning.
+saved = (torch.load(resume_path, map_location="cpu", weights_only=False)
+ if resume_path is not None else None)
+snapshot_id = saved["snapshot_id"] if saved is not None else
training_snapshot_id
+dataset = make_dataset(snapshot_id)
+loader = StatefulDataLoader(dataset, generator=generator, **loader_config)
+epoch, global_step = 0, 0
+
+# Restore model/optimizer/scheduler/scaler before restoring RNGs.
+if saved is not None:
+ if saved["config"] != run_config or saved["snapshot_id"] !=
dataset.snapshot_id:
+ raise ValueError("Training checkpoint configuration changed")
+ model.load_state_dict(saved["model"])
+ optimizer.load_state_dict(saved["optimizer"])
+ loader.load_state_dict(saved["loader"])
+ epoch, global_step = saved["epoch"], saved["global_step"]
+
+iterator = iter(loader) # Dataset/worker state is restored here, not at
load_state_dict.
+if saved is not None:
+ # Iterator construction consumes RNG; restore these after construction.
+ generator.set_state(saved["loader_generator"])
+ random.setstate(saved["python_rng"])
+ np.random.set_state(saved["numpy_rng"])
+ torch.set_rng_state(saved["torch_rng"])
+ if saved["cuda_rng"] is not None:
+ torch.cuda.set_rng_state_all(saved["cuda_rng"])
+
+for batch in iterator:
+ optimizer.zero_grad(set_to_none=True)
+ loss = training_loss(model, batch)
+ loss.backward()
+ optimizer.step()
+ global_step += 1
+ if should_checkpoint(global_step):
+ checkpoint = dict(
+ config=run_config, snapshot_id=dataset.snapshot_id,
+ epoch=epoch, global_step=global_step,
+ model=model.state_dict(), optimizer=optimizer.state_dict(),
+ loader=loader.state_dict(), loader_generator=generator.get_state(),
+ python_rng=random.getstate(), numpy_rng=np.random.get_state(),
+ torch_rng=torch.get_rng_state(),
+ cuda_rng=torch.cuda.get_rng_state_all() if
torch.cuda.is_available() else None,
+ )
+ # Include scheduler/scaler states here when the trainer uses them.
+ torch.save(checkpoint, checkpoint_path + ".tmp")
+ os.replace(checkpoint_path + ".tmp", checkpoint_path)
+```
+
+Alternatively, use the same retained tag in `make_dataset` and verify its
+resolved ID. Move optimizer state to the training device as required by the
+trainer. The example saves inside the optimizer-step loop,
+including its last batch, before the iterator reports exhaustion. With
TorchData
+0.11, saving an already exhausted iterator outside that loop is not equivalent:
+restoring it may immediately start a new epoch and consume another permutation.
+After finishing the restored suffix, start the next epoch with `iter(loader)`.
+
+Keep batch size, `drop_last`, sampler configuration, worker count, prefetch and
+persistent-worker settings unchanged, with `in_order=True`. Keep the same
+PyTorch/TorchData versions; checkpoint compatibility across versions is not
+promised. TorchData may regenerate sampling indices or replay prefetched work
+on restore; CWD does not promise constant-time recovery or zero rereads.
+
+For fixed-topology DDP, use TorchData's `StatefulDistributedSampler`, save one
+loader/generator/RNG state per rank at the same completed global step, and keep
+rank assignment and world size unchanged. Save its seed and epoch in the
training
+configuration and call `sampler.set_epoch(saved_epoch)` before restoring the
+loader; the sampler's state does not itself contain the epoch. Cross-rank
+checkpoint publication, gradient-accumulation boundaries, scheduler/scaler
state
+and deterministic model execution remain the trainer's responsibility.
+
## File Format Metadata Cache
Reusable PyArrow Dataset metadata is cached across reads. Configure its
estimated
diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py
b/paimon-python/pypaimon/multimodal/window_dataset.py
index 4a8209c955..f2240557e3 100644
--- a/paimon-python/pypaimon/multimodal/window_dataset.py
+++ b/paimon-python/pypaimon/multimodal/window_dataset.py
@@ -18,8 +18,11 @@
"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows."""
import copy
+import hashlib
import math
import operator
+import pickle
+import random
from functools import partial
import numpy as np
@@ -41,7 +44,7 @@ from pypaimon.read.query_auth_split import QueryAuthSplit
from pypaimon.schema.data_types import (
is_array_blob_type, is_blob_file_type, is_map_blob_type,
)
-from pypaimon.snapshot.time_travel_util import SCAN_KEYS
+from pypaimon.snapshot.time_travel_util import SCAN_KEYS, TimeTravelUtil
from pypaimon.table.special_fields import SpecialFields
from pypaimon.utils.range import Range
@@ -205,6 +208,7 @@ class ContiguousWindowDataset(Dataset):
for name in self._video_columns
}
self._build_index(index, anchor_row_ids)
+ self._checkpoint_identity = None
@classmethod
def from_query(cls, query, **kwargs):
@@ -256,6 +260,95 @@ class ContiguousWindowDataset(Dataset):
state["_plans"] = {}
return state
+ def state_dict(self):
+ """Return data identity and process state for StatefulDataLoader.
+
+ Rebuild this Dataset with the same retained snapshot, context, anchors,
+ and window settings before loading its state. The loader owns sampler
+ and batch progress; this Dataset never changes its index on restore.
+
+ Python, NumPy and Torch CPU global RNGs are saved for worker
transforms.
+ Callbacks with private state must implement both ``state_dict`` and
+ ``load_state_dict``. Keep callback code and preprocessing configuration
+ unchanged; they cannot be inferred from a callable. The training loop
+ must also checkpoint its model, optimizer and main-process/CUDA RNGs.
+ Do not mutate the Dataset's configuration or index after checkpointing.
+ """
+ return {
+ "version": 1,
+ "identity": self._state_identity(),
+ "callbacks": {
+ # A tuple keeps TorchData from flattening dictionary
subclasses.
+ name: (copy.deepcopy(callback.state_dict()),)
+ for name, callback in self._stateful_callbacks().items()
+ },
+ "python_rng": random.getstate(),
+ "numpy_rng": np.random.get_state(),
+ "torch_rng": torch.get_rng_state(),
+ }
+
+ def load_state_dict(self, state_dict):
+ """Reject checkpoints made from a different CWD data contract."""
+ if state_dict.get("version") != 1:
+ raise ValueError("Unsupported CWD checkpoint version.")
+ if state_dict.get("identity") != self._state_identity():
+ raise ValueError(
+ "Cannot restore state from a different CWD: rebuild the same "
+ "table snapshot, context, anchors and window settings.")
+ callbacks = self._stateful_callbacks()
+ if callbacks.keys() != state_dict["callbacks"].keys():
+ raise ValueError("CWD checkpoint has different stateful
callbacks.")
+ for name, callback in callbacks.items():
+
callback.load_state_dict(copy.deepcopy(state_dict["callbacks"][name][0]))
+ # Callback restoration may itself consume random numbers.
+ random.setstate(state_dict["python_rng"])
+ np.random.set_state(state_dict["numpy_rng"])
+ torch.set_rng_state(state_dict["torch_rng"])
+
+ def _stateful_callbacks(self):
+ callbacks = {"column:" + name: callback
+ for name, callback in self.column_transforms.items()}
+ callbacks["adapter"] = self.adapter
+ stateful = {}
+ for name, callback in callbacks.items():
+ save = callable(getattr(callback, "state_dict", None))
+ restore = callable(getattr(callback, "load_state_dict", None))
+ if save != restore:
+ raise TypeError("CWD callback %s must implement both
state_dict "
+ "and load_state_dict." % name)
+ if save:
+ stateful[name] = callback
+ return stateful
+
+ def _state_identity(self):
+ if self._checkpoint_identity is None:
+ snapshot = (TimeTravelUtil.resolve_snapshot(self._table)
+ if self.snapshot_id is not None else None)
+ masks = sorted(set(
+ tuple(sorted((split.auth_result.column_masking or {}).items()))
+ for split in self._splits if isinstance(split,
QueryAuthSplit)))
+ metadata = (
+ self._table.table_path, self._table.current_branch(),
+ self.snapshot_id,
+ (snapshot.uuid, snapshot.base_manifest_list,
snapshot.delta_manifest_list)
+ if snapshot is not None else None,
+ self._table.table_schema.id,
+ self.group_key, self.order_key, self._group_keys,
+ self.columns, tuple((name, self.frame_offsets[name]) for name
in self.columns),
+ self.window_size, self.stride, self.boundary,
+ _checkpoint_padding(self.pad_values), self.video_backend,
masks,
+ )
+ digest = hashlib.sha256()
+ for value in metadata:
+ digest.update(pickle.dumps(value, protocol=4))
+ for array in (self._row_ids, self._group_starts,
self._group_lengths,
+ self._group_first_orders, self._anchor_groups,
self._anchor_starts):
+ digest.update(pickle.dumps((array.dtype.str, array.shape),
protocol=4))
+ digest.update(memoryview(array))
+ # Worker snapshots are frequent; hash the immutable index only
once.
+ self._checkpoint_identity = digest.hexdigest()
+ return self._checkpoint_identity
+
def close(self):
"""Close video decoders opened by this Dataset process."""
first_error = None
@@ -498,6 +591,23 @@ class ContiguousWindowDataset(Dataset):
return plan
+def _checkpoint_padding(value):
+ # Tensor's default pickle includes storage identity, even for equal values.
+ if isinstance(value, torch.Tensor):
+ raw = value.detach().cpu().contiguous().reshape(-1).view(torch.uint8)
+ return ("tensor", str(value.dtype), tuple(value.shape),
str(value.device),
+ value.requires_grad, raw.numpy().tobytes())
+ if type(value) is dict:
+ return ("dict", tuple(sorted(
+ (_checkpoint_padding(key), _checkpoint_padding(item))
+ for key, item in value.items())))
+ if type(value) in (list, tuple):
+ return (type(value).__name__, tuple(_checkpoint_padding(item) for item
in value))
+ if type(value) is np.ndarray:
+ value = value.copy(order="C")
+ return ("pickle", pickle.dumps(value, protocol=4))
+
+
class _PinnedRowIdPlan:
"""Read a projection from the window context's authorized splits by row ID.
diff --git a/paimon-python/pypaimon/tests/contiguous_window_resume_test.py
b/paimon-python/pypaimon/tests/contiguous_window_resume_test.py
new file mode 100644
index 0000000000..d4d5674661
--- /dev/null
+++ b/paimon-python/pypaimon/tests/contiguous_window_resume_test.py
@@ -0,0 +1,292 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import random
+from collections import defaultdict
+from unittest.mock import patch
+
+import numpy as np
+import pyarrow as pa
+import pytest
+import torch
+
+import pypaimon.multimodal as pmm
+from pypaimon.multimodal import window_dataset
+
+
[email protected]
+def frames(tmp_path):
+ connection = pmm.connect(options={"warehouse": str(tmp_path /
"warehouse")})
+ table = connection.create_table("frames", schema=pa.schema([
+ pa.field("episode", pa.string()),
+ pa.field("step", pa.int64()),
+ pa.field("value", pa.float32()),
+ ]), options={
+ "row-tracking.enabled": "true",
+ "data-evolution.enabled": "true",
+ "file.format": "parquet",
+ "scan.native-plan.enabled": "false",
+ "read.native.enabled": "false",
+ })
+ table.add([
+ {"episode": episode, "step": step, "value": float(step)}
+ for episode in ("a", "b") for step in range(24)
+ ])
+ return table
+
+
+def windows(frames, query=None, **kwargs):
+ options = dict(columns=["value"], group_key="episode", order_key="step",
+ frame_offsets={"value": [0, 1, 2]},
+ column_transforms={"value": torch.tensor})
+ options.update(kwargs)
+ return (query or frames.scan()).to_contiguous_window_dataset(**options)
+
+
+def test_restore_rejects_same_length_different_anchors(frames):
+ original = windows(frames, anchor_row_ids=[0, 1])
+ different = windows(frames, anchor_row_ids=[2, 3])
+ assert len(original) == len(different)
+
+ with pytest.raises(ValueError, match="different CWD"):
+ different.load_state_dict(original.state_dict())
+
+
[email protected]("change", ["context", "offsets", "padding",
"snapshot"])
+def test_restore_rejects_changed_data_contract(frames, change):
+ original = windows(frames, query=frames.scan().where("episode = 'a'"),
boundary="pad")
+ state = original.state_dict()
+ options = {"query": frames.scan().where("episode = 'a'"), "boundary":
"pad"}
+ if change == "context":
+ options["query"] = frames.scan().where("episode = 'b'")
+ elif change == "offsets":
+ options["frame_offsets"] = {"value": [2, 1, 0]}
+ elif change == "padding":
+ options["pad_values"] = {"value": -1}
+ else:
+ frames.add([{"episode": "c", "step": 0, "value": 0.0}])
+ options["query"] = frames.scan().where("episode = 'a'")
+ # A retained old snapshot remains valid after unrelated new commits.
+ pinned = frames.scan(snapshot_id=original.snapshot_id).where("episode
= 'a'")
+ windows(frames, query=pinned, boundary="pad").load_state_dict(state)
+ different = windows(frames, **options)
+ assert len(original) == len(different)
+ with pytest.raises(ValueError, match="different CWD"):
+ different.load_state_dict(state)
+
+
+class _RandomTransform:
+
+ def __init__(self):
+ self.generator = random.Random(37)
+ self.calls = [0]
+
+ def __call__(self, values):
+ self.calls[0] += 1
+ noise = (random.random() + np.random.random() + torch.rand(()).item()
+ + self.generator.random() + self.calls[0])
+ return torch.tensor(values) + noise
+
+ def state_dict(self):
+ return {"generator": self.generator.getstate(), "calls": self.calls}
+
+ def load_state_dict(self, state):
+ self.generator.setstate(state["generator"])
+ self.calls[:] = state["calls"]
+
+
+def test_restore_rng_and_independent_callback_state(frames):
+ original = windows(frames)
+ original.column_transforms["value"] = _RandomTransform()
+ original[0]
+ state = original.state_dict()
+ expected = original[1]["value"]
+ original[2]
+
+ restored = windows(frames)
+ restored.column_transforms["value"] = _RandomTransform()
+ restored.load_state_dict(state)
+ torch.testing.assert_close(restored[1]["value"], expected, rtol=0, atol=0)
+
+
+def test_restore_rejects_missing_callback_and_unknown_version(frames):
+ original = windows(frames, column_transforms={"value": _RandomTransform()})
+ state = original.state_dict()
+ with pytest.raises(ValueError, match="different stateful callbacks"):
+ windows(frames).load_state_dict(state)
+ state["version"] = 2
+ with pytest.raises(ValueError, match="Unsupported CWD checkpoint version"):
+ original.load_state_dict(state)
+
+
+def test_callback_requires_both_state_methods(frames):
+ transform = _RandomTransform()
+ transform.load_state_dict = None
+ dataset = windows(frames, column_transforms={"value": transform})
+ with pytest.raises(TypeError, match="must implement both"):
+ dataset.state_dict()
+
+
+def test_checkpoint_reuses_identity_without_reading_payload(frames):
+ dataset = windows(frames)
+ with patch.object(dataset, "_read_rows",
side_effect=AssertionError("payload read")):
+ state = dataset.state_dict()
+ with patch.object(window_dataset.hashlib, "sha256",
side_effect=AssertionError("rehash")):
+ dataset.load_state_dict(state)
+ assert dataset.state_dict()["identity"] == state["identity"]
+
+
+def test_restore_retained_tag_without_snapshot_file(frames):
+ frames.raw_table.create_tag("train")
+ dataset = windows(frames, query=frames.scan(tag_name="train"))
+ manager = frames.raw_table.snapshot_manager()
+
frames.raw_table.file_io.delete_quietly(manager.get_snapshot_path(dataset.snapshot_id))
+ state = dataset.state_dict()
+ restored = windows(frames, query=frames.scan(tag_name="train"))
+ restored.load_state_dict(state)
+ torch.testing.assert_close(restored[0]["value"], torch.tensor([0.0, 1.0,
2.0]))
+
+
[email protected]("nested", [False, True])
+def test_tensor_padding_identity_is_by_value(frames, nested):
+ def dataset(value):
+ pad = torch.tensor(value)
+ if nested:
+ pad = {"values": [pad]}
+ return windows(frames, boundary="pad", pad_values={"value": pad})
+
+ state = dataset(0.0).state_dict()
+ dataset(0.0).load_state_dict(state)
+ with pytest.raises(ValueError, match="different CWD"):
+ dataset(1.0).load_state_dict(state)
+
+
+def test_numpy_padding_identity_ignores_memory_layout(frames):
+ pad = np.arange(6, dtype=np.float32).reshape(2, 3)
+ original = windows(frames, boundary="pad", pad_values={"value": pad})
+ restored = windows(frames, boundary="pad", pad_values={"value":
np.asfortranarray(pad)})
+ restored.load_state_dict(original.state_dict())
+
+
+def _loader(dataset, workers, persistent=False):
+ loader_type =
pytest.importorskip("torchdata.stateful_dataloader").StatefulDataLoader
+ return loader_type(
+ dataset, batch_size=4, shuffle=True,
+ generator=torch.Generator().manual_seed(91), num_workers=workers,
+ persistent_workers=persistent,
+ **({"multiprocessing_context": "spawn", "prefetch_factor": 2} if
workers else {}))
+
+
+def _batch_values(batch):
+ return (batch["episode"], batch["step"].tolist(), batch["value"].tolist(),
+ batch["value_is_pad"].tolist())
+
+
+class _EpisodeCounter:
+
+ def __init__(self):
+ self.counts = defaultdict(int)
+
+ def __call__(self, sample):
+ episode = sample["episode"]
+ self.counts[episode] += 1
+ return episode, sample["step"], self.counts[episode]
+
+ def state_dict(self):
+ return {"counts": self.counts}
+
+ def load_state_dict(self, state):
+ self.counts = state["counts"]
+
+
+def test_stateful_loader_preserves_callback_dictionary_subclasses(frames,
tmp_path):
+ loader_type =
pytest.importorskip("torchdata.stateful_dataloader").StatefulDataLoader
+
+ def make_loader():
+ return loader_type(
+ windows(frames, adapter=_EpisodeCounter()), batch_size=4,
+ num_workers=1, multiprocessing_context="spawn", prefetch_factor=2)
+
+ def samples(batches):
+ return [(episode, step, count) for episodes, steps, counts in batches
+ for episode, step, count in zip(episodes, steps.tolist(),
counts.tolist())]
+
+ loader = make_loader()
+ iterator = iter(loader)
+ prefix = samples([next(iterator) for _ in range(3)])
+ checkpoint = tmp_path / "loader.pt"
+ torch.save(loader.state_dict(), checkpoint)
+ suffix = samples(iterator)
+ assert prefix + suffix == [(episode, step, step + 1)
+ for episode in ("a", "b") for step in range(22)]
+
+ restored = make_loader()
+ restored.load_state_dict(torch.load(checkpoint, weights_only=False))
+ # Episode b was absent from the saved counters and still needs the default
factory.
+ assert samples(restored) == suffix
+
+
[email protected]("workers,persistent,completed", [
+ (0, False, 3), (2, False, 3), (2, True, 3), (0, False, 11), (2, True, 11),
+])
+def test_stateful_loader_resumes_prefetch_and_next_epoch(
+ frames, tmp_path, workers, persistent, completed):
+ dataset = windows(frames)
+ dataset.column_transforms["value"] = _RandomTransform()
+ loader = _loader(dataset, workers, persistent)
+ iterator = iter(loader)
+ prefix = [_batch_values(next(iterator)) for _ in range(completed)]
+ checkpoint = tmp_path / "loader.pt"
+ torch.save({"loader": loader.state_dict(), "generator":
loader.generator.get_state()}, checkpoint)
+ remaining = [_batch_values(batch) for batch in iterator]
+ next_epoch = [_batch_values(batch) for batch in loader]
+
+ restored_dataset = windows(frames)
+ restored_dataset.column_transforms["value"] = _RandomTransform()
+ restored = _loader(restored_dataset, workers, persistent)
+ state = torch.load(checkpoint, weights_only=False)
+ restored.load_state_dict(state["loader"])
+ restored_iterator = iter(restored)
+ restored.generator.set_state(state["generator"])
+ assert [_batch_values(batch) for batch in restored_iterator] == remaining
+ assert [_batch_values(batch) for batch in restored] == next_epoch
+ assert sum(len(batch[0]) for batch in prefix + remaining) == len(dataset)
+
+
+def test_stateful_distributed_sampler_restores_each_rank(frames):
+ stateful = pytest.importorskip("torchdata.stateful_dataloader")
+ samplers = pytest.importorskip("torchdata.stateful_dataloader.sampler")
+ all_ids = []
+ for rank in range(2):
+ def make_loader():
+ dataset = windows(frames)
+ sampler = samplers.StatefulDistributedSampler(
+ dataset, num_replicas=2, rank=rank, seed=71)
+ sampler.set_epoch(3)
+ return stateful.StatefulDataLoader(dataset, sampler=sampler,
batch_size=4)
+
+ loader = make_loader()
+ iterator = iter(loader)
+ prefix = [_batch_values(next(iterator)) for _ in range(2)]
+ state = loader.state_dict()
+ suffix = [_batch_values(batch) for batch in iterator]
+ restored = make_loader()
+ restored.load_state_dict(state)
+ assert [_batch_values(batch) for batch in restored] == suffix
+ for episodes, steps, _, _ in prefix + suffix:
+ all_ids.extend(zip(episodes, steps))
+ assert len(all_ids) == len(set(all_ids)) == 44
diff --git a/paimon-python/pyproject.toml b/paimon-python/pyproject.toml
index 7d4c1b9ed0..c8761d5856 100644
--- a/paimon-python/pyproject.toml
+++ b/paimon-python/pyproject.toml
@@ -201,6 +201,7 @@ dev = [
torch = [
{include-group = "python-tests"},
"pypaimon[torch,lerobot]",
+ "torchdata>=0.11,<0.12; python_version >= '3.10'",
]
legacy = [
{include-group = "test"},
