From ce4cee184eb6a6893b154cb54980a58e3a2f6fd2 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Tue, 14 Jun 2022 15:23:38 +0900 Subject: [PATCH] exp queue: separate queue implementations into their own modules --- dvc/repo/experiments/__init__.py | 3 +- .../experiments/queue/{local.py => celery.py} | 134 +--------------- dvc/repo/experiments/queue/workspace.py | 143 ++++++++++++++++++ tests/unit/command/test_queue.py | 18 +-- tests/unit/repo/experiments/conftest.py | 2 +- 5 files changed, 157 insertions(+), 143 deletions(-) rename dvc/repo/experiments/queue/{local.py => celery.py} (72%) create mode 100644 dvc/repo/experiments/queue/workspace.py diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 34c432696..8647d6f3c 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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, diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/celery.py similarity index 72% rename from dvc/repo/experiments/queue/local.py rename to dvc/repo/experiments/queue/celery.py index 1f5db766a..b8ef0ded8 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/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 diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py new file mode 100644 index 000000000..477e2676b --- /dev/null +++ b/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 diff --git a/tests/unit/command/test_queue.py b/tests/unit/command/test_queue.py index 332d0c1e4..35d4ae70b 100644 --- a/tests/unit/command/test_queue.py +++ b/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={}, ) diff --git a/tests/unit/repo/experiments/conftest.py b/tests/unit/repo/experiments/conftest.py index 9837c5f1c..d5ae0babc 100644 --- a/tests/unit/repo/experiments/conftest.py +++ b/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