From 7dcf4cfdd036ddfe886607a6b6c2d3292960a6b4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 27 Aug 2026 03:34:49 -0700 Subject: [PATCH 01/10] [python][torch] Support distributed iterable dataset sharding --- docs/docs/pypaimon/pytorch.md | 26 ++ .../pypaimon/read/datasource/torch_dataset.py | 158 +++++-- paimon-python/pypaimon/read/table_read.py | 26 +- .../torch_distributed_sharding_worker.py | 76 ++++ .../pypaimon/tests/torch_read_test.py | 387 +++++++++++++++++- 5 files changed, 641 insertions(+), 32 deletions(-) create mode 100644 paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 9488b20fda1e..a19c4e9546cd 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -60,6 +60,32 @@ when it is false, it will read the full amount of data into memory. **`prefetch_concurrency`** (default: 1): In streaming row mode, controls reader threads per DataLoader worker. It has no effect in non-streaming mode. +### Distributed Sharding + +Streaming datasets can shard splits across both DDP ranks and DataLoader +workers: + +```python +dataset = table_read.to_torch( + splits, + streaming=True, + auto_detect_rank=True, +) +dataloader = DataLoader(dataset, batch_size=32, num_workers=2) +``` + +`auto_detect_rank=False` is the default and preserves the existing worker-only +behavior. Use it when the application already shards `splits` by rank. When +enabled, PyPaimon resolves rank and world size from an initialized +`torch.distributed` process group or `RANK` and `WORLD_SIZE`. Explicit `rank` +and `world_size` arguments take precedence and do not require automatic +detection. PyPaimon first assigns a balanced slice to the rank, then balances +that rank's splits across its DataLoader workers. + +Do not combine `auto_detect_rank=True` with application-side code such as +`splits[rank::world_size]`; remove that slicing or keep automatic detection +disabled to avoid double sharding. + ### Batch Streaming For batch-oriented training, make the streaming dataset yield batches directly: diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index a8c4f4f7cab4..f2a77b1cc809 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -18,6 +18,7 @@ """ Module to read a Paimon table into PyTorch Dataset. """ +import os import queue import random import threading @@ -40,6 +41,76 @@ def _share_epoch_with_torch_workers(value): return torch.tensor(value, dtype=torch.long).share_memory_() +def _validate_distributed_context(rank: int, world_size: int, source: str): + if isinstance(rank, bool) or not isinstance(rank, int): + raise ValueError("%s rank must be an int" % source) + if isinstance(world_size, bool) or not isinstance(world_size, int): + raise ValueError("%s world_size must be an int" % source) + if world_size <= 0: + raise ValueError("%s world_size must be greater than 0" % source) + if rank < 0 or rank >= world_size: + raise ValueError( + "%s rank must satisfy 0 <= rank < world_size" % source + ) + return rank, world_size + + +def _resolve_distributed_context( + auto_detect_rank: bool, + rank: Optional[int], + world_size: Optional[int], +): + if not isinstance(auto_detect_rank, bool): + raise ValueError("auto_detect_rank must be a bool") + if (rank is None) != (world_size is None): + raise ValueError("rank and world_size must be provided together") + + if rank is not None: + return _validate_distributed_context(rank, world_size, "explicit") + + if not auto_detect_rank: + return 0, 1 + + distributed = getattr(torch, "distributed", None) + if ( + distributed is not None + and distributed.is_available() + and distributed.is_initialized() + ): + return _validate_distributed_context( + distributed.get_rank(), + distributed.get_world_size(), + "torch.distributed", + ) + + env_rank = os.environ.get("RANK") + env_world_size = os.environ.get("WORLD_SIZE") + if env_rank is not None or env_world_size is not None: + if env_rank is None or env_world_size is None: + raise ValueError( + "RANK and WORLD_SIZE environment variables must be set together" + ) + try: + parsed_rank = int(env_rank) + parsed_world_size = int(env_world_size) + except ValueError: + raise ValueError( + "RANK and WORLD_SIZE environment variables must be integers" + ) + return _validate_distributed_context( + parsed_rank, parsed_world_size, "environment" + ) + + return 0, 1 + + +def _balanced_slice(values: List[Any], shard_id: int, shard_count: int): + base_size, remainder = divmod(len(values), shard_count) + start = shard_id * base_size + min(shard_id, remainder) + size = base_size + (1 if shard_id < remainder else 0) + return values[start:start + size] + + class TorchDataset(Dataset): """ PyTorch Dataset implementation for reading Paimon table data. @@ -92,10 +163,21 @@ class _BaseTorchIterDataset(IterableDataset): Shared helpers for streaming PyTorch datasets backed by Paimon splits. """ - def __init__(self, table_read: TableRead, splits: List[Split]): + def __init__( + self, + table_read: TableRead, + splits: List[Split], + auto_detect_rank: bool = False, + rank: Optional[int] = None, + world_size: Optional[int] = None, + ): self.table_read = table_read self.splits = splits self.field_names = [field.name for field in table_read.read_type] + self.auto_detect_rank = auto_detect_rank + self.rank, self.world_size = _resolve_distributed_context( + auto_detect_rank, rank, world_size + ) def _row_to_dict(self, offset_row) -> dict: row_dict = {} @@ -135,31 +217,25 @@ def _limit_covers_all_splits(self) -> bool: return False return True - def _worker_splits(self, worker_info) -> List[Split]: - if worker_info is None: - return self.splits + def _assigned_splits(self, worker_info) -> List[Split]: + worker_id = worker_info.id if worker_info is not None else 0 + num_workers = worker_info.num_workers if worker_info is not None else 1 - # DataLoader workers cannot share a limit budget that may truncate. + # Distributed consumers cannot share a limit budget that may truncate. if ( self.table_read.limit is not None and not self._limit_covers_all_splits() ): - return self.splits if worker_info.id == 0 else [] - - worker_id = worker_info.id - num_workers = worker_info.num_workers - total_splits = len(self.splits) - splits_per_worker = total_splits // num_workers - remainder = total_splits % num_workers - - if worker_id < remainder: - start_idx = worker_id * (splits_per_worker + 1) - end_idx = start_idx + splits_per_worker + 1 - else: - start_idx = worker_id * splits_per_worker + remainder - end_idx = start_idx + splits_per_worker + return ( + self.splits + if self.rank == 0 and worker_id == 0 + else [] + ) - return self.splits[start_idx:end_idx] + rank_splits = _balanced_slice( + self.splits, self.rank, self.world_size + ) + return _balanced_slice(rank_splits, worker_id, num_workers) class TorchIterDataset(_BaseTorchIterDataset): @@ -179,7 +255,15 @@ class TorchIterDataset(_BaseTorchIterDataset): _PREFETCH_GET_TIMEOUT_SEC = 300.0 _PREFETCH_JOIN_TIMEOUT_SEC = 5.0 - def __init__(self, table_read: TableRead, splits: List[Split], prefetch_concurrency: int = 1): + def __init__( + self, + table_read: TableRead, + splits: List[Split], + prefetch_concurrency: int = 1, + auto_detect_rank: bool = False, + rank: Optional[int] = None, + world_size: Optional[int] = None, + ): """ Initialize TorchIterDataset. @@ -190,7 +274,9 @@ def __init__(self, table_read: TableRead, splits: List[Split], prefetch_concurre this worker (default 1). When > 1, splits are partitioned across threads to increase read throughput. """ - super().__init__(table_read, splits) + super().__init__( + table_read, splits, auto_detect_rank, rank, world_size + ) self.prefetch_concurrency = max(1, int(prefetch_concurrency)) def __iter__(self): @@ -204,7 +290,7 @@ def __iter__(self): row data of dict type, where keys are column names """ worker_info = torch.utils.data.get_worker_info() - splits_to_process = self._worker_splits(worker_info) + splits_to_process = self._assigned_splits(worker_info) if self.prefetch_concurrency > 1: for row in self._iter_rows(splits_to_process): @@ -393,15 +479,20 @@ def __init__( batch_format: str, batch_size: Optional[int], to_tensor_fn: Optional[Callable[[pa.RecordBatch], Any]] = None, + auto_detect_rank: bool = False, + rank: Optional[int] = None, + world_size: Optional[int] = None, ): - super().__init__(table_read, splits) + super().__init__( + table_read, splits, auto_detect_rank, rank, world_size + ) self.batch_format = batch_format self.batch_size = batch_size self.to_tensor_fn = to_tensor_fn def __iter__(self): worker_info = torch.utils.data.get_worker_info() - splits_to_process = self._worker_splits(worker_info) + splits_to_process = self._assigned_splits(worker_info) raw_batches = self._arrow_batches_for_splits(splits_to_process) batches = _sized_record_batches( @@ -457,8 +548,13 @@ def __init__( seed: int = 0, buffer_size: int = 1000, max_buffer_input_splits: int = 10, + auto_detect_rank: bool = False, + rank: Optional[int] = None, + world_size: Optional[int] = None, ): - super().__init__(table_read, splits) + super().__init__( + table_read, splits, auto_detect_rank, rank, world_size + ) self.seed = self._require_int(seed, "seed") self.buffer_size = self._require_positive_int(buffer_size, "buffer_size") self.max_buffer_input_splits = self._require_positive_int( @@ -497,7 +593,7 @@ def set_epoch(self, epoch: int) -> "TorchShuffledIterDataset": def __iter__(self): worker_info = torch.utils.data.get_worker_info() worker_id = worker_info.id if worker_info is not None else 0 - splits_to_process = self._worker_splits(worker_info) + splits_to_process = self._assigned_splits(worker_info) if self.max_buffer_input_splits == 1: rows = self._iter_ordered_rows(splits_to_process) @@ -559,7 +655,13 @@ def _iter_buffer_shuffled_rows( rows: Iterator[dict], worker_id: int, ) -> Iterator[dict]: - rng = random.Random(self.seed + self.epoch * 1000003 + worker_id) + if self.world_size == 1: + rng_seed = self.seed + self.epoch * 1000003 + worker_id + else: + rng_seed = "%d:%d:%d:%d" % ( + self.seed, self.epoch, self.rank, worker_id + ) + rng = random.Random(rng_seed) buffer = [] for row in rows: if len(buffer) < self.buffer_size: diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index a8fcf92bb333..ddf7109cc26d 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -661,6 +661,9 @@ def to_torch( seed: int = 0, buffer_size: int = 1000, max_buffer_input_splits: int = 10, + auto_detect_rank: bool = False, + rank: Optional[int] = None, + world_size: Optional[int] = None, ) -> "torch.utils.data.Dataset": """Wrap Paimon table data in a PyTorch Dataset. @@ -674,6 +677,10 @@ def to_torch( batch_size: Rows per batch; ``None`` preserves reader batches. to_tensor_fn: Optional RecordBatch converter for Torch batches. shuffle: Whether to shuffle rows; supported only in row format. + auto_detect_rank: Whether to detect DDP rank and world size from + torch.distributed or torchrun environment variables. + rank: Optional explicit DDP rank for distributed sharding. + world_size: Optional explicit DDP world size. """ valid_batch_formats = {"row", "pyarrow", "torch"} if batch_format not in valid_batch_formats: @@ -725,6 +732,9 @@ def to_torch( batch_format=batch_format, batch_size=batch_size, to_tensor_fn=to_tensor_fn, + auto_detect_rank=auto_detect_rank, + rank=rank, + world_size=world_size, ) if shuffle: @@ -739,14 +749,28 @@ def to_torch( seed=seed, buffer_size=buffer_size, max_buffer_input_splits=max_buffer_input_splits, + auto_detect_rank=auto_detect_rank, + rank=rank, + world_size=world_size, ) return dataset if streaming: from pypaimon.read.datasource.torch_dataset import TorchIterDataset - dataset = TorchIterDataset(self, splits, prefetch_concurrency) + dataset = TorchIterDataset( + self, + splits, + prefetch_concurrency, + auto_detect_rank=auto_detect_rank, + rank=rank, + world_size=world_size, + ) return dataset else: + if auto_detect_rank or rank is not None or world_size is not None: + raise ValueError( + "distributed sharding requires streaming=True" + ) from pypaimon.read.datasource.torch_dataset import TorchDataset dataset = TorchDataset(self, splits) return dataset diff --git a/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py new file mode 100644 index 000000000000..d7b65e4dbb7e --- /dev/null +++ b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py @@ -0,0 +1,76 @@ +# 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 json +import os +import sys +from types import SimpleNamespace + +import torch +from torch.utils.data import DataLoader + +from pypaimon.read.datasource.torch_dataset import TorchIterDataset + + +class _OffsetRow: + def __init__(self, values): + self._values = values + + def get_field(self, index): + return self._values[index] + + +class _TableRead: + limit = None + read_type = [ + SimpleNamespace(name="split_id"), + SimpleNamespace(name="rank"), + SimpleNamespace(name="worker"), + ] + + def to_iterator(self, splits): + worker_info = torch.utils.data.get_worker_info() + worker_id = worker_info.id if worker_info is not None else 0 + rank = int(os.environ["RANK"]) + for split_id in splits: + yield _OffsetRow([split_id, rank, worker_id]) + + +def main(): + output_dir = sys.argv[1] + torch.distributed.init_process_group("gloo") + rank = torch.distributed.get_rank() + try: + dataset = TorchIterDataset( + _TableRead(), + list(range(11)), + auto_detect_rank=True, + ) + rows = list(DataLoader(dataset, batch_size=None, num_workers=2)) + with open( + os.path.join(output_dir, "rank-%d.json" % rank), + "w", + encoding="utf-8", + ) as result_file: + json.dump(rows, result_file) + torch.distributed.barrier() + finally: + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 2b3b126b6c2f..e6fa183e2076 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -15,8 +15,12 @@ # specific language governing permissions and limitations # under the License. +import json import os +import pickle import shutil +import subprocess +import sys import tempfile import unittest from types import SimpleNamespace @@ -29,9 +33,345 @@ from pypaimon import CatalogFactory, Schema +from pypaimon.read.datasource.torch_dataset import ( + TorchIterDataset, + TorchShuffledIterDataset, +) from pypaimon.table.file_store_table import FileStoreTable +class TorchDistributedShardingTest(unittest.TestCase): + @staticmethod + def _table_read(limit=None): + return SimpleNamespace(limit=limit, read_type=[]) + + @staticmethod + def _worker(worker_id, num_workers): + return SimpleNamespace(id=worker_id, num_workers=num_workers) + + def _dataset(self, splits, rank=0, world_size=1): + return TorchIterDataset( + self._table_read(), + splits, + rank=rank, + world_size=world_size, + ) + + def _assignments(self, split_count, world_size, num_workers): + splits = list(range(split_count)) + assignments = {} + for rank in range(world_size): + dataset = self._dataset(splits, rank, world_size) + for worker_id in range(num_workers): + assignments[(rank, worker_id)] = dataset._assigned_splits( + self._worker(worker_id, num_workers) + ) + return assignments + + def assertCompleteNonOverlapping(self, assignments, expected): + assigned = [ + split + for splits in assignments.values() + for split in splits + ] + self.assertCountEqual(assigned, expected) + self.assertEqual(len(assigned), len(set(assigned))) + + def test_world_size_one_single_worker_returns_all_splits(self): + splits = list(range(7)) + dataset = self._dataset(splits) + self.assertEqual(dataset._assigned_splits(None), splits) + + def test_world_size_one_preserves_worker_sharding(self): + assignments = self._assignments(10, world_size=1, num_workers=3) + self.assertEqual( + list(assignments.values()), + [list(range(4)), list(range(4, 7)), list(range(7, 10))], + ) + + def test_multiple_ranks_single_worker(self): + assignments = self._assignments(10, world_size=3, num_workers=1) + self.assertCompleteNonOverlapping(assignments, list(range(10))) + self.assertEqual( + list(assignments.values()), + [list(range(4)), list(range(4, 7)), list(range(7, 10))], + ) + + def test_multiple_ranks_and_workers(self): + assignments = self._assignments(17, world_size=3, num_workers=2) + self.assertCompleteNonOverlapping(assignments, list(range(17))) + sizes = [len(splits) for splits in assignments.values()] + self.assertLessEqual(max(sizes) - min(sizes), 1) + + def test_uneven_and_sparse_assignments(self): + uneven = self._assignments(11, world_size=2, num_workers=2) + self.assertCompleteNonOverlapping(uneven, list(range(11))) + self.assertLessEqual( + max(map(len, uneven.values())) - min(map(len, uneven.values())), + 1, + ) + + sparse = self._assignments(3, world_size=2, num_workers=3) + self.assertCompleteNonOverlapping(sparse, list(range(3))) + self.assertTrue(any(not splits for splits in sparse.values())) + + def test_binding_limit_uses_one_distributed_consumer(self): + splits = [SimpleNamespace(row_count=10) for _ in range(4)] + assignments = {} + for rank in range(2): + dataset = TorchIterDataset( + self._table_read(limit=5), + splits, + rank=rank, + world_size=2, + ) + for worker_id in range(2): + assignments[(rank, worker_id)] = dataset._assigned_splits( + self._worker(worker_id, 2) + ) + + self.assertEqual(assignments[(0, 0)], splits) + self.assertTrue( + all( + not assigned + for consumer, assigned in assignments.items() + if consumer != (0, 0) + ) + ) + + def test_explicit_context_has_highest_priority(self): + with patch.dict( + os.environ, {"RANK": "4", "WORLD_SIZE": "5"}, clear=True + ), patch.object( + torch.distributed, "is_available", return_value=True + ), patch.object( + torch.distributed, "is_initialized", return_value=True + ), patch.object( + torch.distributed, "get_rank", return_value=2 + ), patch.object( + torch.distributed, "get_world_size", return_value=3 + ): + dataset = TorchIterDataset( + self._table_read(), + list(range(6)), + auto_detect_rank=True, + rank=1, + world_size=2, + ) + + self.assertEqual((dataset.rank, dataset.world_size), (1, 2)) + + def test_initialized_distributed_context_precedes_environment(self): + with patch.dict( + os.environ, {"RANK": "4", "WORLD_SIZE": "5"}, clear=True + ), patch.object( + torch.distributed, "is_available", return_value=True + ), patch.object( + torch.distributed, "is_initialized", return_value=True + ), patch.object( + torch.distributed, "get_rank", return_value=1 + ), patch.object( + torch.distributed, "get_world_size", return_value=3 + ): + dataset = TorchIterDataset( + self._table_read(), + list(range(6)), + auto_detect_rank=True, + ) + + self.assertEqual((dataset.rank, dataset.world_size), (1, 3)) + + def test_worker_process_can_resolve_torchrun_environment(self): + with patch.dict( + os.environ, {"RANK": "2", "WORLD_SIZE": "4"}, clear=True + ), patch.object( + torch.distributed, "is_available", return_value=True + ), patch.object( + torch.distributed, "is_initialized", return_value=False + ): + dataset = TorchIterDataset( + self._table_read(), + list(range(8)), + auto_detect_rank=True, + ) + + self.assertEqual((dataset.rank, dataset.world_size), (2, 4)) + self.assertEqual(dataset._assigned_splits(None), [4, 5]) + + def test_auto_falls_back_to_single_process(self): + with patch.dict(os.environ, {}, clear=True), patch.object( + torch.distributed, "is_available", return_value=False + ): + dataset = TorchIterDataset( + self._table_read(), + list(range(4)), + auto_detect_rank=True, + ) + + self.assertEqual((dataset.rank, dataset.world_size), (0, 1)) + self.assertEqual(dataset._assigned_splits(None), list(range(4))) + + def test_off_preserves_existing_behavior(self): + splits = list(range(8)) + with patch.dict( + os.environ, {"RANK": "1", "WORLD_SIZE": "2"}, clear=True + ), patch.object( + torch.distributed, "is_available", return_value=True + ), patch.object( + torch.distributed, "is_initialized", return_value=True + ): + dataset = TorchIterDataset(self._table_read(), splits) + + self.assertEqual((dataset.rank, dataset.world_size), (0, 1)) + self.assertEqual( + dataset._assigned_splits(self._worker(1, 2)), + list(range(4, 8)), + ) + + def test_shuffled_dataset_is_reproducible_and_rank_local(self): + splits = list(range(20)) + datasets = [ + TorchShuffledIterDataset( + self._table_read(), + splits, + seed=17, + buffer_size=20, + rank=rank, + world_size=2, + ) + for rank in range(2) + ] + local_splits = [dataset._assigned_splits(None) for dataset in datasets] + self.assertTrue(set(local_splits[0]).isdisjoint(local_splits[1])) + self.assertCountEqual(local_splits[0] + local_splits[1], splits) + restored = pickle.loads(pickle.dumps(datasets[1])) + self.assertEqual((restored.rank, restored.world_size), (1, 2)) + + rows = [{"id": value} for value in range(20)] + first = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0)) + repeat = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0)) + other_rank = list( + datasets[1]._iter_buffer_shuffled_rows(iter(rows), 0) + ) + other_worker = list( + datasets[0]._iter_buffer_shuffled_rows(iter(rows), 1) + ) + self.assertEqual(first, repeat) + self.assertNotEqual(first, other_rank) + self.assertNotEqual(first, other_worker) + + datasets[0].set_epoch(1) + next_epoch = list( + datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0) + ) + self.assertNotEqual(first, next_epoch) + + def test_invalid_distributed_context(self): + invalid_options = [ + ({"auto_detect_rank": "auto"}, "auto_detect_rank"), + ({"rank": 0}, "provided together"), + ( + { + "rank": 0, + "world_size": 0, + }, + "greater than 0", + ), + ( + { + "rank": 2, + "world_size": 2, + }, + "0 <= rank < world_size", + ), + ] + for options, message in invalid_options: + with self.subTest(options=options), self.assertRaisesRegex( + ValueError, message + ): + TorchIterDataset(self._table_read(), [], **options) + + with patch.dict(os.environ, {"RANK": "one"}, clear=True), patch.object( + torch.distributed, "is_available", return_value=False + ), self.assertRaisesRegex(ValueError, "must be set together"): + TorchIterDataset( + self._table_read(), [], auto_detect_rank=True + ) + + with patch.dict( + os.environ, {"RANK": "one", "WORLD_SIZE": "2"}, clear=True + ), patch.object( + torch.distributed, "is_available", return_value=False + ), self.assertRaisesRegex(ValueError, "must be integers"): + TorchIterDataset( + self._table_read(), [], auto_detect_rank=True + ) + + @unittest.skipUnless( + torch.distributed.is_available(), "torch.distributed is unavailable" + ) + def test_torchrun_rank_and_worker_sharding(self): + script = os.path.join( + os.path.dirname(__file__), "torch_distributed_sharding_worker.py" + ) + python_root = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..") + ) + with tempfile.TemporaryDirectory() as output_dir: + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + filter(None, [python_root, env.get("PYTHONPATH")]) + ) + process = subprocess.run( + [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc-per-node=2", + script, + output_dir, + ], + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=180, + ) + self.assertEqual( + process.returncode, + 0, + "torchrun failed:\n%s\n%s" % ( + process.stdout, process.stderr + ), + ) + rows = [] + for rank in range(2): + with open( + os.path.join(output_dir, "rank-%d.json" % rank), + encoding="utf-8", + ) as result_file: + rows.extend(json.load(result_file)) + + split_ids = [row["split_id"] for row in rows] + self.assertCountEqual(split_ids, list(range(11))) + self.assertEqual(len(split_ids), len(set(split_ids))) + assignments = {} + for row in rows: + assignments.setdefault( + (row["rank"], row["worker"]), [] + ).append(row["split_id"]) + self.assertEqual( + {key: sorted(values) for key, values in assignments.items()}, + { + (0, 0): [0, 1, 2], + (0, 1): [3, 4, 5], + (1, 0): [6, 7, 8], + (1, 1): [9, 10], + }, + ) + + class TorchReadTest(unittest.TestCase): @classmethod def setUpClass(cls): @@ -275,7 +615,7 @@ def test_torch_streaming_batches_respect_limit_with_workers(self): batch_size=3, ) self.assertEqual( - dataset._worker_splits(SimpleNamespace(id=1, num_workers=2)), + dataset._assigned_splits(SimpleNamespace(id=1, num_workers=2)), [], ) batches = list(DataLoader( @@ -313,7 +653,7 @@ def test_non_binding_limit_preserves_worker_splits(self): batch_format=batch_format, ) assigned = [ - dataset._worker_splits( + dataset._assigned_splits( SimpleNamespace(id=worker_id, num_workers=2) ) for worker_id in range(2) @@ -335,7 +675,7 @@ def test_non_binding_limit_uses_merged_row_counts(self): dataset = TorchIterDataset(table_read, splits) assigned = [ - dataset._worker_splits( + dataset._assigned_splits( SimpleNamespace(id=worker_id, num_workers=2) ) for worker_id in range(2) @@ -468,6 +808,47 @@ def test_torch_batch_options_validation(self): prefetch_concurrency=invalid, ) + def test_torch_distributed_sharding_public_api(self): + schema = Schema.from_pyarrow_schema( + self.pa_schema, partition_keys=['user_id'] + ) + self.catalog.create_table( + 'default.test_torch_distributed_api', schema, False + ) + table = self.catalog.get_table( + 'default.test_torch_distributed_api' + ) + self._write_test_table(table) + read_builder = table.new_read_builder().with_projection(['user_id']) + splits = read_builder.new_scan().plan().splits() + table_read = read_builder.new_read() + + datasets = [ + table_read.to_torch( + splits, + streaming=True, + batch_format=batch_format, + shuffle=batch_format == 'row' and shuffle, + rank=1, + world_size=2, + ) + for batch_format, shuffle in [ + ('row', False), + ('row', True), + ('pyarrow', False), + ] + ] + for dataset in datasets: + self.assertEqual((dataset.rank, dataset.world_size), (1, 2)) + + with self.assertRaisesRegex(ValueError, 'requires streaming=True'): + table_read.to_torch( + splits, + auto_detect_rank=True, + rank=1, + world_size=2, + ) + def test_blob_torch_read(self): """Test end-to-end blob functionality using blob descriptors.""" import random From a19a0607af3ab1f801f65fbed343147c9a579fa7 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 27 Aug 2026 04:59:16 -0700 Subject: [PATCH 02/10] [python] Clarify upstream rank sharding guidance --- docs/docs/pypaimon/pytorch.md | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index a19c4e9546cd..81d476a4a004 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -82,9 +82,8 @@ and `world_size` arguments take precedence and do not require automatic detection. PyPaimon first assigns a balanced slice to the rank, then balances that rank's splits across its DataLoader workers. -Do not combine `auto_detect_rank=True` with application-side code such as -`splits[rank::world_size]`; remove that slicing or keep automatic detection -disabled to avoid double sharding. +If splits are already sharded by rank upstream, keep `auto_detect_rank=False` +or remove the upstream sharding before enabling it. ### Batch Streaming From 3966610f9d68a17c53280a86fe2027146e21cf46 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 27 Aug 2026 05:07:34 -0700 Subject: [PATCH 03/10] [python][torch] Simplify distributed sharding --- docs/docs/pypaimon/pytorch.md | 18 +-- .../pypaimon/read/datasource/torch_dataset.py | 71 +++------ .../pypaimon/tests/torch_read_test.py | 143 ++++++------------ 3 files changed, 77 insertions(+), 155 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 81d476a4a004..9ee32fc9a668 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -62,8 +62,8 @@ reader threads per DataLoader worker. It has no effect in non-streaming mode. ### Distributed Sharding -Streaming datasets can shard splits across both DDP ranks and DataLoader -workers: +Set `auto_detect_rank=True` to shard streaming reads across DDP ranks and +DataLoader workers: ```python dataset = table_read.to_torch( @@ -74,16 +74,10 @@ dataset = table_read.to_torch( dataloader = DataLoader(dataset, batch_size=32, num_workers=2) ``` -`auto_detect_rank=False` is the default and preserves the existing worker-only -behavior. Use it when the application already shards `splits` by rank. When -enabled, PyPaimon resolves rank and world size from an initialized -`torch.distributed` process group or `RANK` and `WORLD_SIZE`. Explicit `rank` -and `world_size` arguments take precedence and do not require automatic -detection. PyPaimon first assigns a balanced slice to the rank, then balances -that rank's splits across its DataLoader workers. - -If splits are already sharded by rank upstream, keep `auto_detect_rank=False` -or remove the upstream sharding before enabling it. +PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Explicit +`rank` and `world_size` take precedence. Automatic detection is disabled by +default; if splits are already sharded upstream, leave it disabled or remove +the upstream sharding. ### Batch Streaming diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index f2a77b1cc809..f5c8051d69c6 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -41,17 +41,15 @@ def _share_epoch_with_torch_workers(value): return torch.tensor(value, dtype=torch.long).share_memory_() -def _validate_distributed_context(rank: int, world_size: int, source: str): +def _validate_distributed_context(rank: int, world_size: int): if isinstance(rank, bool) or not isinstance(rank, int): - raise ValueError("%s rank must be an int" % source) + raise ValueError("rank must be an int") if isinstance(world_size, bool) or not isinstance(world_size, int): - raise ValueError("%s world_size must be an int" % source) + raise ValueError("world_size must be an int") if world_size <= 0: - raise ValueError("%s world_size must be greater than 0" % source) + raise ValueError("world_size must be greater than 0") if rank < 0 or rank >= world_size: - raise ValueError( - "%s rank must satisfy 0 <= rank < world_size" % source - ) + raise ValueError("rank must satisfy 0 <= rank < world_size") return rank, world_size @@ -66,7 +64,7 @@ def _resolve_distributed_context( raise ValueError("rank and world_size must be provided together") if rank is not None: - return _validate_distributed_context(rank, world_size, "explicit") + return _validate_distributed_context(rank, world_size) if not auto_detect_rank: return 0, 1 @@ -77,11 +75,9 @@ def _resolve_distributed_context( and distributed.is_available() and distributed.is_initialized() ): - return _validate_distributed_context( - distributed.get_rank(), - distributed.get_world_size(), - "torch.distributed", - ) + rank = distributed.get_rank() + world_size = distributed.get_world_size() + return _validate_distributed_context(rank, world_size) env_rank = os.environ.get("RANK") env_world_size = os.environ.get("WORLD_SIZE") @@ -91,15 +87,12 @@ def _resolve_distributed_context( "RANK and WORLD_SIZE environment variables must be set together" ) try: - parsed_rank = int(env_rank) - parsed_world_size = int(env_world_size) + rank, world_size = int(env_rank), int(env_world_size) except ValueError: raise ValueError( "RANK and WORLD_SIZE environment variables must be integers" ) - return _validate_distributed_context( - parsed_rank, parsed_world_size, "environment" - ) + return _validate_distributed_context(rank, world_size) return 0, 1 @@ -174,7 +167,6 @@ def __init__( self.table_read = table_read self.splits = splits self.field_names = [field.name for field in table_read.read_type] - self.auto_detect_rank = auto_detect_rank self.rank, self.world_size = _resolve_distributed_context( auto_detect_rank, rank, world_size ) @@ -217,24 +209,18 @@ def _limit_covers_all_splits(self) -> bool: return False return True - def _assigned_splits(self, worker_info) -> List[Split]: + def _worker_splits(self, worker_info) -> List[Split]: worker_id = worker_info.id if worker_info is not None else 0 num_workers = worker_info.num_workers if worker_info is not None else 1 - # Distributed consumers cannot share a limit budget that may truncate. if ( self.table_read.limit is not None and not self._limit_covers_all_splits() ): - return ( - self.splits - if self.rank == 0 and worker_id == 0 - else [] - ) + # A binding limit cannot be shared safely. + return self.splits if self.rank == 0 and worker_id == 0 else [] - rank_splits = _balanced_slice( - self.splits, self.rank, self.world_size - ) + rank_splits = _balanced_slice(self.splits, self.rank, self.world_size) return _balanced_slice(rank_splits, worker_id, num_workers) @@ -274,9 +260,7 @@ def __init__( this worker (default 1). When > 1, splits are partitioned across threads to increase read throughput. """ - super().__init__( - table_read, splits, auto_detect_rank, rank, world_size - ) + super().__init__(table_read, splits, auto_detect_rank, rank, world_size) self.prefetch_concurrency = max(1, int(prefetch_concurrency)) def __iter__(self): @@ -290,7 +274,7 @@ def __iter__(self): row data of dict type, where keys are column names """ worker_info = torch.utils.data.get_worker_info() - splits_to_process = self._assigned_splits(worker_info) + splits_to_process = self._worker_splits(worker_info) if self.prefetch_concurrency > 1: for row in self._iter_rows(splits_to_process): @@ -483,16 +467,14 @@ def __init__( rank: Optional[int] = None, world_size: Optional[int] = None, ): - super().__init__( - table_read, splits, auto_detect_rank, rank, world_size - ) + super().__init__(table_read, splits, auto_detect_rank, rank, world_size) self.batch_format = batch_format self.batch_size = batch_size self.to_tensor_fn = to_tensor_fn def __iter__(self): worker_info = torch.utils.data.get_worker_info() - splits_to_process = self._assigned_splits(worker_info) + splits_to_process = self._worker_splits(worker_info) raw_batches = self._arrow_batches_for_splits(splits_to_process) batches = _sized_record_batches( @@ -552,9 +534,7 @@ def __init__( rank: Optional[int] = None, world_size: Optional[int] = None, ): - super().__init__( - table_read, splits, auto_detect_rank, rank, world_size - ) + super().__init__(table_read, splits, auto_detect_rank, rank, world_size) self.seed = self._require_int(seed, "seed") self.buffer_size = self._require_positive_int(buffer_size, "buffer_size") self.max_buffer_input_splits = self._require_positive_int( @@ -593,7 +573,7 @@ def set_epoch(self, epoch: int) -> "TorchShuffledIterDataset": def __iter__(self): worker_info = torch.utils.data.get_worker_info() worker_id = worker_info.id if worker_info is not None else 0 - splits_to_process = self._assigned_splits(worker_info) + splits_to_process = self._worker_splits(worker_info) if self.max_buffer_input_splits == 1: rows = self._iter_ordered_rows(splits_to_process) @@ -655,12 +635,9 @@ def _iter_buffer_shuffled_rows( rows: Iterator[dict], worker_id: int, ) -> Iterator[dict]: - if self.world_size == 1: - rng_seed = self.seed + self.epoch * 1000003 + worker_id - else: - rng_seed = "%d:%d:%d:%d" % ( - self.seed, self.epoch, self.rank, worker_id - ) + rng_seed = self.seed + self.epoch * 1000003 + worker_id + if self.world_size > 1: + rng_seed = "%d:%d" % (rng_seed, self.rank) rng = random.Random(rng_seed) buffer = [] for row in rows: diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index e6fa183e2076..c2bc78f6385f 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -36,6 +36,7 @@ from pypaimon.read.datasource.torch_dataset import ( TorchIterDataset, TorchShuffledIterDataset, + _resolve_distributed_context, ) from pypaimon.table.file_store_table import FileStoreTable @@ -63,58 +64,39 @@ def _assignments(self, split_count, world_size, num_workers): for rank in range(world_size): dataset = self._dataset(splits, rank, world_size) for worker_id in range(num_workers): - assignments[(rank, worker_id)] = dataset._assigned_splits( + assignments[(rank, worker_id)] = dataset._worker_splits( self._worker(worker_id, num_workers) ) return assignments def assertCompleteNonOverlapping(self, assignments, expected): assigned = [ - split - for splits in assignments.values() - for split in splits + split for splits in assignments.values() for split in splits ] self.assertCountEqual(assigned, expected) self.assertEqual(len(assigned), len(set(assigned))) - def test_world_size_one_single_worker_returns_all_splits(self): - splits = list(range(7)) - dataset = self._dataset(splits) - self.assertEqual(dataset._assigned_splits(None), splits) - - def test_world_size_one_preserves_worker_sharding(self): - assignments = self._assignments(10, world_size=1, num_workers=3) - self.assertEqual( - list(assignments.values()), - [list(range(4)), list(range(4, 7)), list(range(7, 10))], + @parameterized.expand([ + ("single", 7, 1, 1, [7]), + ("workers", 10, 1, 3, [4, 3, 3]), + ("ranks", 10, 3, 1, [4, 3, 3]), + ("rank_workers", 17, 3, 2, [3, 3, 3, 3, 3, 2]), + ("uneven", 11, 2, 2, [3, 3, 3, 2]), + ("sparse", 3, 2, 3, [1, 1, 0, 1, 0, 0]), + ]) + def test_balanced_assignments( + self, _, split_count, world_size, num_workers, expected_sizes + ): + assignments = self._assignments( + split_count, world_size, num_workers ) - - def test_multiple_ranks_single_worker(self): - assignments = self._assignments(10, world_size=3, num_workers=1) - self.assertCompleteNonOverlapping(assignments, list(range(10))) - self.assertEqual( - list(assignments.values()), - [list(range(4)), list(range(4, 7)), list(range(7, 10))], + self.assertCompleteNonOverlapping( + assignments, list(range(split_count)) ) - - def test_multiple_ranks_and_workers(self): - assignments = self._assignments(17, world_size=3, num_workers=2) - self.assertCompleteNonOverlapping(assignments, list(range(17))) - sizes = [len(splits) for splits in assignments.values()] - self.assertLessEqual(max(sizes) - min(sizes), 1) - - def test_uneven_and_sparse_assignments(self): - uneven = self._assignments(11, world_size=2, num_workers=2) - self.assertCompleteNonOverlapping(uneven, list(range(11))) - self.assertLessEqual( - max(map(len, uneven.values())) - min(map(len, uneven.values())), - 1, + self.assertEqual( + [len(splits) for splits in assignments.values()], expected_sizes ) - sparse = self._assignments(3, world_size=2, num_workers=3) - self.assertCompleteNonOverlapping(sparse, list(range(3))) - self.assertTrue(any(not splits for splits in sparse.values())) - def test_binding_limit_uses_one_distributed_consumer(self): splits = [SimpleNamespace(row_count=10) for _ in range(4)] assignments = {} @@ -126,16 +108,15 @@ def test_binding_limit_uses_one_distributed_consumer(self): world_size=2, ) for worker_id in range(2): - assignments[(rank, worker_id)] = dataset._assigned_splits( + assignments[(rank, worker_id)] = dataset._worker_splits( self._worker(worker_id, 2) ) self.assertEqual(assignments[(0, 0)], splits) self.assertTrue( all( - not assigned - for consumer, assigned in assignments.items() - if consumer != (0, 0) + not value for key, value in assignments.items() + if key != (0, 0) ) ) @@ -151,15 +132,9 @@ def test_explicit_context_has_highest_priority(self): ), patch.object( torch.distributed, "get_world_size", return_value=3 ): - dataset = TorchIterDataset( - self._table_read(), - list(range(6)), - auto_detect_rank=True, - rank=1, - world_size=2, - ) + context = _resolve_distributed_context(True, 1, 2) - self.assertEqual((dataset.rank, dataset.world_size), (1, 2)) + self.assertEqual(context, (1, 2)) def test_initialized_distributed_context_precedes_environment(self): with patch.dict( @@ -173,13 +148,9 @@ def test_initialized_distributed_context_precedes_environment(self): ), patch.object( torch.distributed, "get_world_size", return_value=3 ): - dataset = TorchIterDataset( - self._table_read(), - list(range(6)), - auto_detect_rank=True, - ) + context = _resolve_distributed_context(True, None, None) - self.assertEqual((dataset.rank, dataset.world_size), (1, 3)) + self.assertEqual(context, (1, 3)) def test_worker_process_can_resolve_torchrun_environment(self): with patch.dict( @@ -189,27 +160,19 @@ def test_worker_process_can_resolve_torchrun_environment(self): ), patch.object( torch.distributed, "is_initialized", return_value=False ): - dataset = TorchIterDataset( - self._table_read(), - list(range(8)), - auto_detect_rank=True, - ) + context = _resolve_distributed_context(True, None, None) - self.assertEqual((dataset.rank, dataset.world_size), (2, 4)) - self.assertEqual(dataset._assigned_splits(None), [4, 5]) + self.assertEqual(context, (2, 4)) + dataset = self._dataset(list(range(8)), *context) + self.assertEqual(dataset._worker_splits(None), [4, 5]) def test_auto_falls_back_to_single_process(self): with patch.dict(os.environ, {}, clear=True), patch.object( torch.distributed, "is_available", return_value=False ): - dataset = TorchIterDataset( - self._table_read(), - list(range(4)), - auto_detect_rank=True, - ) + context = _resolve_distributed_context(True, None, None) - self.assertEqual((dataset.rank, dataset.world_size), (0, 1)) - self.assertEqual(dataset._assigned_splits(None), list(range(4))) + self.assertEqual(context, (0, 1)) def test_off_preserves_existing_behavior(self): splits = list(range(8)) @@ -224,7 +187,7 @@ def test_off_preserves_existing_behavior(self): self.assertEqual((dataset.rank, dataset.world_size), (0, 1)) self.assertEqual( - dataset._assigned_splits(self._worker(1, 2)), + dataset._worker_splits(self._worker(1, 2)), list(range(4, 8)), ) @@ -241,7 +204,7 @@ def test_shuffled_dataset_is_reproducible_and_rank_local(self): ) for rank in range(2) ] - local_splits = [dataset._assigned_splits(None) for dataset in datasets] + local_splits = [dataset._worker_splits(None) for dataset in datasets] self.assertTrue(set(local_splits[0]).isdisjoint(local_splits[1])) self.assertCountEqual(local_splits[0] + local_splits[1], splits) restored = pickle.loads(pickle.dumps(datasets[1])) @@ -270,42 +233,30 @@ def test_invalid_distributed_context(self): invalid_options = [ ({"auto_detect_rank": "auto"}, "auto_detect_rank"), ({"rank": 0}, "provided together"), - ( - { - "rank": 0, - "world_size": 0, - }, - "greater than 0", - ), - ( - { - "rank": 2, - "world_size": 2, - }, - "0 <= rank < world_size", - ), + ({"rank": 0, "world_size": 0}, "greater than 0"), + ({"rank": 2, "world_size": 2}, "0 <= rank < world_size"), ] for options, message in invalid_options: with self.subTest(options=options), self.assertRaisesRegex( ValueError, message ): - TorchIterDataset(self._table_read(), [], **options) + _resolve_distributed_context( + options.get("auto_detect_rank", False), + options.get("rank"), + options.get("world_size"), + ) with patch.dict(os.environ, {"RANK": "one"}, clear=True), patch.object( torch.distributed, "is_available", return_value=False ), self.assertRaisesRegex(ValueError, "must be set together"): - TorchIterDataset( - self._table_read(), [], auto_detect_rank=True - ) + _resolve_distributed_context(True, None, None) with patch.dict( os.environ, {"RANK": "one", "WORLD_SIZE": "2"}, clear=True ), patch.object( torch.distributed, "is_available", return_value=False ), self.assertRaisesRegex(ValueError, "must be integers"): - TorchIterDataset( - self._table_read(), [], auto_detect_rank=True - ) + _resolve_distributed_context(True, None, None) @unittest.skipUnless( torch.distributed.is_available(), "torch.distributed is unavailable" @@ -615,7 +566,7 @@ def test_torch_streaming_batches_respect_limit_with_workers(self): batch_size=3, ) self.assertEqual( - dataset._assigned_splits(SimpleNamespace(id=1, num_workers=2)), + dataset._worker_splits(SimpleNamespace(id=1, num_workers=2)), [], ) batches = list(DataLoader( @@ -653,7 +604,7 @@ def test_non_binding_limit_preserves_worker_splits(self): batch_format=batch_format, ) assigned = [ - dataset._assigned_splits( + dataset._worker_splits( SimpleNamespace(id=worker_id, num_workers=2) ) for worker_id in range(2) @@ -675,7 +626,7 @@ def test_non_binding_limit_uses_merged_row_counts(self): dataset = TorchIterDataset(table_read, splits) assigned = [ - dataset._assigned_splits( + dataset._worker_splits( SimpleNamespace(id=worker_id, num_workers=2) ) for worker_id in range(2) From 2b893054f93ce008257907865fb90b92b3b1b5c5 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 27 Aug 2026 05:25:13 -0700 Subject: [PATCH 04/10] [python][torch] Keep rank detection internal --- docs/docs/pypaimon/pytorch.md | 7 +- .../pypaimon/read/datasource/torch_dataset.py | 30 +--- paimon-python/pypaimon/read/table_read.py | 12 +- .../pypaimon/tests/torch_read_test.py | 135 ++++++++---------- 4 files changed, 71 insertions(+), 113 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 9ee32fc9a668..1e486f9ffcff 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -74,10 +74,9 @@ dataset = table_read.to_torch( dataloader = DataLoader(dataset, batch_size=32, num_workers=2) ``` -PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Explicit -`rank` and `world_size` take precedence. Automatic detection is disabled by -default; if splits are already sharded upstream, leave it disabled or remove -the upstream sharding. +PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Detection is +disabled by default; if splits are already sharded upstream, leave it disabled +or remove the upstream sharding. ### Batch Streaming diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index f5c8051d69c6..ee9d9e9e6f75 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -53,19 +53,9 @@ def _validate_distributed_context(rank: int, world_size: int): return rank, world_size -def _resolve_distributed_context( - auto_detect_rank: bool, - rank: Optional[int], - world_size: Optional[int], -): +def _resolve_distributed_context(auto_detect_rank: bool): if not isinstance(auto_detect_rank, bool): raise ValueError("auto_detect_rank must be a bool") - if (rank is None) != (world_size is None): - raise ValueError("rank and world_size must be provided together") - - if rank is not None: - return _validate_distributed_context(rank, world_size) - if not auto_detect_rank: return 0, 1 @@ -161,15 +151,11 @@ def __init__( table_read: TableRead, splits: List[Split], auto_detect_rank: bool = False, - rank: Optional[int] = None, - world_size: Optional[int] = None, ): self.table_read = table_read self.splits = splits self.field_names = [field.name for field in table_read.read_type] - self.rank, self.world_size = _resolve_distributed_context( - auto_detect_rank, rank, world_size - ) + self.rank, self.world_size = _resolve_distributed_context(auto_detect_rank) def _row_to_dict(self, offset_row) -> dict: row_dict = {} @@ -247,8 +233,6 @@ def __init__( splits: List[Split], prefetch_concurrency: int = 1, auto_detect_rank: bool = False, - rank: Optional[int] = None, - world_size: Optional[int] = None, ): """ Initialize TorchIterDataset. @@ -260,7 +244,7 @@ def __init__( this worker (default 1). When > 1, splits are partitioned across threads to increase read throughput. """ - super().__init__(table_read, splits, auto_detect_rank, rank, world_size) + super().__init__(table_read, splits, auto_detect_rank) self.prefetch_concurrency = max(1, int(prefetch_concurrency)) def __iter__(self): @@ -464,10 +448,8 @@ def __init__( batch_size: Optional[int], to_tensor_fn: Optional[Callable[[pa.RecordBatch], Any]] = None, auto_detect_rank: bool = False, - rank: Optional[int] = None, - world_size: Optional[int] = None, ): - super().__init__(table_read, splits, auto_detect_rank, rank, world_size) + super().__init__(table_read, splits, auto_detect_rank) self.batch_format = batch_format self.batch_size = batch_size self.to_tensor_fn = to_tensor_fn @@ -531,10 +513,8 @@ def __init__( buffer_size: int = 1000, max_buffer_input_splits: int = 10, auto_detect_rank: bool = False, - rank: Optional[int] = None, - world_size: Optional[int] = None, ): - super().__init__(table_read, splits, auto_detect_rank, rank, world_size) + super().__init__(table_read, splits, auto_detect_rank) self.seed = self._require_int(seed, "seed") self.buffer_size = self._require_positive_int(buffer_size, "buffer_size") self.max_buffer_input_splits = self._require_positive_int( diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index ddf7109cc26d..09f4421836d7 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -662,8 +662,6 @@ def to_torch( buffer_size: int = 1000, max_buffer_input_splits: int = 10, auto_detect_rank: bool = False, - rank: Optional[int] = None, - world_size: Optional[int] = None, ) -> "torch.utils.data.Dataset": """Wrap Paimon table data in a PyTorch Dataset. @@ -679,8 +677,6 @@ def to_torch( shuffle: Whether to shuffle rows; supported only in row format. auto_detect_rank: Whether to detect DDP rank and world size from torch.distributed or torchrun environment variables. - rank: Optional explicit DDP rank for distributed sharding. - world_size: Optional explicit DDP world size. """ valid_batch_formats = {"row", "pyarrow", "torch"} if batch_format not in valid_batch_formats: @@ -733,8 +729,6 @@ def to_torch( batch_size=batch_size, to_tensor_fn=to_tensor_fn, auto_detect_rank=auto_detect_rank, - rank=rank, - world_size=world_size, ) if shuffle: @@ -750,8 +744,6 @@ def to_torch( buffer_size=buffer_size, max_buffer_input_splits=max_buffer_input_splits, auto_detect_rank=auto_detect_rank, - rank=rank, - world_size=world_size, ) return dataset @@ -762,12 +754,10 @@ def to_torch( splits, prefetch_concurrency, auto_detect_rank=auto_detect_rank, - rank=rank, - world_size=world_size, ) return dataset else: - if auto_detect_rank or rank is not None or world_size is not None: + if auto_detect_rank: raise ValueError( "distributed sharding requires streaming=True" ) diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index c2bc78f6385f..4c957a2848a2 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -50,13 +50,26 @@ def _table_read(limit=None): def _worker(worker_id, num_workers): return SimpleNamespace(id=worker_id, num_workers=num_workers) - def _dataset(self, splits, rank=0, world_size=1): - return TorchIterDataset( - self._table_read(), - splits, - rank=rank, - world_size=world_size, - ) + def _dataset( + self, + splits, + rank=0, + world_size=1, + limit=None, + dataset_type=TorchIterDataset, + **kwargs + ): + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(rank, world_size), + ): + return dataset_type( + self._table_read(limit), + splits, + auto_detect_rank=True, + **kwargs + ) def _assignments(self, split_count, world_size, num_workers): splits = list(range(split_count)) @@ -101,12 +114,7 @@ def test_binding_limit_uses_one_distributed_consumer(self): splits = [SimpleNamespace(row_count=10) for _ in range(4)] assignments = {} for rank in range(2): - dataset = TorchIterDataset( - self._table_read(limit=5), - splits, - rank=rank, - world_size=2, - ) + dataset = self._dataset(splits, rank, 2, limit=5) for worker_id in range(2): assignments[(rank, worker_id)] = dataset._worker_splits( self._worker(worker_id, 2) @@ -120,22 +128,6 @@ def test_binding_limit_uses_one_distributed_consumer(self): ) ) - def test_explicit_context_has_highest_priority(self): - with patch.dict( - os.environ, {"RANK": "4", "WORLD_SIZE": "5"}, clear=True - ), patch.object( - torch.distributed, "is_available", return_value=True - ), patch.object( - torch.distributed, "is_initialized", return_value=True - ), patch.object( - torch.distributed, "get_rank", return_value=2 - ), patch.object( - torch.distributed, "get_world_size", return_value=3 - ): - context = _resolve_distributed_context(True, 1, 2) - - self.assertEqual(context, (1, 2)) - def test_initialized_distributed_context_precedes_environment(self): with patch.dict( os.environ, {"RANK": "4", "WORLD_SIZE": "5"}, clear=True @@ -148,7 +140,7 @@ def test_initialized_distributed_context_precedes_environment(self): ), patch.object( torch.distributed, "get_world_size", return_value=3 ): - context = _resolve_distributed_context(True, None, None) + context = _resolve_distributed_context(True) self.assertEqual(context, (1, 3)) @@ -160,7 +152,7 @@ def test_worker_process_can_resolve_torchrun_environment(self): ), patch.object( torch.distributed, "is_initialized", return_value=False ): - context = _resolve_distributed_context(True, None, None) + context = _resolve_distributed_context(True) self.assertEqual(context, (2, 4)) dataset = self._dataset(list(range(8)), *context) @@ -170,7 +162,7 @@ def test_auto_falls_back_to_single_process(self): with patch.dict(os.environ, {}, clear=True), patch.object( torch.distributed, "is_available", return_value=False ): - context = _resolve_distributed_context(True, None, None) + context = _resolve_distributed_context(True) self.assertEqual(context, (0, 1)) @@ -194,13 +186,13 @@ def test_off_preserves_existing_behavior(self): def test_shuffled_dataset_is_reproducible_and_rank_local(self): splits = list(range(20)) datasets = [ - TorchShuffledIterDataset( - self._table_read(), + self._dataset( splits, + rank, + 2, + dataset_type=TorchShuffledIterDataset, seed=17, buffer_size=20, - rank=rank, - world_size=2, ) for rank in range(2) ] @@ -230,33 +222,31 @@ def test_shuffled_dataset_is_reproducible_and_rank_local(self): self.assertNotEqual(first, next_epoch) def test_invalid_distributed_context(self): - invalid_options = [ - ({"auto_detect_rank": "auto"}, "auto_detect_rank"), - ({"rank": 0}, "provided together"), - ({"rank": 0, "world_size": 0}, "greater than 0"), - ({"rank": 2, "world_size": 2}, "0 <= rank < world_size"), - ] - for options, message in invalid_options: - with self.subTest(options=options), self.assertRaisesRegex( - ValueError, message - ): - _resolve_distributed_context( - options.get("auto_detect_rank", False), - options.get("rank"), - options.get("world_size"), - ) + with self.assertRaisesRegex(ValueError, "auto_detect_rank"): + _resolve_distributed_context("auto") with patch.dict(os.environ, {"RANK": "one"}, clear=True), patch.object( torch.distributed, "is_available", return_value=False ), self.assertRaisesRegex(ValueError, "must be set together"): - _resolve_distributed_context(True, None, None) + _resolve_distributed_context(True) with patch.dict( os.environ, {"RANK": "one", "WORLD_SIZE": "2"}, clear=True ), patch.object( torch.distributed, "is_available", return_value=False ), self.assertRaisesRegex(ValueError, "must be integers"): - _resolve_distributed_context(True, None, None) + _resolve_distributed_context(True) + + for environment, message in [ + ({"RANK": "0", "WORLD_SIZE": "0"}, "greater than 0"), + ({"RANK": "2", "WORLD_SIZE": "2"}, "0 <= rank"), + ]: + with self.subTest(environment=environment), patch.dict( + os.environ, environment, clear=True + ), patch.object( + torch.distributed, "is_available", return_value=False + ), self.assertRaisesRegex(ValueError, message): + _resolve_distributed_context(True) @unittest.skipUnless( torch.distributed.is_available(), "torch.distributed is unavailable" @@ -774,31 +764,30 @@ def test_torch_distributed_sharding_public_api(self): splits = read_builder.new_scan().plan().splits() table_read = read_builder.new_read() - datasets = [ - table_read.to_torch( - splits, - streaming=True, - batch_format=batch_format, - shuffle=batch_format == 'row' and shuffle, - rank=1, - world_size=2, - ) - for batch_format, shuffle in [ - ('row', False), - ('row', True), - ('pyarrow', False), + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(1, 2), + ): + datasets = [ + table_read.to_torch( + splits, + streaming=True, + batch_format=batch_format, + shuffle=batch_format == 'row' and shuffle, + auto_detect_rank=True, + ) + for batch_format, shuffle in [ + ('row', False), + ('row', True), + ('pyarrow', False), + ] ] - ] for dataset in datasets: self.assertEqual((dataset.rank, dataset.world_size), (1, 2)) with self.assertRaisesRegex(ValueError, 'requires streaming=True'): - table_read.to_torch( - splits, - auto_detect_rank=True, - rank=1, - world_size=2, - ) + table_read.to_torch(splits, auto_detect_rank=True) def test_blob_torch_read(self): """Test end-to-end blob functionality using blob descriptors.""" From 1415e0ba4150132348c1e10923317b5ed1251c89 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 27 Aug 2026 05:52:55 -0700 Subject: [PATCH 05/10] [python][torch] Enable rank detection by default --- docs/docs/pypaimon/pytorch.md | 7 ++----- .../pypaimon/read/datasource/torch_dataset.py | 8 ++++---- paimon-python/pypaimon/read/table_read.py | 9 ++------- .../tests/torch_distributed_sharding_worker.py | 1 - paimon-python/pypaimon/tests/torch_read_test.py | 10 +++++----- 5 files changed, 13 insertions(+), 22 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 1e486f9ffcff..d856fe8e4947 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -62,21 +62,18 @@ reader threads per DataLoader worker. It has no effect in non-streaming mode. ### Distributed Sharding -Set `auto_detect_rank=True` to shard streaming reads across DDP ranks and -DataLoader workers: +Streaming reads shard splits across DDP ranks and DataLoader workers: ```python dataset = table_read.to_torch( splits, streaming=True, - auto_detect_rank=True, ) dataloader = DataLoader(dataset, batch_size=32, num_workers=2) ``` PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Detection is -disabled by default; if splits are already sharded upstream, leave it disabled -or remove the upstream sharding. +enabled by default; set `auto_detect_rank=False` to disable rank sharding. ### Batch Streaming diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index ee9d9e9e6f75..2ea9dd77155e 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -150,7 +150,7 @@ def __init__( self, table_read: TableRead, splits: List[Split], - auto_detect_rank: bool = False, + auto_detect_rank: bool = True, ): self.table_read = table_read self.splits = splits @@ -232,7 +232,7 @@ def __init__( table_read: TableRead, splits: List[Split], prefetch_concurrency: int = 1, - auto_detect_rank: bool = False, + auto_detect_rank: bool = True, ): """ Initialize TorchIterDataset. @@ -447,7 +447,7 @@ def __init__( batch_format: str, batch_size: Optional[int], to_tensor_fn: Optional[Callable[[pa.RecordBatch], Any]] = None, - auto_detect_rank: bool = False, + auto_detect_rank: bool = True, ): super().__init__(table_read, splits, auto_detect_rank) self.batch_format = batch_format @@ -512,7 +512,7 @@ def __init__( seed: int = 0, buffer_size: int = 1000, max_buffer_input_splits: int = 10, - auto_detect_rank: bool = False, + auto_detect_rank: bool = True, ): super().__init__(table_read, splits, auto_detect_rank) self.seed = self._require_int(seed, "seed") diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 09f4421836d7..1dad3b7fd5c0 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -661,7 +661,7 @@ def to_torch( seed: int = 0, buffer_size: int = 1000, max_buffer_input_splits: int = 10, - auto_detect_rank: bool = False, + auto_detect_rank: bool = True, ) -> "torch.utils.data.Dataset": """Wrap Paimon table data in a PyTorch Dataset. @@ -675,8 +675,7 @@ def to_torch( batch_size: Rows per batch; ``None`` preserves reader batches. to_tensor_fn: Optional RecordBatch converter for Torch batches. shuffle: Whether to shuffle rows; supported only in row format. - auto_detect_rank: Whether to detect DDP rank and world size from - torch.distributed or torchrun environment variables. + auto_detect_rank: Whether streaming reads detect the DDP context. """ valid_batch_formats = {"row", "pyarrow", "torch"} if batch_format not in valid_batch_formats: @@ -757,10 +756,6 @@ def to_torch( ) return dataset else: - if auto_detect_rank: - raise ValueError( - "distributed sharding requires streaming=True" - ) from pypaimon.read.datasource.torch_dataset import TorchDataset dataset = TorchDataset(self, splits) return dataset diff --git a/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py index d7b65e4dbb7e..c036c39666a4 100644 --- a/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py +++ b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py @@ -58,7 +58,6 @@ def main(): dataset = TorchIterDataset( _TableRead(), list(range(11)), - auto_detect_rank=True, ) rows = list(DataLoader(dataset, batch_size=None, num_workers=2)) with open( diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 4c957a2848a2..47f8b6644d9d 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -166,7 +166,7 @@ def test_auto_falls_back_to_single_process(self): self.assertEqual(context, (0, 1)) - def test_off_preserves_existing_behavior(self): + def test_disabled_preserves_worker_sharding(self): splits = list(range(8)) with patch.dict( os.environ, {"RANK": "1", "WORLD_SIZE": "2"}, clear=True @@ -175,7 +175,9 @@ def test_off_preserves_existing_behavior(self): ), patch.object( torch.distributed, "is_initialized", return_value=True ): - dataset = TorchIterDataset(self._table_read(), splits) + dataset = TorchIterDataset( + self._table_read(), splits, auto_detect_rank=False + ) self.assertEqual((dataset.rank, dataset.world_size), (0, 1)) self.assertEqual( @@ -775,7 +777,6 @@ def test_torch_distributed_sharding_public_api(self): streaming=True, batch_format=batch_format, shuffle=batch_format == 'row' and shuffle, - auto_detect_rank=True, ) for batch_format, shuffle in [ ('row', False), @@ -786,8 +787,7 @@ def test_torch_distributed_sharding_public_api(self): for dataset in datasets: self.assertEqual((dataset.rank, dataset.world_size), (1, 2)) - with self.assertRaisesRegex(ValueError, 'requires streaming=True'): - table_read.to_torch(splits, auto_detect_rank=True) + self.assertIsNotNone(table_read.to_torch(splits)) def test_blob_torch_read(self): """Test end-to-end blob functionality using blob descriptors.""" From 33232c3ae98aa7b2afadab699d2e69dc5851dca4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Thu, 27 Aug 2026 07:14:02 -0700 Subject: [PATCH 06/10] [python][torch] Resolve rank context when iterating --- docs/docs/pypaimon/pytorch.md | 1 + .../pypaimon/read/datasource/torch_dataset.py | 20 ++- .../pypaimon/tests/torch_read_test.py | 139 ++++++++++-------- 3 files changed, 96 insertions(+), 64 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index d856fe8e4947..082c718dde81 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -74,6 +74,7 @@ dataloader = DataLoader(dataset, batch_size=32, num_workers=2) PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Detection is enabled by default; set `auto_detect_rank=False` to disable rank sharding. +A limit that may truncate the input is rejected when multiple ranks are active. ### Batch Streaming diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index 2ea9dd77155e..c945bad92b2e 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -155,7 +155,9 @@ def __init__( self.table_read = table_read self.splits = splits self.field_names = [field.name for field in table_read.read_type] - self.rank, self.world_size = _resolve_distributed_context(auto_detect_rank) + if not isinstance(auto_detect_rank, bool): + raise ValueError("auto_detect_rank must be a bool") + self.auto_detect_rank = auto_detect_rank def _row_to_dict(self, offset_row) -> dict: row_dict = {} @@ -196,17 +198,24 @@ def _limit_covers_all_splits(self) -> bool: return True def _worker_splits(self, worker_info) -> List[Split]: + rank, world_size = _resolve_distributed_context(self.auto_detect_rank) worker_id = worker_info.id if worker_info is not None else 0 num_workers = worker_info.num_workers if worker_info is not None else 1 + if self.table_read.limit == 0: + return [] if ( self.table_read.limit is not None and not self._limit_covers_all_splits() ): + if world_size > 1: + raise ValueError( + "limit is not supported with distributed Torch sharding" + ) # A binding limit cannot be shared safely. - return self.splits if self.rank == 0 and worker_id == 0 else [] + return self.splits if worker_id == 0 else [] - rank_splits = _balanced_slice(self.splits, self.rank, self.world_size) + rank_splits = _balanced_slice(self.splits, rank, world_size) return _balanced_slice(rank_splits, worker_id, num_workers) @@ -615,9 +624,10 @@ def _iter_buffer_shuffled_rows( rows: Iterator[dict], worker_id: int, ) -> Iterator[dict]: + rank, world_size = _resolve_distributed_context(self.auto_detect_rank) rng_seed = self.seed + self.epoch * 1000003 + worker_id - if self.world_size > 1: - rng_seed = "%d:%d" % (rng_seed, self.rank) + if world_size > 1: + rng_seed = "%d:%d" % (rng_seed, rank) rng = random.Random(rng_seed) buffer = [] for row in rows: diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 47f8b6644d9d..68d50dc1a0b1 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -53,33 +53,31 @@ def _worker(worker_id, num_workers): def _dataset( self, splits, - rank=0, - world_size=1, limit=None, dataset_type=TorchIterDataset, **kwargs ): - with patch( - "pypaimon.read.datasource.torch_dataset." - "_resolve_distributed_context", - return_value=(rank, world_size), - ): - return dataset_type( - self._table_read(limit), - splits, - auto_detect_rank=True, - **kwargs - ) + return dataset_type( + self._table_read(limit), + splits, + auto_detect_rank=True, + **kwargs + ) def _assignments(self, split_count, world_size, num_workers): splits = list(range(split_count)) assignments = {} for rank in range(world_size): - dataset = self._dataset(splits, rank, world_size) + dataset = self._dataset(splits) for worker_id in range(num_workers): - assignments[(rank, worker_id)] = dataset._worker_splits( - self._worker(worker_id, num_workers) - ) + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(rank, world_size), + ): + assignments[(rank, worker_id)] = dataset._worker_splits( + self._worker(worker_id, num_workers) + ) return assignments def assertCompleteNonOverlapping(self, assignments, expected): @@ -110,23 +108,24 @@ def test_balanced_assignments( [len(splits) for splits in assignments.values()], expected_sizes ) - def test_binding_limit_uses_one_distributed_consumer(self): + def test_binding_limit_rejects_distributed_sharding(self): splits = [SimpleNamespace(row_count=10) for _ in range(4)] - assignments = {} - for rank in range(2): - dataset = self._dataset(splits, rank, 2, limit=5) - for worker_id in range(2): - assignments[(rank, worker_id)] = dataset._worker_splits( - self._worker(worker_id, 2) - ) + dataset = self._dataset(splits, limit=5) + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(0, 2), + ), self.assertRaisesRegex(ValueError, "limit is not supported"): + dataset._worker_splits(None) - self.assertEqual(assignments[(0, 0)], splits) - self.assertTrue( - all( - not value for key, value in assignments.items() - if key != (0, 0) - ) - ) + def test_zero_limit_returns_no_splits(self): + dataset = self._dataset([SimpleNamespace(row_count=10)], limit=0) + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(1, 2), + ): + self.assertEqual(dataset._worker_splits(None), []) def test_initialized_distributed_context_precedes_environment(self): with patch.dict( @@ -144,7 +143,8 @@ def test_initialized_distributed_context_precedes_environment(self): self.assertEqual(context, (1, 3)) - def test_worker_process_can_resolve_torchrun_environment(self): + def test_context_is_resolved_after_dataset_construction(self): + dataset = TorchIterDataset(self._table_read(), list(range(8))) with patch.dict( os.environ, {"RANK": "2", "WORLD_SIZE": "4"}, clear=True ), patch.object( @@ -153,10 +153,10 @@ def test_worker_process_can_resolve_torchrun_environment(self): torch.distributed, "is_initialized", return_value=False ): context = _resolve_distributed_context(True) + assigned = dataset._worker_splits(None) self.assertEqual(context, (2, 4)) - dataset = self._dataset(list(range(8)), *context) - self.assertEqual(dataset._worker_splits(None), [4, 5]) + self.assertEqual(assigned, [4, 5]) def test_auto_falls_back_to_single_process(self): with patch.dict(os.environ, {}, clear=True), patch.object( @@ -178,54 +178,73 @@ def test_disabled_preserves_worker_sharding(self): dataset = TorchIterDataset( self._table_read(), splits, auto_detect_rank=False ) + assigned = dataset._worker_splits(self._worker(1, 2)) - self.assertEqual((dataset.rank, dataset.world_size), (0, 1)) - self.assertEqual( - dataset._worker_splits(self._worker(1, 2)), - list(range(4, 8)), - ) + self.assertEqual(assigned, list(range(4, 8))) def test_shuffled_dataset_is_reproducible_and_rank_local(self): splits = list(range(20)) datasets = [ self._dataset( splits, - rank, - 2, dataset_type=TorchShuffledIterDataset, seed=17, buffer_size=20, ) - for rank in range(2) + for _ in range(2) ] - local_splits = [dataset._worker_splits(None) for dataset in datasets] + local_splits = [] + for rank, dataset in enumerate(datasets): + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(rank, 2), + ): + local_splits.append(dataset._worker_splits(None)) self.assertTrue(set(local_splits[0]).isdisjoint(local_splits[1])) self.assertCountEqual(local_splits[0] + local_splits[1], splits) restored = pickle.loads(pickle.dumps(datasets[1])) - self.assertEqual((restored.rank, restored.world_size), (1, 2)) + self.assertTrue(restored.auto_detect_rank) rows = [{"id": value} for value in range(20)] - first = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0)) - repeat = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0)) - other_rank = list( - datasets[1]._iter_buffer_shuffled_rows(iter(rows), 0) - ) - other_worker = list( - datasets[0]._iter_buffer_shuffled_rows(iter(rows), 1) - ) + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(0, 2), + ): + first = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0)) + repeat = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0)) + other_worker = list( + datasets[0]._iter_buffer_shuffled_rows(iter(rows), 1) + ) + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(1, 2), + ): + other_rank = list( + datasets[1]._iter_buffer_shuffled_rows(iter(rows), 0) + ) self.assertEqual(first, repeat) self.assertNotEqual(first, other_rank) self.assertNotEqual(first, other_worker) datasets[0].set_epoch(1) - next_epoch = list( - datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0) - ) + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(0, 2), + ): + next_epoch = list( + datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0) + ) self.assertNotEqual(first, next_epoch) def test_invalid_distributed_context(self): with self.assertRaisesRegex(ValueError, "auto_detect_rank"): - _resolve_distributed_context("auto") + TorchIterDataset( + self._table_read(), [], auto_detect_rank="auto" + ) with patch.dict(os.environ, {"RANK": "one"}, clear=True), patch.object( torch.distributed, "is_available", return_value=False @@ -784,8 +803,10 @@ def test_torch_distributed_sharding_public_api(self): ('pyarrow', False), ] ] - for dataset in datasets: - self.assertEqual((dataset.rank, dataset.world_size), (1, 2)) + expected = splits[(len(splits) + 1) // 2:] + for dataset in datasets: + self.assertTrue(dataset.auto_detect_rank) + self.assertEqual(dataset._worker_splits(None), expected) self.assertIsNotNone(table_read.to_torch(splits)) From 0092956f4928fbc73b369e171b98f320836e482c Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Thu, 27 Aug 2026 07:40:05 -0700 Subject: [PATCH 07/10] [python][torch] Preserve rank context for spawned workers --- .../pypaimon/read/datasource/torch_dataset.py | 16 ++++++++++++---- paimon-python/pypaimon/tests/torch_read_test.py | 17 +++++++++++++++++ 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index c945bad92b2e..b68697279566 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -155,9 +155,17 @@ def __init__( self.table_read = table_read self.splits = splits self.field_names = [field.name for field in table_read.read_type] - if not isinstance(auto_detect_rank, bool): - raise ValueError("auto_detect_rank must be a bool") self.auto_detect_rank = auto_detect_rank + self.rank, self.world_size = _resolve_distributed_context(auto_detect_rank) + + def _distributed_context(self): + rank, world_size = _resolve_distributed_context( + self.auto_detect_rank + ) + if world_size == 1 and self.world_size > 1: + return self.rank, self.world_size + self.rank, self.world_size = rank, world_size + return rank, world_size def _row_to_dict(self, offset_row) -> dict: row_dict = {} @@ -198,7 +206,7 @@ def _limit_covers_all_splits(self) -> bool: return True def _worker_splits(self, worker_info) -> List[Split]: - rank, world_size = _resolve_distributed_context(self.auto_detect_rank) + rank, world_size = self._distributed_context() worker_id = worker_info.id if worker_info is not None else 0 num_workers = worker_info.num_workers if worker_info is not None else 1 @@ -624,7 +632,7 @@ def _iter_buffer_shuffled_rows( rows: Iterator[dict], worker_id: int, ) -> Iterator[dict]: - rank, world_size = _resolve_distributed_context(self.auto_detect_rank) + rank, world_size = self._distributed_context() rng_seed = self.seed + self.epoch * 1000003 + worker_id if world_size > 1: rng_seed = "%d:%d" % (rng_seed, rank) diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 68d50dc1a0b1..9a03601f7f0a 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -158,6 +158,23 @@ def test_context_is_resolved_after_dataset_construction(self): self.assertEqual(context, (2, 4)) self.assertEqual(assigned, [4, 5]) + def test_constructor_context_is_preserved_in_spawned_worker(self): + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(1, 2), + ): + dataset = TorchIterDataset(self._table_read(), list(range(8))) + + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(0, 1), + ): + assigned = dataset._worker_splits(None) + + self.assertEqual(assigned, [4, 5, 6, 7]) + def test_auto_falls_back_to_single_process(self): with patch.dict(os.environ, {}, clear=True), patch.object( torch.distributed, "is_available", return_value=False From a1c9ec57926757b1cfceae3e83256f3ad9d1ccc5 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 29 Aug 2026 01:57:38 -0700 Subject: [PATCH 08/10] [python][torch] Scope cached rank to worker processes --- .../pypaimon/read/datasource/torch_dataset.py | 9 ++++- .../pypaimon/tests/torch_read_test.py | 34 +++++++++++++++++-- 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index b68697279566..02935eaf75a2 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -157,14 +157,21 @@ def __init__( self.field_names = [field.name for field in table_read.read_type] self.auto_detect_rank = auto_detect_rank self.rank, self.world_size = _resolve_distributed_context(auto_detect_rank) + self._context_pid = os.getpid() def _distributed_context(self): rank, world_size = _resolve_distributed_context( self.auto_detect_rank ) - if world_size == 1 and self.world_size > 1: + current_pid = os.getpid() + if ( + current_pid != self._context_pid + and world_size == 1 + and self.world_size > 1 + ): return self.rank, self.world_size self.rank, self.world_size = rank, world_size + self._context_pid = current_pid return rank, world_size def _row_to_dict(self, offset_row) -> dict: diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 9a03601f7f0a..c6c07c4e87c7 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -16,6 +16,7 @@ # under the License. import json +import multiprocessing import os import pickle import shutil @@ -41,6 +42,12 @@ from pypaimon.table.file_store_table import FileStoreTable +def _collect_spawned_worker_splits(dataset, output): + os.environ.pop("RANK", None) + os.environ.pop("WORLD_SIZE", None) + output.put(dataset._worker_splits(None)) + + class TorchDistributedShardingTest(unittest.TestCase): @staticmethod def _table_read(limit=None): @@ -166,14 +173,35 @@ def test_constructor_context_is_preserved_in_spawned_worker(self): ): dataset = TorchIterDataset(self._table_read(), list(range(8))) + context = multiprocessing.get_context("spawn") + output = context.Queue() + process = context.Process( + target=_collect_spawned_worker_splits, + args=(dataset, output), + ) + process.start() + process.join(30) + if process.is_alive(): + process.terminate() + process.join() + self.assertEqual(process.exitcode, 0) + self.assertEqual(output.get(timeout=5), [4, 5, 6, 7]) + output.close() + + def test_same_process_uses_latest_context(self): with patch( "pypaimon.read.datasource.torch_dataset." "_resolve_distributed_context", - return_value=(0, 1), + return_value=(1, 2), ): - assigned = dataset._worker_splits(None) + dataset = TorchIterDataset(self._table_read(), list(range(8))) - self.assertEqual(assigned, [4, 5, 6, 7]) + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(0, 1), + ): + self.assertEqual(dataset._worker_splits(None), list(range(8))) def test_auto_falls_back_to_single_process(self): with patch.dict(os.environ, {}, clear=True), patch.object( From 3334cd9508dd5bd1d40f5563c7c911abfcce11c8 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 29 Aug 2026 02:14:28 -0700 Subject: [PATCH 09/10] [python][torch] Handle uneven distributed inputs --- docs/docs/pypaimon/pytorch.md | 6 ++++ .../pypaimon/read/datasource/torch_dataset.py | 8 +++++ .../torch_distributed_sharding_worker.py | 35 +++++++++++++++---- .../pypaimon/tests/torch_read_test.py | 9 +++-- 4 files changed, 50 insertions(+), 8 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 082c718dde81..0a8bc2102735 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -70,10 +70,16 @@ dataset = table_read.to_torch( streaming=True, ) dataloader = DataLoader(dataset, batch_size=32, num_workers=2) + +with model.join(): + for batch in dataloader: + train(batch) ``` PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Detection is enabled by default; set `auto_detect_rank=False` to disable rank sharding. +A rank may receive fewer rows because splits have different sizes; `join()` +keeps DDP collectives aligned while preserving every row without duplication. A limit that may truncate the input is rejected when multiple ranks are active. ### Batch Streaming diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index 02935eaf75a2..784afbadac4e 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -159,6 +159,14 @@ def __init__( self.rank, self.world_size = _resolve_distributed_context(auto_detect_rank) self._context_pid = os.getpid() + def __getstate__(self): + state = self.__dict__.copy() + rank, world_size = _resolve_distributed_context(self.auto_detect_rank) + state["rank"] = rank + state["world_size"] = world_size + state["_context_pid"] = os.getpid() + return state + def _distributed_context(self): rank, world_size = _resolve_distributed_context( self.auto_detect_rank diff --git a/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py index c036c39666a4..24443ff7176c 100644 --- a/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py +++ b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py @@ -22,6 +22,7 @@ import torch from torch.utils.data import DataLoader +from torch.nn.parallel import DistributedDataParallel from pypaimon.read.datasource.torch_dataset import TorchIterDataset @@ -42,24 +43,46 @@ class _TableRead: SimpleNamespace(name="worker"), ] + def __init__(self): + self.rank = None + def to_iterator(self, splits): worker_info = torch.utils.data.get_worker_info() worker_id = worker_info.id if worker_info is not None else 0 - rank = int(os.environ["RANK"]) for split_id in splits: - yield _OffsetRow([split_id, rank, worker_id]) + yield _OffsetRow([split_id, self.rank, worker_id]) def main(): output_dir = sys.argv[1] + rank_env = os.environ.pop("RANK") + world_size_env = os.environ.pop("WORLD_SIZE") + table_read = _TableRead() + dataset = TorchIterDataset(table_read, list(range(11))) + os.environ["RANK"] = rank_env + os.environ["WORLD_SIZE"] = world_size_env torch.distributed.init_process_group("gloo") rank = torch.distributed.get_rank() try: - dataset = TorchIterDataset( - _TableRead(), - list(range(11)), + table_read.rank = rank + os.environ.pop("RANK") + os.environ.pop("WORLD_SIZE") + loader = DataLoader( + dataset, + batch_size=None, + num_workers=2, + multiprocessing_context="spawn", ) - rows = list(DataLoader(dataset, batch_size=None, num_workers=2)) + model = DistributedDataParallel(torch.nn.Linear(1, 1)) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + rows = [] + with model.join(): + for row in loader: + rows.append(row) + value = torch.tensor([[float(row["split_id"])]]) + model(value).sum().backward() + optimizer.step() + optimizer.zero_grad() with open( os.path.join(output_dir, "rank-%d.json" % rank), "w", diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index c6c07c4e87c7..2afb035c1cf3 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -169,7 +169,7 @@ def test_constructor_context_is_preserved_in_spawned_worker(self): with patch( "pypaimon.read.datasource.torch_dataset." "_resolve_distributed_context", - return_value=(1, 2), + return_value=(0, 1), ): dataset = TorchIterDataset(self._table_read(), list(range(8))) @@ -179,7 +179,12 @@ def test_constructor_context_is_preserved_in_spawned_worker(self): target=_collect_spawned_worker_splits, args=(dataset, output), ) - process.start() + with patch( + "pypaimon.read.datasource.torch_dataset." + "_resolve_distributed_context", + return_value=(1, 2), + ): + process.start() process.join(30) if process.is_alive(): process.terminate() From 61e08e5c06621133cca883b74ad1cc49d892c8b4 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Sat, 29 Aug 2026 03:05:45 -0700 Subject: [PATCH 10/10] [python][torch] Make rank sharding opt-in --- docs/docs/pypaimon/pytorch.md | 6 ++++-- .../pypaimon/read/datasource/torch_dataset.py | 8 ++++---- paimon-python/pypaimon/read/table_read.py | 4 ++-- .../torch_distributed_sharding_worker.py | 6 +++++- .../pypaimon/tests/torch_read_test.py | 20 ++++++++++++++++--- 5 files changed, 32 insertions(+), 12 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 0a8bc2102735..806df0f38b00 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -68,6 +68,7 @@ Streaming reads shard splits across DDP ranks and DataLoader workers: dataset = table_read.to_torch( splits, streaming=True, + auto_detect_rank=True, ) dataloader = DataLoader(dataset, batch_size=32, num_workers=2) @@ -76,8 +77,9 @@ with model.join(): train(batch) ``` -PyPaimon checks `torch.distributed`, then `RANK` and `WORLD_SIZE`. Detection is -enabled by default; set `auto_detect_rank=False` to disable rank sharding. +Automatic rank sharding is opt-in. Enable it only when every rank receives the +same ordered, complete splits from one snapshot; leave it disabled for splits +already sharded by the application. A rank may receive fewer rows because splits have different sizes; `join()` keeps DDP collectives aligned while preserving every row without duplication. A limit that may truncate the input is rejected when multiple ranks are active. diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index 784afbadac4e..6fd9da97f0c1 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -150,7 +150,7 @@ def __init__( self, table_read: TableRead, splits: List[Split], - auto_detect_rank: bool = True, + auto_detect_rank: bool = False, ): self.table_read = table_read self.splits = splits @@ -264,7 +264,7 @@ def __init__( table_read: TableRead, splits: List[Split], prefetch_concurrency: int = 1, - auto_detect_rank: bool = True, + auto_detect_rank: bool = False, ): """ Initialize TorchIterDataset. @@ -479,7 +479,7 @@ def __init__( batch_format: str, batch_size: Optional[int], to_tensor_fn: Optional[Callable[[pa.RecordBatch], Any]] = None, - auto_detect_rank: bool = True, + auto_detect_rank: bool = False, ): super().__init__(table_read, splits, auto_detect_rank) self.batch_format = batch_format @@ -544,7 +544,7 @@ def __init__( seed: int = 0, buffer_size: int = 1000, max_buffer_input_splits: int = 10, - auto_detect_rank: bool = True, + auto_detect_rank: bool = False, ): super().__init__(table_read, splits, auto_detect_rank) self.seed = self._require_int(seed, "seed") diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 1dad3b7fd5c0..4491b5dd2997 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -661,7 +661,7 @@ def to_torch( seed: int = 0, buffer_size: int = 1000, max_buffer_input_splits: int = 10, - auto_detect_rank: bool = True, + auto_detect_rank: bool = False, ) -> "torch.utils.data.Dataset": """Wrap Paimon table data in a PyTorch Dataset. @@ -675,7 +675,7 @@ def to_torch( batch_size: Rows per batch; ``None`` preserves reader batches. to_tensor_fn: Optional RecordBatch converter for Torch batches. shuffle: Whether to shuffle rows; supported only in row format. - auto_detect_rank: Whether streaming reads detect the DDP context. + auto_detect_rank: Whether streaming reads shard by DDP rank. """ valid_batch_formats = {"row", "pyarrow", "torch"} if batch_format not in valid_batch_formats: diff --git a/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py index 24443ff7176c..f4c21623fe91 100644 --- a/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py +++ b/paimon-python/pypaimon/tests/torch_distributed_sharding_worker.py @@ -58,7 +58,11 @@ def main(): rank_env = os.environ.pop("RANK") world_size_env = os.environ.pop("WORLD_SIZE") table_read = _TableRead() - dataset = TorchIterDataset(table_read, list(range(11))) + dataset = TorchIterDataset( + table_read, + list(range(11)), + auto_detect_rank=True, + ) os.environ["RANK"] = rank_env os.environ["WORLD_SIZE"] = world_size_env torch.distributed.init_process_group("gloo") diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 2afb035c1cf3..6e62c975a4f7 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -151,7 +151,9 @@ def test_initialized_distributed_context_precedes_environment(self): self.assertEqual(context, (1, 3)) def test_context_is_resolved_after_dataset_construction(self): - dataset = TorchIterDataset(self._table_read(), list(range(8))) + dataset = TorchIterDataset( + self._table_read(), list(range(8)), auto_detect_rank=True + ) with patch.dict( os.environ, {"RANK": "2", "WORLD_SIZE": "4"}, clear=True ), patch.object( @@ -171,7 +173,9 @@ def test_constructor_context_is_preserved_in_spawned_worker(self): "_resolve_distributed_context", return_value=(0, 1), ): - dataset = TorchIterDataset(self._table_read(), list(range(8))) + dataset = TorchIterDataset( + self._table_read(), list(range(8)), auto_detect_rank=True + ) context = multiprocessing.get_context("spawn") output = context.Queue() @@ -199,7 +203,9 @@ def test_same_process_uses_latest_context(self): "_resolve_distributed_context", return_value=(1, 2), ): - dataset = TorchIterDataset(self._table_read(), list(range(8))) + dataset = TorchIterDataset( + self._table_read(), list(range(8)), auto_detect_rank=True + ) with patch( "pypaimon.read.datasource.torch_dataset." @@ -846,6 +852,7 @@ def test_torch_distributed_sharding_public_api(self): streaming=True, batch_format=batch_format, shuffle=batch_format == 'row' and shuffle, + auto_detect_rank=True, ) for batch_format, shuffle in [ ('row', False), @@ -858,6 +865,13 @@ def test_torch_distributed_sharding_public_api(self): self.assertTrue(dataset.auto_detect_rank) self.assertEqual(dataset._worker_splits(None), expected) + pre_sharded = splits[::2] + with patch.dict( + os.environ, {"RANK": "1", "WORLD_SIZE": "2"}, clear=True + ): + dataset = table_read.to_torch(pre_sharded, streaming=True) + self.assertFalse(dataset.auto_detect_rank) + self.assertEqual(dataset._worker_splits(None), pre_sharded) self.assertIsNotNone(table_read.to_torch(splits)) def test_blob_torch_read(self):