From 83795ff3df9f7ce6e863d2b68aa63e6ce285f143 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 19:28:31 +0545 Subject: [PATCH] fs.put_file: add support for uploading fileobjects Also changed the callbacks to inherit from our own so that we can add few more utilities to it. --- dvc/fs/_callback.py | 53 ++++++++++++++++++++++------ dvc/fs/base.py | 53 +++++++++++++--------------- dvc/fs/http.py | 8 ++--- dvc/progress.py | 4 +-- dvc/repo/experiments/executor/ssh.py | 2 +- tests/func/test_fs.py | 20 +++++------ 6 files changed, 84 insertions(+), 56 deletions(-) diff --git a/dvc/fs/_callback.py b/dvc/fs/_callback.py index 54720bda7..58adee2ed 100644 --- a/dvc/fs/_callback.py +++ b/dvc/fs/_callback.py @@ -1,9 +1,51 @@ +from functools import wraps +from typing import IO, TYPE_CHECKING, Optional, TypeVar, cast + import fsspec from dvc.progress import Tqdm +if TYPE_CHECKING: + from typing import Callable + + from typing_extensions import ParamSpec + + _P = ParamSpec("_P") + _R = TypeVar("_R") + class FsspecCallback(fsspec.Callback): + """FsspecCallback usable as a context manager, and a few helper methods.""" + + def wrap_attr(self, fobj: IO, method: str = "read") -> IO: + from tqdm.utils import CallbackIOWrapper + + wrapped = CallbackIOWrapper(self.relative_update, fobj, method) + return cast(IO, wrapped) + + def wrap_fn(self, fn: "Callable[_P, _R]") -> "Callable[_P, _R]": + @wraps(fn) + def wrapped(*args: "_P.args", **kwargs: "_P.kwargs") -> "_R": + res = fn(*args, **kwargs) + self.relative_update() + return res + + return wrapped + + @classmethod + def as_callback( + cls, maybe_callback: Optional["FsspecCallback"] = None + ) -> "FsspecCallback": + if maybe_callback is None: + return DEFAULT_CALLBACK + return maybe_callback + + +class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback): + pass + + +class TqdmCallback(FsspecCallback): def __init__(self, progress_bar): self.progress_bar = progress_bar super().__init__() @@ -22,15 +64,6 @@ class FsspecCallback(fsspec.Callback): self.progress_bar.update_to(value) super().absolute_update(value) - @staticmethod - def wrap_fn(cb, fn): - def wrapped(*args, **kwargs): - res = fn(*args, **kwargs) - cb.relative_update() - return res - - return wrapped - def tdqm_or_callback_wrapped( fobj, method, total, callback=None, **pbar_kwargs @@ -46,4 +79,4 @@ def tdqm_or_callback_wrapped( return Tqdm.wrapattr(fobj, method, total=total, bytes=True, **pbar_kwargs) -DEFAULT_CALLBACK = fsspec.callbacks.NoOpCallback() +DEFAULT_CALLBACK = NoOpCallback() diff --git a/dvc/fs/base.py b/dvc/fs/base.py index a8e496515..7a808f865 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -15,18 +15,19 @@ from typing import ( List, Optional, Union, + cast, overload, ) from funcy import cached_property -from tqdm.utils import CallbackIOWrapper from dvc.exceptions import DvcException -from dvc.fs._callback import DEFAULT_CALLBACK, FsspecCallback from dvc.ui import ui from dvc.utils import tmp_fname from dvc.utils.fs import makedirs, move +from ._callback import DEFAULT_CALLBACK, FsspecCallback + if TYPE_CHECKING: from fsspec.spec import AbstractFileSystem from typing_extensions import Literal @@ -271,19 +272,29 @@ class FileSystem: def put_file( self, - from_file: AnyFSPath, + from_file: Union[AnyFSPath, IO], to_info: AnyFSPath, - callback: Any = DEFAULT_CALLBACK, + callback: FsspecCallback = DEFAULT_CALLBACK, **kwargs, ) -> None: - self.fs.put_file(from_file, to_info, callback=callback, **kwargs) + size = kwargs.get("size") + if size: + callback.set_size(size) + if hasattr(from_file, "read"): + stream = callback.wrap_attr(cast("IO", from_file)) + self.upload_fobj(stream, to_info, size=size) + else: + assert isinstance(from_file, str) + self.fs.put_file( + os.fspath(from_file), to_info, callback=callback, **kwargs + ) self.fs.invalidate_cache(self.path.parent(to_info)) def get_file( self, from_info: AnyFSPath, to_info: AnyFSPath, - callback: Any = DEFAULT_CALLBACK, + callback: FsspecCallback = DEFAULT_CALLBACK, **kwargs, ) -> None: self.fs.get_file(from_info, to_info, callback=callback, **kwargs) @@ -320,15 +331,11 @@ class FileSystem: no_progress_bar: bool = False, **pbar_args: Any, ): - is_file_obj = hasattr(from_info, "read") - method = "upload_fobj" if is_file_obj else "put_file" - if not hasattr(self, method): - raise RemoteActionNotImplemented(method, self.scheme) - - if not is_file_obj: + 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) stack = contextlib.ExitStack() if not callback: @@ -341,24 +348,12 @@ class FileSystem: ) stack.enter_context(pbar) callback = pbar.as_callback() - 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) - - assert isinstance(from_info, str) - logger.debug("Uploading '%s' to '%s'", from_info, to_info) - # pylint: disable=no-member + if total: + callback.set_size(total) return self.put_file( - os.fspath(from_info), to_info, callback=callback + from_info, to_info, callback=callback, size=total ) def download( @@ -406,7 +401,7 @@ class FileSystem: self, from_info: AnyFSPath, to_info: AnyFSPath, - callback=DEFAULT_CALLBACK, + callback: FsspecCallback = DEFAULT_CALLBACK, jobs: int = None, **kwargs, ): @@ -423,7 +418,7 @@ class FileSystem: ) callback.set_size(len(from_infos)) - download_files = FsspecCallback.wrap_fn(callback, self._download_file) + download_files = callback.wrap_fn(self._download_file) max_workers = jobs or self.jobs with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [ diff --git a/dvc/fs/http.py b/dvc/fs/http.py index dbca69263..7f22de11e 100644 --- a/dvc/fs/http.py +++ b/dvc/fs/http.py @@ -1,12 +1,12 @@ import threading -from typing import Any +from typing import IO, Union from funcy import cached_property, memoize, wrap_with from dvc import prompt from dvc.scheme import Schemes -from ._callback import DEFAULT_CALLBACK +from ._callback import DEFAULT_CALLBACK, FsspecCallback from .base import AnyFSPath, FileSystem @@ -141,9 +141,9 @@ class HTTPFileSystem(FileSystem): def put_file( self, - from_file: AnyFSPath, + from_file: Union[AnyFSPath, IO], to_info: AnyFSPath, - callback: Any = DEFAULT_CALLBACK, + callback: FsspecCallback = DEFAULT_CALLBACK, **kwargs, ) -> None: kwargs.setdefault("method", self.upload_method) diff --git a/dvc/progress.py b/dvc/progress.py index cd5fff9a2..3818b1ce5 100644 --- a/dvc/progress.py +++ b/dvc/progress.py @@ -136,9 +136,9 @@ class Tqdm(tqdm): return wrapped def as_callback(self): - from dvc.fs._callback import FsspecCallback + from dvc.fs._callback import TqdmCallback - return FsspecCallback(self) + return TqdmCallback(self) def close(self): self.postfix["info"] = "" diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index e57d3a1ec..29fc79d83 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/dvc/repo/experiments/executor/ssh.py @@ -257,7 +257,7 @@ class SSHExecutor(BaseExecutor): with TemporaryFile(mode="w+", encoding="utf-8") as fobj: json.dump(info.asdict(), fobj) fobj.seek(0) - fs.upload_fobj(fobj, infofile) + fs.put_file(fobj, infofile) cmd = ["source ~/.profile"] script_path = cls._setup_script_path(info.dvc_dir) if fs.exists(posixpath.join(info.root_dir, script_path)): diff --git a/tests/func/test_fs.py b/tests/func/test_fs.py index 6f2adf85c..8988c2237 100644 --- a/tests/func/test_fs.py +++ b/tests/func/test_fs.py @@ -3,10 +3,10 @@ import os from operator import itemgetter from os.path import join -import fsspec import pytest from dvc.fs import get_cloud_fs +from dvc.fs._callback import FsspecCallback from dvc.fs.local import LocalFileSystem from dvc.repo import Repo @@ -201,7 +201,7 @@ def test_walk_dont_ignore_subrepos(tmp_dir, scm, dvc): assert set(get_dirs(next(scm_fs.walk(path)))) == {".dvc", "subdir"} -def test_fs_getsize(dvc, cloud): +def test_fs_size(dvc, cloud): cloud.gen({"data": {"foo": "foo"}, "baz": "baz baz"}) cls, config, path = get_cloud_fs(dvc, **cloud.config) fs = cls(**config) @@ -248,7 +248,7 @@ def test_upload_callback(tmp_dir, dvc, cloud): fs = cls(**config) expected_size = os.path.getsize(tmp_dir / "foo") - callback = fsspec.Callback() + callback = FsspecCallback() fs.upload( (tmp_dir / "foo").fs_path, (cloud / "foo").fs_path, @@ -267,7 +267,7 @@ def test_download_callback(tmp_dir, dvc, cloud, local_cloud): fs.upload((tmp_dir / "to_upload").fs_path, (cloud / "foo").fs_path) expected_size = fs.size((cloud / "foo").fs_path) - callback = fsspec.Callback() + callback = FsspecCallback() fs.download_file( (cloud / "foo").fs_path, (tmp_dir / "foo").fs_path, @@ -284,7 +284,7 @@ def test_download_dir_callback(tmp_dir, dvc, cloud): fs = cls(**config) cloud.gen({"dir": {"foo": "foo", "bar": "bar"}}) - callback = fsspec.Callback() + callback = FsspecCallback() fs.download( (cloud / "dir").fs_path, (tmp_dir / "dir").fs_path, callback=callback ) @@ -303,7 +303,7 @@ def test_download_callbacks_on_dvc_git_fs(tmp_dir, dvc, scm, fs_type): fs = dvc.dvcfs if fs_type == "dvc" else GitFileSystem(scm=scm, rev="HEAD") - callback = fsspec.Callback() + callback = FsspecCallback() fs.download_file( "file", (tmp_dir / "file2").fs_path, @@ -315,7 +315,7 @@ def test_download_callbacks_on_dvc_git_fs(tmp_dir, dvc, scm, fs_type): assert callback.size == size assert callback.value == size - callback = fsspec.Callback() + callback = FsspecCallback() fs.download( "dir", (tmp_dir / "dir2").fs_path, @@ -333,7 +333,7 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm): fs = dvc.repo_fs - callback = fsspec.Callback() + callback = FsspecCallback() fs.download( "dir", (tmp_dir / "dir2").fs_path, @@ -344,7 +344,7 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm): assert callback.size == 2 assert callback.value == 2 - callback = fsspec.Callback() + callback = FsspecCallback() fs.download( os.path.join("dir", "foo"), (tmp_dir / "foo").fs_path, @@ -356,7 +356,7 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm): assert callback.size == size assert callback.value == size - callback = fsspec.Callback() + callback = FsspecCallback() fs.download( os.path.join("dir", "bar"), (tmp_dir / "bar").fs_path,