From 6bd14b21effa836d0d4e1c7958639e4d2e5e164b Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Sat, 17 Sep 2022 10:48:29 +0800 Subject: [PATCH] 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 | 11 +++++++--- dvc/repo/experiments/executor/local.py | 21 +++++++++++++------ dvc/repo/experiments/executor/ssh.py | 7 ++++++- dvc/repo/experiments/queue/base.py | 17 +++++---------- dvc/repo/experiments/queue/tasks.py | 18 +++++++--------- dvc/repo/experiments/queue/tempdir.py | 5 +---- dvc/repo/experiments/queue/workspace.py | 15 +++---------- tests/func/experiments/executor/test_ssh.py | 4 ++-- .../repo/experiments/test_executor_status.py | 19 ++++++++++++----- 9 files changed, 61 insertions(+), 56 deletions(-) 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"])