Make Dataset streaming queries retryable
Opened by rojagtap (open)

### Feature request

Streaming datasets, as intended, do not load the entire dataset in memory or disk. However, while querying the next data chunk from the remote, sometimes it is possible that the service is down or there might be other issues that may cause the query to fail. In such a scenario, it would be nice to make these queries retryable (perhaps with a backoff strategy).

### Motivation

I was working on a model and the model checkpoints after every 1000 steps. At step 1800 I got a 504 HTTP status code error from Huggingface hub for my pytorch `dataloader`. Given the size of my model and data, it took around 2 hours to reach 1800 steps and now it will take about an hour to recover the lost 800. It would be better to get a retryable querying strategy.

### Your contribution

It would be better if someone having experience in this area takes this up as this would require some testing.

Comments:
Hi! The streaming mode also retries requests - `datasets.config.STREAMING_READ_MAX_RETRIES` (20 sec by default) controls the number of retries and `datasets.config.STREAMING_READ_RETRY_INTERVAL` (5 sec) the sleep time between retries.

> At step 1800 I got a 504 HTTP status code error from Huggingface hub for my pytorch dataloader

A minor Hub outage that we experienced yesterday could be the cause.

I wanted something similar. I have a huge dataset I want to process (laion-2b), but after processing several batches, it sometimes fails with this error: `HTTP 502 Bad Gateway for url`. I had the following code to handle it but this way I believe it restarts processing the data from the first batch? How can I set the attribute values you mention above?

```
iterable_dataset = load_dataset("laion/laion2B-multi", streaming=True, split='train')
dataloader = DataLoader(iterable_dataset, batch_size=131072, collate_fn=custom_collate_fn, num_workers=8)

MAX_RETRIES = 5
RETRY_WAIT = 10  # wait 10 seconds before retry

  for retry in range(MAX_RETRIES):
      try:
          for j, batch in enumerate(dataloader):
              < process batch>

      except HfHubHTTPError as e:
          if "502" in str(e) and retry < MAX_RETRIES - 1:
              logging.warning(f"Encountered a 502 error on batch {j}. Waiting for {RETRY_WAIT} seconds before retrying.")
              time.sleep(RETRY_WAIT)
              continue
          else:
              raise

Hey all! Wondering if there's a way of making Datasets streaming mode somewhat robust to Hub outages? Over the weekend, I got two quite cryptic errors, which I reckon were probably from Hub issues:

<details>
<summary> Stack Trace 1 </summary>

```
  File "/home/sanchitgandhi/small-12-4-tpu-timestamped-prob-0.2/run_distillation.py", line 2119, in <module>
    main()
  File "/home/sanchitgandhi/small-12-4-tpu-timestamped-prob-0.2/run_distillation.py", line 1954, in main
    for batch in train_loader:
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1325, in _next_data
    return self._process_data(data)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1371, in _process_data
    data.reraise()
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/_utils.py", line 694, in reraise
    raise exception
ConnectionError: Caught ConnectionError in DataLoader worker process 8.
Original Traceback (most recent call last):
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/aiohttp/connector.py", line 1155, in _create_direct_connection
    hosts = await asyncio.shield(host_resolved)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/aiohttp/connector.py", line 874, in _resolve_host
    addrs = await self._resolver.resolve(host, port, family=self._family)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/aiohttp/resolver.py", line 33, in resolve
    infos = await self._loop.getaddrinfo(
  File "/usr/lib/python3.10/asyncio/base_events.py", line 863, in getaddrinfo
    return await self.run_in_executor(
  File "/usr/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/usr/lib/python3.10/socket.py", line 955, in getaddrinfo
    for res in _socket.getaddrinfo(host, port, family, type, proto, flags):
socket.gaierror: [Errno -3] Temporary failure in name resolution
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
  File "/home/sanchitgandhi/datasets/src/datasets/download/streaming_download_manager.py", line 333, in read_with_retries
    out = read(*args, **kwargs)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/implementations/http.py", line 612, in read
    return super().read(length)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/spec.py", line 1856, in read
    out = self.cache._fetch(self.loc, self.loc + length)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/caching.py", line 439, in _fetch
    new = self.fetcher(self.end, bend)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/asyn.py", line 118, in wrapper
    return sync(self.loop, func, *args, **kwargs)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/asyn.py", line 103, in sync
    raise return_result
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/asyn.py", line 56, in _runner
    result[0] = await coro
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/implementations/http.py", line 660, in async_fetch_range
    r = await self.session.get(
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/aiohttp/client.py", line 562, in _request
    conn = await self._connector.connect(
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/aiohttp/connector.py", line 540, in connect
    proto = await self._create_connection(req, traces, timeout)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/aiohttp/connector.py", line 901, in _create_connection
    _, proto = await self._create_direct_connection(req, traces, timeout)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/aiohttp/connector.py", line 1169, in _create_direct_connection
    raise ClientConnectorError(req.connection_key, exc) from exc
aiohttp.client_exceptions.ClientConnectorError: Cannot connect to host huggingface.co:443 ssl:default [Temporary failure in name resolution]
The above exception was the direct cause of the following exception:
Traceback (most recent call last):
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 32, in fetch
    data.append(next(self.dataset_iter))
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1358, in __iter__
    yield from self._iter_pytorch()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1293, in _iter_pytorch
    for key, example in ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 982, in __iter__
    for x in self.ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 862, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 899, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 862, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 899, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 678, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 740, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 862, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 899, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1114, in __iter__
    for key, example in self.ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 429, in __iter__
    if not iterators[i].hasnext():
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 106, in hasnext
    self._thenext = next(self.it)
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 678, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 740, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1114, in __iter__
    for key, example in self.ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 281, in __iter__
    for key, pa_table in self.generate_tables_fn(**self.kwargs):
  File "/home/sanchitgandhi/.cache/huggingface/modules/datasets_modules/datasets/distil-whisper--switchboard-data/9472ee64cca0e1a7e11909c7033c2354511fa62805f81a2e07616980c765abfe/switchboard-data.py", line 247, in _generate_tables
    for record_batch in pf.iter_batches():
  File "pyarrow/_parquet.pyx", line 1327, in iter_batches
  File "/home/sanchitgandhi/datasets/src/datasets/download/streaming_download_manager.py", line 342, in read_with_retries
    raise ConnectionError("Server Disconnected") from disconnect_err
ConnectionError: Server Disconnected
```

</details>

<details>
<summary> Stack Trace 2 </summary>

```
  File "/home/sanchitgandhi/small-12-2-tpu-v3-timestamped-prob-0.2-bs-512/run_distillation.py", line 2119, in <module>
    main()
  File "/home/sanchitgandhi/small-12-2-tpu-v3-timestamped-prob-0.2-bs-512/run_distillation.py", line 1954, in main
    for batch in train_loader:
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1325, in _next_data
    return self._process_data(data)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 1371, in _process_data
    data.reraise()
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/_utils.py", line 694, in reraise
    raise exception
requests.exceptions.ConnectionError: Caught ConnectionError in DataLoader worker process 13.
Original Traceback (most recent call last):
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/connectionpool.py", line 791, in urlopen
    response = self._make_request(
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/connectionpool.py", line 537, in _make_request
    response = conn.getresponse()
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/connection.py", line 461, in getresponse
    httplib_response = super().getresponse()
  File "/usr/lib/python3.10/http/client.py", line 1375, in getresponse
    response.begin()
  File "/usr/lib/python3.10/http/client.py", line 318, in begin
    version, status, reason = self._read_status()
  File "/usr/lib/python3.10/http/client.py", line 287, in _read_status
    raise RemoteDisconnected("Remote end closed connection without"
http.client.RemoteDisconnected: Remote end closed connection without response
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/requests/adapters.py", line 486, in send
    resp = conn.urlopen(
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/connectionpool.py", line 845, in urlopen
    retries = retries.increment(
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/util/retry.py", line 470, in increment
    raise reraise(type(error), error, _stacktrace)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/util/util.py", line 38, in reraise
    raise value.with_traceback(tb)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/connectionpool.py", line 791, in urlopen
    response = self._make_request(
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/connectionpool.py", line 537, in _make_request
    response = conn.getresponse()
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/urllib3/connection.py", line 461, in getresponse
    httplib_response = super().getresponse()
  File "/usr/lib/python3.10/http/client.py", line 1375, in getresponse
    response.begin()
  File "/usr/lib/python3.10/http/client.py", line 318, in begin
    version, status, reason = self._read_status()
  File "/usr/lib/python3.10/http/client.py", line 287, in _read_status
    raise RemoteDisconnected("Remote end closed connection without"
urllib3.exceptions.ProtocolError: ('Connection aborted.', RemoteDisconnected('Remote end closed connection without response'))
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 32, in fetch
    data.append(next(self.dataset_iter))
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1358, in __iter__
    yield from self._iter_pytorch()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1293, in _iter_pytorch
    for key, example in ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 982, in __iter__
    for x in self.ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 862, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 899, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 862, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 899, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 678, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 740, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 862, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 899, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1114, in __iter__
    for key, example in self.ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 429, in __iter__
    if not iterators[i].hasnext():
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 106, in hasnext
    self._thenext = next(self.it)
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 678, in __iter__
    yield from self._iter()
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 740, in _iter
    for key, example in iterator:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 1114, in __iter__
    for key, example in self.ex_iterable:
  File "/home/sanchitgandhi/datasets/src/datasets/iterable_dataset.py", line 281, in __iter__
    for key, pa_table in self.generate_tables_fn(**self.kwargs):
  File "/home/sanchitgandhi/datasets/src/datasets/packaged_modules/parquet/parquet.py", line 87, in _generate_tables
    for batch_idx, record_batch in enumerate(
  File "pyarrow/_parquet.pyx", line 1327, in iter_batches
  File "/home/sanchitgandhi/datasets/src/datasets/download/streaming_download_manager.py", line 333, in read_with_retries
    out = read(*args, **kwargs)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/spec.py", line 1856, in read
    out = self.cache._fetch(self.loc, self.loc + length)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/fsspec/caching.py", line 189, in _fetch
    self.cache = self.fetcher(start, end)  # new block replaces old
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/huggingface_hub/hf_file_system.py", line 410, in _fetch_range
    r = http_backoff("GET", url, headers=headers)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/huggingface_hub/utils/_http.py", line 258, in http_backoff
    response = session.request(method=method, url=url, **kwargs)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/requests/sessions.py", line 589, in request
    resp = self.send(prep, **send_kwargs)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/requests/sessions.py", line 703, in send
    r = adapter.send(request, **kwargs)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/huggingface_hub/utils/_http.py", line 63, in send
    return super().send(request, *args, **kwargs)
  File "/home/sanchitgandhi/hf/lib/python3.10/site-packages/requests/adapters.py", line 501, in send
    raise ConnectionError(err, request=request)
requests.exceptions.ConnectionError: (ProtocolError('Connection aborted.', RemoteDisconnected('Remote end closed connection without response')), '(Request ID: 5fce9fc2-e22f-41c2-91af-529f13f1d611)')
```

</details>

Having streaming mode fail when the Hub goes down makes using it problematic for long training runs where large amounts of data is involved. However, this is the precise situation for which streaming mode is so appealing!

Wondering if there were a 'common' set of Hub errors that we could catch in `iterable_datasets` and prevent from crashing the script?

cc @lhoestq @mariosasko 

Errors are already caught and requests are already retried.

What you can do is increase the number of retries before an error is raised.

```python
import datasets

datasets.config.STREAMING_READ_MAX_RETRIES = 20  # default
datasets.config.STREAMING_READ_RETRY_INTERVAL = 5  # default
```