From f789cdff89d1991efc0c4683c363665dee2ce925 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Mon, 24 Aug 2020 21:47:25 +0900 Subject: [PATCH] experiments: fix bug where wrong git branch checked-out before pull (#4450) --- dvc/repo/experiments/__init__.py | 16 +++++++++++-- tests/func/experiments/test_experiments.py | 27 ++++++++++++++++++++++ 2 files changed, 41 insertions(+), 2 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 8200812ee..947ed09f5 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -162,13 +162,25 @@ class Experiments: def _scm_checkout(self, rev): self.scm.repo.git.reset(hard=True) if self.scm.repo.head.is_detached: - # switch back to default branch - self.scm.repo.heads[0].checkout() + self._checkout_default_branch() if not Git.is_sha(rev) or not self.scm.has_rev(rev): self.scm.pull() logger.debug("Checking out experiment commit '%s'", rev) self.scm.checkout(rev) + def _checkout_default_branch(self): + # switch to default branch + git_repo = self.scm.repo + origin_refs = git_repo.remotes["origin"].refs + ref = origin_refs["HEAD"].reference + branch_name = ref.name.split("/")[-1] + if branch_name in git_repo.heads: + branch = git_repo.heads[branch_name] + else: + branch = git_repo.create_head(branch_name, ref) + branch.set_tracking_branch(ref) + branch.checkout() + def _stash_exp(self, *args, params: Optional[dict] = None, **kwargs): """Stash changes from the current (parent) workspace as an experiment. diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 33d41c6d2..bc41e62c6 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -22,3 +22,30 @@ def test_new_simple(tmp_dir, scm, dvc, mocker): assert ( tmp_dir / ".dvc" / "experiments" / "metrics.yaml" ).read_text() == "foo: 2" + + +def test_update_with_pull(tmp_dir, scm, dvc, mocker): + 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_revs = [scm.get_rev()] + + tmp_dir.gen("params.yaml", "foo: 2") + dvc.reproduce(stage.addressing, experiment=True) + scm.add(["dvc.yaml", "dvc.lock", "params.yaml", "metrics.yaml"]) + scm.commit("promote experiment") + expected_revs.append(scm.get_rev()) + + tmp_dir.gen("params.yaml", "foo: 3") + dvc.reproduce(stage.addressing, experiment=True) + + exp_scm = dvc.experiments.scm + for rev in expected_revs: + assert exp_scm.has_rev(rev)