plots: add dvc plots -e/--experiment option (#4488)

* plots: add -e/--experiment option for plots diff

* experiments: add method to get baseline for a given experiment

* plots: default to diff against baseline instead of workspace/head for experiments

* use exact match for experiments RE since name-rev may return relative
commits

dvc/command/plots.py

@@ -86,7 +86,12 @@ class CmdPlotsShow(CmdPlots):
 
 class CmdPlotsDiff(CmdPlots):
     def _func(self, *args, **kwargs):
-        return self.repo.plots.diff(*args, revs=self.args.revisions, **kwargs)
+        return self.repo.plots.diff(
+            *args,
+            revs=self.args.revisions,
+            experiment=self.args.experiment,
+            **kwargs,
+        )
 
 
 class CmdPlotsModify(CmdPlots):
@@ -151,6 +156,13 @@ def add_parser(subparsers, parent_parser):
         help="Plots file to visualize. Shows all plots by default.",
         metavar="<path>",
     ).complete = completion.FILE
+    plots_diff_parser.add_argument(
+        "-e",
+        "--experiment",
+        action="store_true",
+        default=False,
+        help=argparse.SUPPRESS,
+    )
     plots_diff_parser.add_argument(
         "revisions", nargs="*", default=None, help="Git commits to plot from",
     )

dvc/repo/experiments/__init__.py

@@ -76,6 +76,9 @@ class Experiments:
     STASH_EXPERIMENT_RE = re.compile(
         r"(?:On \(.*\): )dvc-exp-(?P<baseline_rev>[0-9a-f]+)$"
     )
+    BRANCH_RE = re.compile(
+        r"^(?P<baseline_rev>[a-f0-9]{7})-(?P<exp_sha>[a-f0-9]+)$"
+    )
 
     def __init__(self, repo):
         from dvc.lock import make_lock
@@ -575,6 +578,26 @@ class Experiments:
         git_repo.reset("HEAD")
         git_repo.stash("drop", "stash@{0}")
 
+    @scm_locked
+    def get_baseline(self, rev):
+        """Return the baseline rev for an experiment rev."""
+        from git.exc import GitCommandError
+
+        rev = self.scm.resolve_rev(rev)
+        try:
+            name = self.scm.repo.git.name_rev(rev, name_only=True)
+        except GitCommandError:
+            return None
+        if not name:
+            return None
+        if name in ("undefined", "stash"):
+            _, baseline = self.stash_revs.get(rev, (None, None))
+            return baseline
+        m = self.BRANCH_RE.match(name)
+        if m:
+            return self.scm.resolve_rev(m.group("baseline_rev"))
+        return None
+
     def checkout(self, *args, **kwargs):
         from dvc.repo.experiments.checkout import checkout
 

dvc/repo/experiments/show.py

@@ -1,5 +1,4 @@
 import logging
-import re
 from collections import OrderedDict, defaultdict
 from datetime import datetime
 
@@ -10,9 +9,6 @@ from dvc.repo.params.show import _collect_configs, _read_params
 logger = logging.getLogger(__name__)
 
 
-EXP_RE = re.compile(r"(?P<rev_sha>[a-f0-9]{7})-(?P<exp_sha>[a-f0-9]+)")
-
-
 def _collect_experiment(repo, branch, stash=False):
     res = defaultdict(dict)
     for rev in repo.brancher(revs=[branch]):
@@ -60,9 +56,9 @@ def show(
 
     # collect reproduced experiments
     for exp_branch in repo.experiments.scm.list_branches():
-        m = re.match(EXP_RE, exp_branch)
+        m = repo.experiments.BRANCH_RE.match(exp_branch)
         if m:
-            rev = repo.scm.resolve_rev(m.group("rev_sha"))
+            rev = repo.scm.resolve_rev(m.group("baseline_rev"))
             if rev in revs:
                 exp_rev = repo.experiments.scm.resolve_rev(exp_branch)
                 with repo.experiments.chdir():

dvc/repo/plots/__init__.py

@@ -75,7 +75,7 @@ class Plots:
             for datafile, desc in plots.items()
         }
 
-    def show(self, targets=None, revs=None, props=None):
+    def show(self, targets=None, revs=None, props=None, templates=None):
         from .data import NoMetricInHistoryError
 
         data = self.collect(targets, revs)
@@ -90,7 +90,9 @@ class Plots:
         if not data:
             raise NoPlotsError()
 
-        return self.render(data, revs, props, self.repo.plot_templates)
+        if templates is None:
+            templates = self.repo.plot_templates
+        return self.render(data, revs, props, templates)
 
     def diff(self, *args, **kwargs):
         from .diff import diff

dvc/repo/plots/diff.py

@@ -1,13 +1,23 @@
-def _revisions(revs, is_dirty):
+def _revisions(repo, revs, experiment):
     revisions = revs or []
+    if experiment and len(revisions) == 1:
+        baseline = repo.experiments.get_baseline(revisions[0])
+        if baseline:
+            revisions.append(baseline[:7])
     if len(revisions) <= 1:
-        if len(revisions) == 0 and is_dirty:
+        if len(revisions) == 0 and repo.scm.is_dirty():
             revisions.append("HEAD")
         revisions.append("workspace")
     return revisions
 
 
-def diff(repo, *args, revs=None, **kwargs):
-    return repo.plots.show(
-        *args, revs=_revisions(revs, repo.scm.is_dirty()), **kwargs
+def diff(repo, *args, revs=None, experiment=False, **kwargs):
+    if experiment:
+        # use experiments repo brancher with templates from the main repo
+        kwargs["templates"] = repo.plot_templates
+        plots_repo = repo.experiments.exp_dvc
+    else:
+        plots_repo = repo
+    return plots_repo.plots.show(
+        *args, revs=_revisions(repo, revs, experiment), **kwargs
     )

tests/func/experiments/test_experiments.py

@@ -76,3 +76,27 @@ def test_checkout(tmp_dir, scm, dvc):
     dvc.experiments.checkout(exp_b)
     assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 3"
     assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 3"
+
+
+def test_get_baseline(tmp_dir, scm, dvc):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.run(
+        cmd="python copy.py params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        name="copy-file",
+    )
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("init")
+    expected = scm.get_rev()
+    assert dvc.experiments.get_baseline(expected) is None
+
+    dvc.reproduce(stage.addressing, experiment=True, params=["foo=2"])
+    rev = dvc.experiments.scm.get_rev()
+    assert dvc.experiments.get_baseline(rev) == expected
+
+    dvc.reproduce(
+        stage.addressing, experiment=True, params=["foo=3"], queue=True
+    )
+    assert dvc.experiments.get_baseline("stash@{0}") == expected

tests/unit/command/test_plots.py

@@ -2,7 +2,7 @@ from dvc.cli import parse_args
 from dvc.command.plots import CmdPlotsDiff, CmdPlotsShow
 
 
-def test_metrics_diff(dvc, mocker):
+def test_plots_diff(dvc, mocker):
     cli_args = parse_args(
         [
             "plots",
@@ -24,6 +24,7 @@ def test_metrics_diff(dvc, mocker):
             "x_title",
             "--y-label",
             "y_title",
+            "--experiment",
             "HEAD",
             "tag1",
             "tag2",
@@ -50,10 +51,11 @@ def test_metrics_diff(dvc, mocker):
             "x_label": "x_title",
             "y_label": "y_title",
         },
+        experiment=True,
     )
 
 
-def test_metrics_show(dvc, mocker):
+def test_plots_show(dvc, mocker):
     cli_args = parse_args(
         [
             "plots",

tests/unit/repo/plots/test_diff.py

@@ -13,4 +13,27 @@ from dvc.repo.plots.diff import _revisions
     ],
 )
 def test_revisions(mocker, arg_revisions, is_dirty, expected_revisions):
-    assert _revisions(arg_revisions, is_dirty) == expected_revisions
+    mock_scm = mocker.Mock()
+    mock_scm.configure_mock(**{"is_dirty.return_value": is_dirty})
+    mock_repo = mocker.Mock(scm=mock_scm)
+    assert _revisions(mock_repo, arg_revisions, False) == expected_revisions
+
+
+@pytest.mark.parametrize(
+    "arg_revisions,baseline,expected_revisions",
+    [
+        (["v1"], "v0", ["v1", "v0"]),
+        (["v1"], None, ["v1", "workspace"]),
+        (["v1", "v2"], "v0", ["v1", "v2"]),
+        (["v1", "v2"], None, ["v1", "v2"]),
+    ],
+)
+def test_revisions_experiment(
+    mocker, arg_revisions, baseline, expected_revisions
+):
+    mock_scm = mocker.Mock()
+    mock_scm.configure_mock(**{"is_dirty.return_value": False})
+    mock_experiments = mocker.Mock()
+    mock_experiments.configure_mock(**{"get_baseline.return_value": baseline})
+    mock_repo = mocker.Mock(scm=mock_scm, experiments=mock_experiments)
+    assert _revisions(mock_repo, arg_revisions, True) == expected_revisions