Use task status to replace collected.

fix: #8088

1. Use status to replace collected.
2. Make status show more accurate.

Fix lint

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()

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(

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
 

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)

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(
+            if (
-                executor_info.root_dir
+                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
+        )

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]:

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

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,
+                "running": True if running else False,
-                "executor": info.location,
+                "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
+    if running:
-    assert info.location in cap.out
+        assert "Running" in cap.out
+        assert info.location in cap.out
 
 
 def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker):