diff --git a/dvc/repo/experiments/exceptions.py b/dvc/repo/experiments/exceptions.py index a81b691f9..a609d2adb 100644 --- a/dvc/repo/experiments/exceptions.py +++ b/dvc/repo/experiments/exceptions.py @@ -128,5 +128,9 @@ class UnresolvedQueueExpNamesError(UnresolvedExpNamesError): NAME = "queued experiment name" +class UnresolvedRunningExpNamesError(UnresolvedExpNamesError): + NAME = "running experiment name" + + class ExpQueueEmptyError(DvcException): pass diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index 271b0aa1e..8d2cef4c9 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -21,7 +21,10 @@ from kombu.message import Message from dvc.daemon import daemonize from dvc.exceptions import DvcException -from dvc.repo.experiments.exceptions import UnresolvedQueueExpNamesError +from dvc.repo.experiments.exceptions import ( + UnresolvedQueueExpNamesError, + UnresolvedRunningExpNamesError, +) from dvc.repo.experiments.executor.base import ExecutorInfo, ExecutorResult from dvc.repo.experiments.refs import CELERY_STASH from dvc.repo.experiments.utils import EXEC_TMP_DIR, get_exp_rwlock @@ -317,7 +320,7 @@ class LocalCeleryQueue(BaseStashQueue): self.proc.kill(queue_entry.stash_rev) else: self.proc.interrupt(queue_entry.stash_rev) - logger.debug(f"Task {rev} had been killed.") + ui.write(f"{rev} has been killed.") except ProcessLookupError: fail_to_kill_entries[queue_entry] = rev return fail_to_kill_entries @@ -373,12 +376,12 @@ class LocalCeleryQueue(BaseStashQueue): else: to_kill[queue_entry] = rev - if missing_revs: - raise UnresolvedQueueExpNamesError(missing_revs) - if to_kill: self._kill_entries(to_kill, force) + if missing_revs: + raise UnresolvedRunningExpNamesError(missing_revs) + def shutdown(self, kill: bool = False): self.celery.control.shutdown() if kill: @@ -406,8 +409,8 @@ class LocalCeleryQueue(BaseStashQueue): {rev}, self.iter_active(), self.iter_done() ).get(rev) if queue_entry is None: - if rev in self.match_queue_entry_by_name( - {rev}, self.iter_queued() + if self.match_queue_entry_by_name({rev}, self.iter_queued()).get( + rev ): raise DvcException( f"Experiment '{rev}' is in queue but has not been started" diff --git a/tests/unit/repo/experiments/queue/test_celery.py b/tests/unit/repo/experiments/queue/test_celery.py index 14734ed0b..c867dce32 100644 --- a/tests/unit/repo/experiments/queue/test_celery.py +++ b/tests/unit/repo/experiments/queue/test_celery.py @@ -90,14 +90,6 @@ def test_celery_queue_kill(test_queue, mocker, force): "iter_active", return_value={mock_entry_foo, mock_entry_bar, mock_entry_foobar}, ) - mocker.patch.object( - test_queue, - "match_queue_entry_by_name", - return_value={"bar": None}, - ) - with pytest.raises(UnresolvedExpNamesError): - test_queue.kill("bar") - mocker.patch.object( test_queue, "match_queue_entry_by_name", @@ -151,6 +143,31 @@ def test_celery_queue_kill(test_queue, mocker, force): mark_mocker.assert_called_once_with("bar", None) +@pytest.mark.parametrize("force", [True, False]) +def test_celery_queue_kill_invalid(test_queue, mocker, force): + + mock_entry_foo = mocker.Mock(stash_rev="foo") + mock_entry_bar = mocker.Mock(stash_rev="bar") + + mocker.patch.object( + test_queue, + "match_queue_entry_by_name", + return_value={ + "bar": mock_entry_bar, + "foo": mock_entry_foo, + "foobar": None, + }, + ) + + kill_mock = mocker.patch.object(test_queue, "_kill_entries") + + with pytest.raises(UnresolvedExpNamesError): + test_queue.kill(["bar", "foo", "foobar"], force=force) + assert kill_mock.called_once_with( + {mock_entry_foo: "foo", mock_entry_bar: "bar"}, force + ) + + @pytest.mark.parametrize("status", ["FAILURE", "SUCCESS"]) def test_queue_iter_done_task(test_queue, mocker, status):