From a167dff01959ade268c9bf1fe027a360900b08c2 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Wed, 23 Feb 2022 16:14:48 +0900 Subject: [PATCH] exp: implement celery exp task collection --- dvc/repo/experiments/queue/base.py | 45 ++++++++++++++++++++-- dvc/repo/experiments/queue/local.py | 59 +++++++++++++++++++---------- dvc/repo/experiments/queue/tasks.py | 58 ++++++++++++++++++++++++---- 3 files changed, 129 insertions(+), 33 deletions(-) diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index f54083730..bfd836150 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -18,7 +18,13 @@ from dvc.dependency.param import MissingParamsError from dvc.env import DVCLIVE_RESUME from dvc.exceptions import DvcException -from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR, BaseExecutor +from ..exceptions import CheckpointExistsError, ExperimentExistsError +from ..executor.base import ( + EXEC_PID_DIR, + EXEC_TMP_DIR, + BaseExecutor, + ExecutorResult, +) from ..executor.local import WorkspaceExecutor from ..refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE, ExpRefInfo from ..stash import ExpStash, ExpStashEntry @@ -105,6 +111,10 @@ class BaseStashQueue(ABC): def get(self) -> QueueGetResult: """Pop and return the first item in the queue.""" + @abstractmethod + def reproduce(self) -> Mapping[str, Mapping[str, str]]: + """Reproduce queued experiments sequentially.""" + def _stash_exp( self, *args, @@ -392,6 +402,33 @@ class BaseStashQueue(ABC): f"{name}{BaseExecutor.INFOFILE_EXT}", ) - @abstractmethod - def reproduce(self) -> Mapping[str, Mapping[str, str]]: - """Reproduce queued experiments sequentially.""" + @staticmethod + @scm_locked + def collect_executor( + exp: "Experiments", + executor: BaseExecutor, + exec_result: ExecutorResult, + ) -> Dict[str, str]: + results = {} + + def on_diverged(ref: str, checkpoint: bool): + ref_info = ExpRefInfo.from_ref(ref) + if checkpoint: + raise CheckpointExistsError(ref_info.name) + raise ExperimentExistsError(ref_info.name) + + for ref in executor.fetch_exps( + exp.scm, + force=exec_result.force, + on_diverged=on_diverged, + ): + exp_rev = exp.scm.get_ref(ref) + if exp_rev: + assert exec_result.exp_hash + logger.debug("Collected experiment '%s'.", exp_rev[:7]) + results[exp_rev] = exec_result.exp_hash + + if exec_result.ref_info is not None: + executor.collect_cache(exp.repo, exec_result.ref_info) + + return results diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index e7e6f1d57..d172c3ef7 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -8,15 +8,15 @@ from funcy import cached_property, first from dvc.exceptions import DvcException from ..exceptions import ExpQueueEmptyError -from ..executor.base import EXEC_TMP_DIR, BaseExecutor +from ..executor.base import EXEC_TMP_DIR, BaseExecutor, ExecutorResult from ..executor.local import WorkspaceExecutor from ..refs import EXEC_BRANCH from .base import BaseStashQueue, QueueEntry, QueueGetResult from .tasks import setup_exp if TYPE_CHECKING: - from celery import Celery - + from dvc.repo.experiments import Experiments + from dvc_task.app import FSApp from dvc_task.worker import TemporaryWorker logger = logging.getLogger(__name__) @@ -35,31 +35,42 @@ class LocalCeleryQueue(BaseStashQueue): return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, self.CELERY_DIR) @cached_property - def celery(self) -> "Celery": - from celery import Celery + def celery(self) -> "FSApp": + from dvc_task.app import FSApp - from .celery import get_config - - app = Celery() - app.conf.update(get_config(self.wdir)) + app = FSApp( + "dvc-exp-local", + wdir=self.wdir, + mkdir=True, + include=[ + "dvc.repo.experiments.queue.tasks", + "dvc_task.proc.tasks", + ], + ) return app @cached_property def worker(self) -> "TemporaryWorker": from dvc_task.worker import TemporaryWorker - return TemporaryWorker(self.celery, concurrency=1) + return TemporaryWorker(self.celery, concurrency=1, timeout=10) + + def spawn_worker(self): + from dvc_task.proc.process import ManagedProcess + + logger.debug("Spawning exp queue worker") + ManagedProcess.spawn( + ["dvc", "exp", "queue-worker", "dvc-exp1@localhost"], + wdir=self.wdir, + name="dvc-exp-worker", + ) def put(self, *args, **kwargs) -> QueueEntry: """Stash an experiment and add it to the queue.""" - from celery import chain - - entry = super()._stash_exp(*args, **kwargs) + entry = self._stash_exp(*args, **kwargs) # schedule executor setup # TODO: chain separated git/dvc setup tasks - chain( - self.celery.tasks[setup_exp.name](entry.asdict()), - ).delay() + self.celery.tasks[setup_exp.name].delay(entry.asdict()) return entry # NOTE: Queue consumption should not be done directly. Celery worker(s) @@ -125,7 +136,9 @@ class WorkspaceQueue(BaseStashQueue): ) if exec_result.ref_info: results[rev].update( - self._collect_executor(executor, exec_result) + self.collect_executor( + self.repo.experiments, executor, exec_result + ) ) except CheckpointKilledError: # Checkpoint errors have already been logged @@ -140,12 +153,16 @@ class WorkspaceQueue(BaseStashQueue): executor.cleanup() return results - def _collect_executor( # pylint: disable=unused-argument - self, executor, exec_result + @staticmethod + def collect_executor( # pylint: disable=unused-argument + exp: "Experiments", + executor: BaseExecutor, + exec_result: ExecutorResult, ) -> Dict[str, str]: - results = {} - exp_rev = self.scm.get_ref(EXEC_BRANCH) + 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 diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index e5472ae5a..e04d4f526 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -3,13 +3,16 @@ import os from typing import Any, Dict from celery import chain, shared_task +from celery.utils.log import get_task_logger -from dvc.utils.fs import remove +from dvc.utils.fs import makedirs, remove -from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR +from ..executor.base import EXEC_PID_DIR, EXEC_TMP_DIR, ExecutorInfo from ..executor.local import TempDirExecutor from .base import BaseStashQueue, QueueEntry +logger = get_task_logger(__name__) + @shared_task def setup_exp(entry_dict: Dict[str, Any]) -> None: @@ -23,22 +26,61 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None: executor = BaseStashQueue.setup_executor( repo.experiments, entry, TempDirExecutor ) - infofile = os.path.join(repo.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR) + pid_dir = os.path.join( + repo.tmp_dir, + EXEC_TMP_DIR, + EXEC_PID_DIR, + ) + infofile = os.path.join( + pid_dir, + f"{entry.stash_rev}{TempDirExecutor.INFOFILE_EXT}", + ) + makedirs(os.path.dirname(infofile), exist_ok=True) with open(infofile, "w", encoding="utf-8") as fobj: json.dump(executor.info.asdict(), fobj) cmd = ["dvc", "exp", "exec-run", "--infofile", infofile] # schedule execution + cleanup chain( - run(cmd), # pylint: disable=no-value-for-parameter - cleanup_exp( # pylint: disable=no-value-for-parameter - executor.root_dir, entry.asdict() - ), + run.si(cmd, wdir=pid_dir, name=entry.stash_rev), + collect_exp.s(entry.asdict(), infofile), + cleanup_exp.si(entry.asdict(), executor.root_dir), ).delay() @shared_task +def collect_exp( + proc_dict: Dict[str, Any], + entry_dict: Dict[str, Any], + infofile: str, +) -> None: + from dvc.repo import Repo + from dvc_task.proc.process import ProcessInfo + + proc_info = ProcessInfo.from_dict(proc_dict) + if proc_info.returncode != 0: + # TODO: handle errors, track failed exps separately + pass + + entry = QueueEntry.from_dict(entry_dict) + repo = Repo(entry.dvc_root) + with open(infofile, encoding="utf-8") as fobj: + executor_info = ExecutorInfo.from_dict(json.load(fobj)) + logger.debug("Collecting experiment info '%s'", str(executor_info)) + executor = TempDirExecutor.from_info(executor_info) + exec_result = executor_info.result + if exec_result is not None: + results = BaseStashQueue.collect_executor( + repo.experiments, executor, exec_result + ) + for rev in results: + logger.debug("Collected experiment '%s'", rev[:7]) + else: + logger.debug("Exec result was None") + + +@shared_task(ignore_result=True) def cleanup_exp( # pylint: disable=unused-argument - proc_dict: Dict[str, Any], entry: Dict[str, Any], tmp_dir: str + entry_dict: Dict[str, Any], tmp_dir: str ) -> None: remove(tmp_dir)