diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index 7c53d2749..0ce17aef9 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -73,6 +73,7 @@ class TaskStatus(IntEnum): SUCCESS = 3 FAILED = 4 REVOKED = 5 + FINISHED = 6 @dataclass @@ -84,7 +85,6 @@ class ExecutorInfo: dvc_dir: str name: Optional[str] = None wdir: Optional[str] = None - collected: bool = False result_hash: Optional[str] = None result_ref: Optional[str] = None result_force: bool = False @@ -431,12 +431,9 @@ class BaseExecutor(ABC): exp_ref: Optional["ExpRefInfo"] = None repro_force: bool = False - info.status = TaskStatus.PREPARING - if infofile is not None: - info.dump_json(infofile) - with cls._repro_dvc( info, + infofile, log_errors=log_errors, **kwargs, ) as dvc: @@ -488,9 +485,6 @@ class BaseExecutor(ABC): info.name, repro_force or checkpoint_reset, ) - info.status = TaskStatus.RUNNING - if infofile is not None: - info.dump_json(infofile) stages = dvc_reproduce( dvc, *args, @@ -514,10 +508,6 @@ class BaseExecutor(ABC): info.result_hash = exp_hash info.result_ref = ref info.result_force = repro_force - info.status = TaskStatus.SUCCESS - - if infofile is not None: - info.dump_json(infofile) # ideally we would return stages here like a normal repro() call, but # stages is not currently picklable and cannot be returned across @@ -575,6 +565,7 @@ class BaseExecutor(ABC): def _repro_dvc( cls, info: "ExecutorInfo", + infofile: str = None, log_errors: bool = True, **kwargs, ): @@ -582,6 +573,9 @@ class BaseExecutor(ABC): from dvc.stage.monitor import CheckpointKilledError dvc = Repo(os.path.join(info.root_dir, info.dvc_dir)) + info.status = TaskStatus.RUNNING + if infofile is not None: + info.dump_json(infofile) if cls.QUIET: dvc.scm_context.quiet = cls.QUIET old_cwd = os.getcwd() @@ -593,15 +587,28 @@ class BaseExecutor(ABC): try: logger.debug("Running repro in '%s'", os.getcwd()) yield dvc + info.status = TaskStatus.SUCCESS + if infofile is not None: + info.dump_json(infofile) + except CheckpointKilledError: + info.status = TaskStatus.FAILED + if infofile is not None: + info.dump_json(infofile) raise except DvcException: if log_errors: logger.exception("") + info.status = TaskStatus.FAILED + if infofile is not None: + info.dump_json(infofile) raise except Exception: if log_errors: logger.exception("unexpected error") + info.status = TaskStatus.FAILED + if infofile is not None: + info.dump_json(infofile) raise finally: dvc.close() diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 6bfadd440..f285ea7bf 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -31,6 +31,7 @@ from ..executor.base import ( EXEC_TMP_DIR, BaseExecutor, ExecutorResult, + TaskStatus, ) from ..executor.local import WorkspaceExecutor from ..refs import ExpRefInfo @@ -607,6 +608,7 @@ class BaseStashQueue(ABC): exp: "Experiments", executor: BaseExecutor, exec_result: ExecutorResult, + infofile: str, ) -> Dict[str, str]: results = {} @@ -630,6 +632,10 @@ 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/celery.py b/dvc/repo/experiments/queue/celery.py index 05eee9b28..734d7dcef 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -23,7 +23,12 @@ from dvc.exceptions import DvcException from dvc.ui import ui from ..exceptions import UnresolvedQueueExpNamesError -from ..executor.base import EXEC_TMP_DIR, ExecutorInfo, ExecutorResult +from ..executor.base import ( + EXEC_TMP_DIR, + ExecutorInfo, + ExecutorResult, + TaskStatus, +) from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult from .tasks import run_exp @@ -248,7 +253,8 @@ class LocalCeleryQueue(BaseStashQueue): def _load_collected(rev: str) -> Optional[ExecutorResult]: executor_info = _load_info(rev) - if executor_info.collected: + print("executor_info is", executor_info.status) + if executor_info.status > TaskStatus.SUCCESS: return executor_info.result raise FileNotFoundError diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index ac395371a..43050d425 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -66,7 +66,7 @@ def collect_exp( try: if exec_result is not None: BaseStashQueue.collect_executor( - repo.experiments, executor, exec_result + repo.experiments, executor, exec_result, infofile ) else: logger.debug("Experiment failed (Exec result was None)") @@ -91,15 +91,6 @@ def cleanup_exp( # pylint: disable=unused-argument remove(tmp_dir) -def _set_collected(infofile: str): - try: - executor_info = ExecutorInfo.load_json(infofile) - executor_info.collected = True - executor_info.dump_json(infofile) - except FileNotFoundError: - pass - - @shared_task def run_exp(entry_dict: Dict[str, Any]) -> None: """Run a full experiment. @@ -122,4 +113,3 @@ def run_exp(entry_dict: Dict[str, Any]) -> None: collect_exp.s(proc_dict, entry_dict)() finally: cleanup_exp.s(root_dir, entry_dict)() - _set_collected(infofile) diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index b2cbf3825..dbc459748 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/dvc/repo/experiments/queue/tempdir.py @@ -11,6 +11,7 @@ from ..executor.base import ( BaseExecutor, ExecutorInfo, ExecutorResult, + TaskStatus, ) from ..executor.local import TempDirExecutor from .base import BaseStashQueue, QueueEntry, QueueGetResult @@ -75,8 +76,9 @@ class TempDirQueue(WorkspaceQueue): for stash_rev in self.proc: infofile = self.get_infofile_path(stash_rev) executor_info = ExecutorInfo.load_json(infofile) - if not executor_info.collected and os.path.exists( - executor_info.root_dir + if ( + not executor_info.status <= TaskStatus.SUCCESS + and os.path.exists(executor_info.root_dir) ): yield QueueEntry( self.repo.root_dir, @@ -94,5 +96,8 @@ class TempDirQueue(WorkspaceQueue): exp: "Experiments", executor: BaseExecutor, exec_result: ExecutorResult, + infofile: str, ) -> Dict[str, str]: - return BaseStashQueue.collect_executor(exp, executor, exec_result) + return BaseStashQueue.collect_executor( + exp, executor, exec_result, infofile + ) diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 82f80d241..d36af7bed 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -105,18 +105,15 @@ class WorkspaceQueue(BaseStashQueue): if exec_result.ref_info: results[rev].update( self.collect_executor( - self.repo.experiments, executor, exec_result + self.repo.experiments, executor, exec_result, infofile ) ) except CheckpointKilledError: # Checkpoint errors have already been logged - executor.status = TaskStatus.FAILED return {} except DvcException: - executor.status = TaskStatus.FAILED raise except Exception as exc: - executor.status = TaskStatus.FAILED raise DvcException( f"Failed to reproduce experiment '{rev[:7]}'" ) from exc @@ -131,6 +128,7 @@ 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) @@ -138,6 +136,11 @@ class WorkspaceQueue(BaseStashQueue): assert exec_result.exp_hash 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 c5d930a3d..66584eed4 100644 --- a/tests/func/experiments/executor/test_ssh.py +++ b/tests/func/experiments/executor/test_ssh.py @@ -38,14 +38,14 @@ def test_init_from_stash(tmp_dir, scm, dvc, machine_instance, mocker): @pytest.mark.needs_internet @pytest.mark.parametrize("cloud", [pytest.lazy_fixture("git_ssh")]) -def test_init_git(tmp_dir, scm, cloud): +def test_init_git(tmp_dir, scm, cloud, mocker): tmp_dir.scm_gen({"foo": "foo", "dir": {"bar": "bar"}}, commit="init") baseline_rev = scm.get_rev() - scm.set_ref(EXEC_HEAD, baseline_rev) tmp_dir.gen("foo", "stashed") scm.gitpython.git.stash() rev = scm.resolve_rev("stash@{0}") - scm.set_ref(EXEC_MERGE, rev) + + mock = mocker.Mock(baseline_rev=baseline_rev, head_rev=baseline_rev) root_url = cloud / SSHExecutor.gen_dirname() @@ -58,7 +58,7 @@ def test_init_git(tmp_dir, scm, cloud): username=TEST_SSH_USER, fs_factory=partial(_ssh_factory, cloud), ) - executor.init_git(scm) + executor.init_git(scm, rev, mock) assert root_url.path == executor._repo_abspath fs = cloud._ssh diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py index bf85f8c6b..903e95218 100644 --- a/tests/func/experiments/test_show.py +++ b/tests/func/experiments/test_show.py @@ -12,6 +12,7 @@ from dvc.repo.experiments.executor.base import ( EXEC_TMP_DIR, BaseExecutor, ExecutorInfo, + TaskStatus, ) from dvc.repo.experiments.queue.base import QueueEntry from dvc.repo.experiments.refs import CELERY_STASH, ExpRefInfo @@ -338,9 +339,16 @@ def test_show_sort(tmp_dir, scm, dvc, exp_stage, caplog): @pytest.mark.vscode -def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys): +@pytest.mark.parametrize( + "status, running", [(TaskStatus.RUNNING, True), (TaskStatus.FAILED, False)] +) +def test_show_running_workspace( + tmp_dir, scm, dvc, exp_stage, capsys, status, running +): pid_dir = os.path.join(dvc.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR) - info = make_executor_info(location=BaseExecutor.DEFAULT_LOCATION) + info = make_executor_info( + location=BaseExecutor.DEFAULT_LOCATION, status=status + ) pidfile = os.path.join( pid_dir, "workspace", @@ -349,7 +357,8 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys): makedirs(os.path.dirname(pidfile), True) (tmp_dir / pidfile).dump_json(info.asdict()) - assert dvc.experiments.show()["workspace"] == { + print(dvc.experiments.show()) + assert dvc.experiments.show().get("workspace") == { "baseline": { "data": { "deps": { @@ -363,18 +372,18 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys): "params": {"params.yaml": {"data": {"foo": 1}}}, "outs": {}, "queued": False, - "running": True, - "executor": info.location, + "running": True if running else False, + "executor": info.location if running else None, "timestamp": None, } } } - capsys.readouterr() assert main(["exp", "show", "--csv"]) == 0 cap = capsys.readouterr() - assert "Running" in cap.out - assert info.location in cap.out + if running: + assert "Running" in cap.out + assert info.location in cap.out def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker):