fs.upload: add support for optional callback (#6546)

* 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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

dvc/fs/azure.py

@@ -11,7 +11,8 @@ from dvc.path_info import CloudURLInfo
 from dvc.scheme import Schemes
 from dvc.utils import format_link
 
-from .fsspec_wrapper import ObjectFSWrapper
+from ..progress import DEFAULT_CALLBACK
+from .fsspec_wrapper import CallbackMixin, ObjectFSWrapper
 
 logger = logging.getLogger(__name__)
 _DEFAULT_CREDS_STEPS = (
@@ -159,3 +160,11 @@ class AzureFileSystem(ObjectFSWrapper):
                 " failed.\nLearn more about configuration settings at"
                 f" {format_link('https://man.dvc.org/remote/modify')}"
             ) from e
+
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
+    ):
+        # AzureFileSystem.put_file does not support callbacks yet.
+        return CallbackMixin.put_file_compat(
+            self, from_file, to_info, callback=callback, **kwargs
+        )

dvc/fs/base.py

@@ -1,12 +1,18 @@
+import contextlib
 import logging
+import os
 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 tqdm.utils import CallbackIOWrapper
+
 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,50 +223,65 @@ class BaseFileSystem:
             return False
         return hash_.endswith(cls.CHECKSUM_DIR_SUFFIX)
 
+    @cached_property
+    def _local_fs(self):
+        from dvc.fs import LocalFileSystem
+
+        return LocalFileSystem()
+
     def upload(
         self,
         from_info,
         to_info,
         total=None,
         desc=None,
+        callback=None,
         no_progress_bar=False,
         **pbar_args,
     ):
         is_file_obj = hasattr(from_info, "read")
-        method = "upload_fobj" if is_file_obj else "_upload"
+        method = "upload_fobj" if is_file_obj else "put_file"
         if not hasattr(self, method):
             raise RemoteActionNotImplemented(method, self.scheme)
 
         if to_info.scheme != self.scheme:
             raise NotImplementedError
 
-        logger.debug("Uploading '%s' to '%s'", from_info, to_info)
-        if is_file_obj:
-            with Tqdm.wrapattr(
-                from_info,
-                "read",
+        if not is_file_obj:
+            desc = desc or from_info.name
+
+        stack = contextlib.ExitStack()
+        if not callback:
+            pbar = ui.progress(
+                desc=desc,
                 disable=no_progress_bar,
                 bytes=True,
-                total=total,
-                desc=desc,
+                total=total or -1,
                 **pbar_args,
-            ) as wrapped:
+            )
+            stack.enter_context(pbar)
+            callback = pbar.as_callback(self._local_fs, from_info)
+            if total:
+                callback.set_size(total)
+
+        with stack:
+            if is_file_obj:
+                wrapped = CallbackIOWrapper(
+                    callback.relative_update, from_info
+                )
                 # `size` is used to provide hints to the WebdavFileSystem
                 # for legacy servers.
                 # pylint: disable=no-member
                 return self.upload_fobj(wrapped, to_info, size=total)
 
-        if from_info.scheme != "local":
-            raise NotImplementedError
+            if from_info.scheme != "local":
+                raise NotImplementedError
 
-        name = desc or from_info.name
-        return self._upload(  # noqa, pylint: disable=no-member
-            from_info.fspath,
-            to_info,
-            name=name,
-            no_progress_bar=no_progress_bar,
-            **pbar_args,
-        )
+            logger.debug("Uploading '%s' to '%s'", from_info, to_info)
+            # pylint: disable=no-member
+            return self.put_file(
+                os.fspath(from_info), to_info, callback=callback
+            )
 
     def download(
         self,

dvc/fs/fsspec_wrapper.py

@@ -3,8 +3,9 @@ import shutil
 from functools import lru_cache
 
 from funcy import cached_property
+from tqdm.utils import CallbackIOWrapper
 
-from dvc.progress import Tqdm
+from dvc.progress import DEFAULT_CALLBACK, Tqdm
 
 from .base import BaseFileSystem
 from .local import LocalFileSystem
@@ -130,28 +131,19 @@ class FSSpecWrapper(BaseFileSystem):
             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, **kwargs):
         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:
-            with Tqdm.wrapattr(
-                fobj,
-                "read",
-                disable=no_progress_bar,
-                bytes=True,
-                total=size,
-                desc=name,
-            ) as wrapped:
-                self.upload_fobj(wrapped, to_info, size=size)
-        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
     ):
@@ -263,22 +255,19 @@ 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
+    @staticmethod
+    def put_file_compat(
+        fs, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
     ):
-        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))
+        """Add compatibility support for Callback."""
+        fs.makedirs(to_info.parent)
+        size = os.path.getsize(from_file)
+        with open(from_file, "rb") as fobj:
+            callback.set_size(size)
+            wrapped = CallbackIOWrapper(callback.relative_update, fobj)
+            fs.upload_fobj(wrapped, to_info)
+            # pylint: disable=protected-access
+            fs.fs.invalidate_cache(fs._with_bucket(to_info.parent))
 
     def _download(
         self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args

dvc/fs/gs.py

@@ -6,10 +6,11 @@ from funcy import cached_property, wrap_prop
 from dvc.path_info import CloudURLInfo
 from dvc.scheme import Schemes
 
-from .fsspec_wrapper import ObjectFSWrapper
-
-
 # pylint:disable=abstract-method
+from ..progress import DEFAULT_CALLBACK
+from .fsspec_wrapper import CallbackMixin, ObjectFSWrapper
+
+
 class GSFileSystem(ObjectFSWrapper):
     scheme = Schemes.GS
     PATH_CLS = CloudURLInfo
@@ -35,3 +36,11 @@ class GSFileSystem(ObjectFSWrapper):
         from gcsfs import GCSFileSystem
 
         return GCSFileSystem(**self.fs_args)
+
+    def put_file(
+        self, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
+    ):
+        # GCSFileSystem.put_file does not support callbacks yet.
+        return CallbackMixin.put_file_compat(
+            self, from_file, to_info, callback=callback, **kwargs
+        )

dvc/fs/hdfs.py

@@ -9,7 +9,7 @@ from collections import deque
 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 @@ class HDFSFileSystem(BaseFileSystem):
             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.system import System
 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 @@ class LocalFileSystem(BaseFileSystem):
     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 collections import defaultdict
 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,17 @@ class S3FileSystem(BaseS3FileSystem):  # pylint:disable=abstract-method
         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 import prompt
 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
@@ -124,6 +125,8 @@ class SSHFileSystem(CallbackMixin, FSSpecWrapper):
         with as_atomic(self, to_info) as tmp_file:
             super().upload_fobj(fobj, tmp_file, **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 funcy import cached_property, wrap_prop
 
 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_pbar(pbar, total):
     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 @@ class WebHDFSFileSystem(BaseFileSystem):  # pylint:disable=abstract-method
         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
@@ -182,9 +181,29 @@ class FsspecCallback(fsspec.Callback):
             size = self.fs.info(self.path_info)["size"]
         self.progress_bar.total = size
         self.progress_bar.refresh()
+        super().set_size(size)
 
     def relative_update(self, inc=1):
         self.progress_bar.update(inc)
+        super().relative_update(inc)
 
     def absolute_update(self, value):
         self.progress_bar.update_to(value)
+        super().absolute_update(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,33 +193,39 @@ 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
 
     if os.path.isdir(dest):
         dest = os.path.join(dest, os.path.basename(src))
 
+    if callback:
+        callback.set_size(total)
+
     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)
+
+    if callback:
+        callback.absolute_update(total)
 
 
 def copy_fobj_to_file(fsrc, dest):

tests/func/objects/db/test_index.py

@@ -76,7 +76,7 @@ def test_partial_upload(tmp_dir, dvc, index, mocker):
     tmp_dir.dvc_gen({"foo": "foo content"})
     tmp_dir.dvc_gen({"bar": {"baz": "baz content"}})
 
-    original = LocalFileSystem._upload
+    original = LocalFileSystem.upload
 
     def unreliable_upload(self, from_file, to_info, name=None, **kwargs):
         if "baz" in name:

tests/func/test_fs.py

@@ -3,6 +3,7 @@ import os
 from operator import itemgetter
 from os.path import join
 
+import fsspec
 import pytest
 
 from dvc.fs import get_cloud_fs
@@ -413,3 +414,30 @@ def test_fs_makedirs_on_upload_and_copy(dvc, cloud):
     fs.copy(cloud / "dir" / "foo", cloud / "dir2" / "foo")
     assert fs.isdir(cloud / "dir2")
     assert fs.exists(cloud / "dir2" / "foo")
+
+
+@pytest.mark.needs_internet
+@pytest.mark.parametrize(
+    "cloud",
+    [
+        pytest.lazy_fixture("azure"),
+        pytest.lazy_fixture("gs"),
+        pytest.lazy_fixture("gdrive"),
+        pytest.lazy_fixture("hdfs"),
+        pytest.lazy_fixture("local_cloud"),
+        pytest.lazy_fixture("oss"),
+        pytest.lazy_fixture("s3"),
+        pytest.lazy_fixture("ssh"),
+        pytest.lazy_fixture("webhdfs"),
+    ],
+)
+def test_upload_callback(tmp_dir, dvc, cloud):
+    tmp_dir.gen("foo", "foo")
+    cls, config, _ = get_cloud_fs(dvc, **cloud.config)
+    fs = cls(**config)
+    expected_size = os.path.getsize(tmp_dir / "foo")
+
+    callback = fsspec.Callback()
+    fs.upload(tmp_dir / "foo", cloud / "foo", callback=callback)
+    assert callback.size == expected_size
+    assert callback.value == expected_size

tests/unit/remote/test_remote_tree.py

@@ -91,8 +91,8 @@ def test_walk_files(remote):
 def test_copy_preserve_etag_across_buckets(cloud, dvc):
     cloud.gen(FILE_WITH_CONTENTS)
     rem = _get_odb(dvc, cloud.config)
-    s3 = rem.fs
-    s3.fs.mkdir("another/")
+    s3 = rem.fs.s3
+    s3.create_bucket(Bucket="another")
 
     config = cloud.config.copy()
     config["url"] = "s3://another"