exp: Fix --rev args for ls pull and remove.

Closes #9471

(cherry picked from commit 3ddd4b87fc)

dvc/repo/experiments/ls.py

@@ -1,6 +1,6 @@
 import logging
 from collections import defaultdict
-from typing import Optional
+from typing import List, Optional, Union
 
 from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
@@ -15,14 +15,17 @@ logger = logging.getLogger(__name__)
 @scm_context
 def ls(
     repo,
-    rev: Optional[str] = None,
+    rev: Optional[Union[List[str], str]] = None,
     all_commits: bool = False,
     num: int = 1,
     git_remote: Optional[str] = None,
 ):
     rev_set = None
     if not all_commits:
-        revs = iter_revs(repo.scm, [rev or "HEAD"], num)
+        rev = rev or "HEAD"
+        if isinstance(rev, str):
+            rev = [rev]
+        revs = iter_revs(repo.scm, rev, num)
         rev_set = set(revs.keys())
     ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
 

dvc/repo/experiments/pull.py

@@ -23,7 +23,7 @@ def pull(  # noqa: C901
     git_remote: str,
     exp_names: Union[Iterable[str], str],
     all_commits=False,
-    rev: Optional[str] = None,
+    rev: Optional[Union[List[str], str]] = None,
     num=1,
     force: bool = False,
     pull_cache: bool = False,
@@ -49,7 +49,9 @@ def pull(  # noqa: C901
                 raise UnresolvedExpNamesError(unresolved_exp_names)
 
         if rev:
-            rev_dict = iter_revs(repo.scm, [rev], num)
+            if isinstance(rev, str):
+                rev = [rev]
+            rev_dict = iter_revs(repo.scm, rev, num)
             rev_set = set(rev_dict.keys())
             ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
             for _, ref_info_list in ref_info_dict.items():

dvc/repo/experiments/push.py

@@ -82,9 +82,9 @@ def exp_refs_from_names(scm: "Git", exp_names: List[str]) -> Set["ExpRefInfo"]:
     return exp_ref_set
 
 
-def exp_refs_from_rev(scm: "Git", rev: str, num: int = 1) -> Set["ExpRefInfo"]:
+def exp_refs_from_rev(scm: "Git", rev: List[str], num: int = 1) -> Set["ExpRefInfo"]:
     exp_ref_set = set()
-    rev_dict = iter_revs(scm, [rev], num)
+    rev_dict = iter_revs(scm, rev, num)
     rev_set = set(rev_dict.keys())
     ref_info_dict = exp_refs_by_baseline(scm, rev_set)
     for _, ref_info_list in ref_info_dict.items():
@@ -99,7 +99,7 @@ def push(
     git_remote: str,
     exp_names: Union[List[str], str],
     all_commits: bool = False,
-    rev: Optional[str] = None,
+    rev: Optional[Union[List[str], str]] = None,
     num: int = 1,
     force: bool = False,
     push_cache: bool = False,
@@ -112,6 +112,8 @@ def push(
     if exp_names:
         exp_ref_set.update(exp_refs_from_names(repo.scm, ensure_list(exp_names)))
     if rev:
+        if isinstance(rev, str):
+            rev = [rev]
         exp_ref_set.update(exp_refs_from_rev(repo.scm, rev, num=num))
 
     push_result = _push(repo, git_remote, exp_ref_set, force)

dvc/repo/experiments/remove.py

@@ -24,7 +24,7 @@ logger = logging.getLogger(__name__)
 def remove(  # noqa: C901, PLR0912
     repo: "Repo",
     exp_names: Union[None, str, List[str]] = None,
-    rev: Optional[str] = None,
+    rev: Optional[Union[List[str], str]] = None,
     all_commits: bool = False,
     num: int = 1,
     queue: bool = False,
@@ -60,6 +60,8 @@ def remove(  # noqa: C901, PLR0912
         if remained:
             raise UnresolvedExpNamesError(remained, git_remote=git_remote)
     elif rev:
+        if isinstance(rev, str):
+            rev = [rev]
         exp_ref_dict = _resolve_exp_by_baseline(repo, rev, num, git_remote)
         removed.extend(exp_ref_dict.keys())
         exp_ref_list.extend(exp_ref_dict.values())
@@ -85,14 +87,14 @@ def remove(  # noqa: C901, PLR0912
 
 def _resolve_exp_by_baseline(
     repo: "Repo",
-    rev: str,
+    rev: List[str],
     num: int,
     git_remote: Optional[str] = None,
 ) -> Dict[str, "ExpRefInfo"]:
     assert isinstance(repo.scm, Git)
 
     commit_ref_dict: Dict[str, "ExpRefInfo"] = {}
-    rev_dict = iter_revs(repo.scm, [rev], num)
+    rev_dict = iter_revs(repo.scm, rev, num)
     rev_set = set(rev_dict.keys())
     ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
     for _, ref_info_list in ref_info_dict.items():

dvc/testing/benchmarks/fixtures.py

@@ -7,7 +7,6 @@ import pytest
 from dulwich.porcelain import clone
 from funcy import first
 from packaging import version
-from pytest_virtualenv import VirtualEnv
 
 # pylint: disable=redefined-outer-name,unused-argument
 
@@ -20,6 +19,8 @@ def bench_config(request):
 @pytest.fixture(scope="session")
 def make_dvc_venv(tmp_path_factory):
     def _make_dvc_venv(name):
+        from pytest_virtualenv import VirtualEnv
+
         name = _sanitize_venv_name(name)
         venv_dir = tmp_path_factory.mktemp(f"dvc-venv-{name}")
         return VirtualEnv(workspace=venv_dir)

dvc/utils/hydra.py

@@ -1,3 +1,4 @@
+import logging
 from pathlib import Path
 from typing import TYPE_CHECKING, List
 
@@ -9,6 +10,9 @@ if TYPE_CHECKING:
     from dvc.types import StrPath
 
 
+logger = logging.getLogger(__name__)
+
+
 def compose_and_dump(
     output_file: "StrPath",
     config_dir: str,
@@ -44,6 +48,9 @@ def compose_and_dump(
         dumper(output_file, OmegaConf.to_object(cfg))
     else:
         Path(output_file).write_text(OmegaConf.to_yaml(cfg), encoding="utf-8")
+    logger.trace(  # type: ignore[attr-defined]
+        "Hydra composition enabled. Contents dumped to %s:\n %s", output_file, cfg
+    )
 
 
 def apply_overrides(path: "StrPath", overrides: List[str]) -> None:

tests/func/experiments/test_experiments.py

@@ -369,6 +369,12 @@ def test_list(tmp_dir, scm, dvc, exp_stage):
         baseline_a[:7]: {ref_info_a.name, ref_info_b.name}
     }
 
+    exp_list = dvc.experiments.ls(rev=[baseline_a, scm.get_rev()])
+    assert {key: set(val) for key, val in exp_list.items()} == {
+        baseline_a[:7]: {ref_info_a.name, ref_info_b.name},
+        "master": {ref_info_c.name},
+    }
+
     exp_list = dvc.experiments.ls(all_commits=True)
     assert {key: set(val) for key, val in exp_list.items()} == {
         baseline_a[:7]: {ref_info_a.name, ref_info_b.name},

tests/func/experiments/test_remote.py

@@ -66,6 +66,29 @@ def test_push_args(tmp_dir, scm, dvc, git_upstream, exp_stage, all_, rev, result
     assert git_upstream.tmp_dir.scm.get_ref(str(ref_info3)) == result3
 
 
+def test_push_multi_rev(tmp_dir, scm, dvc, git_upstream, exp_stage):
+    remote = git_upstream.url
+    baseline = scm.get_rev()
+
+    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"])
+    exp2 = first(results)
+    ref_info2 = first(exp_refs_by_rev(scm, exp2))
+
+    scm.commit("new_baseline")
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
+    exp3 = first(results)
+    ref_info3 = first(exp_refs_by_rev(scm, exp3))
+
+    dvc.experiments.push(remote, [], rev=[baseline, scm.get_rev()])
+    assert git_upstream.tmp_dir.scm.get_ref(str(ref_info1)) == exp1
+    assert git_upstream.tmp_dir.scm.get_ref(str(ref_info2)) == exp2
+    assert git_upstream.tmp_dir.scm.get_ref(str(ref_info3)) == exp3
+
+
 def test_push_diverged(tmp_dir, scm, dvc, git_upstream, exp_stage):
     git_upstream.tmp_dir.scm_gen("foo", "foo", commit="init")
     remote_rev = git_upstream.tmp_dir.scm.get_rev()
@@ -228,6 +251,30 @@ def test_pull_args(tmp_dir, scm, dvc, git_downstream, exp_stage, all_, rev, resu
     assert git_downstream.tmp_dir.scm.get_ref(str(ref_info3)) == result3
 
 
+def test_pull_multi_rev(tmp_dir, scm, dvc, git_downstream, exp_stage):
+    baseline = scm.get_rev()
+
+    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"])
+    exp2 = first(results)
+    ref_info2 = first(exp_refs_by_rev(scm, exp2))
+
+    scm.commit("new_baseline")
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
+    exp3 = first(results)
+    ref_info3 = first(exp_refs_by_rev(scm, exp3))
+
+    downstream_exp = git_downstream.tmp_dir.dvc.experiments
+    git_downstream.tmp_dir.scm.fetch_refspecs(str(tmp_dir), ["master:master"])
+    downstream_exp.pull(git_downstream.remote, [], rev=[baseline, scm.get_rev()])
+    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)) == exp3
+
+
 def test_pull_diverged(tmp_dir, scm, dvc, git_downstream, exp_stage):
     git_downstream.tmp_dir.scm_gen("foo", "foo", commit="init")
     remote_rev = git_downstream.tmp_dir.scm.get_rev()

tests/func/experiments/test_remove.py

@@ -156,3 +156,26 @@ def test_remove_experiments_by_rev(tmp_dir, scm, dvc, exp_stage):
     assert "queue2" in queue_revs
     assert scm.get_ref(new_exp_ref) is not None
     assert "queue4" in queue_revs
+
+
+def test_remove_multi_rev(tmp_dir, scm, dvc, exp_stage):
+    baseline = scm.get_rev()
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"])
+    baseline_exp_ref = first(exp_refs_by_rev(scm, first(results)))
+
+    dvc.experiments.run(
+        exp_stage.addressing, params=["foo=2"], queue=True, name="queue2"
+    )
+    scm.commit("new_baseline")
+
+    results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
+    new_exp_ref = first(exp_refs_by_rev(scm, first(results)))
+
+    assert set(dvc.experiments.remove(rev=[baseline, scm.get_rev()])) == {
+        baseline_exp_ref.name,
+        new_exp_ref.name,
+    }
+
+    assert scm.get_ref(str(baseline_exp_ref)) is None
+    assert scm.get_ref(str(new_exp_ref)) is None