exp: implement queue shutdown

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,

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

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

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()