From f39e35445a5e60dcda0378021e27fcc6345776cd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Mon, 25 Apr 2022 23:09:04 +0545 Subject: [PATCH] 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 | 13 +++- dvc/fs/_callback.py | 76 ++++++++++++++----- dvc/fs/base.py | 173 ++++++++++++++------------------------------ dvc/fs/utils.py | 2 +- dvc/progress.py | 5 -- dvc/utils/fs.py | 22 +++--- 6 files changed, 136 insertions(+), 155 deletions(-) diff --git a/dvc/data/stage.py b/dvc/data/stage.py index 67225d4c7..8a9dbaefb 100644 --- a/dvc/data/stage.py +++ b/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( - stream, tmp_info, desc=fs_path.name(from_fs_path), total=size - ) + with FsspecCallback.as_tqdm_callback( + 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) diff --git a/dvc/fs/_callback.py b/dvc/fs/_callback.py index 58adee2ed..23a5ddc5b 100644 --- a/dvc/fs/_callback.py +++ b/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): - self.progress_bar = progress_bar + def __init__(self, progress_bar: "Tqdm" = None, **tqdm_kwargs): + 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 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) + def branch(self, path_1: str, path_2: str, kwargs): + return TqdmCallback(bytes=True, total=-1, desc=path_1, **kwargs) DEFAULT_CALLBACK = NoOpCallback() diff --git a/dvc/fs/base.py b/dvc/fs/base.py index 61d951569..f6271bde6 100644 --- a/dvc/fs/base.py +++ b/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 import tmp_fname -from dvc.utils.fs import makedirs, move +from dvc.utils.fs import as_atomic, makedirs +from dvc.utils.threadpool import ThreadPoolExecutor from ._callback import DEFAULT_CALLBACK, FsspecCallback @@ -325,145 +321,86 @@ class FileSystem: self, from_info: Union[AnyFSPath, IO], to_info: AnyFSPath, - total: int = None, - desc: str = None, - callback=None, - no_progress_bar: bool = False, - **pbar_args: Any, + size: int = None, + callback: FsspecCallback = None, ): + 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() - if not callback: - pbar = ui.progress( - desc=desc, - disable=no_progress_bar, - bytes=True, - total=total or -1, - **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 - ) + with FsspecCallback.as_tqdm_callback( + callback, + desc=desc, + bytes=True, + total=size or -1, + ) as cb: + return self.put_file(from_info, to_info, callback=cb, size=size) def download( self, from_info: AnyFSPath, to_info: AnyFSPath, - name: str = None, - callback=None, - no_progress_bar: bool = False, + callback: "FsspecCallback" = None, 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"): - raise RemoteActionNotImplemented("get_file", self.scheme) - - 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, + pairs = { + info: localfs.path.join( + to_info, *self.path.relparts(info, from_info) ) - stack.enter_context(pbar) - callback = pbar.as_callback() - - 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: + for info in self.find(from_info) + } + if not pairs: makedirs(to_info, exist_ok=True) - return None - - 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) - ] + return + 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): - # NOTE: executor won't let us raise until all futures that - # it has are finished, so we need to cancel them ourselves - # before re-raising. - exc = future.exception() - if exc: - for entry in futures: - entry.cancel() - raise exc + with ThreadPoolExecutor( + max_workers=max_workers, cancel_on_error=True + ) as executor: + list( + executor.imap_unordered( + lambda args: download_files(*args), pairs.items() + ) + ) - def _download_file( + def download_file( self, from_info: AnyFSPath, to_info: AnyFSPath, - callback=DEFAULT_CALLBACK, - ) -> None: + callback: FsspecCallback = None, + ): from .local import localfs - makedirs(localfs.path.parent(to_info), exist_ok=True) - tmp_file = tmp_fname(to_info) - - logger.debug("Downloading '%s' to '%s'", from_info, to_info) - try: - # noqa, pylint: disable=no-member - self.get_file(from_info, tmp_file, callback=callback) - except Exception: # pylint: disable=broad-except - # 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) + with FsspecCallback.as_tqdm_callback( + callback, + total=-1, + desc=self.path.name(from_info), + bytes=True, + ) as cb: + with as_atomic(localfs, to_info, create_parents=True) as tmp_file: + self.get_file(from_info, tmp_file, callback=cb) class ObjectFileSystem(FileSystem): # pylint: disable=abstract-method diff --git a/dvc/fs/utils.py b/dvc/fs/utils.py index c19055b0a..3c96791e4 100644 --- a/dvc/fs/utils.py +++ b/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( diff --git a/dvc/progress.py b/dvc/progress.py index 3818b1ce5..7ca37f05b 100644 --- a/dvc/progress.py +++ b/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 diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py index c12ce4009..7a71fb06e 100644 --- a/dvc/utils/fs.py +++ b/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( - fdest, - "write", - total, - callback=callback, + with FsspecCallback.as_tqdm_callback( + callback, + total=total, + bytes=True, 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: