Skip to content

[python][torch] Support shuffle for Arrow and tensor batch streaming - #10012

Open
TheR1sing3un wants to merge 1 commit into
apache:masterfrom
TheR1sing3un:codex/torch-batch-shuffle
Open

TheR1sing3un wants to merge 1 commit into
apache:masterfrom
TheR1sing3un:codex/torch-batch-shuffle

Conversation

@TheR1sing3un

Copy link
Copy Markdown
Member

Purpose

Torch streaming currently rejects shuffle=True with batch_format="pyarrow" or "torch". Support bounded row shuffling for both batch formats so training can mix samples across reader batches while keeping payloads in Arrow until tensor conversion.

Reuse the existing seed/epoch and rank/worker semantics. Incoming Arrow blocks replace random slots in a rolling buffer, and bounded split interleaving mixes input sources. Filters and sharding precede shuffle; a binding limit retains the existing ordered selection. Output batching and custom tensor conversion run afterward. Close all active readers on completion, early termination, and read/conversion failures.

The default unshuffled path and existing row shuffle order are preserved. Document that this is worker-local buffer shuffling, with additional memory for incoming blocks, gathered output and format readers, and that prefetch_concurrency=1 still applies. Related to #9365.

Tests

  • Torch suite: 69 tests passed initially; the existing standalone torchrun test hit this machine's hostname-resolution timeout, also reproduced on unchanged master. It passed when rerun with PET_LOCAL_ADDR=127.0.0.1 GLOO_SOCKET_IFNAME=lo0.
  • Arrow 16 / NumPy 1.24 compatibility: 20 focused tests passed. Main environment: Python 3.11, Torch 2.8, Arrow 19.
  • Covers Arrow/Tensor order parity, existing row-shuffle parity, nested/null data, bounded input consumption and reader counts, Arrow offset overflow, empty input, filters/limits, vector tensors, deletions, historical snapshots and persistent spawn workers across two explicit ranks.
  • Repository-configured Flake8, license checks, Python 3.6 grammar checks and git diff --check passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant