diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index 248e74738..a21615d1f 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -57,8 +57,6 @@ 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) @@ -333,13 +331,26 @@ class LocalCeleryQueue(BaseStashQueue): if remained_revs: raise CannotKillTasksError(remained_revs) + def _kill_entries(self, entries: Dict[QueueEntry, str]): + logger.debug( + "Found active tasks: '%s' to kill", + list(entries.values()), + ) + inactive_entries: Dict[QueueEntry, str] = self._try_to_kill_tasks( + entries + ) + + if inactive_entries: + self._mark_inactive_tasks_failure(inactive_entries) + def kill(self, revs: Collection[str]) -> None: + name_dict: Dict[ str, Optional[QueueEntry] ] = self.match_queue_entry_by_name(set(revs), self.iter_active()) - to_kill: Dict[QueueEntry, str] = {} missing_revs: List[str] = [] + to_kill: Dict[QueueEntry, str] = {} for rev, queue_entry in name_dict.items(): if queue_entry is None: missing_revs.append(rev) @@ -348,21 +359,17 @@ class LocalCeleryQueue(BaseStashQueue): if missing_revs: raise UnresolvedQueueExpNamesError(missing_revs) - fail_to_kill_entries: Dict[QueueEntry, str] = self._try_to_kill_tasks( - to_kill - ) - - if fail_to_kill_entries: - self._mark_inactive_tasks_failure(fail_to_kill_entries) + if to_kill: + self._kill_entries(to_kill) def shutdown(self, kill: bool = False): self.celery.control.shutdown() if kill: - for _, task_entry in self._iter_active_tasks(): - try: - self.proc.kill(task_entry.stash_rev) - except ProcessLookupError: - continue + to_kill: Dict[QueueEntry, str] = {} + for entry in self.iter_active(): + to_kill[entry] = entry.name or entry.stash_rev + if to_kill: + self._kill_entries(to_kill) def follow( self, diff --git a/tests/unit/repo/experiments/queue/test_celery.py b/tests/unit/repo/experiments/queue/test_celery.py index 6dede4d58..2c389aff6 100644 --- a/tests/unit/repo/experiments/queue/test_celery.py +++ b/tests/unit/repo/experiments/queue/test_celery.py @@ -29,27 +29,25 @@ def test_shutdown(test_queue, mocker): def test_shutdown_with_kill(test_queue, mocker): - sig = _foo.s() - mock_entry = mocker.Mock(stash_rev=_foo.name) - - result = sig.freeze() + mock_entry_foo = mocker.Mock(stash_rev="af12de") + mock_entry_foo.name = "foo" + mock_entry_bar = mocker.Mock(stash_rev="bar") + mock_entry_bar.name = None shutdown_spy = mocker.patch("celery.app.control.Control.shutdown") mocker.patch.object( test_queue, - "_iter_active_tasks", - return_value=[(result, mock_entry)], + "iter_active", + return_value=[mock_entry_foo, mock_entry_bar], ) - kill_spy = mocker.patch.object(test_queue.proc, "kill") + kill_spy = mocker.patch.object(test_queue, "_kill_entries") test_queue.shutdown(kill=True) - sig.delay() - - assert result.get() == "foo" - assert result.id not in test_queue._shutdown_task_ids - kill_spy.assert_called_once_with(mock_entry.stash_rev) shutdown_spy.assert_called_once() + kill_spy.assert_called_once_with( + {mock_entry_foo: "foo", mock_entry_bar: "bar"} + ) def test_post_run_after_kill(test_queue):