From 7c225206abcd1fefdb838354ba3dff24cae56075 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Tue, 18 Oct 2022 18:52:50 +0900 Subject: [PATCH] fetch/push: use dvc-data TransferResult --- dvc/data_cloud.py | 15 +++--- dvc/repo/fetch.py | 120 +++++++++++++++++++++++++-------------------- dvc/repo/push.py | 22 ++++----- dvc/stage/cache.py | 3 +- 4 files changed, 85 insertions(+), 75 deletions(-) diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index ec495b18e..2e1009e2d 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from dvc_data.hashfile.db import HashFileDB from dvc_data.hashfile.hash_info import HashInfo from dvc_data.hashfile.status import CompareStatusResult + from dvc_data.hashfile.transfer import TransferResult logger = logging.getLogger(__name__) @@ -106,14 +107,10 @@ class DataCloud: dest_odb: "HashFileDB", objs: Iterable["HashInfo"], **kwargs, - ): - from dvc.exceptions import FileTransferError - from dvc_data.hashfile.transfer import TransferError, transfer + ) -> "TransferResult": + from dvc_data.hashfile.transfer import transfer - try: - return transfer(src_odb, dest_odb, objs, **kwargs) - except TransferError as exc: - raise FileTransferError(exc.fails) from exc + return transfer(src_odb, dest_odb, objs, **kwargs) def push( self, @@ -121,7 +118,7 @@ class DataCloud: jobs: Optional[int] = None, remote: Optional[str] = None, odb: Optional["HashFileDB"] = None, - ): + ) -> "TransferResult": """Push data items in a cloud-agnostic way. Args: @@ -148,7 +145,7 @@ class DataCloud: jobs: Optional[int] = None, remote: Optional[str] = None, odb: Optional["HashFileDB"] = None, - ): + ) -> "TransferResult": """Pull data items in a cloud-agnostic way. Args: diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index 2bdcb061e..dcbcbaaef 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -3,12 +3,15 @@ from contextlib import suppress from typing import TYPE_CHECKING, Optional from dvc.config import NoRemoteError -from dvc.exceptions import DownloadError, FileTransferError +from dvc.exceptions import DownloadError from dvc.fs import Schemes from . import locked if TYPE_CHECKING: + from dvc.repo import Repo + from dvc.types import TargetType + from dvc_data.hashfile.transfer import TransferResult from dvc_objects.db.base import ObjectDB logger = logging.getLogger(__name__) @@ -63,7 +66,17 @@ def fetch( if _remote.worktree: return _fetch_worktree(self, _remote) - used = self.used_objs( + downloaded_count = 0 + failed_count = 0 + + try: + if run_cache: + self.stage_cache.pull(remote) + except DownloadError as exc: + failed_count += exc.amount + + result = _fetch( + self, targets, all_branches=all_branches, all_tags=all_tags, @@ -74,45 +87,10 @@ def fetch( jobs=jobs, recursive=recursive, revs=revs, + odb=odb, ) - - downloaded = 0 - failed = 0 - - try: - if run_cache: - self.stage_cache.pull(remote) - except DownloadError as exc: - failed += exc.amount - - if odb: - all_ids = set() - for _odb, obj_ids in used.items(): - all_ids.update(obj_ids) - d, f = _fetch( - self, - all_ids, - jobs=jobs, - remote=remote, - odb=odb, - ) - downloaded += d - failed += f - else: - for src_odb, obj_ids in sorted( - used.items(), - key=lambda item: item[0] is not None - and item[0].fs.protocol == Schemes.MEMORY, - ): - d, f = _fetch( - self, - obj_ids, - jobs=jobs, - remote=remote, - odb=src_odb, - ) - downloaded += d - failed += f + downloaded_count = len(result.transferred) + failed_count += len(result.failed) d, f = _fetch_partial_imports( self, @@ -123,23 +101,59 @@ def fetch( recursive=recursive, revs=revs, ) - downloaded += d - failed += f + downloaded_count += d + failed_count += f - if failed: - raise DownloadError(failed) + if failed_count: + raise DownloadError(failed_count) - return downloaded + return downloaded_count -def _fetch(repo, obj_ids, **kwargs): - downloaded = 0 - failed = 0 - try: - downloaded += repo.cloud.pull(obj_ids, **kwargs) - except FileTransferError as exc: - failed += exc.amount - return downloaded, failed +def _fetch( + repo: "Repo", + targets: "TargetType", + remote: Optional[str] = None, + jobs: Optional[int] = None, + odb: Optional["ObjectDB"] = None, + **kwargs, +) -> "TransferResult": + from dvc_data.hashfile.transfer import TransferResult + + result = TransferResult(set(), set()) + used = repo.used_objs( + targets, + remote=remote, + jobs=jobs, + **kwargs, + ) + if odb: + all_ids = set() + for _odb, obj_ids in used.items(): + all_ids.update(obj_ids) + d, f = repo.cloud.pull( + all_ids, + jobs=jobs, + remote=remote, + odb=odb, + ) + result.transferred.update(d) + result.failed.update(f) + else: + for src_odb, obj_ids in sorted( + used.items(), + key=lambda item: item[0] is not None + and item[0].fs.protocol == Schemes.MEMORY, + ): + d, f = repo.cloud.pull( + obj_ids, + jobs=jobs, + remote=remote, + odb=src_odb, + ) + result.transferred.update(d) + result.failed.update(f) + return result def _fetch_partial_imports(repo, targets, **kwargs): diff --git a/dvc/repo/push.py b/dvc/repo/push.py index f8c1e6d5f..905837417 100644 --- a/dvc/repo/push.py +++ b/dvc/repo/push.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING, Optional -from dvc.exceptions import FileTransferError, UploadError +from dvc.exceptions import UploadError from ..utils import glob_targets from . import locked @@ -103,18 +103,18 @@ def push( if not include_imports and dest_odb and dest_odb.read_only: continue all_ids.update(obj_ids) - try: - pushed += self.cloud.push(all_ids, jobs, remote=remote, odb=odb) - except FileTransferError as exc: - raise UploadError(exc.amount) + result = self.cloud.push(all_ids, jobs, remote=remote, odb=odb) + if result.failed: + raise UploadError(len(result.failed)) + pushed += len(result.transferred) else: for dest_odb, obj_ids in used.items(): if dest_odb and dest_odb.read_only: continue - try: - pushed += self.cloud.push( - obj_ids, jobs, remote=remote, odb=odb or dest_odb - ) - except FileTransferError as exc: - raise UploadError(exc.amount) + result = self.cloud.push( + obj_ids, jobs, remote=remote, odb=odb or dest_odb + ) + if result.failed: + raise UploadError(len(result.failed)) + pushed += len(result.transferred) return pushed diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py index 7cd6e46ed..d76e7befc 100644 --- a/dvc/stage/cache.py +++ b/dvc/stage/cache.py @@ -9,7 +9,6 @@ from funcy import cached_property, first from dvc import fs from dvc.exceptions import DvcException from dvc.utils import dict_sha256, relpath -from dvc_data.hashfile.transfer import _log_exceptions if TYPE_CHECKING: from dvc_objects.db import ObjectDB @@ -222,7 +221,7 @@ class StageCache: from_fs = from_odb.fs to_fs = to_odb.fs - func = _log_exceptions(fs.generic.copy) + func = fs.generic.log_exceptions(fs.generic.copy) runs = from_fs.path.join(from_odb.path, "runs") http_odb = next(