Cannot shuffle interleaved IterableDataset with "all_exhausted" stopping strategy

[offchan42]

Describe the bug

Shuffling interleaved IterableDataset with "all_exhausted" strategy yields non-exhaustive sampling.

Steps to reproduce the bug

from datasets import IterableDataset, interleave_datasets

def gen(bias, length):
  for i in range(length):
    yield dict(a=bias+i)

seed = 42
probabilities = [0.2, 0.6, 0.2]
d1 = IterableDataset.from_generator(lambda: gen(0, 3))
d2 = IterableDataset.from_generator(lambda: gen(10, 4))
d3 = IterableDataset.from_generator(lambda: gen(20, 3))
ds = interleave_datasets([d1, d2, d3], probabilities=probabilities, seed=seed, stopping_strategy='all_exhausted')
ds = ds.shuffle(buffer_size=1000)
for x in ds:
  print(x)

This code produces

{'a': 0}
{'a': 22}
{'a': 20}
{'a': 21}
{'a': 10}
{'a': 1}

Expected behavior

It should produce a longer list of examples to exhaust all the datasets.
If you comment out the shuffle line, it will exhaust all the datasets properly.
Here is the output if you comment out shuffling:

{'a': 10}
{'a': 11}
{'a': 20}
{'a': 12}
{'a': 0}
{'a': 21}
{'a': 13}
{'a': 10}
{'a': 1}
{'a': 11}
{'a': 12}
{'a': 22}
{'a': 13}
{'a': 20}
{'a': 10}
{'a': 11}
{'a': 12}
{'a': 2}

Environment info

• datasets version: 2.12.0
• Platform: Linux-5.10.147+-x86_64-with-glibc2.31
• Python version: 3.10.11
• Huggingface_hub version: 0.14.1
• PyArrow version: 9.0.0
• Pandas version: 1.5.3

This was run on Google Colab.