diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index b8c171ac8..c0de27d5b 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/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 diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index 5d64345e0..e797429ef 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/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() diff --git a/dvc/repo/experiments/executor/ssh.py b/dvc/repo/experiments/executor/ssh.py index 0f706bdc3..adccf3aad 100644 --- a/dvc/repo/experiments/executor/ssh.py +++ b/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 .") diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index f285ea7bf..92605e916 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/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( diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index 43050d425..36094d1de 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/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)() diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index dbc459748..1b88a8e12 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/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) diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index d36af7bed..463d41d5e 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/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]: diff --git a/tests/func/experiments/executor/test_ssh.py b/tests/func/experiments/executor/test_ssh.py index 66584eed4..6c2091612 100644 --- a/tests/func/experiments/executor/test_ssh.py +++ b/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() diff --git a/tests/unit/repo/experiments/test_executor_status.py b/tests/unit/repo/experiments/test_executor_status.py index e59dc0ba0..06af2a925 100644 --- a/tests/unit/repo/experiments/test_executor_status.py +++ b/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"])