feat(litdata): Add Support for StreamingRawDataset to Stream Raw Datasets from Cloud Storage

[bhimrazy]

What does this PR do?

Overview

This PR introduces a new StreamingRawDataset class designed for efficient streaming of raw files directly from cloud storage (e.g., S3, GCS).

Unlike optimized formats (e.g., litdata chunks), this class allows users to stream raw datasets (e.g., images, text files, etc) without requiring prior preprocessing or data conversion. It supports scalable, on-the-fly access to data, making it well-suited for training or inference workflows where immediate access to raw files is needed.

Note: This is a beta feature and may be subject to minor changes in future updates based on feedback and pending follow-ups.

Usage Example

from litdata.streaming.raw_dataset import StreamingRawDataset
from torch.utils.data import DataLoader

# Initialize the streaming raw dataset from S3 path
dataset = StreamingRawDataset("s3://bucket/files/") # also accepts transform here

# Access a single file (raw bytes)
first_sample = dataset[0]
print(f"Type of first sample: {type(first_sample)}")  # return bytes

# Access multiple files by indices using __getitems__
indices = [0, 1, 2]
samples = dataset.__getitems__(indices) # returns a list of 'bytes'
print(f"Retrieved {len(samples)} samples. Type of first sample: {type(samples[0])}")

# Use PyTorch DataLoader to batch iterate over the dataset
dataloader = DataLoader(dataset, batch_size=32)

for batch_idx, batch in enumerate(dataloader):
    print(f"Batch {batch_idx} size: {len(batch)}")
    # Each item in batch is raw bytes of a file
    # You can process batch here...

    # For demo, only process first batch
    break
    
 
# output
# Type of first sample: <class 'bytes'>
# Retrieved 3 samples. Type of first sample: <class 'bytes'>
# Batch 0 size: 32

---

ImageNet Benchmarks

# without transform
python benchmarks/stream_raw_imagenet.py --bytes --epoch 1 

# with transform
python benchmarks/stream_raw_imagenet.py --epoch 1 

---

Follow-ups

• Add support for grouping files, resolved by feat(litdata): Add grouping functionality in StreamingRawDataset #665
• Add support for storing index file to the source cloud storage, resolved by feat(litdata/raw): Implement remote and local index caching for StreamingRawDataset #666

  Note from thomas: Can we add an argument to the dataset to re-compute the index. Once the index is computed, can we push it the index to the folder, so next read is fast unless that argument is set to True.

• Switch indexing from fsspec to obstore
• Move the file indexer to its own folder & reorganize, resolved by ref: Move raw dataset code to litdata/raw, expose StreamingRawDataset at top-level #671
• Check for caching files
• Test for azure and other scenarios
• Possible runtime error from eventloop , reported by copilot, resolved by fix: StreamingRawDataset Async Handling #661
• Add support for gcs_folders in resolve_dir : resolved by Add GCP support for directory resolution inresolve_dir #659

  Note from thomas: it would be good to support internal studio path such as /teamspace/gcs_folders/... like we do for StreamingDataset

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

❌ Patch coverage is 78.59649% with 61 lines in your changes missing coverage. Please review.
✅ Project coverage is 83%. Comparing base (aa467f9) to head (9c34bac).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@         Coverage Diff          @@
##           main   #652    +/-   ##
====================================
- Coverage    83%    83%    -0%     
====================================
  Files        49     50     +1     
  Lines      6812   7097   +285     
====================================
+ Hits       5686   5910   +224     
- Misses     1126   1187    +61     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull Request Overview

This PR introduces a new StreamingRawDataset class for efficient streaming of raw files directly from cloud storage (S3, GCS, Azure) without requiring preprocessing. The implementation supports scalable, on-the-fly access to raw datasets with features like fast indexing, local caching, and both synchronous and asynchronous file downloads.

Key changes include:

• New StreamingRawDataset class with support for __getitem__ and __getitems__ methods
• Enhanced downloader classes with download_fileobj and adownload_fileobj methods for direct file streaming
• File indexing system with ZSTD compression for efficient metadata caching

Reviewed Changes

Copilot reviewed 10 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/litdata/streaming/raw_dataset.py	Core implementation of StreamingRawDataset, FileIndexer, and CacheManager classes
src/litdata/streaming/downloader.py	Added download_fileobj and adownload_fileobj methods to all downloader classes
src/litdata/streaming/client.py	Enhanced S3Client to store session reference for obstore integration
src/litdata/constants.py	Added new requirement checks for obstore and asyncio
tests/streaming/test_raw_dataset.py	Comprehensive test suite for the new streaming raw dataset functionality
tests/streaming/test_downloader.py	Tests for new downloader methods
tests/conftest.py	Added obstore mock fixture for testing
requirements/test.txt	Added pytest-asyncio dependency for async test support
requirements/extras.txt	Added asyncio dependency
benchmarks/stream_raw_imagenet.py	Benchmark script for ImageNet streaming performance testing

Comments suppressed due to low confidence (1)

requirements/extras.txt:9

• The 'asyncio' module is part of Python's standard library since Python 3.4 and does not need to be installed as an external dependency. Remove this line as it will cause installation errors.

asyncio

[tchaton]

Looks really neat !