queue: clean local EXEC_HEAD & EXEC_BASELINE after we run exp with queue.

fix: #8031

Currently when we use tmp directory to run a experiment, we set
`EXEC_HEAD` , `EXEC_BASELINE` and `EXEC_MERGE` heads locally but forget
to remove it. This cause some problem when we use `exp show`

1. Clean `EXEC_HEAD` , `EXEC_BASELINE` and `EXEC_MERGE` after we pushed
   them to tmp_dir
2. Add a unit test for it.

Co-authored-by:
This commit is contained in:
karajan1001
2022-07-25 18:17:12 +08:00
committed by Gao
parent cc2d2495c0
commit a4763593c7
2 changed files with 21 additions and 0 deletions
+5
View File
@@ -90,6 +90,11 @@ class TempDirExecutor(BaseLocalExecutor):
head = EXEC_BRANCH if branch else EXEC_HEAD
self.scm.checkout(head, detach=True)
merge_rev = self.scm.get_ref(EXEC_MERGE)
for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
if scm.get_ref(ref):
scm.remove_ref(ref)
try:
self.scm.merge(merge_rev, squash=True, commit=False)
except _SCMError as exc:
@@ -5,6 +5,8 @@ from celery import shared_task
from flaky.flaky_decorator import flaky
from dvc.repo.experiments.exceptions import UnresolvedExpNamesError
from dvc.repo.experiments.executor.local import TempDirExecutor
from dvc.repo.experiments.refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE
def test_shutdown_no_tasks(test_queue, mocker):
@@ -104,3 +106,17 @@ def test_celery_queue_kill(test_queue, mocker):
spy = mocker.patch.object(test_queue.proc, "kill")
test_queue.kill("bar")
assert spy.called_once_with(mock_entry.stash_rev)
def test_queue_clean_workspace_refs(git_dir, tmp_dir):
rev = git_dir.scm.get_rev()
exec_heads = (EXEC_MERGE, EXEC_BASELINE, EXEC_HEAD)
for ref in exec_heads:
git_dir.scm.set_ref(ref, rev)
temp_dir_executor = TempDirExecutor(str(tmp_dir), str(tmp_dir), rev)
temp_dir_executor.init_git(git_dir.scm)
for ref in exec_heads:
assert git_dir.scm.get_ref(ref) is None