exp: use exp baseline for DVC commands that compare to HEAD during exp run (#5818)

* exp: add util method for getting baseline HEAD during run

* exp: use baseline for dvc commands which compare to HEAD during exp run

* add test for new util method

* update plots diff test mock

dvc/repo/diff.py

@@ -7,6 +7,7 @@ from dvc.exceptions import PathMissingError
 from dvc.objects.stage import get_file_hash
 from dvc.objects.stage import stage as ostage
 from dvc.repo import locked
+from dvc.repo.experiments.utils import fix_exp_head
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +29,8 @@ def diff(self, a_rev="HEAD", b_rev=None, targets=None):
 
     repo_fs = RepoFileSystem(self)
 
-    b_rev = b_rev if b_rev else "workspace"
+    a_rev = fix_exp_head(self.scm, a_rev)
+    b_rev = fix_exp_head(self.scm, b_rev) if b_rev else "workspace"
     results = {}
     missing_targets = {}
     for rev in self.brancher(revs=[a_rev, b_rev]):

dvc/repo/experiments/diff.py

@@ -1,5 +1,6 @@
 import logging
 
+from dvc.repo.experiments.utils import fix_exp_head
 from dvc.utils.diff import diff as _diff
 from dvc.utils.diff import format_dict
 
@@ -13,12 +14,16 @@ def diff(repo, *args, a_rev=None, b_rev=None, param_deps=False, **kwargs):
         return {}
 
     if a_rev:
+        a_rev = fix_exp_head(repo.scm, a_rev)
         rev = repo.scm.resolve_rev(a_rev)
         old = _collect_experiment_commit(repo, rev, param_deps=param_deps)
     else:
-        old = _collect_experiment_commit(repo, "HEAD", param_deps=param_deps)
+        old = _collect_experiment_commit(
+            repo, fix_exp_head(repo.scm, "HEAD"), param_deps=param_deps
+        )
 
     if b_rev:
+        b_rev = fix_exp_head(repo.scm, b_rev)
         rev = repo.scm.resolve_rev(b_rev)
         new = _collect_experiment_commit(repo, rev, param_deps=param_deps)
     else:

dvc/repo/experiments/show.py

@@ -5,6 +5,7 @@ from datetime import datetime
 from dvc.exceptions import InvalidArgumentError
 from dvc.repo import locked
 from dvc.repo.experiments.base import ExpRefInfo
+from dvc.repo.experiments.utils import fix_exp_head
 from dvc.repo.metrics.show import _collect_metrics, _read_metrics
 from dvc.repo.params.show import _collect_configs, _read_params
 from dvc.scm.base import SCMError
@@ -100,7 +101,8 @@ def show(
         revs = []
         for n in range(num):
             try:
-                revs.append(repo.scm.resolve_rev(f"HEAD~{n}"))
+                head = fix_exp_head(repo.scm, f"HEAD~{n}")
+                revs.append(repo.scm.resolve_rev(head))
             except SCMError:
                 break
 

dvc/repo/experiments/utils.py

@@ -1,9 +1,14 @@
-from typing import TYPE_CHECKING, Generator, Iterable, Set
+from typing import Generator, Iterable, Optional, Set
 
-from .base import EXEC_NAMESPACE, EXPS_NAMESPACE, EXPS_STASH, ExpRefInfo
+from dvc.scm.git import Git
 
-if TYPE_CHECKING:
-    from dvc.scm.git import Git
+from .base import (
+    EXEC_BASELINE,
+    EXEC_NAMESPACE,
+    EXPS_NAMESPACE,
+    EXPS_STASH,
+    ExpRefInfo,
+)
 
 
 def exp_refs(scm: "Git") -> Generator["ExpRefInfo", None, None]:
@@ -102,3 +107,11 @@ def remove_exp_refs(scm: "Git", ref_infos: Iterable["ExpRefInfo"]):
         if exec_checkpoint and exec_checkpoint == ref:
             scm.remove_ref(EXEC_CHECKPOINT)
         scm.remove_ref(str(ref_info))
+
+
+def fix_exp_head(scm: "Git", ref: Optional[str]) -> Optional[str]:
+    if ref:
+        name, tail = Git.split_ref_pattern(ref)
+        if name == "HEAD" and scm.get_ref(EXEC_BASELINE):
+            return "".join((EXEC_BASELINE, tail))
+    return ref

dvc/repo/metrics/diff.py

@@ -1,4 +1,5 @@
 from dvc.exceptions import MetricsError
+from dvc.repo.experiments.utils import fix_exp_head
 from dvc.utils.diff import diff as _diff
 from dvc.utils.diff import format_dict
 
@@ -18,7 +19,8 @@ def diff(repo, *args, a_rev=None, b_rev=None, **kwargs):
     with_unchanged = kwargs.pop("all", False)
 
     a_rev = a_rev or "HEAD"
-    b_rev = b_rev or "workspace"
+    a_rev = fix_exp_head(repo.scm, a_rev)
+    b_rev = fix_exp_head(repo.scm, b_rev) or "workspace"
 
     metrics = _get_metrics(repo, *args, **kwargs, revs=[a_rev, b_rev])
     old = metrics.get(a_rev, {})

dvc/repo/params/diff.py

@@ -1,3 +1,4 @@
+from dvc.repo.experiments.utils import fix_exp_head
 from dvc.utils.diff import diff as _diff
 from dvc.utils.diff import format_dict
 
@@ -19,7 +20,8 @@ def diff(repo, *args, a_rev=None, b_rev=None, **kwargs):
     with_unchanged = kwargs.pop("all", False)
 
     a_rev = a_rev or "HEAD"
-    b_rev = b_rev or "workspace"
+    a_rev = fix_exp_head(repo.scm, a_rev)
+    b_rev = fix_exp_head(repo.scm, b_rev) or "workspace"
 
     params = _get_params(repo, *args, **kwargs, revs=[a_rev, b_rev])
     old = params.get(a_rev, {})

dvc/repo/plots/diff.py

@@ -1,3 +1,6 @@
+from dvc.repo.experiments.utils import fix_exp_head
+
+
 def _revisions(repo, revs, experiment):
     revisions = revs or []
     if experiment and len(revisions) == 1:
@@ -6,7 +9,7 @@ def _revisions(repo, revs, experiment):
             revisions.append(baseline[:7])
     if len(revisions) <= 1:
         if len(revisions) == 0 and repo.scm.is_dirty():
-            revisions.append("HEAD")
+            revisions.append(fix_exp_head(repo.scm, "HEAD"))
         revisions.append("workspace")
     return revisions
 

dvc/scm/git/__init__.py

@@ -7,7 +7,7 @@ import shlex
 from collections.abc import Mapping
 from contextlib import contextmanager
 from functools import partialmethod
-from typing import Dict, Iterable, List, Optional, Set, Type
+from typing import Dict, Iterable, List, Optional, Set, Tuple, Type
 
 from funcy import cached_property, first
 from pathspec.patterns import GitWildMatchPattern
@@ -77,6 +77,7 @@ class Git(Base):
     GIT_DIR = ".git"
     LOCAL_BRANCH_PREFIX = "refs/heads/"
     RE_HEXSHA = re.compile(r"^[0-9A-Fa-f]{4,40}$")
+    BAD_REF_CHARS_RE = re.compile("[\177\\s~^:?*\\[]")
 
     def __init__(
         self, *args, backends: Optional[Iterable[str]] = None, **kwargs
@@ -123,6 +124,11 @@ class Git(Base):
     def is_sha(cls, rev):
         return rev and cls.RE_HEXSHA.search(rev)
 
+    @classmethod
+    def split_ref_pattern(cls, ref: str) -> Tuple[str, str]:
+        name = cls.BAD_REF_CHARS_RE.split(ref, maxsplit=1)[0]
+        return name, ref[len(name) :]
+
     @staticmethod
     def _get_git_dir(root_dir):
         return os.path.join(root_dir, Git.GIT_DIR)

tests/func/experiments/test_experiments.py

@@ -611,3 +611,16 @@ def test_checkout_targets_deps(tmp_dir, scm, dvc, exp_stage):
     assert (tmp_dir / "foo").exists()
     assert (tmp_dir / "foo").read_text() == "foo"
     assert not (tmp_dir / "bar").exists()
+
+
+@pytest.mark.parametrize("tail", ["", "~1", "^"])
+def test_fix_exp_head(tmp_dir, scm, tail):
+    from dvc.repo.experiments.base import EXEC_BASELINE
+    from dvc.repo.experiments.utils import fix_exp_head
+
+    head = "HEAD" + tail
+    assert head == fix_exp_head(scm, head)
+
+    scm.set_ref(EXEC_BASELINE, "refs/heads/master")
+    assert EXEC_BASELINE + tail == fix_exp_head(scm, head)
+    assert "foo" + tail == fix_exp_head(scm, "foo" + tail)

tests/unit/repo/plots/test_diff.py

@@ -14,7 +14,9 @@ from dvc.repo.plots.diff import _revisions
 )
 def test_revisions(mocker, arg_revisions, is_dirty, expected_revisions):
     mock_scm = mocker.Mock()
-    mock_scm.configure_mock(**{"is_dirty.return_value": is_dirty})
+    mock_scm.configure_mock(
+        **{"is_dirty.return_value": is_dirty, "get_ref.return_value": None}
+    )
     mock_repo = mocker.Mock(scm=mock_scm)
     assert _revisions(mock_repo, arg_revisions, False) == expected_revisions
 
@@ -32,7 +34,9 @@ def test_revisions_experiment(
     mocker, arg_revisions, baseline, expected_revisions
 ):
     mock_scm = mocker.Mock()
-    mock_scm.configure_mock(**{"is_dirty.return_value": False})
+    mock_scm.configure_mock(
+        **{"is_dirty.return_value": False, "get_ref.return_value": None}
+    )
     mock_experiments = mocker.Mock()
     mock_experiments.configure_mock(**{"get_baseline.return_value": baseline})
     mock_repo = mocker.Mock(scm=mock_scm, experiments=mock_experiments)