diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index c57a24221..2a5e18f42 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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[0-9a-f]+)$" + r"(?:On \(.*\): )" + r"dvc-exp:(?P[0-9a-f]+)(:(?P.+))?$" ) BRANCH_RE = re.compile( r"^(?P[a-f0-9]{7})-(?P[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 diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 21ce60142..725c01e5f 100644 --- a/dvc/repo/experiments/executor.py +++ b/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", {}) diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py index 6a2a6d50e..1175d4570 100644 --- a/dvc/repo/experiments/show.py +++ b/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 diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 7ebf71a3f..b11a277f5 100644 --- a/tests/func/experiments/test_experiments.py +++ b/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"