From 686c15dab96aecf12632e70efe33b57bb5541a17 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Fri, 27 May 2022 18:34:32 +0900 Subject: [PATCH] exp: use a single chain for all exp tasks - removes separated scheduling for setup_exp --- dvc/repo/experiments/__init__.py | 1 - dvc/repo/experiments/queue/base.py | 11 +---- dvc/repo/experiments/queue/local.py | 29 +++++++++--- dvc/repo/experiments/queue/tasks.py | 72 ++++++++++++++++++----------- dvc/repo/experiments/show.py | 2 +- 5 files changed, 70 insertions(+), 45 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 69c88d4ec..36586ab9f 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -360,7 +360,6 @@ class Experiments: return exp_baseline raise BaselineMismatchError(exp_baseline, baseline_sha) - @scm_locked def get_baseline(self, rev): """Return the baseline rev for an experiment rev.""" return self._get_baseline(rev) diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 5e58dd581..e95e63248 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -190,15 +190,8 @@ class BaseStashQueue(ABC): """Iterate over items in the queue.""" @abstractmethod - def iter_active( - self, ignore_collected: bool = False - ) -> Generator[QueueEntry, None, None]: - """Iterate over items which are being actively processed. - - Arguments: - ignore_collected: Ignore experiments which have been collected (but - not cleaned up). - """ + def iter_active(self) -> Generator[QueueEntry, None, None]: + """Iterate over items which are being actively processed.""" @abstractmethod def reproduce(self) -> Mapping[str, Mapping[str, str]]: diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index 87953f620..0e38427f7 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -15,6 +15,8 @@ from typing import ( Set, ) +from celery import chain +from celery.canvas import Signature from funcy import cached_property, first from kombu.message import Message @@ -33,7 +35,7 @@ from ..executor.local import WorkspaceExecutor from ..refs import EXEC_BRANCH from ..stash import ExpStashEntry from .base import BaseStashQueue, QueueEntry, QueueGetResult -from .tasks import setup_exp +from .tasks import cleanup_exp, collect_exp, setup_exp if TYPE_CHECKING: from dvc.repo.experiments import Experiments @@ -134,11 +136,26 @@ class LocalCeleryQueue(BaseStashQueue): def put(self, *args, **kwargs) -> QueueEntry: """Stash an experiment and add it to the queue.""" entry = self._stash_exp(*args, **kwargs) - # schedule executor setup - # TODO: chain separated git/dvc setup tasks - self.celery.tasks[setup_exp.name].delay(entry.asdict()) + task_chain = self._chain_exp(entry) + task_chain.freeze() + task_chain.delay() return entry + def _chain_exp(self, entry: QueueEntry) -> Signature: + """Return a complete experiment chain for this queue entry.""" + infofile = self.get_infofile_path(entry.stash_rev) + cmd = ["dvc", "exp", "exec-run", "--infofile", infofile] + return self.celery.signature( + chain( + setup_exp.s(entry.asdict()), + self.proc.run_signature( + cmd, name=entry.stash_rev, immutable=True + ), + collect_exp.s(entry.asdict()), + cleanup_exp.s(entry.asdict()), + ) + ) + # NOTE: Queue consumption should not be done directly. Celery worker(s) # will automatically consume available experiments. def get(self) -> QueueGetResult: @@ -319,9 +336,7 @@ class WorkspaceQueue(BaseStashQueue): entry.name, ) - def iter_active( - self, ignore_collected: bool = False - ) -> Generator[QueueEntry, None, None]: + 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 diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index 1c69a8374..c2912bb64 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -1,6 +1,7 @@ -from typing import Any, Dict +from typing import Any, Dict, List -from celery import chain, shared_task +from celery import shared_task +from celery.signals import task_postrun from celery.utils.log import get_task_logger from dvc.utils.fs import remove @@ -13,18 +14,13 @@ logger = get_task_logger(__name__) @shared_task -def setup_exp(entry_dict: Dict[str, Any]) -> str: - """Setup (queue) an experiment. +def setup_exp(entry_dict: Dict[str, Any]) -> None: + """Setup an experiment. Arguments: entry_dict: Serialized QueueEntry for this experiment. - - Returns: - Celery task ID for the final task in the queued task chain (experiment - cleanup task ID). """ from dvc.repo import Repo - from dvc_task.proc.tasks import run entry = QueueEntry.from_dict(entry_dict) repo = Repo(entry.dvc_root) @@ -33,28 +29,24 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str: executor = BaseStashQueue.setup_executor( repo.experiments, entry, TempDirExecutor ) - proc = repo.experiments.celery_queue.proc infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev) executor.info.dump_json(infofile) - cmd = ["dvc", "exp", "exec-run", "--infofile", infofile] - - # schedule execution + cleanup - exp_chain = chain( - proc.run(cmd, name=entry.stash_rev), - collect_exp.s(entry.asdict(), infofile), - cleanup_exp.si(entry.asdict(), executor.root_dir), - ) - cleanup_task = exp_chain.freeze() - exp_chain.delay() - return cleanup_task.id @shared_task def collect_exp( proc_dict: Dict[str, Any], entry_dict: Dict[str, Any], - infofile: str, -) -> None: +) -> str: + """Collect results for an experiment. + + Arguments: + proc_dict: Serialized ProcessInfo for experiment executor process. + entry_dict: Serialized QueueEntry for this experiment. + + Returns: + Directory to be cleaned up after this experiment. + """ from dvc.repo import Repo from dvc_task.proc.process import ProcessInfo @@ -65,6 +57,7 @@ def collect_exp( entry = QueueEntry.from_dict(entry_dict) repo = Repo(entry.dvc_root) + infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev) executor_info = ExecutorInfo.load_json(infofile) logger.debug("Collecting experiment info '%s'", str(executor_info)) executor = TempDirExecutor.from_info(executor_info) @@ -78,13 +71,38 @@ def collect_exp( logger.debug("Collected experiment '%s'", rev[:7]) else: logger.debug("Exec result was None") - finally: - executor_info.collected = True - executor_info.dump_json(infofile) + except Exception: # pylint: disable=broad-except + # Log exceptions but do not re-raise so that task chain execution + # continues + logger.exception("Failed to collect experiment") + return executor.root_dir @shared_task def cleanup_exp( # pylint: disable=unused-argument - entry_dict: Dict[str, Any], tmp_dir: str + tmp_dir: str, entry_dict: Dict[str, Any] ) -> None: + """Cleanup after an experiment. + + Arguments: + tmp_dir: Temp directory to be removed. + entry_dict: Serialized QueueEntry for this experiment. + """ remove(tmp_dir) + + +@task_postrun.connect(sender=cleanup_exp) +def _cleanup_postrun_handler( + args: List[Any] = None, + **kwargs, +): + from dvc.repo import Repo + + assert args + (_, entry_dict) = args + entry = QueueEntry.from_dict(entry_dict) + repo = Repo(entry.dvc_root) + infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev) + executor_info = ExecutorInfo.load_json(infofile) + executor_info.collected = True + executor_info.dump_json(infofile) diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py index 3e2fbf445..1ed7d76c8 100644 --- a/dvc/repo/experiments/show.py +++ b/dvc/repo/experiments/show.py @@ -195,7 +195,7 @@ def show( ) # collect celery experiments for entry in chain( - repo.experiments.celery_queue.iter_active(ignore_collected=True), + repo.experiments.celery_queue.iter_active(), repo.experiments.celery_queue.iter_queued(), ): stash_rev = entry.stash_rev