Function remove_tasks

dvc/repo/experiments/queue/remove.py:13–58

Remove tasks from task queue. Arguments: queue_entries: An iterable list of task to remove

(  # noqa: C901, PLR0912
    celery_queue: "LocalCeleryQueue",
    queue_entries: Iterable["QueueEntry"],
)

def remove_tasks(  # noqa: C901, PLR0912
    celery_queue: "LocalCeleryQueue",
    queue_entries: Iterable["QueueEntry"],
):
    """Remove tasks from task queue.

    Arguments:
        queue_entries: An iterable list of task to remove
    """
    from celery.result import AsyncResult

    stash_revs: dict[str, ExpStashEntry] = {}
    failed_stash_revs: list[ExpStashEntry] = []
    done_entry_set: set[QueueEntry] = set()
    stash_rev_all = celery_queue.stash.stash_revs
    failed_rev_all: dict[str, ExpStashEntry] = {}
    if celery_queue.failed_stash:
        failed_rev_all = celery_queue.failed_stash.stash_revs
    for entry in queue_entries:
        if entry.stash_rev in stash_rev_all:
            stash_revs[entry.stash_rev] = stash_rev_all[entry.stash_rev]
        else:
            done_entry_set.add(entry)
            if entry.stash_rev in failed_rev_all:
                failed_stash_revs.append(failed_rev_all[entry.stash_rev])

    try:
        for msg, queue_entry in celery_queue._iter_queued():
            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()))

    try:
        for msg, queue_entry in celery_queue._iter_processed():
            if queue_entry not in done_entry_set:
                continue
            task_id = msg.headers["id"]
            result: AsyncResult = AsyncResult(task_id)
            if result is not None:
                result.forget()
            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)


def _get_names(entries: Iterable[Union["QueueEntry", "QueueDoneResult"]]):