From 28848c5407572ef7bc54a4e7acbd54ea6e918f1e Mon Sep 17 00:00:00 2001
From: Fabian Zills <46721498+PythonFZ@users.noreply.github.com>
Date: Mon, 26 Dec 2022 08:49:46 +0100
Subject: [PATCH] `exp run --run-all`: re-enable --jobs concurrency (#8608)

* re-enable --jobs concurrency

* use `start_workers` instead of `spawn_worker`

* make `spawn_worker` a private method.

* update tests: make `spawn_worker` a private method.
---
 dvc/repo/experiments/__init__.py        | 3 +--
 dvc/repo/experiments/queue/celery.py    | 4 ++--
 tests/unit/command/test_queue.py        | 2 +-
 tests/unit/repo/experiments/conftest.py | 2 +-
 4 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index 2ad99285a..a6b19b90e 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -180,8 +180,7 @@ class Experiments:
         if not entries:
             return results
 
-        # TODO: re-enable --jobs concurrency
-        self.celery_queue.spawn_worker()
+        self.celery_queue.start_workers(count=kwargs.get("jobs", 1))
         failed = []
         try:
             ui.write(
diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py
index b8dca00db..248e74738 100644
--- a/dvc/repo/experiments/queue/celery.py
+++ b/dvc/repo/experiments/queue/celery.py
@@ -116,7 +116,7 @@ class LocalCeleryQueue(BaseStashQueue):
             timeout=10,
         )
 
-    def spawn_worker(self, num: int = 1):
+    def _spawn_worker(self, num: int = 1):
         """spawn one single worker to process to queued tasks.
 
         Argument:
@@ -159,7 +159,7 @@ class LocalCeleryQueue(BaseStashQueue):
             if node_name in active_worker:
                 logger.debug(f"Exp queue worker {node_name} already exist")
                 continue
-            self.spawn_worker(num)
+            self._spawn_worker(num)
             started += 1
 
         return started
diff --git a/tests/unit/command/test_queue.py b/tests/unit/command/test_queue.py
index 7f54e7952..610ee3e6f 100644
--- a/tests/unit/command/test_queue.py
+++ b/tests/unit/command/test_queue.py
@@ -114,7 +114,7 @@ def test_experiments_start(dvc, scm, mocker):
 
     cmd = cli_args.func(cli_args)
     m = mocker.patch(
-        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.spawn_worker",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue._spawn_worker",
     )
 
     assert cmd.run() == 0
diff --git a/tests/unit/repo/experiments/conftest.py b/tests/unit/repo/experiments/conftest.py
index 3c5c43f6d..fe4055374 100644
--- a/tests/unit/repo/experiments/conftest.py
+++ b/tests/unit/repo/experiments/conftest.py
@@ -132,7 +132,7 @@ def test_queue(tmp_dir, dvc, scm, mocker) -> LocalCeleryQueue:
     import celery
 
     queue = dvc.experiments.celery_queue
-    mocker.patch.object(queue, "spawn_worker")
+    mocker.patch.object(queue, "_spawn_worker")
 
     f = partial(
         _thread_worker,