exp pull/push: add multi exp name arg support

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/command/experiments/push.py

@@ -88,6 +88,9 @@ def add_parser(experiments_subparsers, parent_parser):
         metavar="<git_remote>",
     )
     experiments_push_parser.add_argument(
-        "experiment", help="Experiment to push.", metavar="<experiment>"
+        "experiment",
+        nargs="+",
+        help="Experiments to push.",
+        metavar="<experiment>",
     ).complete = completion.EXPERIMENT
     experiments_push_parser.set_defaults(func=CmdExperimentsPush)

dvc/repo/experiments/exceptions.py

@@ -0,0 +1,43 @@
+from typing import Iterable, List
+
+from dvc.exceptions import InvalidArgumentError
+
+from .base import ExpRefInfo
+
+
+class AmbiguousExpRefInfo(InvalidArgumentError):
+    def __init__(
+        self,
+        exp_name: str,
+        exp_ref_list: Iterable[ExpRefInfo],
+    ):
+        msg = [
+            (
+                f"Ambiguous name '{exp_name}' refers to multiple experiments."
+                " Use one of the following full refnames instead:"
+            ),
+            "",
+        ]
+        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
+    ):
+        unresolved_names = ";".join(unresolved_list)
+        if not git_remote:
+            if len(unresolved_names) > 1:
+                super().__init__(
+                    f"'{unresolved_names}' are not valid experiment names"
+                )
+            else:
+                super().__init__(
+                    f"'{unresolved_names}' is not a valid experiment name"
+                )
+        else:
+            super().__init__(
+                f"Experiment '{unresolved_names}' does not exist "
+                f"in '{git_remote}'"
+            )

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 @@
 @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,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, push_refspec, resolve_name
 
 logger = logging.getLogger(__name__)
@@ -13,50 +16,66 @@
 @scm_context
 def push(
     repo,
-    git_remote,
-    exp_name: str,
+    git_remote: str,
+    exp_names: Union[Iterable[str], str],
     *args,
-    force=False,
-    push_cache=False,
+    force: bool = False,
+    push_cache: bool = False,
     **kwargs,
 ):
-    exp_ref_dict = resolve_name(repo.scm, exp_name)
-    exp_ref = exp_ref_dict[exp_name]
-    if not exp_ref:
-        raise InvalidArgumentError(
-            f"'{exp_name}' is not a valid experiment name"
-        )
+    if isinstance(exp_names, str):
+        exp_names = [exp_names]
+
+    exp_ref_dict = resolve_name(repo.scm, exp_names)
+    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()
+    _push(repo, git_remote, exp_ref_set, force, push_cache, **kwargs)
 
+
+def _push(
+    repo,
+    git_remote: str,
+    exp_refs,
+    force: bool,
+    push_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 remote experiment "
             "re-run with '--force'."
         )
 
-    refname = str(exp_ref)
-    logger.debug("git push experiment '%s' -> '%s'", exp_ref, git_remote)
+    logger.debug(f"git push experiment '{exp_refs}' -> '{git_remote}'")
 
-    from dvc.scm import TqdmGit
-
-    with TqdmGit(desc="Pushing git refs") as pbar:
-        push_refspec(
-            repo.scm,
-            git_remote,
-            refname,
-            refname,
-            force=force,
-            on_diverged=on_diverged,
-            progress=pbar.update_git,
-        )
+    for exp_ref in exp_refs:
+        with TqdmGit(desc="Pushing git refs") as pbar:
+            push_refspec(
+                repo.scm,
+                git_remote,
+                str(exp_ref),
+                str(exp_ref),
+                force=force,
+                on_diverged=on_diverged,
+                progress=pbar.update_git,
+            )
 
     if push_cache:
-        _push_cache(repo, exp_ref, **kwargs)
+        _push_cache(repo, exp_refs, **kwargs)
 
 
-def _push_cache(repo, exp_ref, dvc_remote=None, jobs=None, run_cache=False):
-    revs = list(exp_commits(repo.scm, [exp_ref]))
-    logger.debug("dvc push experiment '%s'", exp_ref)
+def _push_cache(repo, exp_refs, dvc_remote=None, jobs=None, run_cache=False):
+    revs = list(exp_commits(repo.scm, exp_refs))
+    logger.debug(f"dvc push experiment '{exp_refs}'")
     repo.push(jobs=jobs, remote=dvc_remote, run_cache=run_cache, revs=revs)

dvc/repo/experiments/utils.py

@@ -13,6 +13,7 @@
 from scmrepo.git import Git
 
 from dvc.exceptions import InvalidArgumentError
+from dvc.repo.experiments.exceptions import AmbiguousExpRefInfo
 
 from .base import (
     EXEC_BASELINE,
@@ -23,23 +24,6 @@
 )
 
 
-class AmbiguousExpRefInfo(InvalidArgumentError):
-    def __init__(
-        self,
-        exp_name: str,
-        exp_ref_list: Iterable[ExpRefInfo],
-    ):
-        msg = [
-            (
-                f"Ambiguous name '{exp_name}' refers to multiple experiments."
-                " Use one of the following full refnames instead:"
-            ),
-            "",
-        ]
-        msg.extend([f"\t{info}" for info in exp_ref_list])
-        super().__init__("\n".join(msg))
-
-
 def exp_refs(scm: "Git") -> Generator["ExpRefInfo", None, None]:
     """Iterate over all experiment refs."""
     for ref in scm.iter_refs(base=EXPS_NAMESPACE):