fs.upload/download: simplify, remove unused kwargs, callbacks

1. This PR simplifies `fs.upload`/`fs.download`/`fs.download_file`.
2. Removes unused kwargs.
3. `download` uses `dvc.utils.threadpool.ThreadPoolExecutor` which
   supports canceling futures.
4. `download` for directories supports nested/branched callbacks
   and progress bars.

dvc/data/stage.py

@@ -25,7 +25,8 @@ logger = logging.getLogger(__name__)
 _STAGING_MEMFS_PATH = "dvc-staging"
 
 
-def _upload_file(from_fs_path, fs, odb, upload_odb):
+def _upload_file(from_fs_path, fs, odb, upload_odb, callback=None):
+    from dvc.fs._callback import FsspecCallback
     from dvc.utils import tmp_fname
     from dvc.utils.stream import HashedStreamReader
 
@@ -34,9 +35,13 @@ def _upload_file(from_fs_path, fs, odb, upload_odb):
     with fs.open(from_fs_path, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream:
         stream = HashedStreamReader(stream)
         size = fs.size(from_fs_path)
-        upload_odb.fs.upload(
+        with FsspecCallback.as_tqdm_callback(
-            stream, tmp_info, desc=fs_path.name(from_fs_path), total=size
+            callback,
-        )
+            desc=fs_path.name(from_fs_path),
+            bytes=True,
+            total=size,
+        ) as cb:
+            upload_odb.fs.upload(stream, tmp_info, size=size, callback=cb)
 
     odb.add(tmp_info, upload_odb.fs, stream.hash_info)
     meta = Meta(size=size)

dvc/fs/_callback.py

@@ -1,15 +1,16 @@
 from functools import wraps
-from typing import IO, TYPE_CHECKING, Optional, TypeVar, cast
+from typing import IO, TYPE_CHECKING, Any, Optional, TypeVar, cast
 
 import fsspec
 
-from dvc.progress import Tqdm
-
 if TYPE_CHECKING:
     from typing import Callable
 
     from typing_extensions import ParamSpec
 
+    from dvc.fs.base import FileSystem
+    from dvc.progress import Tqdm
+
     _P = ParamSpec("_P")
     _R = TypeVar("_R")
 
@@ -32,6 +33,39 @@ class FsspecCallback(fsspec.Callback):
 
         return wrapped
 
+    def wrap_and_branch(
+        self, fn: "Callable", fs: "FileSystem" = None
+    ) -> "Callable":
+        """
+        Wraps a function, and pass a new child callback to it.
+        When the function completes, we increment the parent callback by 1.
+        """
+        from .local import localfs
+
+        fs = fs or localfs
+        wrapped = self.wrap_fn(fn)
+
+        def make_callback(path1, path2):
+            # pylint: disable=assignment-from-none
+            child = self.branch(fs.path.name(path1), path2, {})
+            return self.as_callback(child)
+
+        @wraps(fn)
+        def func(path1, path2, **kwargs):
+            with make_callback(path1, path2) as callback:
+                return wrapped(path1, path2, callback=callback, **kwargs)
+
+        return func
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc_args):
+        self.close()
+
+    def close(self):
+        """Handle here on exit."""
+
     @classmethod
     def as_callback(
         cls, maybe_callback: Optional["FsspecCallback"] = None
@@ -40,16 +74,33 @@ class FsspecCallback(fsspec.Callback):
             return DEFAULT_CALLBACK
         return maybe_callback
 
+    @classmethod
+    def as_tqdm_callback(
+        cls,
+        callback: Optional["FsspecCallback"] = None,
+        **tqdm_kwargs: Any,
+    ) -> "FsspecCallback":
+        return callback or TqdmCallback(**tqdm_kwargs)
+
 
 class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback):
     pass
 
 
 class TqdmCallback(FsspecCallback):
-    def __init__(self, progress_bar):
+    def __init__(self, progress_bar: "Tqdm" = None, **tqdm_kwargs):
-        self.progress_bar = progress_bar
+        from dvc.progress import Tqdm
+
+        self.progress_bar = progress_bar or Tqdm(**tqdm_kwargs)
         super().__init__()
 
+    def __enter__(self):
+        self.progress_bar.__enter__()
+        return self
+
+    def close(self):
+        self.progress_bar.close()
+
     def set_size(self, size):
         if size is not None:
             self.progress_bar.total = size
@@ -64,19 +115,8 @@ class TqdmCallback(FsspecCallback):
         self.progress_bar.update_to(value)
         super().absolute_update(value)
 
-
+    def branch(self, path_1: str, path_2: str, kwargs):
-def tdqm_or_callback_wrapped(
+        return TqdmCallback(bytes=True, total=-1, desc=path_1, **kwargs)
-    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 = NoOpCallback()

dvc/fs/base.py

@@ -1,9 +1,6 @@
-import contextlib
 import logging
 import os
 import shutil
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from functools import partialmethod
 from multiprocessing import cpu_count
 from typing import (
     IO,
@@ -22,9 +19,8 @@ from typing import (
 from funcy import cached_property
 
 from dvc.exceptions import DvcException
-from dvc.ui import ui
+from dvc.utils.fs import as_atomic, makedirs
-from dvc.utils import tmp_fname
+from dvc.utils.threadpool import ThreadPoolExecutor
-from dvc.utils.fs import makedirs, move
 
 from ._callback import DEFAULT_CALLBACK, FsspecCallback
 
@@ -325,145 +321,86 @@ class FileSystem:
         self,
         from_info: Union[AnyFSPath, IO],
         to_info: AnyFSPath,
-        total: int = None,
+        size: int = None,
-        desc: str = None,
+        callback: FsspecCallback = None,
-        callback=None,
-        no_progress_bar: bool = False,
-        **pbar_args: Any,
     ):
+        from .local import localfs
+
         if not hasattr(from_info, "read"):
-            from .local import localfs
-
-            desc = desc or localfs.path.name(from_info)
             logger.debug("Uploading '%s' to '%s'", from_info, to_info)
+            desc = localfs.path.name(from_info)
+        else:
+            desc = self.path.name(to_info)
 
-        stack = contextlib.ExitStack()
+        with FsspecCallback.as_tqdm_callback(
-        if not callback:
+            callback,
-            pbar = ui.progress(
+            desc=desc,
-                desc=desc,
+            bytes=True,
-                disable=no_progress_bar,
+            total=size or -1,
-                bytes=True,
+        ) as cb:
-                total=total or -1,
+            return self.put_file(from_info, to_info, callback=cb, size=size)
-                **pbar_args,
-            )
-            stack.enter_context(pbar)
-            callback = pbar.as_callback()
-
-        with stack:
-            if total:
-                callback.set_size(total)
-            return self.put_file(
-                from_info, to_info, callback=callback, size=total
-            )
 
     def download(
         self,
         from_info: AnyFSPath,
         to_info: AnyFSPath,
-        name: str = None,
+        callback: "FsspecCallback" = None,
-        callback=None,
-        no_progress_bar: bool = False,
         jobs: int = None,
-        _only_file: bool = False,
-        **kwargs: Any,
     ):
+        if not self.isdir(from_info):
+            return self.download_file(from_info, to_info, callback=callback)
+
         from .local import localfs
 
-        if not hasattr(self, "get_file"):
+        pairs = {
-            raise RemoteActionNotImplemented("get_file", self.scheme)
+            info: localfs.path.join(
-
+                to_info, *self.path.relparts(info, from_info)
-        download_dir = not _only_file and self.isdir(from_info)
-
-        desc = name or localfs.path.name(to_info)
-        stack = contextlib.ExitStack()
-        if not callback:
-            pbar_kwargs = {"unit": "Files"} if download_dir else {}
-            pbar = ui.progress(
-                total=-1,
-                desc="Downloading directory" if download_dir else desc,
-                bytes=not download_dir,
-                disable=no_progress_bar,
-                **pbar_kwargs,
             )
-            stack.enter_context(pbar)
+            for info in self.find(from_info)
-            callback = pbar.as_callback()
+        }
-
+        if not pairs:
-        with stack:
-            if download_dir:
-                return self._download_dir(
-                    from_info, to_info, callback=callback, jobs=jobs, **kwargs
-                )
-            return self._download_file(from_info, to_info, callback=callback)
-
-    download_file = partialmethod(download, _only_file=True)
-
-    def _download_dir(
-        self,
-        from_info: AnyFSPath,
-        to_info: AnyFSPath,
-        callback: FsspecCallback = DEFAULT_CALLBACK,
-        jobs: int = None,
-        **kwargs,
-    ):
-        from .local import localfs
-
-        from_infos = list(self.find(from_info, **kwargs))
-        if not from_infos:
             makedirs(to_info, exist_ok=True)
-            return None
+            return
-
-        to_infos = (
-            localfs.path.join(to_info, *self.path.relparts(info, from_info))
-            for info in from_infos
-        )
-        callback.set_size(len(from_infos))
-
-        download_files = callback.wrap_fn(self._download_file)
-        max_workers = jobs or self.jobs
-        with ThreadPoolExecutor(max_workers=max_workers) as executor:
-            futures = [
-                executor.submit(download_files, from_info, to_info)
-                for from_info, to_info in zip(from_infos, to_infos)
-            ]
 
+        with FsspecCallback.as_tqdm_callback(
+            callback,
+            total=-1,
+            desc=f"Downloading directory {self.path.name(from_info)}",
+            unit="files",
+        ) as cb:
+            cb.set_size(len(pairs))
+            download_files = cb.wrap_and_branch(self.download_file, fs=self)
+            max_workers = jobs or self.jobs
             # NOTE: unlike pulling/fetching cache, where we need to
             # download everything we can, not raising an error here might
             # turn very ugly, as the user might think that he has
             # downloaded a complete directory, while having a partial one,
             # which might cause unexpected results in his pipeline.
-            for future in as_completed(futures):
+            with ThreadPoolExecutor(
-                # NOTE: executor won't let us raise until all futures that
+                max_workers=max_workers, cancel_on_error=True
-                # it has are finished, so we need to cancel them ourselves
+            ) as executor:
-                # before re-raising.
+                list(
-                exc = future.exception()
+                    executor.imap_unordered(
-                if exc:
+                        lambda args: download_files(*args), pairs.items()
-                    for entry in futures:
+                    )
-                        entry.cancel()
+                )
-                    raise exc
 
-    def _download_file(
+    def download_file(
         self,
         from_info: AnyFSPath,
         to_info: AnyFSPath,
-        callback=DEFAULT_CALLBACK,
+        callback: FsspecCallback = None,
-    ) -> None:
+    ):
         from .local import localfs
 
-        makedirs(localfs.path.parent(to_info), exist_ok=True)
+        with FsspecCallback.as_tqdm_callback(
-        tmp_file = tmp_fname(to_info)
+            callback,
-
+            total=-1,
-        logger.debug("Downloading '%s' to '%s'", from_info, to_info)
+            desc=self.path.name(from_info),
-        try:
+            bytes=True,
-            # noqa, pylint: disable=no-member
+        ) as cb:
-            self.get_file(from_info, tmp_file, callback=callback)
+            with as_atomic(localfs, to_info, create_parents=True) as tmp_file:
-        except Exception:  # pylint: disable=broad-except
+                self.get_file(from_info, tmp_file, callback=cb)
-            # do we need to rollback makedirs for previously not-existing
-            # directories?
-            with contextlib.suppress(FileNotFoundError):
-                os.unlink(tmp_file)
-            raise
-
-        move(tmp_file, to_info)
 
 
 class ObjectFileSystem(FileSystem):  # pylint: disable=abstract-method

dvc/fs/utils.py

@@ -48,7 +48,7 @@ def _copy(
 
     with from_fs.open(from_path, mode="rb") as fobj:
         size = from_fs.size(from_path)
-        return to_fs.upload(fobj, to_path, total=size)
+        return to_fs.upload(fobj, to_path, size=size)
 
 
 def _try_links(

dvc/progress.py

@@ -135,11 +135,6 @@ class Tqdm(tqdm):
 
         return wrapped
 
-    def as_callback(self):
-        from dvc.fs._callback import TqdmCallback
-
-        return TqdmCallback(self)
-
     def close(self):
         self.postfix["info"] = ""
         # remove ETA (either unknown or zero); remove completed bar

dvc/utils/fs.py

@@ -200,17 +200,17 @@ def copyfile(src, dest, callback=None, no_progress_bar=False, name=None):
     try:
         System.reflink(src, dest)
     except OSError:
-        from dvc.fs._callback import tdqm_or_callback_wrapped
+        from dvc.fs._callback import FsspecCallback
 
         with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
-            with tdqm_or_callback_wrapped(
+            with FsspecCallback.as_tqdm_callback(
-                fdest,
+                callback,
-                "write",
+                total=total,
-                total,
+                bytes=True,
-                callback=callback,
                 disable=no_progress_bar,
                 desc=name,
-            ) as wrapped:
+            ) as cb:
+                wrapped = cb.wrap_attr(fdest, "write")
                 while True:
                     buf = fsrc.read(LOCAL_CHUNK_SIZE)
                     if not buf:
@@ -234,10 +234,14 @@ def walk_files(directory):
 
 
 @contextmanager
-def as_atomic(fs, to_info):
+def as_atomic(fs, to_info, create_parents=False):
     from dvc.utils import tmp_fname
 
-    tmp_info = fs.path.join(fs.path.parent(to_info), tmp_fname())
+    parent = fs.path.parent(to_info)
+    if create_parents:
+        fs.makedirs(parent, exist_ok=True)
+
+    tmp_info = fs.path.join(parent, tmp_fname())
     try:
         yield tmp_info
     except BaseException: