diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 2ad99285a1..a6b19b90ef 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -180,8 +180,7 @@ def reproduce_celery( 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 b8dca00db0..248e747385 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -116,7 +116,7 @@ def worker(self) -> "TemporaryWorker": 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 @@ def start_workers(self, count: int) -> int: 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 7f54e7952e..610ee3e6fc 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 3c5c43f6da..fe40553747 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,