diff --git a/dvc/fs/azure.py b/dvc/fs/azure.py index 6ae3e2bc0..518c2cdd8 100644 --- a/dvc/fs/azure.py +++ b/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 + ) diff --git a/dvc/fs/base.py b/dvc/fs/base.py index f53e6bec5..299a0c1a9 100644 --- a/dvc/fs/base.py +++ b/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, diff --git a/dvc/fs/fsspec_wrapper.py b/dvc/fs/fsspec_wrapper.py index 446085f4f..7347e1c07 100644 --- a/dvc/fs/fsspec_wrapper.py +++ b/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 diff --git a/dvc/fs/gs.py b/dvc/fs/gs.py index 5f5b1efe6..c7e24c854 100644 --- a/dvc/fs/gs.py +++ b/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 + ) diff --git a/dvc/fs/hdfs.py b/dvc/fs/hdfs.py index 3812155db..0e988f432 100644 --- a/dvc/fs/hdfs.py +++ b/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( diff --git a/dvc/fs/local.py b/dvc/fs/local.py index 5e1cbef3c..c783d1cc0 100644 --- a/dvc/fs/local.py +++ b/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 diff --git a/dvc/fs/s3.py b/dvc/fs/s3.py index 480b16b1d..bc4b6bb56 100644 --- a/dvc/fs/s3.py +++ b/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 diff --git a/dvc/fs/ssh.py b/dvc/fs/ssh.py index b3b0fbad8..b32ab0950 100644 --- a/dvc/fs/ssh.py +++ b/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) diff --git a/dvc/fs/webhdfs.py b/dvc/fs/webhdfs.py index fdf9de1b5..0c3146083 100644 --- a/dvc/fs/webhdfs.py +++ b/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 diff --git a/dvc/progress.py b/dvc/progress.py index b7bc96ff5..9dc321a24 100644 --- a/dvc/progress.py +++ b/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 diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py index 6c5d10b37..15be5013d 100644 --- a/dvc/utils/fs.py +++ b/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): diff --git a/tests/func/objects/db/test_index.py b/tests/func/objects/db/test_index.py index 37dea8506..7b85aea6a 100644 --- a/tests/func/objects/db/test_index.py +++ b/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: diff --git a/tests/func/test_fs.py b/tests/func/test_fs.py index 04e016fb8..e5a610f2f 100644 --- a/tests/func/test_fs.py +++ b/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 diff --git a/tests/unit/remote/test_remote_tree.py b/tests/unit/remote/test_remote_tree.py index caf7aee8e..3a91018ea 100644 --- a/tests/unit/remote/test_remote_tree.py +++ b/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"