Merge branch 'master' into precommit/ruff-format

Author: Borda

src/lightning/data/processing/data_processor.py

@@ -250,23 +250,35 @@ def _upload_fn(upload_queue: Queue, remove_queue: Queue, cache_dir: str, output_
 
 def _map_items_to_workers_sequentially(num_workers: int, user_items: List[Any]) -> List[List[Any]]:
     num_nodes = _get_num_nodes()
-    current_node_rank = _get_node_rank()
-    node_size = len(user_items) // num_nodes
-    workers_user_items = []
-    for node_rank in range(num_nodes):
-        if node_rank != current_node_rank:
-            continue
-        is_last_node = node_rank == num_nodes - 1
-        start_node = node_rank * node_size
-        end_node = len(user_items) if is_last_node else (node_rank + 1) * node_size
-        node_user_items = user_items[start_node:end_node]
-        worker_size = len(node_user_items) // num_workers
-        for worker_idx in range(num_workers):
-            is_last = worker_idx == num_workers - 1
-            begin = worker_idx * worker_size
-            end = len(node_user_items) if is_last else (worker_idx + 1) * worker_size
-            workers_user_items.append(node_user_items[begin:end])
-    return workers_user_items
+    world_size = (num_nodes * num_workers)
+    num_items_per_worker = len(user_items) // world_size
+
+    num_items_per_worker: List[int] = [num_items_per_worker for _ in range(world_size)]
+    reminder = len(user_items) % world_size
+
+    for worker_idx in range(len(num_items_per_worker) - 1, -1, -1):
+        if reminder == 0:
+            break
+        num_items_per_worker[worker_idx] += 1
+        reminder -= 1
+
+    num_items_cumsum_per_worker = np.cumsum([0] + num_items_per_worker)
+
+    out = []
+    node_rank = _get_node_rank()
+    worker_idx_start = node_rank * num_workers
+    worker_idx_end = (node_rank + 1) * num_workers
+
+    for worker_idx in range(world_size):
+        if worker_idx_start <= worker_idx and worker_idx < worker_idx_end:
+            start = num_items_cumsum_per_worker[worker_idx]
+            end = num_items_cumsum_per_worker[worker_idx + 1]
+            out.append(user_items[start : end])
+
+    if len(out) != num_workers:
+        raise RuntimeError("The items didn't haven't been assigned properly. Please, open an issue on Github.")
+
+    return out
 
 
 def _map_items_to_workers_weighted(

src/lightning/data/streaming/combined.py

@@ -66,6 +66,11 @@ def set_epoch(self, current_epoch: int) -> None:
         for dataset in self._datasets:
             dataset.set_epoch(current_epoch)
 
+    def set_shuffle(self, shuffle: bool) -> None:
+        """Set the current shuffle to the datasets."""
+        for dataset in self._datasets:
+            dataset.set_shuffle(shuffle)
+
     def _check_datasets(self, datasets: List[StreamingDataset]) -> None:
         if any(not isinstance(d, StreamingDataset) for d in datasets):
             raise RuntimeError("The provided datasets should be instances of the StreamingDataset.")

src/lightning/data/streaming/dataloader.py

@@ -541,6 +541,7 @@ def __init__(
         profile_batches: Union[bool, int] = False,
         profile_dir: Optional[str] = None,
         prefetch_factor: Optional[int] = None,
+        shuffle: Optional[bool] = None,
         **kwargs: Any,
     ) -> None:  # pyright: ignore
         if not isinstance(dataset, (StreamingDataset, CombinedStreamingDataset)):
@@ -549,6 +550,11 @@ def __init__(
                 f" Found {dataset}."
             )
 
+        if shuffle is not None:
+            dataset.set_shuffle(shuffle)
+
+        shuffle = None
+
         if profile_batches and not _VIZ_TRACKER_AVAILABLE:
             raise ModuleNotFoundError("To use profile_batches, viztracer is required. Run `pip install viztracer`")
 

src/lightning/data/streaming/dataset.py

@@ -107,6 +107,9 @@ def __init__(
         self.serializers = serializers
         self._state_dict: Optional[Dict[str, Any]] = None
 
+    def set_shuffle(self, shuffle: bool) -> None:
+        self.shuffle = shuffle
+
     def set_epoch(self, current_epoch: int) -> None:
         """Set the current epoch to the dataset on epoch starts.
 

src/lightning/data/utilities/env.py

@@ -31,21 +31,19 @@ def detect(cls) -> "_DistributedEnv":
         if torch.distributed.is_available() and torch.distributed.is_initialized():
             world_size = torch.distributed.get_world_size()
             global_rank = torch.distributed.get_rank()
+            # Note: On multi node CPU, the number of nodes won't be correct.
+            num_nodes = world_size // torch.cuda.device_count() if torch.cuda.is_available() else world_size
+            if torch.cuda.is_available() and world_size % torch.cuda.device_count() != 0:
+                raise RuntimeError("The world size should be divisible by the number of GPUs.")
         else:
             world_size = None
             global_rank = 0
+            num_nodes = 1
 
         if world_size is None or world_size == -1:
             world_size = 1
 
-        # TODO: Add support for other accelerators
-        num_nodes = (world_size // torch.cuda.device_count()) if torch.cuda.is_available() else 1
-
-        if num_nodes > 1:
-            # validate the world size is divisble by the number of GPUs
-            assert world_size % torch.cuda.device_count() == 0
-
-        return cls(world_size=world_size, global_rank=global_rank, num_nodes=max(1, num_nodes))
+        return cls(world_size=world_size, global_rank=global_rank, num_nodes=num_nodes)
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(world_size: {self.world_size}, global_rank: {self.global_rank}\n)"

tests/tests_data/processing/test_data_processor.py

@@ -310,7 +310,7 @@ def test_map_items_to_workers_sequentially(monkeypatch):
     workers_user_items = _map_items_to_workers_sequentially(2, list(range(5)))
     assert workers_user_items == [[0, 1], [2, 3, 4]]
     workers_user_items = _map_items_to_workers_sequentially(3, list(range(5)))
-    assert workers_user_items == [[0], [1], [2, 3, 4]]
+    assert workers_user_items == [[0], [1, 2], [3, 4]]
     workers_user_items = _map_items_to_workers_sequentially(4, list(range(5)))
     assert workers_user_items == [[0], [1], [2], [3, 4]]
 
@@ -335,7 +335,7 @@ def test_map_items_to_workers_sequentially(monkeypatch):
     workers_user_items = _map_items_to_workers_sequentially(2, list(range(32)))
     assert workers_user_items == [[0, 1, 2, 3], [4, 5, 6, 7]]
     workers_user_items = _map_items_to_workers_sequentially(3, list(range(32)))
-    assert workers_user_items == [[0, 1], [2, 3], [4, 5, 6, 7]]
+    assert workers_user_items == [[0, 1], [2, 3], [4, 5]]
     workers_user_items = _map_items_to_workers_sequentially(4, list(range(32)))
     assert workers_user_items == [[0, 1], [2, 3], [4, 5], [6, 7]]
 
@@ -346,7 +346,7 @@ def test_map_items_to_workers_sequentially(monkeypatch):
     workers_user_items = _map_items_to_workers_sequentially(2, list(range(32)))
     assert workers_user_items == [[24, 25, 26, 27], [28, 29, 30, 31]]
     workers_user_items = _map_items_to_workers_sequentially(3, list(range(32)))
-    assert workers_user_items == [[24, 25], [26, 27], [28, 29, 30, 31]]
+    assert workers_user_items == [[23, 24, 25], [26, 27, 28], [29, 30, 31]]
     workers_user_items = _map_items_to_workers_sequentially(4, list(range(32)))
     assert workers_user_items == [[24, 25], [26, 27], [28, 29], [30, 31]]
 

tests/tests_data/streaming/test_dataloader.py

@@ -12,6 +12,10 @@ def __init__(self, size, step):
         self.size = size
         self.step = step
         self.counter = 0
+        self.shuffle = None
+
+    def set_shuffle(self, shuffle):
+        self.shuffle = shuffle
 
     def __len__(self):
         return self.size
@@ -92,3 +96,14 @@ def test_dataloader_profiling(profile, tmpdir, monkeypatch):
         batches.append(batch)
 
     assert os.path.exists(os.path.join(tmpdir, "result.json"))
+
+
+def test_dataloader_shuffle():
+    dataset = TestCombinedStreamingDataset(
+        [TestStatefulDataset(10, 1), TestStatefulDataset(10, -1)], 42, weights=(0.5, 0.5)
+    )
+    assert dataset._datasets[0].shuffle is None
+    assert dataset._datasets[1].shuffle is None
+    StreamingDataLoader(dataset, batch_size=2, num_workers=1, shuffle=True)
+    assert dataset._datasets[0].shuffle
+    assert dataset._datasets[1].shuffle