diff --git a/dvc/repo/experiments/pull.py b/dvc/repo/experiments/pull.py index d2440bf4a..8760e523e 100644 --- a/dvc/repo/experiments/pull.py +++ b/dvc/repo/experiments/pull.py @@ -48,7 +48,16 @@ def pull( _pull_cache(repo, exp_ref, **kwargs) -def _pull_cache(repo, exp_ref, dvc_remote=None, jobs=None, run_cache=False): +def _pull_cache( + repo, + exp_ref, + dvc_remote=None, + jobs=None, + run_cache=False, + odb=None, +): revs = list(exp_commits(repo.scm, [exp_ref])) logger.debug("dvc fetch experiment '%s'", exp_ref) - repo.fetch(jobs=jobs, remote=dvc_remote, run_cache=run_cache, revs=revs) + repo.fetch( + jobs=jobs, remote=dvc_remote, run_cache=run_cache, revs=revs, odb=odb + ) diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index c7af1179b..80acdf2d7 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -1,10 +1,14 @@ import logging +from typing import TYPE_CHECKING, Optional from dvc.exceptions import DownloadError, FileTransferError from dvc.scheme import Schemes from . import locked +if TYPE_CHECKING: + from dvc.objects.db.base import ObjectDB + logger = logging.getLogger(__name__) @@ -21,6 +25,7 @@ def fetch( all_commits=False, run_cache=False, revs=None, + odb: Optional["ObjectDB"] = None, ): """Download data items from a cloud and imported repositories @@ -59,20 +64,34 @@ def fetch( except DownloadError as exc: failed += exc.amount - for odb, obj_ids in sorted( - used.items(), - key=lambda item: item[0] is not None - and item[0].fs.scheme == Schemes.MEMORY, - ): + if odb: + all_ids = set() + for _odb, obj_ids in used.items(): + all_ids.update(obj_ids) d, f = _fetch( self, - obj_ids, + 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.scheme == Schemes.MEMORY, + ): + d, f = _fetch( + self, + obj_ids, + jobs=jobs, + remote=remote, + odb=src_odb, + ) + downloaded += d + failed += f if failed: raise DownloadError(failed) diff --git a/dvc/repo/pull.py b/dvc/repo/pull.py index e0714da3f..b4f472933 100644 --- a/dvc/repo/pull.py +++ b/dvc/repo/pull.py @@ -1,8 +1,12 @@ import logging +from typing import TYPE_CHECKING, Optional from dvc.repo import locked from dvc.utils import glob_targets +if TYPE_CHECKING: + from dvc.objects.db.base import ObjectDB + logger = logging.getLogger(__name__) @@ -20,6 +24,7 @@ def pull( all_commits=False, run_cache=False, glob=False, + odb: Optional["ObjectDB"] = None, ): if isinstance(targets, str): targets = [targets] @@ -36,6 +41,7 @@ def pull( with_deps=with_deps, recursive=recursive, run_cache=run_cache, + odb=odb, ) stats = self.checkout( targets=expanded_targets, diff --git a/dvc/repo/push.py b/dvc/repo/push.py index 7f5320ab3..25a9b3692 100644 --- a/dvc/repo/push.py +++ b/dvc/repo/push.py @@ -24,6 +24,7 @@ def push( revs=None, glob=False, odb: Optional["ObjectDB"] = None, + include_imports=False, ): used_run_cache = ( self.stage_cache.push(remote, odb=odb) if run_cache else [] @@ -49,13 +50,24 @@ def push( ) pushed = len(used_run_cache) - for dest_odb, obj_ids in used.items(): - if dest_odb and dest_odb.read_only: - continue + if odb: + all_ids = set() + for dest_odb, obj_ids in used.items(): + if not include_imports and dest_odb and dest_odb.read_only: + continue + all_ids.update(obj_ids) try: - pushed += self.cloud.push( - obj_ids, jobs, remote=remote, odb=odb or dest_odb - ) + pushed += self.cloud.push(all_ids, jobs, remote=remote, odb=odb) except FileTransferError as exc: raise UploadError(exc.amount) + 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) return pushed