queue kill: we can manually mark problematic tasks as failure

fix: #8461
fix: #8473
For now, for those inactive but uncompleted tasks, they still shown as
Running in `exp show` and `task queue`, we cannot remove them because
they are regarded as running, we cann't kill them because their are no
real process behind them.

1. Now we can examine these false running tasks and mark them as failure
   in `queue kill` command. #8461
2. Better msg handle for initliaizing tasks in `queue kill`#8473
3. Add new unit tests for both cases.
4. bump into new dvc-task 0.1.6

dvc/repo/experiments/queue/celery.py

@@ -15,7 +15,7 @@ from typing import (
 )
 
 from celery.result import AsyncResult
-from funcy import cached_property
+from funcy import cached_property, first
 from kombu.message import Message
 
 from dvc.daemon import daemonize
@@ -25,6 +25,7 @@ from dvc.ui import ui
 from ..exceptions import UnresolvedQueueExpNamesError
 from ..executor.base import EXEC_TMP_DIR, ExecutorInfo, ExecutorResult
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
+from .exceptions import CannotKillTasksError
 from .tasks import run_exp
 from .utils import fetch_running_exp_from_temp_dir
 
@@ -286,24 +287,70 @@ class LocalCeleryQueue(BaseStashQueue):
         # out of the active task list, and needs to be loaded here.
         return self._get_done_result(entry, timeout)
 
+    def _get_running_task_ids(self) -> Set[str]:
+        running_task_ids: Set[str] = set()
+        active_workers = self.worker_status()
+        for _, tasks in active_workers.items():
+            task = first(tasks)
+            if task:
+                running_task_ids.add(task["id"])
+        return running_task_ids
+
+    def _try_to_kill_tasks(
+        self, to_kill: Dict[QueueEntry, str]
+    ) -> Dict[QueueEntry, str]:
+        fail_to_kill_entries: Dict[QueueEntry, str] = {}
+        for queue_entry, rev in to_kill.items():
+            try:
+                self.proc.kill(queue_entry.stash_rev)
+                logger.debug(f"Task {rev} had been killed.")
+            except ProcessLookupError:
+                fail_to_kill_entries[queue_entry] = rev
+        return fail_to_kill_entries
+
+    def _mark_inactive_tasks_failure(self, remained_entries):
+        remained_revs: List[str] = []
+        running_ids = self._get_running_task_ids()
+        logger.debug(f"Current running tasks ids: {running_ids}.")
+        for msg, entry in self._iter_processed():
+            if entry not in remained_entries:
+                continue
+            task_id = msg.headers["id"]
+            if task_id in running_ids:
+                remained_revs.append(remained_entries[entry])
+            else:
+                result: AsyncResult = AsyncResult(task_id)
+                if not result.ready():
+                    logger.debug(
+                        f"Task id {task_id} rev {remained_entries[entry]} "
+                        "marked as failure."
+                    )
+                    self.celery.backend.mark_as_failure(task_id, None)
+
+        if remained_revs:
+            raise CannotKillTasksError(remained_revs)
+
     def kill(self, revs: Collection[str]) -> None:
-        to_kill: Set[QueueEntry] = set()
         name_dict: Dict[
             str, Optional[QueueEntry]
         ] = self.match_queue_entry_by_name(set(revs), self.iter_active())
 
-        missing_rev: List[str] = []
+        to_kill: Dict[QueueEntry, str] = {}
+        missing_revs: List[str] = []
         for rev, queue_entry in name_dict.items():
             if queue_entry is None:
-                missing_rev.append(rev)
+                missing_revs.append(rev)
             else:
-                to_kill.add(queue_entry)
+                to_kill[queue_entry] = rev
+        if missing_revs:
+            raise UnresolvedQueueExpNamesError(missing_revs)
 
-        if missing_rev:
-            raise UnresolvedQueueExpNamesError(missing_rev)
+        fail_to_kill_entries: Dict[QueueEntry, str] = self._try_to_kill_tasks(
+            to_kill
+        )
 
-        for queue_entry in to_kill:
-            self.proc.kill(queue_entry.stash_rev)
+        if fail_to_kill_entries:
+            self._mark_inactive_tasks_failure(fail_to_kill_entries)
 
     def shutdown(self, kill: bool = False):
         self.celery.control.shutdown()
@@ -359,10 +406,11 @@ class LocalCeleryQueue(BaseStashQueue):
         ) as fobj:
             ui.write(fobj.read())
 
-    def worker_status(self) -> Dict:
+    def worker_status(self) -> Dict[str, List[Dict]]:
         """Return the current active celery worker"""
         status = self.celery.control.inspect().active() or {}
         logger.debug(f"Worker status: {status}")
+
         return status
 
     def clear(self, *args, **kwargs):

dvc/repo/experiments/queue/exceptions.py

@@ -0,0 +1,12 @@
+from typing import Collection
+
+from dvc.exceptions import DvcException
+
+
+class CannotKillTasksError(DvcException):
+    def __init__(self, revs: Collection[str]):
+        rev_str = ",".join(revs)
+        super().__init__(
+            f"Task '{rev_str}' is initializing, please wait a few seconds "
+            "until the experiments start running to retry the kill operation."
+        )

pyproject.toml

@@ -52,7 +52,7 @@ dependencies = [
     "typing-extensions>=3.7.4",
     "scmrepo==0.1.3",
     "dvc-render==0.0.14",
-    "dvc-task==0.1.5",
+    "dvc-task==0.1.6",
     "dvclive>=1.0",
     "dvc-data==0.27.0",
     "dvc-http==2.27.2",

tests/unit/repo/experiments/queue/test_celery.py

@@ -2,11 +2,12 @@ import time
 
 import pytest
 from celery import shared_task
-from flaky.flaky_decorator import flaky
+from celery.result import AsyncResult
 
 from dvc.exceptions import DvcException
 from dvc.repo.experiments.exceptions import UnresolvedExpNamesError
 from dvc.repo.experiments.queue.base import QueueDoneResult
+from dvc.repo.experiments.queue.exceptions import CannotKillTasksError
 
 
 def test_shutdown_no_tasks(test_queue, mocker):
@@ -51,16 +52,14 @@ def test_shutdown_with_kill(test_queue, mocker):
     shutdown_spy.assert_called_once()
 
 
-# pytest-celery worker thread may finish the task before we check for PENDING
-@flaky(max_runs=3, min_passes=1)
 def test_post_run_after_kill(test_queue):
 
     from celery import chain
 
     sig_bar = test_queue.proc.run_signature(
-        ["python3", "-c", "import time; time.sleep(5)"], name="bar"
+        ["python3", "-c", "import time; time.sleep(10)"], name="bar"
     )
-    result_bar = sig_bar.freeze()
+    sig_bar.freeze()
     sig_foo = _foo.s()
     result_foo = sig_foo.freeze()
     run_chain = chain(sig_bar, sig_foo)
@@ -69,25 +68,28 @@ def test_post_run_after_kill(test_queue):
     timeout = time.time() + 10
 
     while True:
-        if result_bar.status == "STARTED" or result_bar.ready():
+        try:
+            test_queue.proc.kill("bar")
+            assert result_foo.status == "PENDING"
             break
+        except ProcessLookupError:
+            time.sleep(0.1)
         if time.time() > timeout:
-            raise AssertionError()
-
-    assert result_foo.status == "PENDING"
-    test_queue.proc.kill("bar")
+            raise TimeoutError()
 
     assert result_foo.get(timeout=10) == "foo"
 
 
 def test_celery_queue_kill(test_queue, mocker):
 
-    mock_entry = mocker.Mock(stash_rev=_foo.name)
+    mock_entry_foo = mocker.Mock(stash_rev="foo")
+    mock_entry_bar = mocker.Mock(stash_rev="bar")
+    mock_entry_foobar = mocker.Mock(stash_rev="foobar")
 
     mocker.patch.object(
         test_queue,
         "iter_active",
-        return_value={mock_entry},
+        return_value={mock_entry_foo, mock_entry_bar, mock_entry_foobar},
     )
     mocker.patch.object(
         test_queue,
@@ -100,12 +102,54 @@ def test_celery_queue_kill(test_queue, mocker):
     mocker.patch.object(
         test_queue,
         "match_queue_entry_by_name",
-        return_value={"bar": mock_entry},
+        return_value={
+            "bar": mock_entry_bar,
+            "foo": mock_entry_foo,
+            "foobar": mock_entry_foobar,
+        },
+    )
+    mocker.patch.object(
+        test_queue,
+        "_get_running_task_ids",
+        return_value={"foo", "foobar"},
+    )
+    mocker.patch.object(
+        test_queue,
+        "_iter_processed",
+        return_value=[
+            (mocker.Mock(headers={"id": "foo"}), mock_entry_foo),
+            (mocker.Mock(headers={"id": "bar"}), mock_entry_bar),
+            (mocker.Mock(headers={"id": "foobar"}), mock_entry_foobar),
+        ],
+    )
+    mocker.patch.object(
+        AsyncResult,
+        "ready",
+        return_value=False,
+    )
+    mark_mocker = mocker.patch.object(
+        test_queue.celery.backend,
+        "mark_as_failure",
     )
 
-    spy = mocker.patch.object(test_queue.proc, "kill")
-    test_queue.kill("bar")
-    assert spy.called_once_with(mock_entry.stash_rev)
+    def kill_function(rev):
+        if rev == "foo":
+            return True
+        raise ProcessLookupError
+
+    kill_mock = mocker.patch.object(
+        test_queue.proc,
+        "kill",
+        side_effect=mocker.MagicMock(side_effect=kill_function),
+    )
+    with pytest.raises(
+        CannotKillTasksError, match="Task 'foobar' is initializing,"
+    ):
+        test_queue.kill(["bar", "foo", "foobar"])
+    assert kill_mock.called_once_with(mock_entry_foo.stash_rev)
+    assert kill_mock.called_once_with(mock_entry_bar.stash_rev)
+    assert kill_mock.called_once_with(mock_entry_foobar.stash_rev)
+    mark_mocker.assert_called_once_with("bar", None)
 
 
 @pytest.mark.parametrize("status", ["FAILURE", "SUCCESS"])