exp queue: separate queue implementations into their own modules

dvc/repo/experiments/__init__.py

@@ -18,7 +18,8 @@ from .exceptions import (
 )
 from .executor.base import BaseExecutor, ExecutorInfo
 from .queue.base import BaseStashQueue, QueueEntry
-from .queue.local import LocalCeleryQueue, WorkspaceQueue
+from .queue.celery import LocalCeleryQueue
+from .queue.workspace import WorkspaceQueue
 from .refs import (
     CELERY_FAILED_STASH,
     CELERY_STASH,

dvc/repo/experiments/queue/celery.py

@@ -2,7 +2,6 @@ import hashlib
 import locale
 import logging
 import os
-from collections import defaultdict
 from typing import (
     TYPE_CHECKING,
     Collection,
@@ -15,29 +14,25 @@ from typing import (
     Set,
 )
 
-from funcy import cached_property, first
+from funcy import cached_property
 from kombu.message import Message
 
 from dvc.daemon import daemonize
 from dvc.exceptions import DvcException
 from dvc.ui import ui
 
-from ..exceptions import ExpQueueEmptyError, UnresolvedQueueExpNamesError
+from ..exceptions import UnresolvedQueueExpNamesError
 from ..executor.base import (
     EXEC_PID_DIR,
     EXEC_TMP_DIR,
-    BaseExecutor,
     ExecutorInfo,
     ExecutorResult,
 )
-from ..executor.local import WorkspaceExecutor
-from ..refs import EXEC_BRANCH
 from ..stash import ExpStashEntry
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
 from .tasks import run_exp
 
 if TYPE_CHECKING:
-    from dvc.repo.experiments import Experiments
     from dvc_task.app import FSApp
     from dvc_task.proc.manager import ProcessManager
     from dvc_task.worker import TemporaryWorker
@@ -352,128 +347,3 @@ class LocalCeleryQueue(BaseStashQueue):
         status = self.celery.control.inspect().active() or {}
         logger.debug(f"Worker status: {status}")
         return status
-
-
-class WorkspaceQueue(BaseStashQueue):
-    def put(self, *args, **kwargs) -> QueueEntry:
-        return self._stash_exp(*args, **kwargs)
-
-    def get(self) -> QueueGetResult:
-        revs = self.stash.stash_revs
-        if not revs:
-            raise ExpQueueEmptyError("No experiments in the queue.")
-        stash_rev, stash_entry = first(revs.items())
-        entry = QueueEntry(
-            self.repo.root_dir,
-            self.scm.root_dir,
-            self.ref,
-            stash_rev,
-            stash_entry.baseline_rev,
-            stash_entry.branch,
-            stash_entry.name,
-            stash_entry.head_rev,
-        )
-        executor = self.setup_executor(self.repo.experiments, entry)
-        return QueueGetResult(entry, executor)
-
-    def iter_queued(self) -> Generator[QueueEntry, None, None]:
-        for rev, entry in self.stash.stash_revs:
-            yield QueueEntry(
-                self.repo.root_dir,
-                self.scm.root_dir,
-                self.ref,
-                rev,
-                entry.baseline_rev,
-                entry.branch,
-                entry.name,
-                entry.head_rev,
-            )
-
-    def iter_active(self) -> Generator[QueueEntry, None, None]:
-        # Workspace run state is reflected in the workspace itself and does not
-        # need to be handled via the queue
-        raise NotImplementedError
-
-    def iter_done(self) -> Generator[QueueDoneResult, None, None]:
-        raise NotImplementedError
-
-    def reproduce(self) -> Dict[str, Dict[str, str]]:
-        results: Dict[str, Dict[str, str]] = defaultdict(dict)
-        try:
-            while True:
-                entry, executor = self.get()
-                results.update(self._reproduce_entry(entry, executor))
-        except ExpQueueEmptyError:
-            pass
-        return results
-
-    def _reproduce_entry(
-        self, entry: QueueEntry, executor: BaseExecutor
-    ) -> Dict[str, Dict[str, str]]:
-        from dvc.stage.monitor import CheckpointKilledError
-
-        results: Dict[str, Dict[str, str]] = defaultdict(dict)
-        exec_name = "workspace"
-        infofile = self.get_infofile_path(exec_name)
-        try:
-            rev = entry.stash_rev
-            exec_result = executor.reproduce(
-                info=executor.info,
-                rev=rev,
-                infofile=infofile,
-                log_level=logger.getEffectiveLevel(),
-                log_errors=not isinstance(executor, WorkspaceExecutor),
-            )
-            if not exec_result.exp_hash:
-                raise DvcException(
-                    f"Failed to reproduce experiment '{rev[:7]}'"
-                )
-            if exec_result.ref_info:
-                results[rev].update(
-                    self.collect_executor(
-                        self.repo.experiments, executor, exec_result
-                    )
-                )
-        except CheckpointKilledError:
-            # Checkpoint errors have already been logged
-            return {}
-        except DvcException:
-            raise
-        except Exception as exc:
-            raise DvcException(
-                f"Failed to reproduce experiment '{rev[:7]}'"
-            ) from exc
-        finally:
-            executor.cleanup()
-        return results
-
-    @staticmethod
-    def collect_executor(  # pylint: disable=unused-argument
-        exp: "Experiments",
-        executor: BaseExecutor,
-        exec_result: ExecutorResult,
-    ) -> Dict[str, str]:
-        results: Dict[str, str] = {}
-        exp_rev = exp.scm.get_ref(EXEC_BRANCH)
-        if exp_rev:
-            assert exec_result.exp_hash
-            logger.debug("Collected experiment '%s'.", exp_rev[:7])
-            results[exp_rev] = exec_result.exp_hash
-        return results
-
-    def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:
-        raise NotImplementedError
-
-    def kill(self, revs: Collection[str]) -> None:
-        raise NotImplementedError
-
-    def shutdown(self, kill: bool = False):
-        raise NotImplementedError
-
-    def logs(
-        self,
-        rev: str,
-        encoding: Optional[str] = None,
-        follow: bool = False,
-    ):
-        raise NotImplementedError

dvc/repo/experiments/queue/workspace.py

@@ -0,0 +1,143 @@
+import logging
+from collections import defaultdict
+from typing import TYPE_CHECKING, Collection, Dict, Generator, Optional
+
+from funcy import first
+
+from dvc.exceptions import DvcException
+
+from ..exceptions import ExpQueueEmptyError
+from ..executor.base import BaseExecutor, ExecutorResult
+from ..executor.local import WorkspaceExecutor
+from ..refs import EXEC_BRANCH
+from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
+
+if TYPE_CHECKING:
+    from dvc.repo.experiments import Experiments
+
+logger = logging.getLogger(__name__)
+
+
+class WorkspaceQueue(BaseStashQueue):
+    def put(self, *args, **kwargs) -> QueueEntry:
+        return self._stash_exp(*args, **kwargs)
+
+    def get(self) -> QueueGetResult:
+        revs = self.stash.stash_revs
+        if not revs:
+            raise ExpQueueEmptyError("No experiments in the queue.")
+        stash_rev, stash_entry = first(revs.items())
+        entry = QueueEntry(
+            self.repo.root_dir,
+            self.scm.root_dir,
+            self.ref,
+            stash_rev,
+            stash_entry.baseline_rev,
+            stash_entry.branch,
+            stash_entry.name,
+            stash_entry.head_rev,
+        )
+        executor = self.setup_executor(self.repo.experiments, entry)
+        return QueueGetResult(entry, executor)
+
+    def iter_queued(self) -> Generator[QueueEntry, None, None]:
+        for rev, entry in self.stash.stash_revs:
+            yield QueueEntry(
+                self.repo.root_dir,
+                self.scm.root_dir,
+                self.ref,
+                rev,
+                entry.baseline_rev,
+                entry.branch,
+                entry.name,
+                entry.head_rev,
+            )
+
+    def iter_active(self) -> Generator[QueueEntry, None, None]:
+        # Workspace run state is reflected in the workspace itself and does not
+        # need to be handled via the queue
+        raise NotImplementedError
+
+    def iter_done(self) -> Generator[QueueDoneResult, None, None]:
+        raise NotImplementedError
+
+    def reproduce(self) -> Dict[str, Dict[str, str]]:
+        results: Dict[str, Dict[str, str]] = defaultdict(dict)
+        try:
+            while True:
+                entry, executor = self.get()
+                results.update(self._reproduce_entry(entry, executor))
+        except ExpQueueEmptyError:
+            pass
+        return results
+
+    def _reproduce_entry(
+        self, entry: QueueEntry, executor: BaseExecutor
+    ) -> Dict[str, Dict[str, str]]:
+        from dvc.stage.monitor import CheckpointKilledError
+
+        results: Dict[str, Dict[str, str]] = defaultdict(dict)
+        exec_name = "workspace"
+        infofile = self.get_infofile_path(exec_name)
+        try:
+            rev = entry.stash_rev
+            exec_result = executor.reproduce(
+                info=executor.info,
+                rev=rev,
+                infofile=infofile,
+                log_level=logger.getEffectiveLevel(),
+                log_errors=not isinstance(executor, WorkspaceExecutor),
+            )
+            if not exec_result.exp_hash:
+                raise DvcException(
+                    f"Failed to reproduce experiment '{rev[:7]}'"
+                )
+            if exec_result.ref_info:
+                results[rev].update(
+                    self.collect_executor(
+                        self.repo.experiments, executor, exec_result
+                    )
+                )
+        except CheckpointKilledError:
+            # Checkpoint errors have already been logged
+            return {}
+        except DvcException:
+            raise
+        except Exception as exc:
+            raise DvcException(
+                f"Failed to reproduce experiment '{rev[:7]}'"
+            ) from exc
+        finally:
+            executor.cleanup()
+        return results
+
+    @staticmethod
+    def collect_executor(  # pylint: disable=unused-argument
+        exp: "Experiments",
+        executor: BaseExecutor,
+        exec_result: ExecutorResult,
+    ) -> Dict[str, str]:
+        results: Dict[str, str] = {}
+        exp_rev = exp.scm.get_ref(EXEC_BRANCH)
+        if exp_rev:
+            assert exec_result.exp_hash
+            logger.debug("Collected experiment '%s'.", exp_rev[:7])
+            results[exp_rev] = exec_result.exp_hash
+        return results
+
+    def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:
+        raise NotImplementedError
+
+    def kill(self, revs: Collection[str]) -> None:
+        raise NotImplementedError
+
+    def shutdown(self, kill: bool = False):
+        raise NotImplementedError
+
+    def logs(
+        self,
+        rev: str,
+        encoding: Optional[str] = None,
+        follow: bool = False,
+    ):
+        raise NotImplementedError

tests/unit/command/test_queue.py

@@ -21,7 +21,7 @@ def test_experiments_remove(dvc, scm, mocker):
 
     cmd = cli_args.func(cli_args)
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.clear",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.clear",
         return_value={},
     )
 
@@ -40,7 +40,7 @@ def test_experiments_remove(dvc, scm, mocker):
 
     cmd = cli_args.func(cli_args)
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.remove",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.remove",
         return_value={},
     )
 
@@ -61,7 +61,7 @@ def test_experiments_kill(dvc, scm, mocker):
 
     cmd = cli_args.func(cli_args)
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.kill",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.kill",
         return_value={},
     )
 
@@ -75,7 +75,7 @@ def test_experiments_start(dvc, scm, mocker):
 
     cmd = cli_args.func(cli_args)
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.spawn_worker",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.spawn_worker",
     )
 
     assert cmd.run() == 0
@@ -94,7 +94,7 @@ def test_experiments_stop(dvc, scm, mocker):
 
     cmd = cli_args.func(cli_args)
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.shutdown",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.shutdown",
     )
 
     assert cmd.run() == 0
@@ -134,11 +134,11 @@ def test_worker_status(dvc, scm, worker_status, output, mocker, capsys):
 
     cmd = cli_args.func(cli_args)
     mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.status",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.status",
         return_value=[],
     )
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.worker_status",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.worker_status",
         return_value=worker_status,
     )
 
@@ -176,7 +176,7 @@ def test_experiments_status(dvc, scm, mocker, capsys):
         },
     ]
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.status",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.status",
         return_value=status_result,
     )
 
@@ -194,7 +194,7 @@ def test_queue_logs(dvc, scm, mocker):
 
     cmd = cli_args.func(cli_args)
     m = mocker.patch(
-        "dvc.repo.experiments.queue.local.LocalCeleryQueue.logs",
+        "dvc.repo.experiments.queue.celery.LocalCeleryQueue.logs",
         return_value={},
     )
 

tests/unit/repo/experiments/conftest.py

@@ -2,7 +2,7 @@ from contextlib import contextmanager
 
 import pytest
 
-from dvc.repo.experiments.queue.local import LocalCeleryQueue
+from dvc.repo.experiments.queue.celery import LocalCeleryQueue
 
 
 @contextmanager