queue: cleanup flaky tests

dvc/repo/experiments/queue/celery.py

@@ -246,10 +246,14 @@ class LocalCeleryQueue(BaseStashQueue):
             infofile = self.get_infofile_path(rev)
             return ExecutorInfo.load_json(infofile)
 
-        try:
-            executor_info = _load_info(entry.stash_rev)
+        def _load_collected(rev: str) -> Optional[ExecutorResult]:
+            executor_info = _load_info(rev)
             if executor_info.collected:
                 return executor_info.result
+            raise FileNotFoundError
+
+        try:
+            return _load_collected(entry.stash_rev)
         except FileNotFoundError:
             # Infofile will not be created until execution begins
             pass
@@ -269,7 +273,15 @@ class LocalCeleryQueue(BaseStashQueue):
                     ) from exc
                 executor_info = _load_info(entry.stash_rev)
                 return executor_info.result
-        raise DvcException("Invalid experiment.")
+
+        # 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.
+        try:
+            return _load_collected(entry.stash_rev)
+        except FileNotFoundError:
+            pass
+        raise DvcException("Invalid experiment '{entry.stash_rev[:7]}'.")
 
     def kill(self, revs: Collection[str]) -> None:
         to_kill: Set[QueueEntry] = set()

tests/func/experiments/test_experiments.py

@@ -199,8 +199,7 @@ def test_apply(tmp_dir, scm, dvc, exp_stage):
     assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 3"
 
 
-def test_apply_queued(tmp_dir, scm, dvc, exp_stage, test_queue):
-    # from dvc.exceptions import InvalidArgumentError
+def test_apply_queued(tmp_dir, scm, dvc, exp_stage):
     from dvc.repo.experiments.exceptions import ApplyConflictError
 
     metrics_original = (tmp_dir / "metrics.yaml").read_text().strip()

tests/unit/repo/experiments/queue/test_local.py

@@ -2,6 +2,7 @@ import time
 
 import pytest
 from celery import shared_task
+from flaky.flaky_decorator import flaky
 
 from dvc.repo.experiments.exceptions import UnresolvedExpNamesError
 
@@ -48,6 +49,8 @@ def test_shutdown_with_kill(test_queue, mocker):
     shutdown_spy.assert_called_once()
 
 
+# pytest-celery worker thread may finish the task before we check for PENDING
+@flaky(max_runs=3, min_passes=1)
 def test_post_run_after_kill(test_queue):
 
     from celery import chain