Title: `drop_last_batch` does not drop the last batch using IterableDataset + interleave_datasets + multi_worker
Author: memray (NONE)
State: open
Created: 2025-03-08T10:28:44Z

### 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

Comments:
Hi @memray, I’d like to help fix the issue with `drop_last_batch` not working when `num_workers > 1`. I’ll investigate and propose a solution. Thanks!


Thank you very much for offering to help! I also noticed a problem related to a previous issue and left a comment [here](https://github.com/huggingface/datasets/issues/6565#issuecomment-2708169303) (the code checks the validity before certain columns removed). Can you take a look as well?