GitHub issue #5374: Using too many threads results in: Got disconnected from remote data host. Retrying in 5sec
State: closed | Created: 2022-12-18T11:38:58Z

### Describe the bug

`streaming_download_manager` seems to disconnect if too many runs access the same underlying dataset 🧐
The code works fine for me if I have ~100 runs in parallel, but disconnects once scaling to 200.

Possibly related:
- https://github.com/huggingface/datasets/pull/3100
- https://github.com/huggingface/datasets/pull/3050

### Steps to reproduce the bug

Running
```python
c4 = datasets.load_dataset("c4", "en", split="train", streaming=True).skip(args.start).take(args.end-args.start)
df = pd.DataFrame(c4, index=None)
```
with different start & end arguments on 200 CPUs in parallel yields:
```
WARNING:datasets.load:Using the latest cached version of the module from /users/muennighoff/.cache/huggingface/modules/datasets_modules/datasets/c4/df532b158939272d032cc63ef19cd5b83e9b4d00c922b833e4cb18b2e9869b01 (last modified on Mon Dec 12 10:45:02 2022) since it couldn't be found locally at c4.
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [1/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [2/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [3/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [4/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [5/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [6/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [7/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [8/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [9/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [10/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [11/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [12/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [13/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [14/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [15/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [16/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [17/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [18/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [19/20]
WARNING:datasets.download.streaming_download_manager:Got disconnected from remote data host. Retrying in 5sec [20/20]
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /pfs/lustrep4/scratch/project_462000119/muennighoff/dec-2022-tasky/inference │
│ _c4.py:68 in <module>                                                        │
│                                                                              │
│   65 │   model.eval()                                                        │
│   66 │                                                                       │
│   67 │   c4 = datasets.load_dataset("c4", "en", split="train", streaming=Tru │
│ ❱ 68 │   df = pd.DataFrame(c4, index=None)                                   │
│   69 │   texts = df["text"].to_list()                                        │
│   70 │   preds = batch_inference(texts, batch_size=args.batch_size)          │
│   71                                                                         │
│                                                                              │
│ /opt/cray/pe/python/3.9.12.1/lib/python3.9/site-packages/pandas/core/frame.p │
│ y:684 in __init__                                                            │
│                                                                              │
│     681 │   │   # For data is list-like, or Iterable (will consume into list │
│     682 │   │   elif is_list_like(data):                                     │
│     683 │   │   │   if not isinstance(data, (abc.Sequence, ExtensionArray)): │
│ ❱   684 │   │   │   │   data = list(data)                                    │
│     685 │   │   │   if len(data) > 0:                                        │
│     686 │   │   │   │   if is_dataclass(data[0]):                            │
│     687 │   │   │   │   │   data = dataclasses_to_dicts(data)                │
│                                                                              │
│ /pfs/lustrep4/scratch/project_462000119/muennighoff/nov-2022-bettercom/venv/ │
│ lib/python3.9/site-packages/datasets/iterable_dataset.py:751 in __iter__     │
│                                                                              │
│    748 │   │   yield from ex_iterable.shard_data_sources(shard_idx)          │
│    749 │                                                                     │
│    750 │   def __iter__(self):                                               │
│ ❱  751 │   │   for key, example in self._iter():                             │
│    752 │   │   │   if self.features:                                         │
│    753 │   │   │   │   # `IterableDataset` automatically fills missing colum │
│    754 │   │   │   │   # This is done with `_apply_feature_types`.           │
│                                                                              │
│ /pfs/lustrep4/scratch/project_462000119/muennighoff/nov-2022-bettercom/venv/ │
│ lib/python3.9/site-packages/datasets/iterable_dataset.py:741 in _iter        │
│                                                                              │
│    738 │   │   │   ex_iterable = self._ex_iterable.shuffle_data_sources(self │
│    739 │   │   else:                                                         │
│    740 │   │   │   ex_iterable = self._ex_iterable                           │
│ ❱  741 │   │   yield from ex_iterable                                        │
│    742 │                                                                     │
│    743 │   def _iter_shard(self, shard_idx: int):                            │
│    744 │   │   if self._shuffling:                                           │
│                                                                              │
│ /pfs/lustrep4/scratch/project_462000119/muennighoff/nov-2022-bettercom/venv/ │
│ lib/python3.9/site-packages/datasets/iterable_dataset.py:617 in __iter__     │
│                                                                              │
│    614 │   │   self.n = n                                                    │
│    615 │                                                                     │
│    616 │   def __iter__(self):                                               │
│ ❱  617 │   │   yield from islice(self.ex_iterable, self.n)                   │
│    618 │                                                                     │
│    619 │   def shuffle_data_sources(self, generator: np.random.Generator) -> │
│    620 │   │   """Doesn't shuffle the wrapped examples iterable since it wou │
│                                                                              │
│ /pfs/lustrep4/scratch/project_462000119/muennighoff/nov-2022-bettercom/venv/ │
│ lib/python3.9/site-packages/datasets/iterable_dataset.py:594 in __iter__     │
│                                                                              │
│    591 │                                                                     │
│    592 │   def __iter__(self):                                               │
│    593 │   │   #ex_iterator = iter(self.ex_iterable)                         │
│ ❱  594 │   │   yield from islice(self.ex_iterable, self.n, None)             │
│    595 │   │   #for _ in range(self.n):                                      │
│    596 │   │   #    next(ex_iterator)                                        │
│    597 │   │   #yield from islice(ex_iterator, self.n, None)                 │
│                                                                              │
│ /pfs/lustrep4/scratch/project_462000119/muennighoff/nov-2022-bettercom/venv/ │
│ lib/python3.9/site-packages/datasets/iterable_dataset.py:106 in __iter__     │
│                                                                              │
│    103 │   │   self.kwargs = kwargs                                          │
│    104 │                                                                     │
│    105 │   def __iter__(self):                                               │
│ ❱  106 │   │   yield from self.generate_examples_fn(**self.kwargs)           │
│    107 │                                                                     │
│    108 │   def shuffle_data_sources(self, generator: np.random.Generator) -> │
│    109 │   │   return ShardShuffledExamplesIterable(self.generate_examples_f │
│                                                                              │
│ /users/muennighoff/.cache/huggingface/modules/datasets_modules/datasets/c4/d │
│ f532b158939272d032cc63ef19cd5b83e9b4d00c922b833e4cb18b2e9869b01/c4.py:89 in  │
│ _generate_examples                                                           │
│                                                                              │
│   86 │   │   for filepath in filepaths:                                      │
│   87 │   │   │   logger.info("generating examples from = %s", filepath)      │
│   88 │   │   │   with gzip.open(open(filepath, "rb"), "rt", encoding="utf-8" │
│ ❱ 89 │   │   │   │   for line in f:                                          │
│   90 │   │   │   │   │   if line:                                            │
│   91 │   │   │   │   │   │   example = json.loads(line)                      │
│   92 │   │   │   │   │   │   yield id_, example                              │
│                                                                              │
│ /opt/cray/pe/python/3.9.12.1/lib/python3.9/gzip.py:313 in read1              │
│                                                                              │
│   310 │   │                                                                  │
│   311 │   │   if size < 0:                                                   │
│   312 │   │   │   size = io.DEFAULT_BUFFER_SIZE                              │
│ ❱ 313 │   │   return self._buffer.read1(size)                                │
│   314 │                                                                      │
│   315 │   def peek(self, n):                                                 │
│   316 │   │   self._check_not_closed()                                       │
│                                                                              │
│ /opt/cray/pe/python/3.9.12.1/lib/python3.9/_compression.py:68 in readinto    │
│                                                                              │
│    65 │                                                                      │
│    66 │   def readinto(self, b):                                             │
│    67 │   │   with memoryview(b) as view, view.cast("B") as byte_view:       │
│ ❱  68 │   │   │   data = self.read(len(byte_view))                           │
│    69 │   │   │   byte_view[:len(data)] = data                               │
│    70 │   │   return len(data)                                               │
│    71                                                                        │
│                                                                              │
│ /opt/cray/pe/python/3.9.12.1/lib/python3.9/gzip.py:493 in read               │
│                                                                              │
│   490 │   │   │   │   self._new_member = False                               │
│   491 │   │   │                                                              │
│   492 │   │   │   # Read a chunk of data from the file                       │
│ ❱ 493 │   │   │   buf = self._fp.read(io.DEFAULT_BUFFER_SIZE)                │
│   494 │   │   │                                                              │
│   495 │   │   │   uncompress = self._decompressor.decompress(buf, size)      │
│   496 │   │   │   if self._decompressor.unconsumed_tail != b"":              │
│                                                                              │
│ /opt/cray/pe/python/3.9.12.1/lib/python3.9/gzip.py:96 in read                │
│                                                                              │
│    93 │   │   │   read = self._read                                          │
│    94 │   │   │   self._read = None                                          │
│    95 │   │   │   return self._buffer[read:] + \                             │
│ ❱  96 │   │   │   │      self.file.read(size-self._length+read)              │
│    97 │                                                                      │
│    98 │   def prepend(self, prepend=b''):                                    │
│    99 │   │   if self._read is None:                                         │
│                                                                              │
│ /pfs/lustrep4/scratch/project_462000119/muennighoff/nov-2022-bettercom/venv/ │
│ lib/python3.9/site-packages/datasets/download/streaming_download_manager.py: │
│ 365 in read_with_retries                                                     │
│                                                                              │
│   362 │   │   │   │   )                                                      │
│   363 │   │   │   │   time.sleep(config.STREAMING_READ_RETRY_INTERVAL)       │
│   364 │   │   else:                                                          │
│ ❱ 365 │   │   │   raise ConnectionError("Server Disconnected")               │
│   366 │   │   return out                                                     │
│   367 │                                                                      │
│   368 │   file_obj.read = read_with_retries                                  │
╰──────────────────────────────────────────────────────────────────────────────╯
ConnectionError: Server Disconnected

```

### Expected behavior

There should be no disconnect I think.

### Environment info

```
datasets=2.7.0
Python 3.9.12
```