Solve run_all flaky tests

1. Modify run all to include currently running exps.
2. bump dvc-task to 0.1.5

Author: karajan1001

dvc/repo/experiments/__init__.py

@@ -4,7 +4,7 @@
 import time
 from typing import Dict, Iterable, Optional
 
-from funcy import cached_property, first
+from funcy import cached_property, chain, first
 
 from dvc.exceptions import DvcException
 from dvc.ui import ui
@@ -171,7 +171,18 @@ def reproduce_celery(
     ) -> Dict[str, str]:
         results: Dict[str, str] = {}
         if entries is None:
-            entries = list(self.celery_queue.iter_queued())
+            entries = list(
+                chain(
+                    self.celery_queue.iter_active(),
+                    self.celery_queue.iter_queued(),
+                )
+            )
+
+        logger.debug(
+            "reproduce all these entries '%s'",
+            entries,
+        )
+
         if not entries:
             return results
 

dvc/repo/experiments/queue/celery.py

@@ -182,6 +182,7 @@ def _iter_queued(self) -> Generator[_MessageEntry, None, None]:
                 continue
             args, kwargs, _embed = msg.decode()
             entry_dict = kwargs.get("entry_dict", args[0])
+            logger.debug("Found queued task %s", entry_dict["stash_rev"])
             yield _MessageEntry(msg, QueueEntry.from_dict(entry_dict))
 
     def _iter_processed(self) -> Generator[_MessageEntry, None, None]:
@@ -198,6 +199,7 @@ def _iter_active_tasks(self) -> Generator[_TaskEntry, None, None]:
             task_id = msg.headers["id"]
             result: AsyncResult = AsyncResult(task_id)
             if not result.ready():
+                logger.debug("Found active task %s", entry.stash_rev)
                 yield _TaskEntry(result, entry)
 
     def _iter_done_tasks(self) -> Generator[_TaskEntry, None, None]:
@@ -206,6 +208,7 @@ def _iter_done_tasks(self) -> Generator[_TaskEntry, None, None]:
             task_id = msg.headers["id"]
             result: AsyncResult = AsyncResult(task_id)
             if result.ready():
+                logger.debug("Found done task %s", entry.stash_rev)
                 yield _TaskEntry(result, entry)
 
     def iter_active(self) -> Generator[QueueEntry, None, None]:

dvc/repo/experiments/run.py

@@ -29,8 +29,7 @@ def run(
     of `repro` for that experiment.
     """
     if run_all:
-        entries = list(repo.experiments.celery_queue.iter_queued())
-        return repo.experiments.reproduce_celery(entries, jobs=jobs)
+        return repo.experiments.reproduce_celery(jobs=jobs)
 
     hydra_sweep = None
     if params:

tests/func/experiments/test_experiments.py

@@ -536,7 +536,7 @@ def test_run_celery(tmp_dir, scm, dvc, exp_stage, mocker):
     repro_spy = mocker.spy(dvc.experiments, "reproduce_celery")
     results = dvc.experiments.run(run_all=True)
     assert len(results) == 2
-    repro_spy.assert_called_once_with(entries=mocker.ANY, jobs=1)
+    repro_spy.assert_called_once_with(jobs=1)
 
     expected = {"foo: 2", "foo: 3"}
     metrics = set()

tests/func/experiments/test_show.py

@@ -646,6 +646,9 @@ def _get_rev_isotimestamp(rev):
     result1 = dvc.experiments.run(exp_stage.addressing, params=["foo=2"])
     rev1 = first(result1)
     ref_info1 = first(exp_refs_by_rev(scm, rev1))
+
+    # at least 1 second gap between these experiments to make sure
+    # the previous experiment to be regarded as branch_base
     time.sleep(1)
     result2 = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
     rev2 = first(result2)