Merge pull request #2809 from casperdcl/gs-progress

GS progress for push & pull

dvc/remote/gs.py

@@ -1,14 +1,17 @@
-from __future__ import unicode_literals
+from __future__ import unicode_literals, division
 
 import logging
 from datetime import timedelta
 from functools import wraps
+import io
+import os.path
 
 from funcy import cached_property
 
 from dvc.config import Config
 from dvc.exceptions import DvcException
 from dvc.path_info import CloudURLInfo
+from dvc.progress import Tqdm
 from dvc.remote.base import RemoteBASE
 from dvc.scheme import Schemes
 from dvc.utils.compat import FileNotFoundError  # skipcq: PYL-W0622
@@ -20,25 +23,24 @@ def dynamic_chunk_size(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         import requests
-        from google.cloud.storage.blob import Blob, _DEFAULT_CHUNKSIZE
+        from google.cloud.storage.blob import Blob
 
-        # Default chunk size for gs is 100M, which might be too much for
-        # particular network (see [1]). So if we are getting ConnectionError,
-        # we should try lowering the chunk size until we reach the minimum
-        # allowed chunk size of 256K. Also note that `chunk_size` must be a
-        # multiple of 256K per the API specification.
+        # `ConnectionError` may be due to too large `chunk_size`
+        # (see [#2572]) so try halving on error.
+        # Note: start with 40 * [default: 256K] = 10M.
+        # Note: must be multiple of 256K.
         #
-        # [1] https://github.com/iterative/dvc/issues/2572
+        # [#2572]: https://github.com/iterative/dvc/issues/2572
 
         # skipcq: PYL-W0212
-        multiplier = int(_DEFAULT_CHUNKSIZE / Blob._CHUNK_SIZE_MULTIPLE)
+        multiplier = 40
         while True:
             try:
                 # skipcq: PYL-W0212
                 chunk_size = Blob._CHUNK_SIZE_MULTIPLE * multiplier
                 return func(*args, chunk_size=chunk_size, **kwargs)
             except requests.exceptions.ConnectionError:
-                multiplier = int(multiplier / 2)
+                multiplier //= 2
                 if not multiplier:
                     raise
 
@@ -46,9 +48,32 @@ def dynamic_chunk_size(func):
 
 
 @dynamic_chunk_size
-def _upload_to_bucket(bucket, from_file, to_info, **kwargs):
-    blob = bucket.blob(to_info.path, **kwargs)
-    blob.upload_from_filename(from_file)
+def _upload_to_bucket(
+    bucket,
+    from_file,
+    to_info,
+    chunk_size=None,
+    name=None,
+    no_progress_bar=True,
+):
+    blob = bucket.blob(to_info.path, chunk_size=chunk_size)
+    with Tqdm(
+        desc=name or to_info.path,
+        total=os.path.getsize(from_file),
+        bytes=True,
+        disable=no_progress_bar,
+    ) as pbar:
+        with io.open(from_file, mode="rb") as fobj:
+            raw_read = fobj.read
+
+            def read(size=chunk_size):
+                res = raw_read(size)
+                if res:
+                    pbar.update(len(res))
+                return res
+
+            fobj.read = read
+            blob.upload_from_file(fobj)
 
 
 class RemoteGS(RemoteBASE):
@@ -123,14 +148,34 @@ class RemoteGS(RemoteBASE):
         paths = set(self._list_paths(path_info.bucket, path_info.path))
         return any(path_info.path == path for path in paths)
 
-    def _upload(self, from_file, to_info, **_kwargs):
+    def _upload(self, from_file, to_info, name=None, no_progress_bar=True):
         bucket = self.gs.bucket(to_info.bucket)
-        _upload_to_bucket(bucket, from_file, to_info)
+        _upload_to_bucket(
+            bucket,
+            from_file,
+            to_info,
+            name=name,
+            no_progress_bar=no_progress_bar,
+        )
 
-    def _download(self, from_info, to_file, **_kwargs):
+    def _download(self, from_info, to_file, name=None, no_progress_bar=True):
         bucket = self.gs.bucket(from_info.bucket)
         blob = bucket.get_blob(from_info.path)
-        blob.download_to_filename(to_file)
+        with Tqdm(
+            desc=name or from_info.path,
+            total=blob.size,
+            bytes=True,
+            disable=no_progress_bar,
+        ) as pbar:
+            with io.open(to_file, mode="wb") as fobj:
+                raw_write = fobj.write
+
+                def write(byte_string):
+                    raw_write(byte_string)
+                    pbar.update(len(byte_string))
+
+                fobj.write = write
+                blob.download_to_file(fobj)
 
     def _generate_download_url(self, path_info, expires=3600):
         expiration = timedelta(seconds=int(expires))

tests/unit/remote/test_gs.py

@@ -53,14 +53,4 @@ def test_dynamic_chunk_size():
     with pytest.raises(requests.exceptions.ConnectionError):
         upload()
 
-    assert chunk_sizes == [
-        104857600,
-        52428800,
-        26214400,
-        13107200,
-        6553600,
-        3145728,
-        1572864,
-        786432,
-        262144,
-    ]
+    assert chunk_sizes == [10485760, 5242880, 2621440, 1310720, 524288, 262144]