From e796a822e01232149cbe1df44c1550708cdf00db Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Sat, 25 Jul 2026 00:42:36 +0000 Subject: [PATCH 1/3] chore(tests): speed up streaming pytest suite Collapse redundant num_workers=2 parallel-matrix cells onto a single-process path (keep one MP cell per behavior), reuse shared Cache dirs across params, shrink dataloader fixtures, replace fixed sleeps with polls/backdated clocks, and raise Ubuntu xdist workers to 4. Co-authored-by: Cursor --- .github/workflows/ci-testing.yml | 9 +- tests/streaming/test_client.py | 11 +- tests/streaming/test_dataloader.py | 56 +++---- tests/streaming/test_dataset.py | 14 +- tests/streaming/test_parallel.py | 250 ++++++++++++++++++----------- 5 files changed, 205 insertions(+), 135 deletions(-) diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml index f5e0d1209..84971bda7 100644 --- a/.github/workflows/ci-testing.yml +++ b/.github/workflows/ci-testing.yml @@ -61,10 +61,17 @@ jobs: - name: Run fast tests in parallel run: | + # Ubuntu runners can take more xdist workers; keep macOS/Windows at 2 to + # avoid spawn/resource pressure from DataLoader multiprocessing tests. + if [ "${{ runner.os }}" = "Linux" ]; then + XDIST_WORKERS=4 + else + XDIST_WORKERS=2 + fi pytest tests \ --ignore=tests/processing \ --ignore=tests/raw \ - -n 2 --dist=loadgroup --cov=litdata --durations=0 --timeout=120 --capture=no --verbose + -n "${XDIST_WORKERS}" --dist=loadgroup --cov=litdata --durations=0 --timeout=120 --capture=no --verbose - name: Run processing tests sequentially run: | diff --git a/tests/streaming/test_client.py b/tests/streaming/test_client.py index e98f8974e..d165813a1 100644 --- a/tests/streaming/test_client.py +++ b/tests/streaming/test_client.py @@ -1,5 +1,5 @@ import sys -from time import sleep, time +from time import time from unittest import mock import pytest @@ -92,11 +92,12 @@ def test_s3_client_with_cloud_space_id(use_shared_credentials, monkeypatch): assert s3.client assert s3.client boto3_session().client.assert_called_once() - sleep(1 - (time() - s3._last_time)) + # Backdate last fetch so the next property access refreshes without sleeping. + s3._last_time = time() - s3._refetch_interval - 0.01 assert s3.client assert s3.client assert len(boto3_session().client._mock_mock_calls) == 6 - sleep(1 - (time() - s3._last_time)) + s3._last_time = time() - s3._refetch_interval - 0.01 assert s3.client assert s3.client assert len(boto3_session().client._mock_mock_calls) == 9 @@ -422,8 +423,8 @@ def test_r2_client_property_refreshes_expired_credentials(monkeypatch): r2_client.client first_call_count = boto3_session().client.call_count - # Wait for credentials to expire - sleep(1.1) + # Expire credentials without sleeping through the refetch interval. + r2_client._last_time = time() - r2_client._refetch_interval - 0.01 # Second access should refresh credentials r2_client.client diff --git a/tests/streaming/test_dataloader.py b/tests/streaming/test_dataloader.py index 1458476fc..b110ad7f8 100644 --- a/tests/streaming/test_dataloader.py +++ b/tests/streaming/test_dataloader.py @@ -203,7 +203,7 @@ def test_custom_collate_multiworker(): def test_dataloader_no_workers(tmpdir): cache = Cache(input_dir=str(tmpdir), chunk_bytes="64MB") - for i in range(1000): + for i in range(100): cache[i] = i cache.done() @@ -211,15 +211,15 @@ def test_dataloader_no_workers(tmpdir): dataset = StreamingDataset(str(tmpdir), shuffle=True) dataloader = StreamingDataLoader(dataset) - assert len(dataset) == 1000 - assert len(dataloader) == 1000 - assert len(dataset) == 1000 + assert len(dataset) == 100 + assert len(dataloader) == 100 + assert len(dataset) == 100 @pytest.mark.timeout(120) def test_dataloader_with_loading_states(tmpdir): cache = Cache(input_dir=str(tmpdir), chunk_bytes="64MB") - for i in range(100): + for i in range(40): cache[i] = i cache.done() cache.merge() @@ -231,16 +231,16 @@ def test_dataloader_with_loading_states(tmpdir): dataloader.load_state_dict(dataloader.state_dict()) batch = next(iter(dataloader)) assert len(batch) == 4, "Batch size should be 4" - assert len(dataloader) == 25, "Dataloader length should be 25 (100 items / batch size 4)" + assert len(dataloader) == 10, "Dataloader length should be 10 (40 items / batch size 4)" # Test dataloader with num workers dataloader = StreamingDataLoader(dataset, batch_size=4, num_workers=2) - assert len(dataloader) == 25, "Dataloader length should be 25 (100 items / batch size 4)" + assert len(dataloader) == 10, "Dataloader length should be 10 (40 items / batch size 4)" # Verify dataloader state after partial iteration for batch_idx, batch in enumerate(dataloader): assert dataloader.current_epoch == 1, "Current epoch should be 1" - if batch_idx == 10: + if batch_idx == 4: break dataloader.load_state_dict(dataloader.state_dict()) assert dataloader.restore @@ -249,8 +249,8 @@ def test_dataloader_with_loading_states(tmpdir): for _ in dataloader: assert dataloader.current_epoch == 1, "Current epoch should be 1" count += 1 - # we consumed 11 batches (batch_idx==10) before. - assert count == 14, "There should be at least 14 batches remaining in the first epoch" + # we consumed 5 batches (batch_idx==4) before. + assert count == 5, "There should be 5 batches remaining in the first epoch" assert not dataloader.restore # Verify batches in the second epoch @@ -258,7 +258,7 @@ def test_dataloader_with_loading_states(tmpdir): for _ in dataloader: assert dataloader.current_epoch == 2, "Current epoch should be 2" count += 1 - assert count >= 25, "There should be at least 25 batches in the second epoch" + assert count >= 10, "There should be at least 10 batches in the second epoch" # Verify that the datalaoder can resume after complete last epoch dataloader.load_state_dict(dataloader.state_dict()) @@ -267,13 +267,13 @@ def test_dataloader_with_loading_states(tmpdir): for _ in dataloader: assert dataloader.current_epoch == 3, "Current epoch should be 3" count += 1 - assert count >= 25, "There should be at least 25 batches in the third epoch" + assert count >= 10, "There should be at least 10 batches in the third epoch" @pytest.mark.timeout(120) def test_dataloader_states_with_persistent_workers(tmpdir): cache = Cache(input_dir=str(tmpdir), chunk_bytes="64MB") - for i in range(100): + for i in range(40): cache[i] = i cache.done() cache.merge() @@ -281,12 +281,12 @@ def test_dataloader_states_with_persistent_workers(tmpdir): dataset = StreamingDataset(str(tmpdir), shuffle=True) dataloader = StreamingDataLoader(dataset, batch_size=4, num_workers=2) - assert len(dataloader) == 25, "Dataloader length should be 25 (100 items / batch size 4)" + assert len(dataloader) == 10, "Dataloader length should be 10 (40 items / batch size 4)" # Verify dataloader state after partial iteration for batch_idx, batch in enumerate(dataloader): assert dataloader.current_epoch == 1, "Current epoch should be 1" - if batch_idx == 10: + if batch_idx == 4: break prev_dataloader_state = dataloader.state_dict() @@ -299,8 +299,8 @@ def test_dataloader_states_with_persistent_workers(tmpdir): for _ in dataloader: assert dataloader.current_epoch == 1, "Current epoch should be 1" count += 1 - # batch_idx==10 means we consumed 11 batches before. - assert count == 14, "There should be at least 14 batches remaining in the first epoch" + # batch_idx==4 means we consumed 5 batches before. + assert count == 5, "There should be 5 batches remaining in the first epoch" assert not dataloader.restore # Verify batches in the second epoch @@ -308,7 +308,7 @@ def test_dataloader_states_with_persistent_workers(tmpdir): for _ in dataloader: assert dataloader.current_epoch == 2, "Current epoch should be 2" count += 1 - assert count >= 25, "There should be at least 25 batches in the second epoch" + assert count >= 10, "There should be at least 10 batches in the second epoch" # Verify that the datalaoder can resume after complete last epoch dataloader.load_state_dict(dataloader.state_dict()) @@ -317,7 +317,7 @@ def test_dataloader_states_with_persistent_workers(tmpdir): for _ in dataloader: assert dataloader.current_epoch == 3, "Current epoch should be 3" count += 1 - assert count >= 25, "There should be at least 25 batches in the third epoch" + assert count >= 10, "There should be at least 10 batches in the third epoch" @pytest.mark.timeout(90) @@ -326,7 +326,7 @@ def test_resume_dataloader_with_new_dataset(tmpdir): dataset_2_path = tmpdir.join("dataset_2") for dataset in [dataset_1_path, dataset_2_path]: cache = Cache(input_dir=str(dataset), chunk_bytes="64MB") - for i in range(50): + for i in range(20): cache[i] = i cache.done() cache.merge() @@ -423,16 +423,16 @@ def test_dataloader_dataset_transform(tmpdir, shuffle): os.makedirs(cache_dir) os.makedirs(data_dir) - # Create a dataset with 100 items, 20 items per chunk - cache = Cache(str(data_dir), chunk_size=20) - for i in range(100): + # Create a dataset with 40 items, 10 items per chunk + cache = Cache(str(data_dir), chunk_size=10) + for i in range(40): cache[i] = i cache.done() cache.merge() dataset = StreamingDataset(data_dir, cache_dir=str(cache_dir), shuffle=shuffle, transform=transform_fn) dataset_length = len(dataset) - assert dataset_length == 100 + assert dataset_length == 40 # ACT dl = StreamingDataLoader(dataset, batch_size=10, num_workers=2, shuffle=shuffle) @@ -472,16 +472,16 @@ def test_dataloader_dataset_transform_inheritance(tmpdir, shuffle): os.makedirs(cache_dir) os.makedirs(data_dir) - # Create a dataset with 100 items, 20 items per chunk - cache = Cache(str(data_dir), chunk_size=20) - for i in range(100): + # Create a dataset with 40 items, 10 items per chunk + cache = Cache(str(data_dir), chunk_size=10) + for i in range(40): cache[i] = i cache.done() cache.merge() dataset = StreamingDatasetWithTransform(data_dir, cache_dir=str(cache_dir), shuffle=shuffle) dataset_length = len(dataset) - assert dataset_length == 100 + assert dataset_length == 40 # ACT dl = StreamingDataLoader(dataset, batch_size=10, num_workers=2, shuffle=shuffle) diff --git a/tests/streaming/test_dataset.py b/tests/streaming/test_dataset.py index 74cca7796..4c261d5ac 100644 --- a/tests/streaming/test_dataset.py +++ b/tests/streaming/test_dataset.py @@ -18,7 +18,7 @@ import shutil import sys from functools import partial -from time import sleep +from time import perf_counter, sleep from typing import Any from unittest import mock from unittest.mock import patch @@ -132,7 +132,13 @@ def test_optimize_dataset( keep_data_ordered=keep_data_ordered, ) - sleep(2) # wait for the cache to be created + # optimize writes index.json when the dataset is ready; poll instead of a fixed sleep. + index_path = os.path.join(data_dir, "index.json") + deadline = perf_counter() + 5.0 + while not os.path.exists(index_path): + if perf_counter() > deadline: + raise TimeoutError(f"Timed out waiting for {index_path}") + sleep(0.05) ds = StreamingDataset(input_dir=data_dir) @@ -1187,8 +1193,6 @@ def fn(remote_chunkpath: str, local_chunkpath: str): dataloader_iter = iter(dataloader) next(dataloader_iter) - sleep(1) - state_dict = dataset.state_dict(0, 1, 2) dataset.load_state_dict(state_dict) @@ -1324,8 +1328,6 @@ def fn(remote_chunkpath: str, local_chunkpath: str): dataloader_iter = iter(dataloader) next(dataloader_iter) - sleep(1) - state_dict = dataset.state_dict(0, 1, 2) dataset.load_state_dict(state_dict) diff --git a/tests/streaming/test_parallel.py b/tests/streaming/test_parallel.py index cb17d592a..455f1b544 100644 --- a/tests/streaming/test_parallel.py +++ b/tests/streaming/test_parallel.py @@ -1,13 +1,16 @@ import functools +import os import sys from copy import deepcopy from dataclasses import dataclass +from pathlib import Path from unittest.mock import ANY, MagicMock import pytest import torch from torch.utils.data import IterableDataset +from litdata.constants import _INDEX_FILENAME from litdata.streaming.cache import Cache from litdata.streaming.dataloader import StreamingDataLoader from litdata.streaming.dataset import StreamingDataset @@ -333,10 +336,20 @@ def rng_transform(_, rngs, which): return rngs[which].random() -@pytest.mark.parametrize("length", [None, 7]) -@pytest.mark.parametrize("num_workers", [0, 2]) -@pytest.mark.parametrize("which", ["random", "numpy", "torch"]) -@pytest.mark.parametrize("reset_rngs", [False, True]) +@pytest.mark.parametrize( + ("length", "num_workers", "which", "reset_rngs"), + [ + # Combinatorial coverage stays on the cheap single-process path. + *[ + (length, 0, which, reset_rngs) + for length in (None, 7) + for which in ("random", "numpy", "torch") + for reset_rngs in (False, True) + ], + # One multi-worker cell proves spawn + RNG seeding still works. + (None, 2, "torch", True), + ], +) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_rng(length, num_workers, which, reset_rngs): transform = functools.partial(rng_transform, which=which) @@ -422,17 +435,60 @@ def test_dataloader_shuffle(tmp_path, shuffle): assert shuffle ^ all(torch.equal(x, y) for x, y in zip(epoch_1_batches[:3], epoch_2_batches[-3:])) -def prepare_parallel_dataset_and_dataloder( - tmp_path_factory, parlen, len1=48, len2=56, num_workers=0, batch_size=4, shuffle=True, resume=True, tmpdir=None -): - tmpdir = tmp_path_factory.mktemp("data") if tmpdir is None else tmpdir +def _ensure_parallel_cache_pair(tmpdir: Path, len1: int, len2: int) -> list[str]: + """Build (or reuse) two Cache datasets under ``tmpdir`` for parallel-dataset tests.""" datasets = [str(tmpdir / f"dataset_{i}") for i in range(2)] for dataset, num_items in zip(datasets, [len1, len2]): + if os.path.exists(os.path.join(dataset, _INDEX_FILENAME)): + continue + os.makedirs(dataset, exist_ok=True) cache = Cache(input_dir=dataset, chunk_size=10) for i in range(num_items): cache[i] = i cache.done() cache.merge() + return datasets + + +def _parallel_sample_order(n_items: int, num_workers: int) -> list[int]: + """Return the sample index order for equal-sized shards under ``num_workers``. + + With ``num_workers<=1`` this is sequential. With multiple workers, LitData/PyTorch + assigns contiguous shards then the DataLoader round-robins worker results + (e.g. ``n=10, workers=2`` → ``[0, 5, 1, 6, ...]``). + """ + if num_workers <= 1: + return list(range(n_items)) + sizes = [n_items // num_workers + (1 if i < n_items % num_workers else 0) for i in range(num_workers)] + parts: list[list[int]] = [] + start = 0 + for size in sizes: + parts.append(list(range(start, start + size))) + start += size + order: list[int] = [] + for i in range(max(len(p) for p in parts)): + for part in parts: + if i < len(part): + order.append(part[i]) + return order + + +def _pair_batches(order: list[int]) -> list[list[torch.Tensor]]: + return [[torch.tensor([i]), torch.tensor([i])] for i in order] + + +def prepare_parallel_dataset_and_dataloder( + tmp_path_factory, parlen, len1=48, len2=56, num_workers=0, batch_size=4, shuffle=True, resume=True, tmpdir=None +): + # Reuse a shared cache root per (len1, len2) within this xdist worker so parametrized + # cells do not pay Cache write/merge cost on every case. + if tmpdir is None: + shared = Path(tmp_path_factory.getbasetemp()) / "parallel_shared_caches" / f"l{len1}_{len2}" + shared.mkdir(parents=True, exist_ok=True) + tmpdir = shared + else: + tmpdir = Path(tmpdir) + datasets = _ensure_parallel_cache_pair(tmpdir, len1, len2) dset1 = StreamingDataset(datasets[0], shuffle=shuffle) dset2 = StreamingDataset(datasets[1], shuffle=shuffle) pardset = ParallelStreamingDataset(datasets=[dset1, dset2], length=parlen, resume=resume) @@ -451,13 +507,18 @@ def test_parallel_dataset_dataloader_states_without_any_iterations(tmp_path_fact @pytest.mark.timeout(120) -@pytest.mark.parametrize("length", [None, 24]) -@pytest.mark.parametrize("num_workers", [0, 2]) +@pytest.mark.parametrize( + ("length", "num_workers"), + [ + (None, 0), + (24, 0), + # One multi-worker cell covers spawn + multi-epoch state transitions. + (24, 2), + ], +) @pytest.mark.parametrize("batch_size", [2]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_dataloader_states_complete_iterations(tmp_path_factory, length, num_workers, batch_size): - print(f"Testing with num_workers={num_workers}") - _, _, parallel_dataset, dataloader, _ = prepare_parallel_dataset_and_dataloder( tmp_path_factory, length, @@ -514,16 +575,21 @@ def test_parallel_dataset_dataloader_states_complete_iterations(tmp_path_factory @pytest.mark.timeout(300) -@pytest.mark.parametrize("length", [None, 20, 48]) -@pytest.mark.parametrize("num_workers", [0, 2]) +@pytest.mark.parametrize( + ("length", "num_workers", "break_at"), + [ + (None, 0, 3), + (48, 0, 3), + (48, 0, 7), + # One multi-worker cell covers restore after early break under spawn. + (None, 2, 3), + ], +) @pytest.mark.parametrize("batch_size", [2]) -@pytest.mark.parametrize("break_at", [3, 7]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_dataloader_states_partial_iterations( tmp_path_factory, length, num_workers, batch_size, break_at ): - print(f"Testing with num_workers={num_workers}, break_at={break_at}") - _, _, parallel_dataset, dataloader, _ = prepare_parallel_dataset_and_dataloder( tmp_path_factory, length, batch_size=batch_size, num_workers=num_workers, shuffle=True ) @@ -888,22 +954,42 @@ def test_parallel_infinite_restore_survives_early_break_without_hang(tmp_path_fa assert dloader.restore -@pytest.mark.parametrize("length", [None, 16, float("inf")]) -@pytest.mark.parametrize("resume", [False, True]) -@pytest.mark.parametrize("shuffle", [False, True]) +@pytest.mark.parametrize( + ("length", "resume", "shuffle", "num_workers"), + [ + *[ + (length, resume, shuffle, 0) + for length in (None, 16) + for resume in (False, True) + for shuffle in (False, True) + ], + # Multi-worker coverage (float("inf") is covered by the dedicated infinite restore test). + (16, True, False, 2), + (None, False, True, 2), + ], +) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") -def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, resume, shuffle): +def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, resume, shuffle, num_workers): _, _, pardset, dloader, tmpdir = prepare_parallel_dataset_and_dataloder( - tmp_path_factory, parlen=length, len1=10, len2=10, batch_size=1, num_workers=2, shuffle=shuffle, resume=resume + tmp_path_factory, + parlen=length, + len1=10, + len2=10, + batch_size=1, + num_workers=num_workers, + shuffle=shuffle, + resume=resume, ) assert pardset.is_cycling() or length is None break_at = 3 - expected_1 = [ - [torch.tensor([0]), torch.tensor([0])], - [torch.tensor([5]), torch.tensor([5])], - [torch.tensor([1]), torch.tensor([1])], - [torch.tensor([6]), torch.tensor([6])], - ] + order = _parallel_sample_order(10, num_workers) + # Three partial windows of 4 batches each along the epoch order (and wrap when cycling). + window = _pair_batches(order) + expected_1 = window[:4] + expected_2 = window[4:8] + expected_3 = window[8:10] + window[:2] + expected_4 = window[2:6] + expected_5 = window[6:10] batches_1 = [] for i, batch in enumerate(dloader): if not shuffle: @@ -911,12 +997,6 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res batches_1.append(batch) if i == break_at: break - expected_2 = [ - [torch.tensor([2]), torch.tensor([2])], - [torch.tensor([7]), torch.tensor([7])], - [torch.tensor([3]), torch.tensor([3])], - [torch.tensor([8]), torch.tensor([8])], - ] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -928,12 +1008,6 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res if i == break_at: break state_dict_after_2 = dloader.state_dict() - expected_3 = [ - [torch.tensor([4]), torch.tensor([4])], - [torch.tensor([9]), torch.tensor([9])], - [torch.tensor([0]), torch.tensor([0])], - [torch.tensor([5]), torch.tensor([5])], - ] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -952,7 +1026,7 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res len1=10, len2=10, batch_size=1, - num_workers=2, + num_workers=num_workers, shuffle=shuffle, resume=resume, tmpdir=tmpdir, @@ -976,14 +1050,9 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res # worker index assignment may differ from the previous session. This is expected PyTorch behavior: # each worker processes a subset of indices, but the order workers deliver batches can vary. # We adjust expected values to match the actual (deterministic but different) worker ordering. - expected_2 = [expected_2[i + 1] if i % 2 == 0 else expected_2[i - 1] for i in range(len(expected_2))] - batches_2 = [batches_2[i + 1] if i % 2 == 0 else batches_2[i - 1] for i in range(len(batches_2))] - expected_4 = [ - [torch.tensor([1]), torch.tensor([1])], - [torch.tensor([6]), torch.tensor([6])], - [torch.tensor([2]), torch.tensor([2])], - [torch.tensor([7]), torch.tensor([7])], - ] + if num_workers > 1: + expected_2 = [expected_2[i + 1] if i % 2 == 0 else expected_2[i - 1] for i in range(len(expected_2))] + batches_2 = [batches_2[i + 1] if i % 2 == 0 else batches_2[i - 1] for i in range(len(batches_2))] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -994,12 +1063,6 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res assert all(torch.equal(x, y) for x, y in zip(batch, batches_2[i])) if i == break_at: break - expected_5 = [ - [torch.tensor([3]), torch.tensor([3])], - [torch.tensor([8]), torch.tensor([8])], - [torch.tensor([4]), torch.tensor([4])], - [torch.tensor([9]), torch.tensor([9])], - ] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -1012,34 +1075,47 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res break -@pytest.mark.parametrize("length", [None, 4]) -@pytest.mark.parametrize("resume", [False, True]) -@pytest.mark.parametrize("shuffle", [False, True]) +@pytest.mark.parametrize( + ("length", "resume", "shuffle", "num_workers"), + [ + *[(length, resume, shuffle, 0) for length in (None, 4) for resume in (False, True) for shuffle in (False, True)], + (4, True, False, 2), + (None, False, True, 2), + ], +) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") -def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, resume, shuffle): +def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, resume, shuffle, num_workers): _, _, pardset, dloader, tmpdir = prepare_parallel_dataset_and_dataloder( - tmp_path_factory, parlen=length, len1=6, len2=6, batch_size=1, num_workers=2, shuffle=shuffle, resume=resume + tmp_path_factory, + parlen=length, + len1=6, + len2=6, + batch_size=1, + num_workers=num_workers, + shuffle=shuffle, + resume=resume, ) assert pardset.is_cycling() or length is None - expected_1 = [ - [torch.tensor([0]), torch.tensor([0])], - [torch.tensor([3]), torch.tensor([3])], - [torch.tensor([1]), torch.tensor([1])], - [torch.tensor([4]), torch.tensor([4])], - [torch.tensor([2]), torch.tensor([2])], - [torch.tensor([5]), torch.tensor([5])], - ] + order = _parallel_sample_order(6, num_workers) + n = len(order) + + def _window(start: int, window: int) -> list[list[torch.Tensor]]: + return _pair_batches([order[(start + i) % n] for i in range(window)]) + + if length is None: + expected_1 = _pair_batches(order) + expected_2 = expected_3 = expected_4 = expected_5 = expected_1 + else: + expected_1 = _window(0, length) + expected_2 = _window(length, length) + expected_3 = _window(2 * length, length) + expected_4 = _window(3 * length, length) + expected_5 = _window(4 * length, length) batches_1 = [] for i, batch in enumerate(dloader): if not shuffle: assert all(torch.equal(x, y) for x, y in zip(batch, expected_1[i])) batches_1.append(batch) - expected_2 = [ - [torch.tensor([2]), torch.tensor([2])], - [torch.tensor([5]), torch.tensor([5])], - [torch.tensor([0]), torch.tensor([0])], - [torch.tensor([3]), torch.tensor([3])], - ] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -1049,12 +1125,6 @@ def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, re elif not resume and length is not None: assert all(torch.equal(x, y) for x, y in zip(batch, batches_1[i])) state_dict_after_2 = dloader.state_dict() - expected_3 = [ - [torch.tensor([1]), torch.tensor([1])], - [torch.tensor([4]), torch.tensor([4])], - [torch.tensor([2]), torch.tensor([2])], - [torch.tensor([5]), torch.tensor([5])], - ] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -1070,7 +1140,7 @@ def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, re len1=6, len2=6, batch_size=1, - num_workers=2, + num_workers=num_workers, shuffle=shuffle, resume=resume, tmpdir=tmpdir, @@ -1088,12 +1158,6 @@ def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, re ) elif not resume and length is not None: assert all(torch.equal(x, y) for x, y in zip(batch, batches_1[i])) - expected_4 = [ - [torch.tensor([0]), torch.tensor([0])], - [torch.tensor([3]), torch.tensor([3])], - [torch.tensor([1]), torch.tensor([1])], - [torch.tensor([4]), torch.tensor([4])], - ] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -1102,12 +1166,6 @@ def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, re ) elif not resume and length is not None: assert all(torch.equal(x, y) for x, y in zip(batch, batches_1[i])) - expected_5 = [ - [torch.tensor([2]), torch.tensor([2])], - [torch.tensor([5]), torch.tensor([5])], - [torch.tensor([0]), torch.tensor([0])], - [torch.tensor([3]), torch.tensor([3])], - ] for i, batch in enumerate(dloader): if not shuffle: assert all( @@ -1123,8 +1181,9 @@ def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, re @pytest.mark.parametrize("shuffle", [False, True]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_partial_iteration_resume_without_dataloader(tmp_path_factory, length, resume, shuffle): + # Dataset-only iteration: no DataLoader workers are used, so keep num_workers=0. _, _, pardset, _, _ = prepare_parallel_dataset_and_dataloder( - tmp_path_factory, parlen=length, len1=10, len2=10, batch_size=1, num_workers=2, shuffle=shuffle, resume=resume + tmp_path_factory, parlen=length, len1=10, len2=10, batch_size=1, num_workers=0, shuffle=shuffle, resume=resume ) assert pardset.is_cycling() or length is None break_at = 3 @@ -1155,8 +1214,9 @@ def test_parallel_dataset_partial_iteration_resume_without_dataloader(tmp_path_f @pytest.mark.parametrize("shuffle", [False, True]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_complete_iteration_resume_without_dataloader(tmp_path_factory, length, resume, shuffle): + # Dataset-only iteration: no DataLoader workers are used, so keep num_workers=0. _, _, pardset, _, _ = prepare_parallel_dataset_and_dataloder( - tmp_path_factory, parlen=length, len1=4, len2=4, batch_size=1, num_workers=2, shuffle=shuffle, resume=resume + tmp_path_factory, parlen=length, len1=4, len2=4, batch_size=1, num_workers=0, shuffle=shuffle, resume=resume ) assert pardset.is_cycling() or length is None expected = [ From aa5f5df7a2f6394a3633b419c9b277d5ac08be93 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 00:42:47 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/streaming/test_parallel.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/streaming/test_parallel.py b/tests/streaming/test_parallel.py index 455f1b544..1c618375b 100644 --- a/tests/streaming/test_parallel.py +++ b/tests/streaming/test_parallel.py @@ -1078,7 +1078,12 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res @pytest.mark.parametrize( ("length", "resume", "shuffle", "num_workers"), [ - *[(length, resume, shuffle, 0) for length in (None, 4) for resume in (False, True) for shuffle in (False, True)], + *[ + (length, resume, shuffle, 0) + for length in (None, 4) + for resume in (False, True) + for shuffle in (False, True) + ], (4, True, False, 2), (None, False, True, 2), ], From 6545ef68056dfb6d59318f452f519cacb300922f Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Sat, 25 Jul 2026 00:47:39 +0000 Subject: [PATCH 3/3] test(parallel): restore full multi-worker resume coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep every resume/rng/state matrix cell on num_workers=2 (or the original 0/2 grid). Speedups remain shared Cache reuse, smaller dataloader fixtures, sleep→poll/clock backdating, and Ubuntu xdist=4 — not coverage cuts. Co-authored-by: Cursor --- tests/streaming/test_parallel.py | 83 ++++++++------------------------ 1 file changed, 21 insertions(+), 62 deletions(-) diff --git a/tests/streaming/test_parallel.py b/tests/streaming/test_parallel.py index 1c618375b..eb6bde45a 100644 --- a/tests/streaming/test_parallel.py +++ b/tests/streaming/test_parallel.py @@ -336,20 +336,10 @@ def rng_transform(_, rngs, which): return rngs[which].random() -@pytest.mark.parametrize( - ("length", "num_workers", "which", "reset_rngs"), - [ - # Combinatorial coverage stays on the cheap single-process path. - *[ - (length, 0, which, reset_rngs) - for length in (None, 7) - for which in ("random", "numpy", "torch") - for reset_rngs in (False, True) - ], - # One multi-worker cell proves spawn + RNG seeding still works. - (None, 2, "torch", True), - ], -) +@pytest.mark.parametrize("length", [None, 7]) +@pytest.mark.parametrize("num_workers", [0, 2]) +@pytest.mark.parametrize("which", ["random", "numpy", "torch"]) +@pytest.mark.parametrize("reset_rngs", [False, True]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_rng(length, num_workers, which, reset_rngs): transform = functools.partial(rng_transform, which=which) @@ -507,15 +497,8 @@ def test_parallel_dataset_dataloader_states_without_any_iterations(tmp_path_fact @pytest.mark.timeout(120) -@pytest.mark.parametrize( - ("length", "num_workers"), - [ - (None, 0), - (24, 0), - # One multi-worker cell covers spawn + multi-epoch state transitions. - (24, 2), - ], -) +@pytest.mark.parametrize("length", [None, 24]) +@pytest.mark.parametrize("num_workers", [0, 2]) @pytest.mark.parametrize("batch_size", [2]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_dataloader_states_complete_iterations(tmp_path_factory, length, num_workers, batch_size): @@ -575,17 +558,10 @@ def test_parallel_dataset_dataloader_states_complete_iterations(tmp_path_factory @pytest.mark.timeout(300) -@pytest.mark.parametrize( - ("length", "num_workers", "break_at"), - [ - (None, 0, 3), - (48, 0, 3), - (48, 0, 7), - # One multi-worker cell covers restore after early break under spawn. - (None, 2, 3), - ], -) +@pytest.mark.parametrize("length", [None, 20, 48]) +@pytest.mark.parametrize("num_workers", [0, 2]) @pytest.mark.parametrize("batch_size", [2]) +@pytest.mark.parametrize("break_at", [3, 7]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") def test_parallel_dataset_dataloader_states_partial_iterations( tmp_path_factory, length, num_workers, batch_size, break_at @@ -954,22 +930,13 @@ def test_parallel_infinite_restore_survives_early_break_without_hang(tmp_path_fa assert dloader.restore -@pytest.mark.parametrize( - ("length", "resume", "shuffle", "num_workers"), - [ - *[ - (length, resume, shuffle, 0) - for length in (None, 16) - for resume in (False, True) - for shuffle in (False, True) - ], - # Multi-worker coverage (float("inf") is covered by the dedicated infinite restore test). - (16, True, False, 2), - (None, False, True, 2), - ], -) +@pytest.mark.parametrize("length", [None, 16, float("inf")]) +@pytest.mark.parametrize("resume", [False, True]) +@pytest.mark.parametrize("shuffle", [False, True]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") -def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, resume, shuffle, num_workers): +def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, resume, shuffle): + # Keep num_workers=2 for every cell: resume + worker priming is the regression surface. + num_workers = 2 _, _, pardset, dloader, tmpdir = prepare_parallel_dataset_and_dataloder( tmp_path_factory, parlen=length, @@ -1075,21 +1042,13 @@ def test_parallel_dataset_partial_iteration_resume(tmp_path_factory, length, res break -@pytest.mark.parametrize( - ("length", "resume", "shuffle", "num_workers"), - [ - *[ - (length, resume, shuffle, 0) - for length in (None, 4) - for resume in (False, True) - for shuffle in (False, True) - ], - (4, True, False, 2), - (None, False, True, 2), - ], -) +@pytest.mark.parametrize("length", [None, 4]) +@pytest.mark.parametrize("resume", [False, True]) +@pytest.mark.parametrize("shuffle", [False, True]) @pytest.mark.skipif(sys.platform in ("win32", "darwin"), reason="too slow in CI") -def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, resume, shuffle, num_workers): +def test_parallel_dataset_complete_iteration_resume(tmp_path_factory, length, resume, shuffle): + # Keep num_workers=2 for every cell: full-epoch resume under spawn is the regression surface. + num_workers = 2 _, _, pardset, dloader, tmpdir = prepare_parallel_dataset_and_dataloder( tmp_path_factory, parlen=length,