experiments: support committing to existing experiment branches (#4573)

* experiments: stash optional branch to extend

* experiments: support commit to existing experiment branch

dvc/repo/experiments/__init__.py

@@ -2,6 +2,7 @@ import logging
 import os
 import re
 import tempfile
+from collections import namedtuple
 from collections.abc import Mapping
 from concurrent.futures import (
     ProcessPoolExecutor,
@@ -72,14 +73,17 @@ class Experiments:
 
     EXPERIMENTS_DIR = "experiments"
     PACKED_ARGS_FILE = "repro.dat"
-    STASH_MSG_PREFIX = "dvc-exp-"
+    STASH_MSG_PREFIX = "dvc-exp:"
     STASH_EXPERIMENT_RE = re.compile(
-        r"(?:On \(.*\): )dvc-exp-(?P<baseline_rev>[0-9a-f]+)$"
+        r"(?:On \(.*\): )"
+        r"dvc-exp:(?P<baseline_rev>[0-9a-f]+)(:(?P<branch>.+))?$"
     )
     BRANCH_RE = re.compile(
         r"^(?P<baseline_rev>[a-f0-9]{7})-(?P<exp_sha>[a-f0-9]+)$"
     )
 
+    StashEntry = namedtuple("StashEntry", ["index", "baseline_rev", "branch"])
+
     def __init__(self, repo):
         from dvc.lock import make_lock
 
@@ -144,7 +148,9 @@ class Experiments:
         for i, entry in enumerate(self.stash_reflog):
             m = self.STASH_EXPERIMENT_RE.match(entry.message)
             if m:
-                revs[entry.newhexsha] = (i, m.group("baseline_rev"))
+                revs[entry.newhexsha] = self.StashEntry(
+                    i, m.group("baseline_rev"), m.group("branch")
+                )
         return revs
 
     def _init_clone(self):
@@ -184,30 +190,54 @@ class Experiments:
             branch.set_tracking_branch(ref)
         branch.checkout()
 
-    def _stash_exp(self, *args, params: Optional[dict] = None, **kwargs):
+    def _stash_exp(
+        self,
+        *args,
+        params: Optional[dict] = None,
+        branch: Optional[str] = None,
+        allow_unchanged: Optional[bool] = False,
+        apply_workspace: Optional[bool] = True,
+        **kwargs,
+    ):
         """Stash changes from the current (parent) workspace as an experiment.
 
         Args:
             params: Optional dictionary of parameter values to be used.
                 Values take priority over any parameters specified in the
                 user's workspace.
+            branch: Optional experiment branch name. If specified, the
+                experiment will be added to `branch` instead of creating
+                a new branch.
+            allow_unchanged: Force experiment reproduction even if params are
+                unchanged from the baseline.
+            apply_workspace: Apply changes from the user workspace to the
+                experiment workspace.
         """
-        rev = self.scm.get_rev()
+        if branch:
+            rev = self.scm.resolve_rev(branch)
+        else:
+            rev = self.scm.get_rev()
 
-        # patch user's workspace into experiments clone
-        tmp = tempfile.NamedTemporaryFile(delete=False).name
-        try:
-            self.repo.scm.repo.git.diff(
-                patch=True, full_index=True, binary=True, output=tmp
-            )
-            if os.path.getsize(tmp):
-                logger.debug("Patching experiment workspace")
-                self.scm.repo.git.apply(tmp)
-            elif not params:
-                # experiment matches original baseline
-                raise UnchangedExperimentError(rev)
-        finally:
-            remove(tmp)
+        if apply_workspace:
+            # patch user's workspace into experiments clone
+            #
+            # TODO: patching changes into an extended branch will require
+            # propagating changes down the full branch, since the workspace
+            # patch will be based on the original branch point, not the tip
+            # of the experiment branch
+            tmp = tempfile.NamedTemporaryFile(delete=False).name
+            try:
+                self.repo.scm.repo.git.diff(
+                    patch=True, full_index=True, binary=True, output=tmp
+                )
+                if os.path.getsize(tmp):
+                    logger.debug("Patching experiment workspace")
+                    self.scm.repo.git.apply(tmp)
+                elif not params and not allow_unchanged:
+                    # experiment matches original baseline
+                    raise UnchangedExperimentError(rev)
+            finally:
+                remove(tmp)
 
         # update experiment params from command line
         if params:
@@ -219,10 +249,15 @@ class Experiments:
         # save experiment as a stash commit w/message containing baseline rev
         # (stash commits are merge commits and do not contain a parent commit
         # SHA)
-        msg = f"{self.STASH_MSG_PREFIX}{rev}"
+        msg = self._stash_msg(rev, branch)
         self.scm.repo.git.stash("push", "-m", msg)
         return self.scm.resolve_rev("stash@{0}")
 
+    def _stash_msg(self, rev, branch=None):
+        if branch:
+            return f"{self.STASH_MSG_PREFIX}{rev}:{branch}"
+        return f"{self.STASH_MSG_PREFIX}{rev}"
+
     def _pack_args(self, *args, **kwargs):
         ExperimentExecutor.pack_repro_args(self.args_file, *args, **kwargs)
         self.scm.add(self.args_file)
@@ -252,19 +287,21 @@ class Experiments:
             with modify_data(path, tree=self.exp_dvc.tree) as data:
                 _update(data, params[params_fname])
 
-    def _commit(self, exp_hash, check_exists=True, branch=True):
+    def _commit(self, exp_hash, check_exists=True, create_branch=True):
         """Commit stages as an experiment and return the commit SHA."""
         if not self.scm.is_dirty():
             raise UnchangedExperimentError(self.scm.get_rev())
 
         rev = self.scm.get_rev()
         exp_name = f"{rev[:7]}-{exp_hash}"
-        if branch:
+        if create_branch:
             if check_exists and exp_name in self.scm.list_branches():
                 logger.debug("Using existing experiment branch '%s'", exp_name)
                 return self.scm.resolve_rev(exp_name)
             self.scm.checkout(exp_name, create_new=True)
-        logger.debug("Commit new experiment branch '%s'", exp_name)
+            logger.debug("Commit new experiment branch '%s'", exp_name)
+        else:
+            logger.debug("Commit to current experiment branch")
         self.scm.repo.git.add(A=True)
         self.scm.commit(f"Add experiment {exp_name}")
         return self.scm.get_rev()
@@ -302,7 +339,14 @@ class Experiments:
         Experiment will be reproduced and checked out into the user's
         workspace.
         """
-        rev = self.repo.scm.get_rev()
+        branch = kwargs.get("branch")
+        if branch:
+            rev = self.scm.resolve_rev(branch)
+            logger.debug(
+                "Using '%s' (tip of branch '%s') as baseline", rev, branch
+            )
+        else:
+            rev = self.repo.scm.get_rev()
         self._scm_checkout(rev)
         try:
             stash_rev = self._stash_exp(*args, **kwargs)
@@ -337,13 +381,12 @@ class Experiments:
         # commit) and baseline_rev is the baseline to compare output against.
         # The final experiment commit will be branched from baseline_rev.
         if revs is None:
-            to_run = {
-                rev: baseline_rev
-                for rev, (_, baseline_rev) in stash_revs.items()
-            }
+            to_run = dict(stash_revs)
         else:
             to_run = {
-                rev: stash_revs[rev][1] if rev in stash_revs else rev
+                rev: stash_revs[rev]
+                if rev in stash_revs
+                else self.StashEntry(None, rev, None)
                 for rev in revs
             }
 
@@ -355,12 +398,13 @@ class Experiments:
         # setup executors - unstash experiment, generate executor, upload
         # contents of (unstashed) exp workspace to the executor tree
         executors = {}
-        for rev, baseline_rev in to_run.items():
-            self._scm_checkout(baseline_rev)
+        for rev, item in to_run.items():
+            self._scm_checkout(item.baseline_rev)
             self.scm.repo.git.stash("apply", rev)
             repro_args, repro_kwargs = self._unpack_args()
             executor = LocalExecutor(
-                baseline_rev,
+                item.baseline_rev,
+                branch=item.branch,
                 repro_args=repro_args,
                 repro_kwargs=repro_kwargs,
                 dvc_dir=self.dvc_dir,
@@ -417,7 +461,10 @@ class Experiments:
                 exc = future.exception()
                 if exc is None:
                     exp_hash = future.result()
-                    self._scm_checkout(executor.baseline_rev)
+                    if executor.branch:
+                        self._scm_checkout(executor.branch)
+                    else:
+                        self._scm_checkout(executor.baseline_rev)
                     try:
                         self._collect_output(executor)
                     except DownloadError:
@@ -431,7 +478,8 @@ class Experiments:
                             remove(self.args_file)
 
                     try:
-                        exp_rev = self._commit(exp_hash)
+                        branch = not executor.branch
+                        exp_rev = self._commit(exp_hash, create_branch=branch)
                     except UnchangedExperimentError:
                         logger.debug(
                             "Experiment '%s' identical to baseline '%s'",
@@ -517,12 +565,12 @@ class Experiments:
 
         from dvc.repo.checkout import checkout as dvc_checkout
 
-        self._check_baseline(rev)
+        baseline_rev = self._check_baseline(rev)
         self._scm_checkout(rev)
 
         tmp = tempfile.NamedTemporaryFile(delete=False).name
         self.scm.repo.head.commit.diff(
-            "HEAD~1", patch=True, full_index=True, binary=True, output=tmp
+            baseline_rev, patch=True, full_index=True, binary=True, output=tmp
         )
 
         dirty = self.repo.scm.is_dirty()
@@ -549,11 +597,14 @@ class Experiments:
 
     def _check_baseline(self, exp_rev):
         baseline_sha = self.repo.scm.get_rev()
-        exp_commit = self.scm.repo.rev_parse(exp_rev)
-        parent = first(exp_commit.parents)
-        if parent is not None and parent.hexsha == baseline_sha:
-            return
-        raise BaselineMismatchError(parent, baseline_sha)
+        exp_baseline = self._get_baseline(exp_rev)
+        if exp_baseline is None:
+            # if we can't tell from branch name, fall back to parent commit
+            exp_commit = self.scm.repo.rev_parse(exp_rev)
+            exp_baseline = first(exp_commit.parents).hexsha
+        if exp_baseline == baseline_sha:
+            return exp_baseline
+        raise BaselineMismatchError(exp_baseline, baseline_sha)
 
     def _unstash_workspace(self):
         # Essentially we want `git stash pop` with `-X ours` merge strategy
@@ -581,6 +632,9 @@ class Experiments:
     @scm_locked
     def get_baseline(self, rev):
         """Return the baseline rev for an experiment rev."""
+        return self._get_baseline(rev)
+
+    def _get_baseline(self, rev):
         from git.exc import GitCommandError
 
         rev = self.scm.resolve_rev(rev)
@@ -588,16 +642,28 @@ class Experiments:
             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
+            entry = self.stash_revs.get(rev)
+            if entry:
+                return entry.baseline_rev
+            return None
         m = self.BRANCH_RE.match(name)
         if m:
             return self.scm.resolve_rev(m.group("baseline_rev"))
         return None
 
+    def _get_branch_containing(self, rev):
+        from git.exc import GitCommandError
+
+        if self.scm.repo.head.is_detached:
+            self._checkout_default_branch()
+        try:
+            name = self.scm.repo.git.branch(contains=rev)
+            return name.rsplit("/")[-1].strip()
+        except GitCommandError:
+            pass
+        return None
+
     def checkout(self, *args, **kwargs):
         from dvc.repo.experiments.checkout import checkout
 

dvc/repo/experiments/executor.py

@@ -2,7 +2,7 @@ import logging
 import os
 import pickle
 from tempfile import TemporaryDirectory
-from typing import Iterable
+from typing import Iterable, Optional
 
 from funcy import cached_property
 
@@ -22,12 +22,16 @@ class ExperimentExecutor:
         baseline_rev: baseline revision that this experiment is derived from.
 
     Optional keyword args:
+        branch: Existing git branch for this experiment.
         repro_args: Args to be passed into reproduce.
         repro_kwargs: Keyword args to be passed into reproduce.
     """
 
-    def __init__(self, baseline_rev: str, **kwargs):
+    def __init__(
+        self, baseline_rev: str, branch: Optional[str] = None, **kwargs,
+    ):
         self.baseline_rev = baseline_rev
+        self.branch = branch
         self.repro_args = kwargs.pop("repro_args", [])
         self.repro_kwargs = kwargs.pop("repro_kwargs", {})
 

dvc/repo/experiments/show.py

@@ -87,12 +87,12 @@ def show(
                 res[rev][exp_rev] = experiment
 
     # collect queued (not yet reproduced) experiments
-    for stash_rev, (_, baseline_rev) in repo.experiments.stash_revs.items():
-        if baseline_rev in revs:
+    for stash_rev, entry in repo.experiments.stash_revs.items():
+        if entry.baseline_rev in revs:
             with repo.experiments.chdir():
                 experiment = _collect_experiment(
                     repo.experiments.exp_dvc, stash_rev, stash=True
                 )
-            res[baseline_rev][stash_rev] = experiment
+            res[entry.baseline_rev][stash_rev] = experiment
 
     return res

tests/func/experiments/test_experiments.py

@@ -100,3 +100,39 @@ def test_get_baseline(tmp_dir, scm, dvc):
         stage.addressing, experiment=True, params=["foo=3"], queue=True
     )
     assert dvc.experiments.get_baseline("stash@{0}") == expected
+
+
+def test_extend_branch(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")
+
+    dvc.reproduce(stage.addressing, experiment=True, params=["foo=2"])
+    exp_a = dvc.experiments.scm.get_rev()
+    exp_branch = dvc.experiments._get_branch_containing(exp_a)
+
+    dvc.reproduce(
+        stage.addressing,
+        experiment=True,
+        params=["foo=3"],
+        branch=exp_branch,
+        apply_workspace=False,
+    )
+    exp_b = dvc.experiments.scm.get_rev()
+
+    assert dvc.experiments._get_branch_containing(exp_b) == exp_branch
+
+    dvc.experiments.checkout(exp_a)
+    assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 2"
+    assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 2"
+
+    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"