Title: Is it possible to shard a single-sharded IterableDataset?
Author: jonathanasdf (NONE)
State: closed
Created: 2025-03-21T04:33:37Z

I thought https://github.com/huggingface/datasets/pull/7252 might be applicable but looking at it maybe not.

Say we have a process, eg. a database query, that can return data in slightly different order each time. So, the initial query needs to be run by a single thread (not to mention running multiple times incurs more cost too). But the results are also big enough that we don't want to materialize it entirely and instead stream it with an IterableDataset.

But after we have the results we want to split it up across workers to parallelize processing.

Is something like this possible to do?

Here's a failed attempt. The end result should be that each of the shards has unique data, but unfortunately with this attempt the generator gets run once in each shard and the results end up with duplicates...

```
import random
import datasets


def gen():
  print('RUNNING GENERATOR!')
  items = list(range(10))
  random.shuffle(items)
  yield from items


ds = datasets.IterableDataset.from_generator(gen)

print('dataset contents:')
for item in ds:
  print(item)
print()

print('dataset contents (2):')
for item in ds:
  print(item)
print()


num_shards = 3


def sharded(shard_id):
  for i, example in enumerate(ds):
    if i % num_shards in shard_id:
      yield example


ds1 = datasets.IterableDataset.from_generator(
  sharded, gen_kwargs={'shard_id': list(range(num_shards))}
)

for shard in range(num_shards):
  print('shard', shard)
  for item in ds1.shard(num_shards, shard):
    print(item)
```

Comments:
Hi ! Maybe you can look for an option in your dataset to partition your data based on a deterministic filter ? For example each worker could stream the data based on `row.id % num_shards` or something like that ?

So the recommendation is to start out with multiple shards initially and re-sharding after is not expected to work? :(

Would something like the following work? Some DiskCachingIterableDataset, where worker 0 streams from the datasource, but also writes to disk, and all of the other workers read from what worker 0 wrote? Then that would produce a stream with a deterministic order and we can subsample.

To be honest it would be cool to support native multiprocessing in `IterableDataset.map` so you can parallelize any specific processing step without having to rely on a torch Dataloader. What do you think ?

related: https://github.com/huggingface/datasets/issues/7193 https://github.com/huggingface/datasets/issues/3444 
original issue: https://github.com/huggingface/datasets/issues/2642

Alternatively the DiskCachingIterableDataset idea works, just note that to make it work with a torch Dataloader with num_workers>0 you'll need:
1. to make your own `torch.utils.data.IterableDataset` and have rank=0 stream the data and share them with the other workers (either via disk as suggested or IPC)
2. take into account that`datasets.IterableDataset` will yield 0 examples for ranks with id>0 if there is only one shard, but in your case it's ok since you'd only stream from rank=0

Ohh that would be pretty cool!

Thanks for the suggestions, as there's no actionable items for this repo I'm going to close this issue now.

Another usecase for this resharding:

If we have a bunch of jsonl files, and we load it as an IterableDataset with multiple dataloader workers, each file gets naively assigned to a worker.

If the files were not carefully produced to be equally sized, eg if the very last file is significantly shorter, containing just a few examples, and it gets assigned onto a dataloader worker by itself, then the examples in that file will be significantly oversampled.

It would be nice if datasets had an internal way to rebalance this without requiring offline reprocessing of the data files