Fix some celery queue realted ci failure

fix: #8403
1. remove some of the flaky mark
2. In `get_result` make sure the celery task is completed.

dvc/repo/experiments/__init__.py

@@ -189,7 +189,10 @@ class Experiments:
                     time.sleep(1)
                 self.celery_queue.follow(entry)
                 # wait for task collection to complete
-                result = self.celery_queue.get_result(entry)
+                try:
+                    result = self.celery_queue.get_result(entry)
+                except FileNotFoundError:
+                    result = None
                 if result is None or result.exp_hash is None:
                     name = entry.name or entry.stash_rev[:7]
                     failed.append(name)

dvc/repo/experiments/executor/base.py

@@ -602,29 +602,22 @@ class BaseExecutor(ABC):
             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:
+            if infofile is not None:
+                info.dump_json(infofile)
             dvc.close()
             os.chdir(old_cwd)
 

dvc/repo/experiments/executor/local.py

@@ -47,9 +47,9 @@ class BaseLocalExecutor(BaseExecutor):
         return SCM(self.root_dir)
 
     def cleanup(self, infofile: str):
-        super().cleanup(infofile)
         self.scm.close()
         del self.scm
+        super().cleanup(infofile)
 
     def collect_cache(
         self, repo: "Repo", exp_ref: "ExpRefInfo", run_cache: bool = True

dvc/repo/experiments/queue/celery.py

@@ -23,12 +23,7 @@ from dvc.exceptions import DvcException
 from dvc.ui import ui
 
 from ..exceptions import UnresolvedQueueExpNamesError
-from ..executor.base import (
-    EXEC_TMP_DIR,
-    ExecutorInfo,
-    ExecutorResult,
-    TaskStatus,
-)
+from ..executor.base import EXEC_TMP_DIR, ExecutorInfo, ExecutorResult
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
 from .tasks import run_exp
 from .utils import fetch_running_exp_from_temp_dir
@@ -243,48 +238,50 @@ class LocalCeleryQueue(BaseStashQueue):
     def reproduce(self) -> Mapping[str, Mapping[str, str]]:
         raise NotImplementedError
 
-    def get_result(
+    def _load_info(self, rev: str) -> ExecutorInfo:
+        infofile = self.get_infofile_path(rev)
+        return ExecutorInfo.load_json(infofile)
+
+    def _get_done_result(
         self, entry: QueueEntry, timeout: Optional[float] = None
-    ) -> Optional[ExecutorResult]:
+    ) -> Optional["ExecutorResult"]:
         from celery.exceptions import TimeoutError as _CeleryTimeout
 
-        def _load_info(rev: str) -> ExecutorInfo:
-            infofile = self.get_infofile_path(rev)
-            return ExecutorInfo.load_json(infofile)
-
-        def _load_collected(rev: str) -> Optional[ExecutorResult]:
-            executor_info = _load_info(rev)
-            if executor_info.status > TaskStatus.SUCCESS:
+        for msg, processed_entry in self._iter_processed():
+            if entry.stash_rev == processed_entry.stash_rev:
+                task_id = msg.headers["id"]
+                result: AsyncResult = AsyncResult(task_id)
+                if not result.ready():
+                    logger.debug(
+                        "Waiting for exp task '%s' to complete", result.id
+                    )
+                    try:
+                        result.get(timeout=timeout)
+                    except _CeleryTimeout as exc:
+                        raise DvcException(
+                            "Timed out waiting for exp to finish."
+                        ) from exc
+                executor_info = self._load_info(entry.stash_rev)
                 return executor_info.result
-            raise FileNotFoundError
+        raise FileNotFoundError
+
+    def get_result(
+        self, entry: QueueEntry, timeout: Optional[float] = None
+    ) -> Optional["ExecutorResult"]:
 
         try:
-            return _load_collected(entry.stash_rev)
+            return self._get_done_result(entry, timeout)
         except FileNotFoundError:
-            # Infofile will not be created until execution begins
             pass
 
         for queue_entry in self.iter_queued():
             if entry.stash_rev == queue_entry.stash_rev:
                 raise DvcException("Experiment has not been started.")
-        for result, active_entry in self._iter_active_tasks():
-            if entry.stash_rev == active_entry.stash_rev:
-                logger.debug(
-                    "Waiting for exp task '%s' to complete", result.id
-                )
-                try:
-                    result.get(timeout=timeout)
-                except _CeleryTimeout as exc:
-                    raise DvcException(
-                        "Timed out waiting for exp to finish."
-                    ) from exc
-                executor_info = _load_info(entry.stash_rev)
-                return executor_info.result
 
         # NOTE: It's possible for an exp to complete while iterating through
         # other queued and active tasks, in which case the exp will get moved
         # out of the active task list, and needs to be loaded here.
-        return _load_collected(entry.stash_rev)
+        return self._get_done_result(entry, timeout)
 
     def kill(self, revs: Collection[str]) -> None:
         to_kill: Set[QueueEntry] = set()

dvc/repo/experiments/show.py

@@ -398,7 +398,6 @@ def show(
             repo,
             found_revs,
             running,
-            sha_only=sha_only,
             param_deps=param_deps,
             onerror=onerror,
         )

tests/func/experiments/test_queue.py

@@ -11,64 +11,54 @@ def to_dict(tasks):
     return status_dict
 
 
-@pytest.fixture
-def success_tasks(tmp_dir, dvc, scm, test_queue, exp_stage):
-    queue_length = 3
-    name_list = []
-    for i in range(queue_length):
-        name = f"success{i}"
-        name_list.append(name)
-        dvc.experiments.run(
-            exp_stage.addressing, params=[f"foo={i}"], queue=True, name=name
-        )
-    dvc.experiments.run(run_all=True)
-    return name_list
-
-
-@pytest.fixture
-def failed_tasks(tmp_dir, dvc, scm, test_queue, failed_exp_stage):
-    queue_length = 3
-    name_list = []
-    for i in range(queue_length):
-        name = f"failed{i}"
-        name_list.append(name)
-        dvc.experiments.run(
-            failed_exp_stage.addressing,
-            params=[f"foo={i+queue_length}"],
-            queue=True,
-            name=name,
-        )
-    dvc.experiments.run(run_all=True)
-    return name_list
-
-
-@pytest.mark.xfail(strict=False, reason="pytest-celery flaky")
 @pytest.mark.parametrize("follow", [True, False])
 def test_celery_logs(
     tmp_dir,
     scm,
     dvc,
     failed_exp_stage,
-    test_queue,
     follow,
     capsys,
 ):
+    celery_queue = dvc.experiments.celery_queue
     dvc.experiments.run(failed_exp_stage.addressing, queue=True)
     dvc.experiments.run(run_all=True)
 
-    queue = dvc.experiments.celery_queue
-    done_result = first(queue.iter_done())
+    done_result = first(celery_queue.iter_done())
+
     name = done_result.entry.stash_rev
     captured = capsys.readouterr()
-    queue.logs(name, follow=follow)
+    celery_queue.logs(name, follow=follow)
     captured = capsys.readouterr()
     assert "failed to reproduce 'failed-copy-file'" in captured.out
 
 
-@pytest.mark.xfail(strict=False, reason="pytest-celery flaky")
-def test_queue_remove_done(dvc, failed_tasks, success_tasks):
-    assert len(dvc.experiments.celery_queue.failed_stash) == 3
-    status = to_dict(dvc.experiments.celery_queue.status())
+def test_queue_remove_done(
+    dvc,
+    exp_stage,
+    failed_exp_stage,
+):
+    queue_length = 3
+    success_tasks = []
+    failed_tasks = []
+    celery_queue = dvc.experiments.celery_queue
+    for i in range(queue_length):
+        name = f"success{i}"
+        success_tasks.append(name)
+        dvc.experiments.run(
+            exp_stage.addressing, params=[f"foo={i}"], queue=True, name=name
+        )
+        name_fail = f"failed{i}"
+        failed_tasks.append(name_fail)
+        dvc.experiments.run(
+            failed_exp_stage.addressing,
+            params=[f"foo={i+queue_length}"],
+            queue=True,
+            name=name_fail,
+        )
+    dvc.experiments.run(run_all=True)
+    assert len(celery_queue.failed_stash) == 3
+    status = to_dict(celery_queue.status())
     assert len(status) == 6
     for name in failed_tasks:
         assert status[name] == "Failed"
@@ -76,27 +66,21 @@ def test_queue_remove_done(dvc, failed_tasks, success_tasks):
         assert status[name] == "Success"
 
     with pytest.raises(InvalidArgumentError):
-        dvc.experiments.celery_queue.remove(failed_tasks[:2] + ["non-exist"])
-    assert len(dvc.experiments.celery_queue.status()) == 6
+        celery_queue.remove(failed_tasks[:2] + ["non-exist"])
+    assert len(celery_queue.status()) == 6
 
     to_remove = [failed_tasks[0], success_tasks[2]]
-    assert set(dvc.experiments.celery_queue.remove(to_remove)) == set(
-        to_remove
-    )
+    assert set(celery_queue.remove(to_remove)) == set(to_remove)
 
-    assert len(dvc.experiments.celery_queue.failed_stash) == 2
-    status = to_dict(dvc.experiments.celery_queue.status())
+    assert len(celery_queue.failed_stash) == 2
+    status = to_dict(celery_queue.status())
     assert set(status) == set(failed_tasks[1:] + success_tasks[:2])
 
-    assert dvc.experiments.celery_queue.clear(failed=True) == failed_tasks[1:]
+    assert set(celery_queue.clear(failed=True)) == set(failed_tasks[1:])
 
-    assert len(dvc.experiments.celery_queue.failed_stash) == 0
-    assert set(to_dict(dvc.experiments.celery_queue.status())) == set(
-        success_tasks[:2]
-    )
+    assert len(celery_queue.failed_stash) == 0
+    assert set(to_dict(celery_queue.status())) == set(success_tasks[:2])
 
-    assert (
-        dvc.experiments.celery_queue.clear(success=True) == success_tasks[:2]
-    )
+    assert set(celery_queue.clear(success=True)) == set(success_tasks[:2])
 
-    assert dvc.experiments.celery_queue.status() == []
+    assert celery_queue.status() == []

tests/func/experiments/test_show.py

@@ -174,7 +174,6 @@ def test_show_queued(tmp_dir, scm, dvc, exp_stage):
 
 
 @pytest.mark.vscode
-@pytest.mark.xfail(strict=False, reason="pytest-celery flaky")
 def test_show_failed_experiment(tmp_dir, scm, dvc, failed_exp_stage):
     baseline_rev = scm.get_rev()
     timestamp = datetime.fromtimestamp(

tests/unit/repo/experiments/test_executor_status.py

@@ -82,21 +82,34 @@ def test_workspace_executor_success_status(dvc, scm, exp_stage, queue_type):
         assert not os.path.exists(infofile)
 
 
-@pytest.mark.parametrize("queue_type", ["workspace_queue", "tempdir_queue"])
+@pytest.mark.parametrize(
+    "queue_type",
+    ["workspace_queue", "tempdir_queue"],
+)
 def test_workspace_executor_failed_status(
     dvc, scm, failed_exp_stage, queue_type
 ):
-    workspace_queue = getattr(dvc.experiments, queue_type)
-    queue_entry = workspace_queue.put(
+    queue = getattr(dvc.experiments, queue_type)
+    queue.put(
         params={"params.yaml": ["foo=1"]},
         targets=failed_exp_stage.addressing,
         name="failed",
     )
-    name = workspace_queue._EXEC_NAME or queue_entry.stash_rev
+    entry, executor = queue.get()
+    name = queue._EXEC_NAME or entry.stash_rev
+    infofile = queue.get_infofile_path(name)
+    rev = entry.stash_rev
 
-    infofile = workspace_queue.get_infofile_path(name)
     with pytest.raises(DvcException):
-        workspace_queue.reproduce()
+        executor.reproduce(
+            info=executor.info,
+            rev=rev,
+            infofile=infofile,
+        )
+    executor_info = ExecutorInfo.load_json(infofile)
+    assert executor_info.status == TaskStatus.FAILED
+
+    cleanup_exp.s(executor, infofile)()
     if queue_type == "workspace_queue":
         assert not os.path.exists(infofile)
     else: