huggingface / huggingface/datasets

`drop_last_batch` does not drop the last batch using IterableDataset + interleave_datasets + multi_worker

Open
#7,441 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
22k
Forks
3.4k
Avg merge
5d 7h
Merged PRs (30d)
17

Description

### Describe the bug

See the script below
`drop_last_batch=True` is defined using map() for each dataset.
The last batch for each dataset is expected to be dropped, id 21-25.
The code behaves as expected when num_workers=0 or 1.
When using num_workers>1, 'a-11', 'b-11', 'a-12', 'b-12' are gone and instead 21 and 22 are sampled.

### Steps to reproduce the bug

```
from datasets import Dataset
from datasets import interleave_datasets
from torch.utils.data import DataLoader

def convert_to_str(batch, dataset_name):
batch['a'] = [f"{dataset_name}-{e}" for e in batch['a']]
return batch

def gen1():
for ii in range(1, 25):
yield {"a": ii}

def gen2():
for ii in range(1, 25):
yield {"a": ii}

# https://github.com/huggingface/datasets/issues/6565
if __name__ == '__main__':
dataset1 = Dataset.from_generator(gen1).to_iterable_dataset(num_shards=2)
dataset2 = Dataset.from_generator(gen2).to_iterable_dataset(num_shards=2)
dataset1 = dataset1.map(lambda x: convert_to_str(x, dataset_name="a"), batched=True, batch_size=10, drop_last_batch=True)
dataset2 = dataset2.map(lambda x: convert_to_str(x, dataset_name="b"), batched=True, batch_size=10, drop_last_batch=True)

interleaved = interleave_datasets([dataset1, dataset2], stopping_strategy="all_exhausted")

print(f"num_workers=0")
loader = DataLoader(interleaved, batch_size=5, num_workers=0)
i = 0
for b in loader:
print(i, b['a'])
i += 1

print('=-' * 20)
print(f"num_workers=1")
loader = DataLoader(interleaved, batch_size=5, num_workers=1)
i = 0
for b in loader:
print(i, b['a'])
i += 1

print('=-' * 20)
print(f"num_workers=2")
loader = DataLoader(interleaved, batch_size=5, num_workers=2)
i = 0
for b in loader:
print(i, b['a'])
i += 1

print('=-' * 20)
print(f"num_workers=3")
loader = DataLoader(interleaved, batch_size=5, num_workers=3)
i = 0
for b in loader:
print(i, b['a'])
i += 1
```

output is:
```
num_workers=0
0 ['a-1', 'b-1', 'a-2', 'b-2', 'a-3']
1 ['b-3', 'a-4', 'b-4', 'a-5', 'b-5']
2 ['a-6', 'b-6', 'a-7', 'b-7', 'a-8']
3 ['b-8', 'a-9', 'b-9', 'a-10', 'b-10']
4 ['a-11', 'b-11', 'a-12', 'b-12', 'a-13']
5 ['b-13', 'a-14', 'b-14', 'a-15', 'b-15']
6 ['a-16', 'b-16', 'a-17', 'b-17', 'a-18']
7 ['b-18', 'a-19', 'b-19', 'a-20', 'b-20']
=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
num_workers=1
0 ['a-1', 'b-1', 'a-2', 'b-2', 'a-3']
1 ['b-3', 'a-4', 'b-4', 'a-5', 'b-5']
2 ['a-6', 'b-6', 'a-7', 'b-7', 'a-8']
3 ['b-8', 'a-9', 'b-9', 'a-10', 'b-10']
4 ['a-11', 'b-11', 'a-12', 'b-12', 'a-13']
5 ['b-13', 'a-14', 'b-14', 'a-15', 'b-15']
6 ['a-16', 'b-16', 'a-17', 'b-17', 'a-18']
7 ['b-18', 'a-19', 'b-19', 'a-20', 'b-20']
=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
num_workers=2
0 ['a-1', 'b-1', 'a-2', 'b-2', 'a-3']
1 ['a-13', 'b-13', 'a-14', 'b-14', 'a-15']
2 ['b-3', 'a-4', 'b-4', 'a-5', 'b-5']
3 ['b-15', 'a-16', 'b-16', 'a-17', 'b-17']
4 ['a-6', 'b-6', 'a-7', 'b-7', 'a-8']
5 ['a-18', 'b-18', 'a-19', 'b-19', 'a-20']
6 ['b-8', 'a-9', 'b-9', 'a-10', 'b-10']
7 ['b-20', 'a-21', 'b-21', 'a-22', 'b-22']
=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-
num_workers=3
Too many dataloader workers: 3 (max is dataset.num_shards=2). Stopping 1 dataloader workers.
0 ['a-1', 'b-1', 'a-2', 'b-2', 'a-3']
1 ['a-13', 'b-13', 'a-14', 'b-14', 'a-15']
2 ['b-3', 'a-4', 'b-4', 'a-5', 'b-5']
3 ['b-15', 'a-16', 'b-16', 'a-17', 'b-17']
4 ['a-6', 'b-6', 'a-7', 'b-7', 'a-8']
5 ['a-18', 'b-18', 'a-19', 'b-19', 'a-20']
6 ['b-8', 'a-9', 'b-9', 'a-10', 'b-10']
7 ['b-20', 'a-21', 'b-21', 'a-22', 'b-22']

```

### Expected behavior

`'a-21', 'b-21', 'a-22', 'b-22'` should be dropped

### Environment info

- `datasets` version: 3.3.2
- Platform: Linux-5.15.0-1056-aws-x86_64-with-glibc2.31
- Python version: 3.10.16
- `huggingface_hub` version: 0.28.0
- PyArrow version: 19.0.0
- Pandas version: 2.2.3
- `fsspec` version: 2024.6.1

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.