fs.upload: add support for optional callback

Also removes _upload methods in other fses, and replaces them
with fsspec-compatible put_file which supports fsspec callback
even if they are not fsspec based filesystem.

Author: skshetry

dvc/fs/base.py

@@ -1,12 +1,16 @@
+import contextlib
 import logging
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from functools import partial, partialmethod
 from multiprocessing import cpu_count
 from typing import Any, ClassVar, Dict, FrozenSet, Optional
 
+from funcy import cached_property
+
 from dvc.exceptions import DvcException
 from dvc.path_info import URLInfo
 from dvc.progress import Tqdm
+from dvc.ui import ui
 from dvc.utils import tmp_fname
 from dvc.utils.fs import makedirs, move
 
@@ -217,8 +221,22 @@ def is_dir_hash(cls, hash_):
             return False
         return hash_.endswith(cls.CHECKSUM_DIR_SUFFIX)
 
-    def upload(self, from_info, to_info, name=None, no_progress_bar=False):
-        if not hasattr(self, "_upload"):
+    @cached_property
+    def _local_fs(self):
+        from dvc.fs import LocalFileSystem
+
+        return LocalFileSystem()
+
+    def upload(
+        self,
+        from_info,
+        to_info,
+        name=None,
+        callback=None,
+        no_progress_bar=False,
+        **pbar_kw,
+    ):
+        if not hasattr(self, "put_file"):
             raise RemoteActionNotImplemented("upload", self.scheme)
 
         if to_info.scheme != self.scheme:
@@ -231,12 +249,23 @@ def upload(self, from_info, to_info, name=None, no_progress_bar=False):
 
         name = name or from_info.name
 
-        self._upload(  # noqa, pylint: disable=no-member
-            from_info.fspath,
-            to_info,
-            name=name,
-            no_progress_bar=no_progress_bar,
-        )
+        stack = contextlib.ExitStack()
+        fs = self._local_fs
+        if not callback:
+            pbar = ui.progress(
+                desc=name,
+                disable=no_progress_bar,
+                bytes=True,
+                total=-1,
+                **pbar_kw,
+            )
+            stack.enter_context(pbar)
+            callback = pbar.as_callback(fs, from_info)
+
+        with stack:
+            src = from_info.fspath
+            # pylint: disable=no-member
+            return self.put_file(src, to_info, callback=callback)
 
     def upload_fobj(
         self, fobj, to_info, no_progress_bar=False, size=None, **pbar_args

dvc/fs/fsspec_wrapper.py

@@ -4,7 +4,7 @@
 
 from funcy import cached_property
 
-from dvc.progress import Tqdm
+from dvc.progress import DEFAULT_CALLBACK, Tqdm
 
 from .base import BaseFileSystem
 from .local import LocalFileSystem
@@ -130,26 +130,19 @@ def makedirs(self, path_info, **kwargs):
             self._with_bucket(path_info), exist_ok=kwargs.pop("exist_ok", True)
         )
 
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
+    ):
+        self.fs.put_file(
+            from_file, self._with_bucket(to_info), callback=callback, **kwargs
+        )
+        self.fs.invalidate_cache(self._with_bucket(to_info.parent))
+
     def _upload_fobj(self, fobj, to_info, size=None):
         self.makedirs(to_info.parent)
         with self.open(to_info, "wb") as fdest:
             shutil.copyfileobj(fobj, fdest, length=fdest.blocksize)
 
-    def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **kwargs
-    ):
-        self.makedirs(to_info.parent)
-        size = os.path.getsize(from_file)
-        with open(from_file, "rb") as fobj:
-            self.upload_fobj(
-                fobj,
-                to_info,
-                size=size,
-                desc=name,
-                no_progress_bar=no_progress_bar,
-            )
-        self.fs.invalidate_cache(self._with_bucket(to_info.parent))
-
     def _download(
         self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
     ):
@@ -261,23 +254,6 @@ class CallbackMixin:
     """Use the native ``get_file()``/``put_file()`` APIs
     if the target filesystem supports callbacks."""
 
-    def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **pbar_args
-    ):
-        with Tqdm(
-            desc=name,
-            disable=no_progress_bar,
-            bytes=True,
-            total=-1,
-            **pbar_args,
-        ) as pbar:
-            self.fs.put_file(
-                os.fspath(from_file),
-                self._with_bucket(to_info),
-                callback=pbar.as_callback(_LOCAL_FS, from_file),
-            )
-        self.fs.invalidate_cache(self._with_bucket(to_info.parent))
-
     def _download(
         self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
     ):

dvc/fs/hdfs.py

@@ -9,7 +9,7 @@
 from contextlib import closing, contextmanager
 
 from dvc.hash_info import HashInfo
-from dvc.progress import Tqdm
+from dvc.progress import DEFAULT_CALLBACK, Tqdm
 from dvc.scheme import Schemes
 from dvc.utils import fix_env, tmp_fname
 
@@ -249,22 +249,22 @@ def _upload_fobj(self, fobj, to_info, **kwargs):
             with hdfs.open_output_stream(to_info.path) as fdest:
                 shutil.copyfileobj(fobj, fdest, self.BLOCK_SIZE)
 
-    def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
     ):
+        from tqdm.utils import CallbackIOWrapper
+
         with self.hdfs(to_info) as hdfs:
+            hdfs.create_dir(to_info.parent.path)
+
             tmp_file = tmp_fname(to_info.path)
-            total = os.path.getsize(from_file)
+            total: int = os.path.getsize(from_file)
+            callback.set_size(total)
+
             with open(from_file, "rb") as fobj:
-                with Tqdm.wrapattr(
-                    fobj,
-                    "read",
-                    desc=name,
-                    total=total,
-                    disable=no_progress_bar,
-                ) as wrapped:
-                    with hdfs.open_output_stream(tmp_file) as sobj:
-                        shutil.copyfileobj(wrapped, sobj, self.BLOCK_SIZE)
+                wrapped = CallbackIOWrapper(callback.relative_update, fobj)
+                with hdfs.open_output_stream(tmp_file) as sobj:
+                    shutil.copyfileobj(wrapped, sobj, self.BLOCK_SIZE)
             hdfs.move(tmp_file, to_info.path)
 
     def _download(

dvc/fs/local.py

@@ -7,6 +7,7 @@
 from dvc.utils import is_exec, tmp_fname
 from dvc.utils.fs import copy_fobj_to_file, copyfile, makedirs, move, remove
 
+from ..progress import DEFAULT_CALLBACK
 from .base import BaseFileSystem
 
 logger = logging.getLogger(__name__)
@@ -158,15 +159,12 @@ def reflink(self, from_info, to_info):
     def info(self, path_info):
         return self.fs.info(path_info)
 
-    def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
     ):
         makedirs(to_info.parent, exist_ok=True)
-
         tmp_file = tmp_fname(to_info)
-        copyfile(
-            from_file, tmp_file, name=name, no_progress_bar=no_progress_bar
-        )
+        copyfile(from_file, tmp_file, callback=callback)
         os.replace(tmp_file, to_info)
 
     @staticmethod

dvc/fs/s3.py

@@ -6,7 +6,7 @@
 from funcy import cached_property, wrap_prop
 
 from dvc.path_info import CloudURLInfo
-from dvc.progress import Tqdm
+from dvc.progress import DEFAULT_CALLBACK, Tqdm
 from dvc.scheme import Schemes
 
 from .fsspec_wrapper import ObjectFSWrapper
@@ -204,24 +204,18 @@ def _get_obj(self, path_info):
         return bucket.Object(path_info.path)
 
     @_translate_exceptions
-    def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **pbar_args
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
     ):
-        total = os.path.getsize(from_file)
-        with Tqdm(
-            disable=no_progress_bar,
-            total=total,
-            bytes=True,
-            desc=name,
-            **pbar_args,
-        ) as pbar:
-            obj = self._get_obj(to_info)
-            obj.upload_file(
-                from_file,
-                Callback=pbar.update,
-                ExtraArgs=self.fs_args.get("s3_additional_kwargs"),
-                Config=self._transfer_config,
-            )
+        callback.set_size(os.path.getsize(from_file))
+
+        obj = self._get_obj(to_info)
+        obj.upload_file(
+            from_file,
+            Callback=callback.relative_update,
+            ExtraArgs=self.fs_args.get("s3_additional_kwargs"),
+            Config=self._transfer_config,
+        )
         self.fs.invalidate_cache(self._with_bucket(to_info.parent))
 
     @_translate_exceptions

dvc/fs/ssh.py

@@ -8,6 +8,7 @@
 from dvc.scheme import Schemes
 from dvc.utils.fs import as_atomic
 
+from ..progress import DEFAULT_CALLBACK
 from .fsspec_wrapper import CallbackMixin, FSSpecWrapper
 
 _SSH_TIMEOUT = 60 * 30
@@ -123,6 +124,8 @@ def _upload_fobj(self, fobj, to_info, *args, **kwargs):
         with as_atomic(self, to_info) as tmp_file:
             super()._upload_fobj(fobj, tmp_file, *args, **kwargs)
 
-    def _upload(self, from_file, to_info, *args, **kwargs):
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
+    ):
         with as_atomic(self, to_info) as tmp_file:
-            super()._upload(from_file, tmp_file, *args, **kwargs)
+            super().put_file(from_file, tmp_file, callback=callback, **kwargs)

dvc/fs/webhdfs.py

@@ -9,7 +9,7 @@
 
 from dvc.hash_info import HashInfo
 from dvc.path_info import CloudURLInfo
-from dvc.progress import Tqdm
+from dvc.progress import DEFAULT_CALLBACK, Tqdm
 from dvc.scheme import Schemes
 
 from .base import BaseFileSystem
@@ -29,6 +29,15 @@ def update(_, bytes_transfered):
     return update
 
 
+def update_callback(callback, total):
+    def update(_, bytes_transfered):
+        if bytes_transfered == -1:
+            return callback.absolute_update(total)
+        return callback.relative_update(bytes_transfered)
+
+    return update
+
+
 class WebHDFSFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
     scheme = Schemes.WEBHDFS
     PATH_CLS = CloudURLInfo
@@ -145,19 +154,19 @@ def _upload_fobj(self, fobj, to_info, **kwargs):
         with self.hdfs_client.write(to_info.path) as fdest:
             shutil.copyfileobj(fobj, fdest)
 
-    def _upload(
-        self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
     ):
         total = os.path.getsize(from_file)
-        with Tqdm(
-            desc=name, total=total, disable=no_progress_bar, bytes=True
-        ) as pbar:
-            self.hdfs_client.upload(
-                to_info.path,
-                from_file,
-                overwrite=True,
-                progress=update_pbar(pbar, total),
-            )
+        callback.set_size(total)
+
+        self.hdfs_client.makedirs(to_info.parent.path)
+        return self.hdfs_client.upload(
+            to_info.path,
+            from_file,
+            overwrite=True,
+            progress=update_callback(callback, total),
+        )
 
     def _download(
         self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs

dvc/progress.py

@@ -1,5 +1,4 @@
 """Manages progress bars for DVC repo."""
-
 import logging
 import sys
 from threading import RLock
@@ -188,3 +187,20 @@ def relative_update(self, inc=1):
 
     def absolute_update(self, value):
         self.progress_bar.update_to(value)
+
+
+def tdqm_or_callback_wrapped(
+    fobj, method, total, callback=None, **pbar_kwargs
+):
+    if callback:
+        from funcy import nullcontext
+        from tqdm.utils import CallbackIOWrapper
+
+        callback.set_size(total)
+        wrapper = CallbackIOWrapper(callback.relative_update, fobj, method)
+        return nullcontext(wrapper)
+
+    return Tqdm.wrapattr(fobj, method, total=total, bytes=True, **pbar_kwargs)
+
+
+DEFAULT_CALLBACK = fsspec.callbacks.NoOpCallback

dvc/utils/fs.py

@@ -193,10 +193,8 @@ def makedirs(path, exist_ok=False, mode=None):
         logger.trace("failed to chmod '%o' '%s'", mode, path, exc_info=True)
 
 
-def copyfile(src, dest, no_progress_bar=False, name=None):
+def copyfile(src, dest, callback=None, no_progress_bar=False, name=None):
     """Copy file with progress bar"""
-    from dvc.progress import Tqdm
-
     name = name if name else os.path.basename(dest)
     total = os.stat(src).st_size
 
@@ -206,20 +204,22 @@ def copyfile(src, dest, no_progress_bar=False, name=None):
     try:
         System.reflink(src, dest)
     except DvcException:
+        from dvc.progress import tdqm_or_callback_wrapped
+
         with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
-            with Tqdm.wrapattr(
+            with tdqm_or_callback_wrapped(
                 fdest,
                 "write",
-                desc=name,
+                total,
+                callback=callback,
                 disable=no_progress_bar,
-                total=total,
-                bytes=True,
-            ) as fdest_wrapped:
+                desc=name,
+            ) as wrapped:
                 while True:
                     buf = fsrc.read(LOCAL_CHUNK_SIZE)
                     if not buf:
                         break
-                    fdest_wrapped.write(buf)
+                    wrapped.write(buf)
 
 
 def copy_fobj_to_file(fsrc, dest):