Repository: iterative/dvc

Issue:
experiments: provide access to experiment plots
While thinking about #4448 it occurred to me that there's currently no way to use `dvc plots` with experiment revisions. Implementation-wise, we should only need a `dvc exp plots ...` subcommand that runs regular plots in the experiments workspace, the same way that show/diff work for experiments.

One concern would be how to handle situations where someone wants to save an experiments plot - do we need functionality to git add & commit plots output into the internal experiment branches? Or can we just say for experiments, generated plots should always be considered temporary, and not be committed anywhere?

@dmpetrov @pared 

Discussion:

Patch:
diff --git a/dvc/command/plots.py b/dvc/command/plots.py
index 8aaf07d8f..2fa779669 100644
--- a/dvc/command/plots.py
+++ b/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",
     )
diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index 1d65a743a..ca5663f37 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/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
 
diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py
index 36f82d227..55594d951 100644
--- a/dvc/repo/experiments/show.py
+++ b/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():
diff --git a/dvc/repo/plots/__init__.py b/dvc/repo/plots/__init__.py
index a3222b395..3fac7fb1d 100644
--- a/dvc/repo/plots/__init__.py
+++ b/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
diff --git a/dvc/repo/plots/diff.py b/dvc/repo/plots/diff.py
index 4c1fa0ae5..a5daf535f 100644
--- a/dvc/repo/plots/diff.py
+++ b/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
     )


Test patch:
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index 52330df0a..7ebf71a3f 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/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
diff --git a/tests/unit/command/test_plots.py b/tests/unit/command/test_plots.py
index 1cc7543ee..219b9d543 100644
--- a/tests/unit/command/test_plots.py
+++ b/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",
diff --git a/tests/unit/repo/plots/test_diff.py b/tests/unit/repo/plots/test_diff.py
index f3525384d..9a0e80404 100644
--- a/tests/unit/repo/plots/test_diff.py
+++ b/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