YannByron opened a new pull request, #10628: URL: https://github.com/apache/paimon/pull/10628
### 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. 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 - **19 new tests pass**: changed snapshot/context/anchors/offsets/padding rejection, retained tags, value-based Tensor/NumPy padding identity, callback/RNG restoration, cached identity, shuffled resume with 0/2 workers, spawn/prefetch/persistent workers, the last yielded batch and next epoch, and per-rank `StatefulDistributedSampler` progress. - Broader CWD/window-transform/TorchDataset regression run: **143 passed, 2 failed**. Both failures reproduce on unmodified base `220f0fe5fc`: the old native wheel lacks `Plan.stats()` (tracked separately in #10620), and local macOS `torchrun --standalone` times out during rendezvous. The final padding refinements are covered by the final 19-test run. - 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. -- 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]
