push/fetch: support push/fetch to specific ODB

- prerequisite for remote execution
- allow explicitly pushing imports which would normally be skipped
  (internal-use only, does not apply to `dvc push`)

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
+    )

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)

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,

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