IterableDataset raises exception instead of retrying

[bauwenst]

Describe the bug

In light of the recent server outages, I decided to look into whether I could somehow wrap my IterableDataset streams to retry rather than error out immediately. To my surprise, datasets already supports retries. Since a commit by @lhoestq last week, that code lives here:

https://github.com/huggingface/datasets/blob/fe2bea6a4b09b180bd23b88fe96dfd1a11191a4f/src/datasets/utils/file_utils.py#L1097C1-L1111C19

If GitHub code snippets still aren't working, here's a copy:

    def read_with_retries(*args, **kwargs):
        disconnect_err = None
        for retry in range(1, max_retries + 1):
            try:
                out = read(*args, **kwargs)
                break
            except (ClientError, TimeoutError) as err:
                disconnect_err = err
                logger.warning(
                    f"Got disconnected from remote data host. Retrying in {config.STREAMING_READ_RETRY_INTERVAL}sec [{retry}/{max_retries}]"
                )
                time.sleep(config.STREAMING_READ_RETRY_INTERVAL)
        else:
            raise ConnectionError("Server Disconnected") from disconnect_err
        return out

With the latest outage, the end of my stack trace looked like this:

...
  File "/miniconda3/envs/draft/lib/python3.11/site-packages/datasets/download/streaming_download_manager.py", line 342, in read_with_retries
    out = read(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/draft/lib/python3.11/gzip.py", line 301, in read
    return self._buffer.read(size)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/draft/lib/python3.11/_compression.py", line 68, in readinto
    data = self.read(len(byte_view))
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/draft/lib/python3.11/gzip.py", line 505, in read
    buf = self._fp.read(io.DEFAULT_BUFFER_SIZE)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/draft/lib/python3.11/gzip.py", line 88, in read
    return self.file.read(size)
           ^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/draft/lib/python3.11/site-packages/fsspec/spec.py", line 1856, in read
    out = self.cache._fetch(self.loc, self.loc + length)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/draft/lib/python3.11/site-packages/fsspec/caching.py", line 189, in _fetch
    self.cache = self.fetcher(start, end)  # new block replaces old
                 ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/miniconda3/envs/draft/lib/python3.11/site-packages/huggingface_hub/hf_file_system.py", line 626, in _fetch_range
    hf_raise_for_status(r)
  File "/miniconda3/envs/draft/lib/python3.11/site-packages/huggingface_hub/utils/_errors.py", line 333, in hf_raise_for_status
    raise HfHubHTTPError(str(e), response=response) from e
huggingface_hub.utils._errors.HfHubHTTPError: 504 Server Error: Gateway Time-out for url: https://huggingface.co/datasets/allenai/c4/resolve/1588ec454efa1a09f29cd18ddd04fe05fc8653a2/en/c4-train.00346-of-01024.json.gz

Indeed, the code for retries only catches ClientErrors and TimeoutErrors, and all other exceptions, including HuggingFace's own custom HTTP error class, are not caught. Nothing is retried, and instead the exception is propagated upwards immediately.

Steps to reproduce the bug

Not sure how you reproduce this. Maybe unplug your Ethernet cable while streaming a dataset; the issue is pretty clear from the stack trace.

Expected behavior

All HTTP errors while iterating a streamable dataset should cause retries.

Environment info

Output from datasets-cli env:

• datasets version: 2.18.0
• Platform: Linux-4.18.0-513.24.1.el8_9.x86_64-x86_64-with-glibc2.28
• Python version: 3.11.7
• huggingface_hub version: 0.20.3
• PyArrow version: 15.0.0
• Pandas version: 2.2.0
• fsspec version: 2023.10.0

[mariosasko]

Thanks for reporting! I've opened a PR with a fix.

[bauwenst]

Thanks, @mariosasko! Related question (although I guess this is a feature request): could we have some kind of exponential back-off for these retries? Here's my reasoning:

• If a one-time accidental error happens, you should retry immediately and will succeed immediately.
• If the Hub has a small outage on the order of minutes, you don't want to retry on the order of hours.
• If the Hub has a prologned outage of several hours, we don't want to keep retrying on the order of minutes.

There actually already exists an implementation for (clipped) exponential backoff in the HuggingFace suite (here), but I don't think it is used here.

The requirements are basically that you have an initial minimum waiting time and a maximum waiting time, and with each retry, the waiting time is doubled. We don't want to overload your servers with needless retries, especially when they're down :sweat_smile:

[mariosasko]

Oh, I've just remembered that we added retries to the HfFileSystem in huggingface_hub 0.21.0 (see this), so I'll close the linked PR as we don't want to retry the retries :).

I agree with the exponential backoff suggestion, so I'll open another PR.

[bauwenst]

@mariosasko The call you linked indeed points to the implementation I linked in my previous comment, yes, but it has no configurability. Arguably, you want to have this hidden backoff under the hood that catches small network disturbances on the time scale of seconds -- perhaps even with hardcoded limits as is the case currently -- but you also still want to have a separate backoff on top of that with the configurability as suggested by @lhoestq in the comment I linked.

My particular use-case is that I'm streaming a dataset while training on a university cluster with a very long scheduling queue. This means that when the backoff runs out of retries (which happens in under 30 seconds with the call you linked), I lose my spot on the cluster and have to queue for a whole day or more. Ideally, I should be able to specify that I want to retry for 2 to 3 hours but with more and more time between requests, so that I can smooth over hours-long outages without a setback of days.

[haydn-jones]

I also have my runs crash a surprising amount due to the dataloader crashing because of the hub, some way to address this would be nice.