From 350e0fb1f7ddb39bc540143070c7d81c45fe57f8 Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Fri, 18 Nov 2022 09:59:41 +0800 Subject: [PATCH] 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 | 68 ++++++++++++++--- dvc/repo/experiments/queue/exceptions.py | 12 +++ pyproject.toml | 2 +- tests/unit/repo/experiments/queue/__init__.py | 0 .../queue/{test_local.py => test_celery.py} | 76 +++++++++++++++---- 5 files changed, 131 insertions(+), 27 deletions(-) create mode 100644 dvc/repo/experiments/queue/exceptions.py create mode 100644 tests/unit/repo/experiments/queue/__init__.py rename tests/unit/repo/experiments/queue/{test_local.py => test_celery.py} (70%) diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index 12cf3f592..d151ce200 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/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): diff --git a/dvc/repo/experiments/queue/exceptions.py b/dvc/repo/experiments/queue/exceptions.py new file mode 100644 index 000000000..7bc29ad4a --- /dev/null +++ b/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." + ) diff --git a/pyproject.toml b/pyproject.toml index c3ee86cd3..626baa3fd 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/unit/repo/experiments/queue/__init__.py b/tests/unit/repo/experiments/queue/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/repo/experiments/queue/test_local.py b/tests/unit/repo/experiments/queue/test_celery.py similarity index 70% rename from tests/unit/repo/experiments/queue/test_local.py rename to tests/unit/repo/experiments/queue/test_celery.py index 19e1dc1de..6dede4d58 100644 --- a/tests/unit/repo/experiments/queue/test_local.py +++ b/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"])