Lightning-AI / Lightning-AI/litData
Improve CombinedStreamingDataset to handle multiple subdatasets efficiently
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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