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"])