Move infofile operation to executor and

Currently infofile dump was spread in both executor and queue classes.
put them all into executor to make them more managable. Move all ending
dump operations to the cleanup function, because the collect_result progress
might failed because of git operations, in which can guarantee them to
be run in a `finally` scope.

dvc/repo/experiments/executor/base.py

@@ -72,7 +72,7 @@ class TaskStatus(IntEnum):
     RUNNING = 2
     SUCCESS = 3
     FAILED = 4
-    REVOKED = 5
+    CANCELED = 5
     FINISHED = 6
 
 
@@ -173,6 +173,7 @@ class BaseExecutor(ABC):
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
+        infofile: Optional[str],
         branch: Optional[str] = None,
     ):
         """Init git repo and populate it using exp refs from the specified
@@ -281,8 +282,12 @@ class BaseExecutor(ABC):
                 exp_data.update(to_lockfile(stage))
         return dict_sha256(exp_data)
 
-    def cleanup(self):
-        pass
+    def cleanup(self, infofile: str):
+        if infofile is not None:
+            info = ExecutorInfo.load_json(infofile)
+            if info.status < TaskStatus.FAILED:
+                info.status = TaskStatus.FINISHED
+            info.dump_json(infofile)
 
     # TODO: come up with better way to stash repro arguments
     @staticmethod

dvc/repo/experiments/executor/local.py

@@ -46,8 +46,8 @@ class BaseLocalExecutor(BaseExecutor):
     def scm(self):
         return SCM(self.root_dir)
 
-    def cleanup(self):
-        super().cleanup()
+    def cleanup(self, infofile: str):
+        super().cleanup(infofile)
         self.scm.close()
         del self.scm
 
@@ -72,6 +72,7 @@ class TempDirExecutor(BaseLocalExecutor):
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
+        infofile: Optional[str],
         branch: Optional[str] = None,
     ):
         from dulwich.repo import Repo as DulwichRepo
@@ -81,6 +82,8 @@ class TempDirExecutor(BaseLocalExecutor):
         DulwichRepo.init(os.fspath(self.root_dir))
 
         self.status = TaskStatus.PREPARING
+        if infofile:
+            self.info.dump_json(infofile)
 
         with self.set_exec_refs(scm, stash_rev, entry):
             refspec = f"{EXEC_NAMESPACE}/"
@@ -119,8 +122,8 @@ class TempDirExecutor(BaseLocalExecutor):
         """Initialize DVC cache."""
         self._config(repo.odb.repo.path)
 
-    def cleanup(self):
-        super().cleanup()
+    def cleanup(self, infofile: str):
+        super().cleanup(infofile)
         logger.debug("Removing tmpdir '%s'", self.root_dir)
         remove(self.root_dir)
 
@@ -169,8 +172,13 @@ class WorkspaceExecutor(BaseLocalExecutor):
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
+        infofile: Optional[str],
         branch: Optional[str] = None,
     ):
+        self.status = TaskStatus.PREPARING
+        if infofile:
+            self.info.dump_json(infofile)
+
         scm.set_ref(EXEC_HEAD, entry.head_rev)
         scm.set_ref(EXEC_MERGE, stash_rev)
         scm.set_ref(EXEC_BASELINE, entry.baseline_rev)
@@ -194,7 +202,9 @@ class WorkspaceExecutor(BaseLocalExecutor):
     def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True):
         pass
 
-    def cleanup(self):
+    def cleanup(self, infofile: str):
+        super().cleanup(infofile)
+        remove(os.path.dirname(infofile))
         with self._detach_stack:
             self.scm.remove_ref(EXEC_BASELINE)
             self.scm.remove_ref(EXEC_MERGE)
@@ -203,4 +213,3 @@ class WorkspaceExecutor(BaseLocalExecutor):
             checkpoint = self.scm.get_ref(EXEC_CHECKPOINT)
             if checkpoint and checkpoint != self._orig_checkpoint:
                 self.scm.set_ref(EXEC_APPLY, checkpoint)
-        super().cleanup()

dvc/repo/experiments/executor/ssh.py

@@ -15,7 +15,7 @@ from ..refs import (
     EXEC_MERGE,
     EXEC_NAMESPACE,
 )
-from .base import BaseExecutor, ExecutorInfo, ExecutorResult
+from .base import BaseExecutor, ExecutorInfo, ExecutorResult, TaskStatus
 
 if TYPE_CHECKING:
     from queue import Queue
@@ -122,10 +122,15 @@ class SSHExecutor(BaseExecutor):
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
+        infofile: Optional[str],
         branch: Optional[str] = None,
     ):
         from ..utils import push_refspec
 
+        self.status = TaskStatus.PREPARING
+        if infofile:
+            self.info.dump_json(infofile)
+
         with self.sshfs() as fs:
             fs.makedirs(self.root_dir)
             self._ssh_cmd(fs, "git init .")

dvc/repo/experiments/queue/base.py

@@ -31,7 +31,6 @@ from ..executor.base import (
     EXEC_TMP_DIR,
     BaseExecutor,
     ExecutorResult,
-    TaskStatus,
 )
 from ..executor.local import WorkspaceExecutor
 from ..refs import ExpRefInfo
@@ -582,15 +581,14 @@ class BaseStashQueue(ABC):
         )
 
         infofile = exp.celery_queue.get_infofile_path(stash_rev)
-        executor.info.dump_json(infofile)
-
         executor.init_git(
-            exp.repo.scm, stash_rev, stash_entry, branch=stash_entry.branch
+            exp.repo.scm,
+            stash_rev,
+            stash_entry,
+            infofile,
+            branch=stash_entry.branch,
         )
 
-        infofile = exp.celery_queue.get_infofile_path(stash_rev)
-        executor.info.dump_json(infofile)
-
         executor.init_cache(exp.repo, stash_rev)
 
         return executor
@@ -608,7 +606,6 @@ class BaseStashQueue(ABC):
         exp: "Experiments",
         executor: BaseExecutor,
         exec_result: ExecutorResult,
-        infofile: str,
     ) -> Dict[str, str]:
         results = {}
 
@@ -632,10 +629,6 @@ class BaseStashQueue(ABC):
         if exec_result.ref_info is not None:
             executor.collect_cache(exp.repo, exec_result.ref_info)
 
-        executor.status = TaskStatus.FINISHED
-        if infofile is not None:
-            executor.info.dump_json(infofile)
-
         return results
 
     def match_queue_entry_by_name(

dvc/repo/experiments/queue/tasks.py

@@ -3,8 +3,6 @@ from typing import Any, Dict
 from celery import shared_task
 from celery.utils.log import get_task_logger
 
-from dvc.utils.fs import remove
-
 from ..executor.base import ExecutorInfo
 from ..executor.local import TempDirExecutor
 from .base import BaseStashQueue, QueueEntry
@@ -13,7 +11,7 @@ logger = get_task_logger(__name__)
 
 
 @shared_task
-def setup_exp(entry_dict: Dict[str, Any]) -> str:
+def setup_exp(entry_dict: Dict[str, Any]) -> TempDirExecutor:
     """Setup an experiment.
 
     Arguments:
@@ -36,7 +34,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str:
     )
     infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev)
     executor.info.dump_json(infofile)
-    return executor.root_dir
+    return executor
 
 
 @shared_task
@@ -66,7 +64,7 @@ def collect_exp(
     try:
         if exec_result is not None:
             BaseStashQueue.collect_executor(
-                repo.experiments, executor, exec_result, infofile
+                repo.experiments, executor, exec_result
             )
         else:
             logger.debug("Experiment failed (Exec result was None)")
@@ -79,16 +77,14 @@ def collect_exp(
 
 
 @shared_task
-def cleanup_exp(  # pylint: disable=unused-argument
-    tmp_dir: str, entry_dict: Dict[str, Any]
-) -> None:
+def cleanup_exp(executor: TempDirExecutor, infofile: str) -> None:
     """Cleanup after an experiment.
 
     Arguments:
         tmp_dir: Temp directory to be removed.
         entry_dict: Serialized QueueEntry for this experiment.
     """
-    remove(tmp_dir)
+    executor.cleanup(infofile)
 
 
 @shared_task
@@ -106,10 +102,10 @@ def run_exp(entry_dict: Dict[str, Any]) -> None:
     repo = Repo(entry.dvc_root)
     queue = repo.experiments.celery_queue
     infofile = queue.get_infofile_path(entry.stash_rev)
-    root_dir = setup_exp.s(entry_dict)()
+    executor = setup_exp.s(entry_dict)()
     try:
         cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
         proc_dict = queue.proc.run_signature(cmd, name=entry.stash_rev)()
         collect_exp.s(proc_dict, entry_dict)()
     finally:
-        cleanup_exp.s(root_dir, entry_dict)()
+        cleanup_exp.s(executor, infofile)()

dvc/repo/experiments/queue/tempdir.py

@@ -96,8 +96,5 @@ class TempDirQueue(WorkspaceQueue):
         exp: "Experiments",
         executor: BaseExecutor,
         exec_result: ExecutorResult,
-        infofile: str,
     ) -> Dict[str, str]:
-        return BaseStashQueue.collect_executor(
-            exp, executor, exec_result, infofile
-        )
+        return BaseStashQueue.collect_executor(exp, executor, exec_result)

dvc/repo/experiments/queue/workspace.py

@@ -1,15 +1,13 @@
 import logging
-import os
 from collections import defaultdict
 from typing import TYPE_CHECKING, Collection, Dict, Generator, Optional
 
 from funcy import first
 
 from dvc.exceptions import DvcException
-from dvc.utils.fs import remove
 
 from ..exceptions import ExpQueueEmptyError
-from ..executor.base import BaseExecutor, ExecutorResult, TaskStatus
+from ..executor.base import BaseExecutor, ExecutorResult
 from ..executor.local import WorkspaceExecutor
 from ..refs import EXEC_BRANCH
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
@@ -105,7 +103,7 @@ class WorkspaceQueue(BaseStashQueue):
             if exec_result.ref_info:
                 results[rev].update(
                     self.collect_executor(
-                        self.repo.experiments, executor, exec_result, infofile
+                        self.repo.experiments, executor, exec_result
                     )
                 )
         except CheckpointKilledError:
@@ -118,9 +116,7 @@ class WorkspaceQueue(BaseStashQueue):
                 f"Failed to reproduce experiment '{rev[:7]}'"
             ) from exc
         finally:
-            if self._EXEC_NAME == exec_name:
-                remove(os.path.join(self.pid_dir, exec_name))
-            executor.cleanup()
+            executor.cleanup(infofile)
         return results
 
     @staticmethod
@@ -128,7 +124,6 @@ class WorkspaceQueue(BaseStashQueue):
         exp: "Experiments",
         executor: BaseExecutor,
         exec_result: ExecutorResult,
-        infofile: str,
     ) -> Dict[str, str]:
         results: Dict[str, str] = {}
         exp_rev = exp.scm.get_ref(EXEC_BRANCH)
@@ -137,10 +132,6 @@ class WorkspaceQueue(BaseStashQueue):
             logger.debug("Collected experiment '%s'.", exp_rev[:7])
             results[exp_rev] = exec_result.exp_hash
 
-        executor.status = TaskStatus.FINISHED
-        if infofile is not None:
-            executor.info.dump_json(infofile)
-
         return results
 
     def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:

tests/func/experiments/executor/test_ssh.py

@@ -58,7 +58,8 @@ def test_init_git(tmp_dir, scm, cloud, mocker):
         username=TEST_SSH_USER,
         fs_factory=partial(_ssh_factory, cloud),
     )
-    executor.init_git(scm, rev, mock)
+    infofile = str((root_url / "foo.run").path)
+    executor.init_git(scm, rev, mock, infofile=infofile)
     assert root_url.path == executor._repo_abspath
 
     fs = cloud._ssh
@@ -116,7 +117,6 @@ def test_reproduce(tmp_dir, scm, dvc, cloud, exp_stage, mocker):
     SSHExecutor.reproduce(
         info,
         rev,
-        infofile=infofile,
         fs_factory=partial(_ssh_factory, cloud),
     )
     assert mock_execute.called_once()

tests/unit/repo/experiments/test_executor_status.py

@@ -4,7 +4,11 @@ import pytest
 
 from dvc.exceptions import DvcException
 from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus
-from dvc.repo.experiments.queue.tasks import collect_exp, setup_exp
+from dvc.repo.experiments.queue.tasks import (
+    cleanup_exp,
+    collect_exp,
+    setup_exp,
+)
 
 
 def test_celery_queue_success_status(dvc, scm, test_queue, exp_stage):
@@ -14,7 +18,7 @@ def test_celery_queue_success_status(dvc, scm, test_queue, exp_stage):
         name="success",
     )
     infofile = test_queue.get_infofile_path(queue_entry.stash_rev)
-    setup_exp.s(queue_entry.asdict())()
+    executor = setup_exp.s(queue_entry.asdict())()
     executor_info = ExecutorInfo.load_json(infofile)
     assert executor_info.status == TaskStatus.PREPARING
 
@@ -27,6 +31,7 @@ def test_celery_queue_success_status(dvc, scm, test_queue, exp_stage):
     assert executor_info.status == TaskStatus.SUCCESS
 
     collect_exp.s(proc_dict, queue_entry.asdict())()
+    cleanup_exp.s(executor, infofile)()
     executor_info = ExecutorInfo.load_json(infofile)
     assert executor_info.status == TaskStatus.FINISHED
 
@@ -66,11 +71,15 @@ def test_workspace_executor_success_status(dvc, scm, exp_stage, queue_type):
     assert executor_info.status == TaskStatus.SUCCESS
     if exec_result.ref_info:
         workspace_queue.collect_executor(
-            dvc.experiments, executor, exec_result, infofile
+            dvc.experiments, executor, exec_result
         )
+    executor.cleanup(infofile)
 
-    executor_info = ExecutorInfo.load_json(infofile)
-    assert executor_info.status == TaskStatus.FINISHED
+    if queue_type == "tempdir_queue":
+        executor_info = ExecutorInfo.load_json(infofile)
+        assert executor_info.status == TaskStatus.FINISHED
+    else:
+        assert not os.path.exists(infofile)
 
 
 @pytest.mark.parametrize("queue_type", ["workspace_queue", "tempdir_queue"])