Repository: iterative/dvc

Issue:
ci: celery queue remove/status race condition
see:
- https://github.com/iterative/dvc/actions/runs/4371541807/jobs/7647506453
- https://github.com/iterative/dvc/actions/runs/4485105768/jobs/7886311287


Discussion:

Resolving patch:
diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -1,7 +1,6 @@
 import logging
 import os
 import re
-import time
 from typing import TYPE_CHECKING, Dict, Iterable, Optional
 
 from funcy import chain, first
@@ -205,8 +204,7 @@ def reproduce_celery(  # noqa: C901
             )
             for entry in entries:
                 # wait for task execution to start
-                while not self.celery_queue.proc.get(entry.stash_rev):
-                    time.sleep(1)
+                self.celery_queue.wait_for_start(entry, sleep_interval=1)
                 self.celery_queue.follow(entry)
                 # wait for task collection to complete
                 try:
diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py
--- a/dvc/repo/experiments/queue/celery.py
+++ b/dvc/repo/experiments/queue/celery.py
@@ -295,6 +295,28 @@ def get_result(
         # out of the active task list, and needs to be loaded here.
         return self._get_done_result(entry, timeout)
 
+    def wait(self, revs: Collection[str], **kwargs) -> None:
+        """Block until the specified tasks have completed."""
+        revs = [revs] if isinstance(revs, str) else revs
+        results = self.match_queue_entry_by_name(
+            revs, self.iter_queued(), self.iter_done(), self.iter_failed()
+        )
+        for entry in results.values():
+            if not entry:
+                continue
+            self.wait_for_start(entry, **kwargs)
+            try:
+                self.get_result(entry)
+            except FileNotFoundError:
+                pass
+
+    def wait_for_start(self, entry: QueueEntry, sleep_interval: float = 0.001) -> None:
+        """Block until the specified task has been started."""
+        import time
+
+        while not self.proc.get(entry.stash_rev):
+            time.sleep(sleep_interval)
+
     def _get_running_task_ids(self) -> Set[str]:
         running_task_ids: Set[str] = set()
         active_workers = self.worker_status()


Test patch:
diff --git a/tests/func/experiments/test_apply.py b/tests/func/experiments/test_apply.py
--- a/tests/func/experiments/test_apply.py
+++ b/tests/func/experiments/test_apply.py
@@ -35,7 +35,7 @@ def test_apply_failed(tmp_dir, scm, dvc, failed_exp_stage, session_queue):
     dvc.experiments.run(
         failed_exp_stage.addressing, params=["foo=3"], queue=True, name="foo"
     )
-    dvc.experiments.run(run_all=True)
+    session_queue.wait(["foo"])
 
     with pytest.raises(InvalidArgumentError):
         dvc.experiments.apply("bar")
diff --git a/tests/func/experiments/test_queue.py b/tests/func/experiments/test_queue.py
--- a/tests/func/experiments/test_queue.py
+++ b/tests/func/experiments/test_queue.py
@@ -22,8 +22,9 @@ def test_celery_logs(
     test_queue,
 ):
     celery_queue = dvc.experiments.celery_queue
-    dvc.experiments.run(failed_exp_stage.addressing, queue=True)
+    dvc.experiments.run(failed_exp_stage.addressing, queue=True, name="foo")
     dvc.experiments.run(run_all=True)
+    test_queue.wait(["foo"])
 
     done_result = first(celery_queue.iter_done())
 
@@ -34,10 +35,6 @@ def test_celery_logs(
     assert "failed to reproduce 'failed-copy-file'" in captured.out
 
 
-@pytest.mark.xfail(
-    strict=False,
-    reason="https://github.com/iterative/dvc/issues/9143",
-)
 def test_queue_remove_done(
     dvc,
     exp_stage,
@@ -63,6 +60,7 @@ def test_queue_remove_done(
             name=name_fail,
         )
     dvc.experiments.run(run_all=True)
+    test_queue.wait(success_tasks + failed_tasks)
     assert len(celery_queue.failed_stash) == 3
     status = to_dict(celery_queue.status())
     assert len(status) == 6
diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py
--- a/tests/func/experiments/test_show.py
+++ b/tests/func/experiments/test_show.py
@@ -156,15 +156,18 @@ def test_show_queued(tmp_dir, scm, dvc, exp_stage):
 
 
 @pytest.mark.vscode
-def test_show_failed_experiment(tmp_dir, scm, dvc, failed_exp_stage):
+def test_show_failed_experiment(tmp_dir, scm, dvc, failed_exp_stage, test_queue):
     baseline_rev = scm.get_rev()
     timestamp = datetime.fromtimestamp(
         scm.gitpython.repo.rev_parse(baseline_rev).committed_date
     )
 
-    dvc.experiments.run(failed_exp_stage.addressing, params=["foo=2"], queue=True)
+    dvc.experiments.run(
+        failed_exp_stage.addressing, params=["foo=2"], queue=True, name="show-failed"
+    )
     exp_rev = dvc.experiments.scm.resolve_rev(f"{CELERY_STASH}@{{0}}")
     dvc.experiments.run(run_all=True)
+    test_queue.wait(["show-failed"])
     experiments = dvc.experiments.show()[baseline_rev]
 
     expected_baseline = {
@@ -542,15 +545,20 @@ def test_show_running_celery(tmp_dir, scm, dvc, exp_stage, mocker):
     assert results["workspace"]["baseline"]["data"]["status"] == "Success"
 
 
-def test_show_running_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, mocker):
+def test_show_running_checkpoint(
+    tmp_dir, scm, dvc, checkpoint_stage, mocker, test_queue
+):
     from dvc.repo.experiments.executor.local import TempDirExecutor
 
     baseline_rev = scm.get_rev()
-    dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"], queue=True)
+    dvc.experiments.run(
+        checkpoint_stage.addressing, params=["foo=2"], queue=True, name="foo"
+    )
     queue = dvc.experiments.celery_queue
     entries = list(queue.iter_queued())
 
     run_results = dvc.experiments.run(run_all=True)
+    test_queue.wait(["foo"])
     checkpoint_rev = first(run_results)
     exp_ref = first(exp_refs_by_rev(scm, checkpoint_rev))