Repository: pytorch/vision

Issue:
Downloads from Google Drive return empty files / are still broken
## 🐛 Bug

All downloaded files are empty when downloading from google Drive via `torchvision.utils.download_file_from_google_drive` (or methods that resolve to this function, e.g. `download_url`)

## To Reproduce

Steps to reproduce the behavior:

1. install torchvision (tested with `master` branch at 959666891589c35e6d225943253f523cffbae4cc)
2. Run code below
```
from torchvision.datasets.utils import download_url, download_file_from_google_drive

try:
    download_url(
        "http://www.vision.caltech.edu/Image_Datasets/Caltech101/101_ObjectCategories.tar.gz",
        "./caltech101",
        filename="101_ObjectCategories.tar.gz",
        md5="b224c7392d521a49829488ab0f1120d9")
except:
    pass
finally:
    folder = './miniimagenet'
    gdrive_id = '16V_ZlkW4SsnNDtnGmaBRq2OoPmUOc5mY'
    gz_filename = 'mini-imagenet.tar.gz'
    gz_md5 = 'b38f1eb4251fb9459ecc8e7febf9b2eb'
    download_file_from_google_drive(gdrive_id, folder, gz_filename, md5=gz_md5)
```
3. Afterwards, both `101_ObjectCategories.tar.gz` and `mini-imagenet.tar.gz` are empty files

## Expected behavior

Download fails explicitly if google drive quota is exceeded, and succeeds otherwise.

## Additional context & Error source

Related Issues #3708 #2992
Related PRs: #3710 #3035

The issue stems from the fact that google drive issues a quota on downloads and that the returned `response.status_code` cannot be used to check the quota consistently (refer #2992). As a workaround, we therefore need to check the payload for the corresponding string via `_quota_exceeded`. Before #3710, this required parsing the whole payload/content, which was infeasible and therefore disabled in #3035.

However, the proposed solution in #3710 breaks `torchvision.utils.download_file_from_google_drive`. The reason is that one should only iterate once over the content of a Response, refer the [requests documentation](https://docs.python-requests.org/en/latest/user/advanced/#body-content-workflow) for streaming content.

Opposed to this, we currently construct iterators twice from a streaming Response: first [here](https://github.com/pytorch/vision/blob/master/torchvision/datasets/utils.py#L189), second [here](https://github.com/pytorch/vision/blob/master/torchvision/datasets/utils.py#L252).

As a result, the second iterator has length of 0 and therefore the files written to disk are empty.

## Proposed solution

Since we do not want to issue the same request twice to google drive (the response may be different between first request and 2nd one from google drive), I suggest to 
1. construct the Iterator + extract its first chunk inside `download_file_from_google_drive`
2. pass only the first chunk to `_quota_exceeded`
3. pass the first chunk + partially consumed Iterator to `_save_response_content` if the quota check is passed

I have begun working on such a solution here: 

https://github.com/ORippler/vision/tree/fix_google_drive_quotacheck

@pmeier 


Discussion:

Resolving patch:
diff --git a/torchvision/datasets/utils.py b/torchvision/datasets/utils.py
--- a/torchvision/datasets/utils.py
+++ b/torchvision/datasets/utils.py
@@ -5,15 +5,15 @@
 import gzip
 import re
 import tarfile
-from typing import Any, Callable, List, Iterable, Optional, TypeVar, Dict, IO, Tuple
+from typing import Any, Callable, List, Iterable, Optional, TypeVar, Dict, IO, Tuple, Iterator
 from urllib.parse import urlparse
 import zipfile
 import lzma
-import contextlib
 import urllib
 import urllib.request
 import urllib.error
 import pathlib
+import itertools
 
 import torch
 from torch.utils.model_zoo import tqdm
@@ -184,11 +184,10 @@ def list_files(root: str, suffix: str, prefix: bool = False) -> List[str]:
     return files
 
 
-def _quota_exceeded(response: "requests.models.Response") -> bool:  # type: ignore[name-defined]
+def _quota_exceeded(first_chunk: bytes) -> bool:  # type: ignore[name-defined]
     try:
-        start = next(response.iter_content(chunk_size=128, decode_unicode=True))
-        return isinstance(start, str) and "Google Drive - Quota exceeded" in start
-    except StopIteration:
+        return "Google Drive - Quota exceeded" in first_chunk.decode()
+    except UnicodeDecodeError:
         return False
 
 
@@ -224,7 +223,16 @@ def download_file_from_google_drive(file_id: str, root: str, filename: Optional[
             params = {'id': file_id, 'confirm': token}
             response = session.get(url, params=params, stream=True)
 
-        if _quota_exceeded(response):
+        # Ideally, one would use response.status_code to check for quota limits, but google drive is not consistent
+        # with their own API, refer https://github.com/pytorch/vision/issues/2992#issuecomment-730614517.
+        # Should this be fixed at some place in future, one could refactor the following to no longer rely on decoding
+        # the first_chunk of the payload
+        response_content_generator = response.iter_content(32768)
+        first_chunk = None
+        while not first_chunk:  # filter out keep-alive new chunks
+            first_chunk = next(response_content_generator)
+
+        if _quota_exceeded(first_chunk):
             msg = (
                 f"The daily quota of the file {filename} is exceeded and it "
                 f"can't be downloaded. This is a limitation of Google Drive "
@@ -232,7 +240,8 @@ def download_file_from_google_drive(file_id: str, root: str, filename: Optional[
             )
             raise RuntimeError(msg)
 
-        _save_response_content(response, fpath)
+        _save_response_content(itertools.chain((first_chunk, ), response_content_generator), fpath)
+        response.close()
 
 
 def _get_confirm_token(response: "requests.models.Response") -> Optional[str]:  # type: ignore[name-defined]
@@ -244,12 +253,13 @@ def _get_confirm_token(response: "requests.models.Response") -> Optional[str]:
 
 
 def _save_response_content(
-    response: "requests.models.Response", destination: str, chunk_size: int = 32768,  # type: ignore[name-defined]
+    response_gen: Iterator[bytes], destination: str,  # type: ignore[name-defined]
 ) -> None:
     with open(destination, "wb") as f:
         pbar = tqdm(total=None)
         progress = 0
-        for chunk in response.iter_content(chunk_size):
+
+        for chunk in response_gen:
             if chunk:  # filter out keep-alive new chunks
                 f.write(chunk)
                 progress += len(chunk)


Test patch: