Repository: iterative/dvc

Issue:
gs: support progress callback
I may have missed an issue concerning this, but when I am `dvc push`ing to a google cloud remote, the progress bars that are currently displayed don't get updated progressively (They go from 0 when starting to 100 when finished).

It would be nice to have a dynamic progressbar with upload speed and remaining time estimation, especially when push huge sets of data.


Discussion:
Hi @mhham !

Indeed, callback for google cloud storate push/pull is not implemented yet. Let's keep this issue open as a reminder. Thank you for the feedback! 🙂 
Problem is there's no current API.

`google.cloud.storage.bucket.Bucket.`:
- [`copy_blob`](https://googleapis.github.io/google-cloud-python/latest/storage/buckets.html#google.cloud.storage.bucket.Bucket.copy_blob)
- [`upload_from_filename`](https://googleapis.github.io/google-cloud-python/latest/storage/blobs.html#google.cloud.storage.blob.Blob.upload_from_filename)
- [`download_to_filename`](https://googleapis.github.io/google-cloud-python/latest/storage/blobs.html#google.cloud.storage.blob.Blob.download_to_filename)

Possibly could work around this by using [`upload_from_file`](https://googleapis.github.io/google-cloud-python/latest/storage/blobs.html#google.cloud.storage.blob.Blob.upload_from_file) and [`download_to_file`](https://googleapis.github.io/google-cloud-python/latest/storage/blobs.html#google.cloud.storage.blob.Blob.download_to_file) with our own custom file objects which have progress for `read()` and `write()`
Looks good, let me look at it.
Let's focus on the downloading first.. .
> with our own custom file objects

What do you mean @casperdcl ? Could you provide a link to that? I only found this
https://github.com/iterative/dvc/blob/38c210068ba5aeb70ad92a8c22e85832640d8ebd/dvc/api.py#L56
which doesn't seems to be helpful.

I think the more appropriate approach would be to use this [google-resumable-media](https://googleapis.github.io/google-resumable-media-python/latest/google.resumable_media.requests.html#chunked-downloads) as it's recommended 

> For more fine-grained control over the download process, check out google-resumable-media. For example, this library allows downloading parts of a blob rather than the whole thing.

[here](https://googleapis.dev/python/storage/latest/blobs.html#google.cloud.storage.blob.Blob.download_to_file).

I will check if using `download_to_file` with `start` & `end` gives the similar result to `google-resumable-media`. If so, we could avoid adding a new dependency. 

> .. with our own custom file objects which have progress for read() and write()

Ok, got it. I think @casperdcl meant the custom file object is to be written as a part of this task.
Based on the resumable media link it seems like the proposed `download_to_file` approach would be just as complicated to implement really. If you have issues though feel free to open an incomplete PR and I can take a look.
My main problem is testing - @efiop how would you recommended setting up local testing for google cloud storage push/pull?
@casperdcl I'm not aware of any local implementation that you could run on your machine 🙁 So far I was just using real GS to test things out.
@casperdcl in meantime i started to play a bit with your concept of custom file object and did this very limited PoC (since it's PoC I used a simple file in the root of the dir &  `unittests`)
https://github.com/iterative/dvc/compare/master...xliiv:progress?expand=1

It's for `StringIO`, so some work needed to adjust it to our needs, but looks promising.

Maybe `tqdm` should expose objects like the `ProgressStringIO`?
@xliiv A side note: Please use pytest-style tests for new tests. We do have some unittest tests as a legacy, but all new ones are written with pytest.

Also, why StringIO and not BytestIO?
@efiop sure, i'll be using `pytest`.

> Also, why StringIO and not BytestIO?

It is PoC and I felt it would be simpler to use `StringIO`. I'll convert it to `BytesIO`.
> So far I was just using real GS to test things out.

Ok @efiop what's the easiest way to set up real GS? I mean example scripts to

1. add a remote (ideally one which already exists and doesn't need credentials if possible)
2. push
3. pull
4. repeat steps 2 & 3 each time a code change is made while fixing this issue



Resolving patch:
diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py
--- a/dvc/remote/gs.py
+++ b/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 wrapper(*args, **kwargs):
 
 
 @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 @@ def exists(self, path_info):
         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))


Test patch:
diff --git a/tests/unit/remote/test_gs.py b/tests/unit/remote/test_gs.py
--- a/tests/unit/remote/test_gs.py
+++ b/tests/unit/remote/test_gs.py
@@ -53,14 +53,4 @@ def upload(chunk_size=None):
     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]