Lightning-AI / Lightning-AI/litData

Improve CombinedStreamingDataset to handle multiple subdatasets efficiently

Open
#386 2 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
614
Forks
106
Avg merge
15h 8m
Merged PRs (30d)
22

Description

**Issue Title:** Improve CombinedStreamingDataset to handle multiple subdatasets efficiently

**Description:**

As raised by Emile Clastres, there is a performance issue when using `CombinedStreamingDataset` on hundreds of sub-datasets. The dataset crashes after a few batches, making it impractical for scenarios where multiple subdatasets need to be combined.

**[Initial User Inquiry:](https://discord.com/channels/1077906959069626439/1290640326779736157)**

> "I have an S3 bucket in which I have hundreds of directories containing the output of LitData's `optimize` on some subset of my full data. This structure comes from the fact that each subdataset was processed independently on different machines and called its own `optimize`. I was very happy about this since it also allows me to combine subdatasets flexibly to create train/val/test splits.
>
> However, it seems that `CombinedStreamingDataset` has terrible performance when used on hundreds of sub-datasets. It even crashes after a few batches have been yielded.

**Suggested Solution:**

As per @tchaton ' suggestion, We could also re-think the CombinedDataset to virtually re-combine the different index.json into one and decide how to fetch the chunks accordingly using the right remote path. This would be pretty interesting. Same performance as normal StreamingDataset but no need to combine them and create data copy.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by locating CombinedStreamingDataset and reproducing the failure with hundreds of subdatasets, as described in the issue. Compare its behavior with a normal StreamingDataset and inspect how the separate index.json files and remote paths are handled. Done means large combinations no longer crash after a few batches and have practical performance.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.