Pull: accept multiple as input(#7196)

fix: #7196

New feature: now `dvc exp pull` can accept a list of experiment names as input.

2. add support to list input for the `exp pull`
3. add function tests for the multi exp pull
4. add unit test for the arg parse

dvc/command/experiments/pull.py

@@ -87,6 +87,9 @@ def add_parser(experiments_subparsers, parent_parser):
         metavar="<git_remote>",
     )
     experiments_pull_parser.add_argument(
-        "experiment", help="Experiment to pull.", metavar="<experiment>"
+        "experiment",
+        nargs="+",
+        help="Experiments to pull.",
+        metavar="<experiment>",
     )
     experiments_pull_parser.set_defaults(func=CmdExperimentsPull)

dvc/repo/experiments/exceptions.py

@@ -1,5 +1,6 @@
+from typing import Iterable, List
+
 from dvc.exceptions import InvalidArgumentError
-from dvc.types import Iterable, List, Optional
 
 from .base import ExpRefInfo
 
@@ -20,6 +21,7 @@ class AmbiguousExpRefInfo(InvalidArgumentError):
         msg.extend([f"\t{info}" for info in exp_ref_list])
         super().__init__("\n".join(msg))
 
+
 class UnresolvedExpNamesError(InvalidArgumentError):
     def __init__(
         self, unresolved_list: List[str], *args, git_remote: str = None

dvc/repo/experiments/pull.py

@@ -1,9 +1,12 @@
 import logging
+from typing import Iterable, Union
 
-from dvc.exceptions import DvcException, InvalidArgumentError
+from dvc.exceptions import DvcException
 from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
+from dvc.scm import TqdmGit
 
+from .exceptions import UnresolvedExpNamesError
 from .utils import exp_commits, resolve_name
 
 logger = logging.getLogger(__name__)
@@ -12,40 +15,61 @@ logger = logging.getLogger(__name__)
 @locked
 @scm_context
 def pull(
-    repo, git_remote, exp_name, *args, force=False, pull_cache=False, **kwargs
+    repo,
+    git_remote: str,
+    exp_names: Union[Iterable[str], str],
+    *args,
+    force: bool = False,
+    pull_cache: bool = False,
+    **kwargs,
 ):
-    exp_ref_dict = resolve_name(repo.scm, exp_name, git_remote)
-    exp_ref = exp_ref_dict[exp_name]
-    if not exp_ref:
-        raise InvalidArgumentError(
-            f"Experiment '{exp_name}' does not exist in '{git_remote}'"
-        )
+    if isinstance(exp_names, str):
+        exp_names = [exp_names]
+    exp_ref_dict = resolve_name(repo.scm, exp_names, git_remote)
+    unresolved_exp_names = [
+        exp_name
+        for exp_name, exp_ref in exp_ref_dict.items()
+        if exp_ref is None
+    ]
+    if unresolved_exp_names:
+        raise UnresolvedExpNamesError(unresolved_exp_names)
 
+    exp_ref_set = exp_ref_dict.values()
+    _pull(repo, git_remote, exp_ref_set, force, pull_cache, **kwargs)
+
+
+def _pull(
+    repo,
+    git_remote: str,
+    exp_refs,
+    force: bool,
+    pull_cache: bool,
+    **kwargs,
+):
     def on_diverged(refname: str, rev: str) -> bool:
         if repo.scm.get_ref(refname) == rev:
             return True
+        exp_name = refname.split("/")[-1]
         raise DvcException(
             f"Local experiment '{exp_name}' has diverged from remote "
             "experiment with the same name. To override the local experiment "
             "re-run with '--force'."
         )
 
-    refspec = f"{exp_ref}:{exp_ref}"
-    logger.debug("git pull experiment '%s' -> '%s'", git_remote, refspec)
-
-    from dvc.scm import TqdmGit
+    refspec_list = [f"{exp_ref}:{exp_ref}" for exp_ref in exp_refs]
+    logger.debug(f"git pull experiment '{git_remote}' -> '{refspec_list}'")
 
     with TqdmGit(desc="Fetching git refs") as pbar:
         repo.scm.fetch_refspecs(
             git_remote,
-            [refspec],
+            refspec_list,
             force=force,
             on_diverged=on_diverged,
             progress=pbar.update_git,
         )
 
     if pull_cache:
-        _pull_cache(repo, exp_ref, **kwargs)
+        _pull_cache(repo, exp_refs, **kwargs)
 
 
 def _pull_cache(
@@ -56,8 +80,8 @@ def _pull_cache(
     run_cache=False,
     odb=None,
 ):
-    revs = list(exp_commits(repo.scm, [exp_ref]))
-    logger.debug("dvc fetch experiment '%s'", exp_ref)
+    revs = list(exp_commits(repo.scm, exp_ref))
+    logger.debug(f"dvc fetch experiment '{exp_ref}'")
     repo.fetch(
         jobs=jobs, remote=dvc_remote, run_cache=run_cache, revs=revs, odb=odb
     )

dvc/repo/experiments/push.py

@@ -1,10 +1,10 @@
 import logging
+from typing import Iterable, Union
 
 from dvc.exceptions import DvcException
 from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
 from dvc.scm import TqdmGit
-from dvc.types import IterableStr, Union
 
 from .exceptions import UnresolvedExpNamesError
 from .utils import exp_commits, push_refspec, resolve_name
@@ -17,7 +17,7 @@ logger = logging.getLogger(__name__)
 def push(
     repo,
     git_remote: str,
-    exp_names: Union[IterableStr, str],
+    exp_names: Union[Iterable[str], str],
     *args,
     force: bool = False,
     push_cache: bool = False,

dvc/repo/experiments/utils.py

@@ -13,6 +13,7 @@ from typing import (
 from scmrepo.git import Git
 
 from dvc.exceptions import InvalidArgumentError
+from dvc.repo.experiments.exceptions import AmbiguousExpRefInfo
 
 from .base import (
     EXEC_BASELINE,
@@ -21,8 +22,6 @@ from .base import (
     EXPS_STASH,
     ExpRefInfo,
 )
-from dvc.repo.experiments.exceptions import AmbiguousExpRefInfo
-
 
 
 def exp_refs(scm: "Git") -> Generator["ExpRefInfo", None, None]:

dvc/types.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 if TYPE_CHECKING:
     from os import PathLike
@@ -10,4 +10,3 @@ OptStr = Optional[str]
 TargetType = Union[List[str], str]
 DictStrAny = Dict[str, Any]
 DictAny = Dict[Any, Any]
-IterableStr = Iterable[str]

tests/func/experiments/test_remote.py

@@ -151,19 +151,29 @@ def test_pull(tmp_dir, scm, dvc, git_downstream, exp_stage, use_url):
     remote = git_downstream.url if use_url else git_downstream.remote
     downstream_exp = git_downstream.tmp_dir.dvc.experiments
     with pytest.raises(InvalidArgumentError):
-        downstream_exp.pull(remote, "foo")
+        downstream_exp.pull(remote, ["foo"])
 
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"])
+    exp1 = first(results)
+    ref_info1 = first(exp_refs_by_rev(scm, exp1))
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"])
-    exp = first(results)
-    ref_info = first(exp_refs_by_rev(scm, exp))
+    exp2 = first(results)
+    ref_info2 = first(exp_refs_by_rev(scm, exp2))
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
+    exp3 = first(results)
+    ref_info3 = first(exp_refs_by_rev(scm, exp3))
 
-    downstream_exp.pull(remote, ref_info.name)
-    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info)) == exp
+    downstream_exp.pull(
+        git_downstream.remote, [ref_info1.name, ref_info2.name], force=True
+    )
+    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info1)) == exp1
+    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info2)) == exp2
+    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info3)) is None
 
-    git_downstream.tmp_dir.scm.remove_ref(str(ref_info))
+    git_downstream.tmp_dir.scm.remove_ref(str(ref_info1))
 
-    downstream_exp.pull(remote, str(ref_info))
-    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info)) == exp
+    downstream_exp.pull(remote, [str(ref_info1)])
+    assert git_downstream.tmp_dir.scm.get_ref(str(ref_info1)) == exp1
 
 
 def test_pull_diverged(tmp_dir, scm, dvc, git_downstream, exp_stage):
@@ -193,7 +203,7 @@ def test_pull_checkpoint(tmp_dir, scm, dvc, git_downstream, checkpoint_stage):
     ref_info_a = first(exp_refs_by_rev(scm, exp_a))
 
     downstream_exp = git_downstream.tmp_dir.dvc.experiments
-    downstream_exp.pull(git_downstream.remote, ref_info_a.name, force=True)
+    downstream_exp.pull(git_downstream.remote, [ref_info_a.name], force=True)
     assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_a)) == exp_a
 
     results = dvc.experiments.run(
@@ -202,7 +212,7 @@ def test_pull_checkpoint(tmp_dir, scm, dvc, git_downstream, checkpoint_stage):
     exp_b = first(results)
     ref_info_b = first(exp_refs_by_rev(scm, exp_b))
 
-    downstream_exp.pull(git_downstream.remote, ref_info_b.name, force=True)
+    downstream_exp.pull(git_downstream.remote, [ref_info_b.name], force=True)
     assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_b)) == exp_b
 
 
@@ -225,13 +235,13 @@ def test_pull_ambiguous_name(tmp_dir, scm, dvc, git_downstream, exp_stage):
     remote = git_downstream.remote
     downstream_exp = git_downstream.tmp_dir.dvc.experiments
     with pytest.raises(InvalidArgumentError):
-        downstream_exp.pull(remote, "foo")
+        downstream_exp.pull(remote, ["foo"])
 
-    downstream_exp.pull(remote, str(ref_info_b))
+    downstream_exp.pull(remote, [str(ref_info_b)])
     assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_b)) == exp_b
 
     with git_downstream.tmp_dir.scm.detach_head(ref_info_a.baseline_sha):
-        downstream_exp.pull(remote, "foo")
+        downstream_exp.pull(remote, ["foo"])
     assert git_downstream.tmp_dir.scm.get_ref(str(ref_info_a)) == exp_a
 
 
@@ -257,7 +267,7 @@ def test_push_pull_cache(
 
     remove(dvc.odb.local.cache_dir)
 
-    dvc.experiments.pull(remote, ref_info.name, pull_cache=True)
+    dvc.experiments.pull(remote, [ref_info.name], pull_cache=True)
     for x in range(2, checkpoint_stage.iterations + 1):
         hash_ = digest(str(x))
         path = os.path.join(dvc.odb.local.cache_dir, hash_[:2], hash_[2:])
@@ -282,7 +292,7 @@ def test_auth_error_pull(tmp_dir, scm, dvc, http_auth_patch):
         GitAuthError,
         match=f"HTTP Git authentication is not supported: '{http_auth_patch}'",
     ):
-        dvc.experiments.pull(http_auth_patch, "foo")
+        dvc.experiments.pull(http_auth_patch, ["foo"])
 
 
 def test_auth_error_push(tmp_dir, scm, dvc, exp_stage, http_auth_patch):

tests/unit/command/test_experiments.py

@@ -277,7 +277,7 @@ def test_experiments_pull(dvc, scm, mocker):
     m.assert_called_once_with(
         cmd.repo,
         "origin",
-        "experiment",
+        ["experiment"],
         force=True,
         pull_cache=False,
         dvc_remote="my-remote",