[python] Support CWD checkpoints with StatefulDataLoader

[YannByron]

Purpose

Allow ContiguousWindowDataset training to resume through TorchData's StatefulDataLoader, including shuffled, prefetched batches and random transforms. Saving only the model and a sampler cursor can resume from the wrong batch because workers have already prefetched later samples.

CWD now implements state_dict() / load_state_dict() to validate the pinned table/snapshot, actual context and anchor index, and window configuration, and to restore Python/NumPy/Torch CPU RNGs and stateful transform/adapter callbacks. Callback states are wrapped as opaque tuple leaves so TorchData's dictionary flattening preserves dictionary subclasses and their metadata. The immutable index identity is cached after its first checkpoint. Reader caches and decoder handles are rebuilt.

This reuses TorchData's loader and samplers. TorchData is added only to the torch test dependency group. The documentation shows saving model/optimizer/loader/RNG state at a completed optimizer step, including the iterator-before-generator restoration order needed for the next epoch's shuffle.

The initial contract requires the same retained snapshot, preprocessing, loader configuration, framework versions and worker/rank topology. Cross-rank checkpoint coordination and model execution determinism remain the trainer's responsibility.

Tests

• 20 checkpoint tests pass: changed snapshot/context/anchors/offsets/padding rejection, retained tags, value-based Tensor/NumPy padding identity, callback/RNG restoration, dictionary-subclass preservation through worker snapshots, cached identity, shuffled resume with 0/2 workers, spawn/prefetch/persistent workers, the last yielded batch and next epoch, and per-rank StatefulDistributedSampler progress.
• 73 tests pass across the CWD checkpoint, dataset and window-transform suites. The new regression checkpoints after three batches with one spawn worker and restores a per-episode defaultdict(int) counter; the original code fails with KeyError on the first unseen episode.
• Real RoboMIND JPEG smoke test in three independent processes: continuous 20-step CPU Conv/Dropout/AdamW training versus exit after step 7 and fresh-process resume. All batch tensors, losses, per-step model/optimizer states and final RNG states match exactly, including the next epoch. This is a recovery correctness check, not a GPU/DDP, ACT or throughput benchmark.
• Flake8 with the repository configuration and git diff --check pass. Local integration environment: Python 3.13, Torch 2.10.0, TorchData 0.11.0.