From 2b0c365f35c90042cc4bc816ed8cec2a863d90b8 Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Fri, 3 Feb 2023 18:26:27 +0800 Subject: [PATCH] Fix some mypy error. --- dvc/repo/experiments/queue/celery.py | 5 ++++- dvc/repo/experiments/queue/remove.py | 5 +++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index d46662d2c..07598b2a6 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -336,7 +336,10 @@ class LocalCeleryQueue(BaseStashQueue): task_id, remained_entries[entry], ) - self.celery.backend.mark_as_failure(task_id, None) + backend = self.celery.backend + backend.mark_as_failure( # type: ignore[attr-defined] + task_id, None + ) if remained_revs: raise CannotKillTasksError(remained_revs) diff --git a/dvc/repo/experiments/queue/remove.py b/dvc/repo/experiments/queue/remove.py index ef857280c..34f6f106c 100644 --- a/dvc/repo/experiments/queue/remove.py +++ b/dvc/repo/experiments/queue/remove.py @@ -40,7 +40,7 @@ def remove_tasks( # noqa: C901 msg, queue_entry, ) in celery_queue._iter_queued(): # pylint: disable=protected-access - if queue_entry.stash_rev in stash_revs: + if queue_entry.stash_rev in stash_revs and msg.delivery_tag: celery_queue.celery.reject(msg.delivery_tag) finally: celery_queue.stash.remove_revs(list(stash_revs.values())) @@ -56,7 +56,8 @@ def remove_tasks( # noqa: C901 result: AsyncResult = AsyncResult(task_id) if result is not None: result.forget() - celery_queue.celery.purge(msg.delivery_tag) + if msg.delivery_tag: + celery_queue.celery.purge(msg.delivery_tag) finally: if celery_queue.failed_stash: celery_queue.failed_stash.remove_revs(failed_stash_revs)