Skip to content

chore(tests): speed up streaming pytest suite #856

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
tchaton wants to merge 3 commits into
base: main
Choose a base branch
from
Open

chore(tests): speed up streaming pytest suite #856

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e796a82
chore(tests): speed up streaming pytest suite
tchaton Jul 25, 2026
aa5f5df
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 25, 2026
6545ef6
test(parallel): restore full multi-worker resume coverage
tchaton Jul 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion .github/workflows/ci-testing.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

@deependujha deependujha Jul 25, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we can also try standalone script for running tests


- name: Run processing tests sequentially
run: |
Expand Down
11 changes: 6 additions & 5 deletions tests/streaming/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import sys
from time import sleep, time
from time import time
from unittest import mock

import pytest
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
56 changes: 28 additions & 28 deletions tests/streaming/test_dataloader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,23 +203,23 @@ 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()
cache.merge()

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()
Expand All @@ -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
Expand All @@ -249,16 +249,16 @@ 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
count = 0
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())
Expand All @@ -267,26 +267,26 @@ 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()

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()
Expand All @@ -299,16 +299,16 @@ 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
count = 0
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())
Expand All @@ -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)
Expand All @@ -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()
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
14 changes: 8 additions & 6 deletions tests/streaming/test_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
Loading
Loading