Fix streaming: catch Timeout error
Opened by borisdayma (closed)

Catches Timeout error during streaming.

fix #3049 

Comments:
I'm running a large test.
Let's see if I get any error within a few days.

This time it stopped after 8h but correctly raised `ConnectionError: Server Disconnected`.

Traceback:
```
Traceback (most recent call last):                                                                                                                                                                               
  File "/home/koush/dalle-mini/dev/seq2seq/run_seq2seq_flax.py", line 1027, in <module>                                                                                                                          
    main()                                                                                                                                                                                                       
  File "/home/koush/dalle-mini/dev/seq2seq/run_seq2seq_flax.py", line 991, in main                                                                                                                               
    for batch in tqdm(                                                                                                                                                                                           
  File "/home/koush/.pyenv/versions/dev/lib/python3.9/site-packages/tqdm/std.py", line 1180, in __iter__                                                                                                         
    for obj in iterable:                                                                                                                                                                                         
  File "/home/koush/dalle-mini/dev/seq2seq/run_seq2seq_flax.py", line 376, in data_loader_streaming
    for item in dataset:
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 341, in __iter__
    for key, example in self._iter():
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 338, in _iter
    yield from ex_iterable
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 179, in __iter__
    key_examples_list = [(key, example)] + [
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 179, in <listcomp>
    key_examples_list = [(key, example)] + [
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 176, in __iter__
    for key, example in iterator:
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 225, in __iter__
    for x in self.ex_iterable:
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 99, in __iter__
    for key, example in self.generate_examples_fn(**kwargs_with_shuffled_shards):
  File "/home/koush/datasets/src/datasets/iterable_dataset.py", line 287, in wrapper
    for key, table in generate_tables_fn(**kwargs):
  File "/home/koush/datasets/src/datasets/packaged_modules/json/json.py", line 107, in _generate_tables
    batch = f.read(self.config.chunksize)
  File "/home/koush/datasets/src/datasets/utils/streaming_download_manager.py", line 136, in read_with_retries
    raise ConnectionError("Server Disconnected")
ConnectionError: Server Disconnected
```

Right before this error, the warnings were correctly raised:

```
10/10/2021 06:02:26 - WARNING - datasets.utils.streaming_download_manager - Got disconnected from remote data host. Retrying in 1sec [1/3]
10/10/2021 06:02:27 - WARNING - datasets.utils.streaming_download_manager - Got disconnected from remote data host. Retrying in 1sec [2/3]                                                                       
10/10/2021 06:02:28 - WARNING - datasets.utils.streaming_download_manager - Got disconnected from remote data host. Retrying in 1sec [3/3
```

I'm going to see what happens if I change the max retries to 20 and the interval to 5.

Also maybe we can raise the Server Disconnected error with more info about what kind of error caused it (client error, time out, etc.)

I have 2 runs:
* [run 1](https://wandb.ai/dalle-mini/dalle-mini/runs/1nj161cl?workspace=user-borisd13) with [this data](https://huggingface.co/datasets/dalle-mini/encoded) that I will remove soon because I now use the 2nd one
* [run 2](https://wandb.ai/dalle-mini/dalle-mini/runs/he9rrc3q?workspace=user-borisd13) with [this data](https://huggingface.co/datasets/dalle-mini/encoded-vqgan_imagenet_f16_16384)
* `load_dataset(dataset_repo, data_files={'train':'data/train/*.jsonl', 'validation':'data/valid/*.jsonl'}, streaming=True)`

They have now been running by a bit more than a day for one run and 15h for the other.

The error logs are not shown in wandb because the script use `pylogging` (not sure why, I should change it) but basically so far with the new settings I had one timeout in each with successful reconnect afterwards.

So I think it's a good idea to have:
* `STREAMING_READ_RETRY_INTERVAL = 5` since before my runs would get 3 errors in a row (with the default 1 second pause)
* `STREAMING_READ_MAX_RETRIES` should also be increased. Since this type of error does not happen a lot, I would still have a large number (at least 10) because a stopped training run may be a big issue if checkpointing/restart is not well implemented which is not always trivial

I agree ! Feel free to open a PR to increase both values