diff --git a/dvc/fs/_callback.py b/dvc/fs/_callback.py index 394b6b05b..668242fc0 100644 --- a/dvc/fs/_callback.py +++ b/dvc/fs/_callback.py @@ -1,5 +1,5 @@ from functools import wraps -from typing import IO, TYPE_CHECKING, Any, Optional, TypeVar, cast +from typing import IO, TYPE_CHECKING, Any, Dict, Optional, TypeVar, cast import fsspec @@ -48,8 +48,7 @@ class FsspecCallback(fsspec.Callback): def make_callback(path1, path2): # pylint: disable=assignment-from-none - child = self.branch(fs.path.name(path1), path2, {}) - return self.as_callback(child) + return self.branch(fs.path.name(path1), path2, {}) @wraps(fn) def func(path1, path2, **kwargs): @@ -89,6 +88,11 @@ class FsspecCallback(fsspec.Callback): ): return callback or RichCallback(**rich_kwargs) + def branch( + self, path_1: str, path_2: str, kwargs: Dict[str, Any] + ) -> "FsspecCallback": + return DEFAULT_CALLBACK + class NoOpCallback(FsspecCallback, fsspec.callbacks.NoOpCallback): pass diff --git a/dvc/fs/base.py b/dvc/fs/base.py index da96b2e39..e16de68a5 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -19,7 +19,6 @@ from typing import ( from funcy import cached_property from dvc.exceptions import DvcException -from dvc.utils.fs import as_atomic, makedirs from dvc.utils.threadpool import ThreadPoolExecutor from ._callback import DEFAULT_CALLBACK, FsspecCallback @@ -360,20 +359,25 @@ class FileSystem: def size(self, path: AnyFSPath) -> Optional[int]: return self.fs.size(path) - # pylint: enable=unused-argument - - def download( + def get( self, from_info: AnyFSPath, to_info: AnyFSPath, - callback: "FsspecCallback" = None, - jobs: int = None, + callback: "FsspecCallback" = DEFAULT_CALLBACK, + batch_size: int = None, ): - if not self.isdir(from_info): - return self.download_file(from_info, to_info, callback=callback) - from .local import localfs + def get_file(from_info, to_info): + localfs.makedirs(localfs.path.parent(to_info), exist_ok=True) + branch = callback.branch(from_info, to_info, {}) + self.get_file(from_info, to_info, callback=branch) + callback.relative_update() + + if not self.isdir(from_info): + callback.set_size(1) + return get_file(from_info, to_info) + pairs = { info: localfs.path.join( to_info, *self.path.relparts(info, from_info) @@ -381,48 +385,24 @@ class FileSystem: for info in self.find(from_info) } if not pairs: - makedirs(to_info, exist_ok=True) + localfs.makedirs(to_info, exist_ok=True) 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. - with ThreadPoolExecutor( - max_workers=max_workers, cancel_on_error=True - ) as executor: - list( - executor.imap_unordered( - lambda args: download_files(*args), pairs.items() - ) + callback.set_size(len(pairs)) + max_workers = batch_size 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. + with ThreadPoolExecutor( + max_workers=max_workers, cancel_on_error=True + ) as executor: + list( + executor.imap_unordered( + lambda args: get_file(*args), pairs.items() ) - - def download_file( - self, - from_info: AnyFSPath, - to_info: AnyFSPath, - callback: FsspecCallback = None, - ): - from .local import localfs - - 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 d1cffe994..f379cef4b 100644 --- a/dvc/fs/utils.py +++ b/dvc/fs/utils.py @@ -3,6 +3,8 @@ import logging import os from typing import TYPE_CHECKING, List, Optional +from dvc.utils.fs import as_atomic + from ._callback import DEFAULT_CALLBACK, FsspecCallback from .base import RemoteActionNotImplemented from .local import LocalFileSystem @@ -45,7 +47,10 @@ def copy( return to_fs.put_file(from_path, to_path, callback=callback) if isinstance(to_fs, LocalFileSystem): - return from_fs.download_file(from_path, to_path, callback=callback) + from .local import localfs + + with as_atomic(localfs, to_path, create_parents=True) as tmp_file: + return from_fs.get_file(from_path, tmp_file, callback=callback) with from_fs.open(from_path, mode="rb") as fobj: size = from_fs.size(from_path) diff --git a/dvc/output.py b/dvc/output.py index a532f9860..c064f10b9 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -695,7 +695,14 @@ class Output: raise DvcException(msg.format(self.fs_path, name)) def download(self, to, jobs=None): - self.fs.download(self.fs_path, to.fs_path, jobs=jobs) + from dvc.fs._callback import FsspecCallback + + with FsspecCallback.as_tqdm_callback( + total=-1, + desc=f"Downloading {self.fs.path.name(self.fs_path)}", + unit="files", + ) as cb: + self.fs.get(self.fs_path, to.fs_path, batch_size=jobs, callback=cb) def get_obj(self, filter_info=None, **kwargs): if self.obj: diff --git a/dvc/repo/get.py b/dvc/repo/get.py index f1690a40f..55a4f218a 100644 --- a/dvc/repo/get.py +++ b/dvc/repo/get.py @@ -21,6 +21,7 @@ def get(url, path, out=None, rev=None, jobs=None): from dvc.dvcfile import is_valid_filename from dvc.external_repo import external_repo + from dvc.fs._callback import FsspecCallback out = resolve_output(path, out) @@ -49,6 +50,14 @@ def get(url, path, out=None, rev=None, jobs=None): with external_repo( url=url, rev=rev, cache_dir=tmp_dir, cache_types=cache_types ) as repo: - repo.repo_fs.download(path, os.path.abspath(out), jobs=jobs) + fs = repo.repo_fs + with FsspecCallback.as_tqdm_callback( + total=-1, + desc=f"Downloading {fs.path.name(path)}", + unit="files", + ) as cb: + fs.get( + path, os.path.abspath(out), batch_size=jobs, callback=cb + ) finally: remove(tmp_dir) diff --git a/tests/func/test_external_repo.py b/tests/func/test_external_repo.py index 3a5cb2a52..b97095635 100644 --- a/tests/func/test_external_repo.py +++ b/tests/func/test_external_repo.py @@ -95,7 +95,7 @@ def test_pull_subdir_file(tmp_dir, erepo_dir): dest = tmp_dir / "file" with external_repo(os.fspath(erepo_dir)) as repo: - repo.repo_fs.download( + repo.repo_fs.get( "subdir/file", os.fspath(dest), ) @@ -191,7 +191,7 @@ def test_subrepos_are_ignored(tmp_dir, erepo_dir): subrepo.dvc_gen({"file": "file"}, commit="add files on subrepo") with external_repo(os.fspath(erepo_dir)) as repo: - repo.repo_fs.download( + repo.repo_fs.get( "dir", os.fspath(tmp_dir / "out"), ) @@ -237,7 +237,7 @@ def test_subrepos_are_ignored_for_git_tracked_dirs(tmp_dir, erepo_dir): subrepo.dvc_gen({"file": "file"}, commit="add files on subrepo") with external_repo(os.fspath(erepo_dir)) as repo: - repo.repo_fs.download( + repo.repo_fs.get( "dir", os.fspath(tmp_dir / "out"), ) diff --git a/tests/func/test_fs.py b/tests/func/test_fs.py index d40227012..5610e433c 100644 --- a/tests/func/test_fs.py +++ b/tests/func/test_fs.py @@ -3,8 +3,6 @@ from io import BytesIO from operator import itemgetter from os.path import join -import pytest - from dvc.fs import get_cloud_fs from dvc.fs._callback import FsspecCallback from dvc.fs.local import LocalFileSystem @@ -258,33 +256,13 @@ def test_upload_callback(tmp_dir, dvc, cloud): assert callback.value == expected_size -def test_download_callback(tmp_dir, dvc, cloud, local_cloud): - cls, config, _ = get_cloud_fs(dvc, **cloud.config) - fs = cls(**config) - - (tmp_dir / "to_upload").write_text("foo") - fs.put_file((tmp_dir / "to_upload").fs_path, (cloud / "foo").fs_path) - expected_size = fs.size((cloud / "foo").fs_path) - - callback = FsspecCallback() - fs.download_file( - (cloud / "foo").fs_path, - (tmp_dir / "foo").fs_path, - callback=callback, - ) - - assert callback.size == expected_size - assert callback.value == expected_size - assert (tmp_dir / "foo").read_text() == "foo" - - -def test_download_dir_callback(tmp_dir, dvc, cloud): +def test_get_dir_callback(tmp_dir, dvc, cloud): cls, config, _ = get_cloud_fs(dvc, **cloud.config) fs = cls(**config) cloud.gen({"dir": {"foo": "foo", "bar": "bar"}}) callback = FsspecCallback() - fs.download( + fs.get( (cloud / "dir").fs_path, (tmp_dir / "dir").fs_path, callback=callback ) @@ -293,47 +271,14 @@ def test_download_dir_callback(tmp_dir, dvc, cloud): assert (tmp_dir / "dir").read_text() == {"foo": "foo", "bar": "bar"} -@pytest.mark.parametrize("fs_type", ["git", "dvc"]) -def test_download_callbacks_on_dvc_git_fs(tmp_dir, dvc, scm, fs_type): - from dvc.fs.git import GitFileSystem - - gen = tmp_dir.scm_gen if fs_type == "git" else tmp_dir.dvc_gen - gen({"dir": {"foo": "foo", "bar": "bar"}, "file": "file"}, commit="gen") - - fs = dvc.dvcfs if fs_type == "dvc" else GitFileSystem(scm=scm, rev="HEAD") - - callback = FsspecCallback() - fs.download_file( - "file", - (tmp_dir / "file2").fs_path, - callback=callback, - ) - - size = os.path.getsize(tmp_dir / "file") - assert (tmp_dir / "file2").read_text() == "file" - assert callback.size == size - assert callback.value == size - - callback = FsspecCallback() - fs.download( - "dir", - (tmp_dir / "dir2").fs_path, - callback=callback, - ) - - assert (tmp_dir / "dir2").read_text() == {"foo": "foo", "bar": "bar"} - assert callback.size == 2 - assert callback.value == 2 - - -def test_callback_on_repo_fs(tmp_dir, dvc, scm): +def test_callback_on_repo_fs(tmp_dir, dvc, scm, mocker): tmp_dir.dvc_gen({"dir": {"bar": "bar"}}, commit="dvc") tmp_dir.scm_gen({"dir": {"foo": "foo"}}, commit="git") fs = dvc.repo_fs callback = FsspecCallback() - fs.download( + fs.get( "dir", (tmp_dir / "dir2").fs_path, callback=callback, @@ -344,7 +289,8 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm): assert callback.value == 2 callback = FsspecCallback() - fs.download( + branch = mocker.spy(callback, "branch") + fs.get( os.path.join("dir", "foo"), (tmp_dir / "foo").fs_path, callback=callback, @@ -352,11 +298,18 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm): size = os.path.getsize(tmp_dir / "dir" / "foo") assert (tmp_dir / "foo").read_text() == "foo" - assert callback.size == size - assert callback.value == size + assert callback.size == 1 + assert callback.value == 1 + + assert branch.call_count == 1 + assert branch.spy_return.size == size + assert branch.spy_return.value == size + + branch.reset_mock() callback = FsspecCallback() - fs.download( + branch = mocker.spy(callback, "branch") + fs.get( os.path.join("dir", "bar"), (tmp_dir / "bar").fs_path, callback=callback, @@ -364,5 +317,9 @@ def test_callback_on_repo_fs(tmp_dir, dvc, scm): size = os.path.getsize(tmp_dir / "dir" / "bar") assert (tmp_dir / "bar").read_text() == "bar" - assert callback.size == size - assert callback.value == size + assert callback.size == 1 + assert callback.value == 1 + + assert branch.call_count == 1 + assert branch.spy_return.size == size + assert branch.spy_return.value == size diff --git a/tests/func/test_repro.py b/tests/func/test_repro.py index 4426e9aa2..ced595dea 100644 --- a/tests/func/test_repro.py +++ b/tests/func/test_repro.py @@ -877,10 +877,10 @@ class TestReproAlreadyCached(TestRepro): ret = main(["import-url", self.FOO, self.BAR]) self.assertEqual(ret, 0) - patch_download = patch.object( + patch_get = patch.object( LocalFileSystem, - "download", - side_effect=LocalFileSystem.download, + "get", + side_effect=LocalFileSystem.get, autospec=True, ) @@ -891,7 +891,7 @@ class TestReproAlreadyCached(TestRepro): autospec=True, ) - with patch_download as mock_download: + with patch_get as mock_download: with patch_checkout as mock_checkout: assert main(["unfreeze", "bar.dvc"]) == 0 ret = main(["repro", "--force", "bar.dvc"])