Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

fs.download: add callback support (#6610)

Browse Source 
* fs.download: add callback support

* fix lint issues

* compat for gdrive, use fs to upload text for tests

* fix typing issues

* skip gdrive test

* Update tests/func/test_fs.py
This commit is contained in:
Saugat Pachhai
2021-09-14 15:53:16 +05:45
committed by GitHub
parent 788b4d3bb5
commit 237e9d4354
18 changed files with 195 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/fs/azure.py
+1 -10
View File
@@ -11,7 +11,6 @@ from dvc.path_info import CloudURLInfo
from dvc.scheme import Schemes
from dvc.utils import format_link
from ..progress import DEFAULT_CALLBACK
from .fsspec_wrapper import CallbackMixin, ObjectFSWrapper
logger = logging.getLogger(__name__)
@@ -40,7 +39,7 @@ def _az_config():
# pylint:disable=abstract-method
class AzureFileSystem(ObjectFSWrapper):
class AzureFileSystem(CallbackMixin, ObjectFSWrapper):
scheme = Schemes.AZURE
PATH_CLS = CloudURLInfo
PARAM_CHECKSUM = "etag"
@@ -160,11 +159,3 @@ 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
+54 -16
View File
@@ -11,7 +11,6 @@ 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
@@ -288,13 +287,14 @@ class BaseFileSystem:
from_info,
to_info,
name=None,
callback=None,
no_progress_bar=False,
jobs=None,
_only_file=False,
**kwargs,
):
if not hasattr(self, "_download"):
raise RemoteActionNotImplemented("download", self.scheme)
if not hasattr(self, "get_file"):
raise RemoteActionNotImplemented("get_file", self.scheme)
if from_info.scheme != self.scheme:
raise NotImplementedError
@@ -308,14 +308,31 @@ class BaseFileSystem:
if not _only_file and self.isdir(from_info):
return self._download_dir(
from_info, to_info, name, no_progress_bar, jobs, **kwargs
from_info,
to_info,
desc=name,
no_progress_bar=no_progress_bar,
jobs=jobs,
**kwargs,
)
return self._download_file(from_info, to_info, name, no_progress_bar)
return self._download_file(
from_info,
to_info,
callback=callback,
desc=name,
no_progress_bar=no_progress_bar,
)
download_file = partialmethod(download, _only_file=True)
def _download_dir(
self, from_info, to_info, name, no_progress_bar, jobs, **kwargs
self,
from_info,
to_info,
desc=None,
no_progress_bar=False,
jobs=None,
**kwargs,
):
from_infos = list(self.walk_files(from_info, **kwargs))
if not from_infos:
@@ -325,14 +342,14 @@ class BaseFileSystem:
to_info / info.relative_to(from_info) for info in from_infos
)
with Tqdm(
with ui.progress(
total=len(from_infos),
desc="Downloading directory",
unit="Files",
disable=no_progress_bar,
) as pbar:
download_files = pbar.wrap_fn(
partial(self._download_file, name=name, no_progress_bar=True)
partial(self._download_file, desc=desc, no_progress_bar=True)
)
max_workers = jobs or self.jobs
with ThreadPoolExecutor(max_workers=max_workers) as executor:
@@ -356,16 +373,37 @@ class BaseFileSystem:
entry.cancel()
raise exc
def _download_file(self, from_info, to_info, name, no_progress_bar):
def _download_file(
self,
from_info,
to_info,
desc=None,
no_progress_bar=False,
callback=None,
):
makedirs(to_info.parent, exist_ok=True)
logger.debug("Downloading '%s' to '%s'", from_info, to_info)
name = name or to_info.name
tmp_file = tmp_fname(to_info)
self._download( # noqa, pylint: disable=no-member
from_info, tmp_file, name=name, no_progress_bar=no_progress_bar
)
stack = contextlib.ExitStack()
if not callback:
pbar = ui.progress(
desc=desc or to_info.name,
disable=no_progress_bar,
bytes=True,
total=-1,
)
stack.enter_context(pbar)
callback = pbar.as_callback(self, from_info)
logger.debug("Downloading '%s' to '%s'", from_info, to_info)
try:
with stack:
# 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?
os.unlink(tmp_file)
raise
move(tmp_file, to_info)
dvc/fs/dvc.py
+6 -3
View File
@@ -6,6 +6,7 @@ from dvc.exceptions import OutputNotFoundError
from dvc.path_info import PathInfo
from dvc.utils import relpath
from ..progress import DEFAULT_CALLBACK
from ._metadata import Metadata
from .base import BaseFileSystem
@@ -256,10 +257,12 @@ class DvcFileSystem(BaseFileSystem): # pylint:disable=abstract-method
return ret
def _download(self, from_info, to_file, **kwargs):
def get_file(
self, from_info, to_file, callback=DEFAULT_CALLBACK, **kwargs
):
fs, path = self._get_fs_path(from_info)
fs._download( # pylint: disable=protected-access
path, to_file, **kwargs
fs.get_file( # pylint: disable=protected-access
path, to_file, callback=callback, **kwargs
)
def checksum(self, path_info):
dvc/fs/fsspec_wrapper.py
+33 -41
View File
@@ -5,7 +5,7 @@ from functools import lru_cache
from funcy import cached_property
from tqdm.utils import CallbackIOWrapper
from dvc.progress import DEFAULT_CALLBACK, Tqdm
from dvc.progress import DEFAULT_CALLBACK
from .base import BaseFileSystem
from .local import LocalFileSystem
@@ -139,28 +139,18 @@ class FSSpecWrapper(BaseFileSystem):
)
self.fs.invalidate_cache(self._with_bucket(to_info.parent))
def get_file(
self, from_info, to_info, callback=DEFAULT_CALLBACK, **kwargs
):
self.fs.get_file(
self._with_bucket(from_info), to_info, callback=callback, **kwargs
)
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 _download(
self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
):
total = self.getsize(self._with_bucket(from_info))
with self.open(from_info, "rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name,
disable=no_progress_bar,
bytes=True,
total=total,
**pbar_args,
) as wrapped:
with open(to_file, "wb") as fdest:
shutil.copyfileobj(wrapped, fdest, length=fobj.blocksize)
# pylint: disable=abstract-method
class ObjectFSWrapper(FSSpecWrapper):
@@ -252,35 +242,37 @@ _LOCAL_FS = LocalFileSystem()
class CallbackMixin:
"""Use the native ``get_file()``/``put_file()`` APIs
if the target filesystem supports callbacks."""
"""Provides callback support for the filesystem that don't support yet."""
@staticmethod
def put_file_compat(
fs, from_file, to_info, callback=DEFAULT_CALLBACK, **kwargs
def put_file(
self,
from_file,
to_info,
callback=DEFAULT_CALLBACK,
**kwargs,
):
"""Add compatibility support for Callback."""
fs.makedirs(to_info.parent)
# pylint: disable=protected-access
self.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))
self.upload_fobj(wrapped, to_info)
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
def get_file(
self,
from_info,
to_info,
callback=DEFAULT_CALLBACK,
**kwargs,
):
with Tqdm(
desc=name,
disable=no_progress_bar,
bytes=True,
total=-1,
**pbar_args,
) as pbar:
self.fs.get_file(
self._with_bucket(from_info),
os.fspath(to_file),
callback=pbar.as_callback(self, from_info),
)
# pylint: disable=protected-access
total: int = self.getsize(from_info)
if total:
callback.set_size(total)
with self.open(from_info, "rb") as fobj, open(to_info, "wb") as fdest:
wrapped = CallbackIOWrapper(callback.relative_update, fobj)
shutil.copyfileobj(wrapped, fdest, length=fobj.blocksize)
dvc/fs/gdrive.py
+2 -4
View File
@@ -12,7 +12,7 @@ from dvc.path_info import CloudURLInfo
from dvc.scheme import Schemes
from dvc.utils import format_link, tmp_fname
from .fsspec_wrapper import CallbackMixin, FSSpecWrapper
from .fsspec_wrapper import FSSpecWrapper
logger = logging.getLogger(__name__)
FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"
@@ -79,9 +79,7 @@ class GDriveURLInfo(CloudURLInfo):
self._spath = re.sub("/{2,}", "/", self._spath.rstrip("/"))
class GDriveFileSystem(
CallbackMixin, FSSpecWrapper
): # pylint:disable=abstract-method
class GDriveFileSystem(FSSpecWrapper): # pylint:disable=abstract-method
scheme = Schemes.GDRIVE
PATH_CLS = GDriveURLInfo
PARAM_CHECKSUM = "checksum"
dvc/fs/git.py
+14 -9
View File
@@ -1,8 +1,11 @@
import errno
import os
from tqdm.utils import CallbackIOWrapper
from dvc.utils import is_exec, relpath
from ..progress import DEFAULT_CALLBACK
from .base import BaseFileSystem
@@ -116,16 +119,18 @@ class GitFileSystem(BaseFileSystem): # pylint:disable=abstract-method
# NOTE: os.path.join is ~5.5 times slower
yield f"{root}{os.sep}{file}"
def _download(
self, from_info, to_file, name=None, no_progress_bar=False, **kwargs
def get_file(
self, from_info, to_file, callback=DEFAULT_CALLBACK, **kwargs
):
import shutil
from dvc.progress import Tqdm
total = self.getsize(from_info)
if total:
callback.set_size(total)
with open(to_file, "wb+") as to_fobj:
with Tqdm.wrapattr(
to_fobj, "write", desc=name, disable=no_progress_bar
) as wrapped:
with self.open(from_info, "rb", **kwargs) as from_fobj:
shutil.copyfileobj(from_fobj, wrapped)
with self.open(from_info, "rb", **kwargs) as from_fobj:
with open(to_file, "wb+") as to_fobj:
wrapped = CallbackIOWrapper(
callback.relative_update, from_fobj
)
shutil.copyfileobj(wrapped, to_fobj)
dvc/fs/gs.py
+1 -10
View File
@@ -7,11 +7,10 @@ from dvc.path_info import CloudURLInfo
from dvc.scheme import Schemes
# pylint:disable=abstract-method
from ..progress import DEFAULT_CALLBACK
from .fsspec_wrapper import CallbackMixin, ObjectFSWrapper
class GSFileSystem(ObjectFSWrapper):
class GSFileSystem(CallbackMixin, ObjectFSWrapper):
scheme = Schemes.GS
PATH_CLS = CloudURLInfo
REQUIRES = {"gcsfs": "gcsfs"}
@@ -36,11 +35,3 @@ 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
+12 -15
View File
@@ -8,8 +8,10 @@ import subprocess
from collections import deque
from contextlib import closing, contextmanager
from tqdm.utils import CallbackIOWrapper
from dvc.hash_info import HashInfo
from dvc.progress import DEFAULT_CALLBACK, Tqdm
from dvc.progress import DEFAULT_CALLBACK
from dvc.scheme import Schemes
from dvc.utils import fix_env, tmp_fname
@@ -252,8 +254,6 @@ class HDFSFileSystem(BaseFileSystem):
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)
@@ -267,19 +267,16 @@ class HDFSFileSystem(BaseFileSystem):
shutil.copyfileobj(wrapped, sobj, self.BLOCK_SIZE)
hdfs.move(tmp_file, to_info.path)
def _download(
self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
def get_file(
self, from_info, to_file, callback=DEFAULT_CALLBACK, **kwargs
):
with self.hdfs(from_info) as hdfs:
file_info = hdfs.get_file_info(from_info.path)
total = file_info.size
with open(to_file, "wb+") as fobj:
with Tqdm.wrapattr(
fobj,
"write",
desc=name,
total=total,
disable=no_progress_bar,
) as wrapped:
with hdfs.open_input_stream(from_info.path) as sobj:
shutil.copyfileobj(sobj, wrapped, self.BLOCK_SIZE)
if total:
callback.set_size(total)
with hdfs.open_input_stream(from_info.path) as sobj:
with open(to_file, "wb+") as fobj:
wrapped = CallbackIOWrapper(callback.relative_update, sobj)
shutil.copyfileobj(wrapped, fobj, self.BLOCK_SIZE)
dvc/fs/http.py
+2 -2
View File
@@ -6,7 +6,7 @@ from dvc import prompt
from dvc.path_info import HTTPURLInfo
from dvc.scheme import Schemes
from .fsspec_wrapper import CallbackMixin, FSSpecWrapper, NoDirectoriesMixin
from .fsspec_wrapper import FSSpecWrapper, NoDirectoriesMixin
@wrap_with(threading.Lock())
@@ -32,7 +32,7 @@ def make_context(ssl_verify):
# pylint: disable=abstract-method
class HTTPFileSystem(CallbackMixin, NoDirectoriesMixin, FSSpecWrapper):
class HTTPFileSystem(NoDirectoriesMixin, FSSpecWrapper):
scheme = Schemes.HTTP
PATH_CLS = HTTPURLInfo
PARAM_CHECKSUM = "checksum"
dvc/fs/local.py
+3 -6
View File
@@ -167,10 +167,7 @@ class LocalFileSystem(BaseFileSystem):
copyfile(from_file, tmp_file, callback=callback)
os.replace(tmp_file, to_info)
@staticmethod
def _download(
from_info, to_file, name=None, no_progress_bar=False, **_kwargs
def get_file(
self, from_info, to_file, callback=DEFAULT_CALLBACK, **kwargs
):
copyfile(
from_info, to_file, no_progress_bar=no_progress_bar, name=name
)
copyfile(from_info, to_file, callback=callback)
dvc/fs/oss.py
+2 -2
View File
@@ -7,13 +7,13 @@ from funcy import cached_property, wrap_prop
from dvc.path_info import CloudURLInfo
from dvc.scheme import Schemes
from .fsspec_wrapper import CallbackMixin, ObjectFSWrapper
from .fsspec_wrapper import ObjectFSWrapper
logger = logging.getLogger(__name__)
# pylint:disable=abstract-method
class OSSFileSystem(CallbackMixin, ObjectFSWrapper):
class OSSFileSystem(ObjectFSWrapper):
scheme = Schemes.OSS
PATH_CLS = CloudURLInfo
REQUIRES = {"ossfs": "ossfs"}
dvc/fs/repo.py
+8 -5
View File
@@ -9,6 +9,7 @@ from funcy import lfilter, wrap_with
from dvc.path_info import PathInfo
from ..progress import DEFAULT_CALLBACK
from .base import BaseFileSystem
from .dvc import DvcFileSystem
@@ -448,19 +449,21 @@ class RepoFileSystem(BaseFileSystem): # pylint:disable=abstract-method
for fname in files:
yield PathInfo(root) / fname
def _download(self, from_info, to_file, **kwargs):
def get_file(
self, from_info, to_file, callback=DEFAULT_CALLBACK, **kwargs
):
fs, dvc_fs = self._get_fs_pair(from_info)
try:
fs._download( # pylint: disable=protected-access
from_info, to_file, **kwargs
fs.get_file( # pylint: disable=protected-access
from_info, to_file, callback=callback, **kwargs
)
return
except FileNotFoundError:
if not dvc_fs:
raise
dvc_fs._download( # pylint: disable=protected-access
from_info, to_file, **kwargs
dvc_fs.get_file( # pylint: disable=protected-access
from_info, to_file, callback=callback, **kwargs
)
def metadata(self, path):
dvc/fs/s3.py
+5 -11
View File
@@ -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 DEFAULT_CALLBACK, Tqdm
from dvc.progress import DEFAULT_CALLBACK
from dvc.scheme import Schemes
from .fsspec_wrapper import ObjectFSWrapper
@@ -218,15 +218,9 @@ class S3FileSystem(BaseS3FileSystem): # pylint:disable=abstract-method
self.fs.invalidate_cache(self._with_bucket(to_info.parent))
@_translate_exceptions
def _download(
self, from_info, to_file, name=None, no_progress_bar=False, **pbar_args
def get_file(
self, from_info, to_info, callback=DEFAULT_CALLBACK, **kwargs
):
obj = self._get_obj(from_info)
with Tqdm(
disable=no_progress_bar,
total=obj.content_length,
bytes=True,
desc=name,
**pbar_args,
) as pbar:
obj.download_file(to_file, Callback=pbar.update)
callback.set_size(obj.content_length)
obj.download_file(to_info, Callback=callback.relative_update)
dvc/fs/ssh.py
+2 -2
View File
@@ -9,7 +9,7 @@ from dvc.scheme import Schemes
from dvc.utils.fs import as_atomic
from ..progress import DEFAULT_CALLBACK
from .fsspec_wrapper import CallbackMixin, FSSpecWrapper
from .fsspec_wrapper import FSSpecWrapper
_SSH_TIMEOUT = 60 * 30
_SSH_CONFIG_FILE = os.path.expanduser(os.path.join("~", ".ssh", "config"))
@@ -27,7 +27,7 @@ def ask_password(host, user, port):
# pylint:disable=abstract-method
class SSHFileSystem(CallbackMixin, FSSpecWrapper):
class SSHFileSystem(FSSpecWrapper):
scheme = Schemes.SSH
REQUIRES = {"sshfs": "sshfs"}
dvc/fs/webhdfs.py
+9 -9
View File
@@ -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 DEFAULT_CALLBACK, Tqdm
from dvc.progress import DEFAULT_CALLBACK
from dvc.scheme import Schemes
from .base import BaseFileSystem
@@ -168,13 +168,13 @@ class WebHDFSFileSystem(BaseFileSystem): # pylint:disable=abstract-method
progress=update_callback(callback, total),
)
def _download(
self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs
def get_file(
self, from_info, to_file, callback=DEFAULT_CALLBACK, **kwargs
):
total = self.getsize(from_info)
with Tqdm(
desc=name, total=total, disable=no_progress_bar, bytes=True
) as pbar:
self.hdfs_client.download(
from_info.path, to_file, progress=update_pbar(pbar, total)
)
if total:
callback.set_size(total)
self.hdfs_client.download(
from_info.path, to_file, progress=update_callback(callback, total)
)
dvc/progress.py
+4 -1
View File
@@ -191,6 +191,9 @@ class FsspecCallback(fsspec.Callback):
self.progress_bar.update_to(value)
super().absolute_update(value)
def branch(self, path_1, path_2, kwargs):
return FsspecCallback(self.fs, path_1, Tqdm(**kwargs))
def tdqm_or_callback_wrapped(
fobj, method, total, callback=None, **pbar_kwargs
@@ -206,4 +209,4 @@ def tdqm_or_callback_wrapped(
return Tqdm.wrapattr(fobj, method, total=total, bytes=True, **pbar_kwargs)
DEFAULT_CALLBACK = fsspec.callbacks.NoOpCallback
DEFAULT_CALLBACK = fsspec.callbacks.NoOpCallback()
tests/func/test_fs.py
+35
View File
@@ -439,5 +439,40 @@ def test_upload_callback(tmp_dir, dvc, cloud):
callback = fsspec.Callback()
fs.upload(tmp_dir / "foo", cloud / "foo", callback=callback)
assert callback.size == expected_size
assert callback.value == expected_size
@pytest.mark.needs_internet
@pytest.mark.parametrize(
"cloud",
[
pytest.lazy_fixture("azure"),
pytest.lazy_fixture("gs"),
pytest.param(
pytest.lazy_fixture("gdrive"),
marks=pytest.mark.xfail(
reason="https://github.com/iterative/PyDrive2/issues/136"
),
),
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_download_callback(tmp_dir, dvc, cloud):
cls, config, _ = get_cloud_fs(dvc, **cloud.config)
fs = cls(**config)
fs.upload(io.BytesIO(b"foo"), cloud / "foo")
expected_size = fs.getsize(cloud / "foo")
callback = fsspec.Callback()
fs.download_file(cloud / "foo", tmp_dir / "foo", callback=callback)
assert callback.size == expected_size
assert callback.value == expected_size
assert (tmp_dir / "foo").read_text() == "foo"
tests/unit/remote/test_http.py
+2 -2
View File
@@ -6,11 +6,11 @@ from mock import patch
from dvc.fs.http import HTTPFileSystem
def test_download_fails_on_error_code(dvc, http):
def test_download_fails_on_error_code(tmp_dir, http):
fs = HTTPFileSystem(**http.config)
with pytest.raises(FileNotFoundError):
fs._download(http / "missing.txt", "missing.txt")
fs.download_file(http / "missing.txt", tmp_dir / "missing.txt")
def test_public_auth_method(dvc):