From d1aa2de205f72ba344c733622d4930294d7e5499 Mon Sep 17 00:00:00 2001 From: Yann Date: Sat, 17 Oct 2026 16:47:45 +0800 Subject: [PATCH 1/2] feat(python): support CWD checkpoints with StatefulDataLoader Validate the pinned dataset contract and restore worker RNG and callback state through TorchData, with checkpoint integration tests and trainer guidance. Co-Authored-By: Codex AI-Model: gpt-6 Co-Authored-By: Codex AI-Contributed/Feature: 228/228 AI-Contributed/UT: 247/247 --- docs/docs/pypaimon/pytorch.md | 116 ++++++++ .../pypaimon/multimodal/window_dataset.py | 111 +++++++- .../tests/contiguous_window_resume_test.py | 247 ++++++++++++++++++ paimon-python/pyproject.toml | 1 + 4 files changed, 474 insertions(+), 1 deletion(-) create mode 100644 paimon-python/pypaimon/tests/contiguous_window_resume_test.py diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 5deedfe8dacb..cc3b94e08bc9 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 4a8209c955ee..c1658e8d9aa7 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.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 @@ def __init__( 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,94 @@ def __getstate__(self): 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": { + 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])) + # 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 +590,23 @@ def _plan_for(self, columns): 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 000000000000..0b75b211c713 --- /dev/null +++ b/paimon-python/pypaimon/tests/contiguous_window_resume_test.py @@ -0,0 +1,247 @@ +# 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 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 + + +@pytest.fixture +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()) + + +@pytest.mark.parametrize("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])) + + +@pytest.mark.parametrize("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()) + + +@pytest.mark.parametrize("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 7d4c1b9ed0b5..c8761d58561f 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"}, From f34c48e6b035b11952487ffe8558bde36f47d818 Mon Sep 17 00:00:00 2001 From: Yann Date: Mon, 19 Oct 2026 22:50:57 +0800 Subject: [PATCH 2/2] fix(python): preserve callback dictionaries during CWD resume Keep callback state opaque to TorchData worker dictionary flattening, so dictionary subclasses retain their metadata after loader restoration. Add a spawn-worker regression for per-episode defaultdict counters. Co-Authored-By: Codex AI-Model: gpt-6 Co-Authored-By: Codex AI-Contributed/Feature: 5/5 AI-Contributed/UT: 45/45 --- .../pypaimon/multimodal/window_dataset.py | 5 ++- .../tests/contiguous_window_resume_test.py | 45 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index c1658e8d9aa7..f2240557e3de 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -278,7 +278,8 @@ def state_dict(self): "version": 1, "identity": self._state_identity(), "callbacks": { - name: copy.deepcopy(callback.state_dict()) + # 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(), @@ -298,7 +299,7 @@ def load_state_dict(self, state_dict): 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])) + 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"]) diff --git a/paimon-python/pypaimon/tests/contiguous_window_resume_test.py b/paimon-python/pypaimon/tests/contiguous_window_resume_test.py index 0b75b211c713..d4d567466160 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_resume_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_resume_test.py @@ -16,6 +16,7 @@ # under the License. import random +from collections import defaultdict from unittest.mock import patch import numpy as np @@ -195,6 +196,50 @@ def _batch_values(batch): 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 + + @pytest.mark.parametrize("workers,persistent,completed", [ (0, False, 3), (2, False, 3), (2, True, 3), (0, False, 11), (2, True, 11), ])