JingsongLi commented on code in PR #10628:
URL: https://github.com/apache/paimon/pull/10628#discussion_r4302962029
##########
paimon-python/pypaimon/multimodal/window_dataset.py:
##########
@@ -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()
+ },
Review Comment:
[P2] Preserve callback state through TorchData's dictionary flattening
Passing the callback's state dictionary through unchanged is not lossless
with TorchData 0.11: its worker snapshot code recursively flattens dictionaries
and reconstructs them as plain `dict`, discarding dictionary subclasses and
their metadata. I reproduced this with an adapter that stores per-episode
counters in `defaultdict(int)`, returns `{"counts": self.counts}` from
`state_dict()`, and restores it in `load_state_dict()`. A direct CWD roundtrip
succeeds, but checkpointing a `StatefulDataLoader(num_workers=1,
multiprocessing_context="spawn", batch_size=4)` after three batches and
restoring it raises `KeyError: 'b'` on the first unseen episode; the
uninterrupted suffix completes successfully. Both callback methods are
implemented and the dataset/configuration are unchanged.
Please encapsulate each callback state as an opaque value that TorchData
will not recursively flatten, restore the original state before invoking the
callback's `load_state_dict()`, and add a loader-level regression test covering
this case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
