diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 620da7808..1a29575c6 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -170,6 +170,16 @@ class BaseStashQueue(ABC): This method blocks until the specified item has been collected. """ + @abstractmethod + def shutdown(self, kill: bool = False): + """Shutdown the queue worker. + + Arguments: + kill: If True, the any active experiments will be killed and the + worker will shutdown immediately. If False, the worker will + finish any active experiments before shutting down. + """ + def _stash_exp( self, *args, diff --git a/dvc/repo/experiments/queue/local.py b/dvc/repo/experiments/queue/local.py index 7d598d66d..213994904 100644 --- a/dvc/repo/experiments/queue/local.py +++ b/dvc/repo/experiments/queue/local.py @@ -11,6 +11,7 @@ from typing import ( Mapping, NamedTuple, Optional, + Set, ) from funcy import cached_property, first @@ -47,6 +48,11 @@ class _MessageEntry(NamedTuple): entry: QueueEntry +class _TaskEntry(NamedTuple): + task_id: str + entry: QueueEntry + + class LocalCeleryQueue(BaseStashQueue): """DVC experiment queue. @@ -55,6 +61,8 @@ class LocalCeleryQueue(BaseStashQueue): CELERY_DIR = "celery" + _shutdown_task_ids: Set[str] = set() + @cached_property def wdir(self) -> str: return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, self.CELERY_DIR) @@ -139,19 +147,26 @@ class LocalCeleryQueue(BaseStashQueue): entry_dict = kwargs.get("entry_dict", args[0]) yield _MessageEntry(msg, QueueEntry.from_dict(entry_dict)) - def _iter_processed(self) -> Generator[QueueEntry, None, None]: + def _iter_processed(self) -> Generator[_MessageEntry, None, None]: for msg in self.celery.iter_processed(): if msg.headers.get("task") != setup_exp.name: continue args, kwargs, _embed = msg.decode() entry_dict = kwargs.get("entry_dict", args[0]) - yield QueueEntry.from_dict(entry_dict) + yield _MessageEntry(msg, QueueEntry.from_dict(entry_dict)) + + def _iter_active_tasks(self) -> Generator[_TaskEntry, None, None]: + from celery.result import AsyncResult + + for msg, entry in self._iter_processed(): + task_id = msg.headers["id"] + result = AsyncResult(task_id) + if not result.ready(): + yield _TaskEntry(task_id, entry) def iter_active(self) -> Generator[QueueEntry, None, None]: - for entry in self._iter_processed(): - proc_info = self.proc.get(entry.stash_rev) - if proc_info is not None and proc_info.returncode is None: - yield entry + for _, entry in self._iter_active_tasks(): + yield entry def reproduce(self) -> Mapping[str, Mapping[str, str]]: raise NotImplementedError @@ -168,6 +183,25 @@ class LocalCeleryQueue(BaseStashQueue): pass time.sleep(1) + def _shutdown_handler(self, task_id: str = None, **kwargs): + if task_id in self._shutdown_task_ids: + self._shutdown_task_ids.remove(task_id) + if not self._shutdown_task_ids: + self.celery.control.shutdown() + + def shutdown(self, kill: bool = False): + from celery.signals import task_postrun + + if kill: + raise NotImplementedError + tasks = list(self._iter_active_tasks()) + if tasks: + for task_id, _ in tasks: + self._shutdown_task_ids.add(task_id) + task_postrun.connect()(self._shutdown_handler) + else: + self.celery.control.shutdown() + class WorkspaceQueue(BaseStashQueue): def put(self, *args, **kwargs) -> QueueEntry: @@ -284,3 +318,6 @@ class WorkspaceQueue(BaseStashQueue): def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]: raise NotImplementedError + + def shutdown(self, kill: bool = False): + raise NotImplementedError diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index 9104e39f2..e3fc60dae 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/dvc/repo/experiments/queue/tasks.py @@ -13,7 +13,15 @@ logger = get_task_logger(__name__) @shared_task -def setup_exp(entry_dict: Dict[str, Any]) -> None: +def setup_exp(entry_dict: Dict[str, Any]) -> str: + """Setup (queue) an experiment. + + Arguments: + entry_dict: Serialized QueueEntry for this experiment. + + Returns: + Celery task ID for the queued experiment task chain. + """ from dvc.repo import Repo from dvc_task.proc.tasks import run @@ -30,11 +38,14 @@ def setup_exp(entry_dict: Dict[str, Any]) -> None: cmd = ["dvc", "exp", "exec-run", "--infofile", infofile] # schedule execution + cleanup - chain( + exp_chain = chain( proc.run(cmd, name=entry.stash_rev), collect_exp.s(entry.asdict(), infofile), cleanup_exp.si(entry.asdict(), executor.root_dir), - ).delay() + ) + exp_chain.freeze() + exp_chain.delay() + return exp_chain.id @shared_task diff --git a/tests/unit/repo/experiments/queue/test_local.py b/tests/unit/repo/experiments/queue/test_local.py new file mode 100644 index 000000000..b8d6337aa --- /dev/null +++ b/tests/unit/repo/experiments/queue/test_local.py @@ -0,0 +1,37 @@ +from celery import shared_task + + +def test_shutdown_no_tasks(test_queue, mocker): + shutdown_spy = mocker.spy(test_queue.celery.control, "shutdown") + test_queue.shutdown() + shutdown_spy.assert_called_once() + + +@shared_task +def _foo(): + return "foo" + + +def test_shutdown_active_tasks(test_queue, mocker): + + shutdown_spy = mocker.spy(test_queue.celery.control, "shutdown") + # use frozen signature to get an assigned task ID and AsyncResult object + # before running the task + sig = _foo.s() + result = sig.freeze() + mocker.patch.object( + test_queue, "_iter_active_tasks", return_value=[(result.id, None)] + ) + + # shutdown signal handler should be registered but not executed since task + # has not completed + test_queue.shutdown() + assert result.id in test_queue._shutdown_task_ids + assert not result.ready() + shutdown_spy.assert_not_called() + + sig.delay() + # block until task completes + assert "foo" == result.get() + assert result.id not in test_queue._shutdown_task_ids + shutdown_spy.assert_called_once()