From fa54c1af7c0970af2b42eba8cd48d446c0a00926 Mon Sep 17 00:00:00 2001
From: karajan1001 <mishanyo1001@gmail.com>
Date: Wed, 12 Oct 2022 11:17:27 +0800
Subject: [PATCH] Solve run_all flaky tests

1. Modify run all to include currently running exps.
2. bump dvc-task to 0.1.5
---
 dvc/repo/experiments/__init__.py           | 15 +++++++++++++--
 dvc/repo/experiments/queue/celery.py       |  3 +++
 dvc/repo/experiments/run.py                |  3 +--
 tests/func/experiments/test_experiments.py |  2 +-
 tests/func/experiments/test_show.py        |  3 +++
 5 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index e003acc93..287490fcb 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -4,7 +4,7 @@ import re
 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 @@ class Experiments:
     ) -> 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
 
diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py
index 6018e01a9..12cf3f592 100644
--- a/dvc/repo/experiments/queue/celery.py
+++ b/dvc/repo/experiments/queue/celery.py
@@ -182,6 +182,7 @@ class LocalCeleryQueue(BaseStashQueue):
                 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 @@ class LocalCeleryQueue(BaseStashQueue):
             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 @@ class LocalCeleryQueue(BaseStashQueue):
             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]:
diff --git a/dvc/repo/experiments/run.py b/dvc/repo/experiments/run.py
index 38361f063..d03f2ea27 100644
--- a/dvc/repo/experiments/run.py
+++ b/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:
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index c0b2665c7..3eb7d6419 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/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()
diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py
index 490fff646..92f3aabbb 100644
--- a/tests/func/experiments/test_show.py
+++ b/tests/func/experiments/test_show.py
@@ -646,6 +646,9 @@ def test_show_csv(tmp_dir, scm, dvc, exp_stage, capsys):
     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)