From eadd0da39431c35f2476d26cfaf61dc96880ec0a Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Fri, 16 Dec 2022 17:54:18 +0800 Subject: [PATCH] Refactor `get_ref_and_entry_by_names` 1. Move get_ref_and_entry_by_names from `exp remove` to `celery` for other commands usage. --- dvc/repo/experiments/queue/base.py | 5 ++ dvc/repo/experiments/queue/celery.py | 40 ++++++++++- dvc/repo/experiments/remove.py | 81 +++++----------------- tests/unit/repo/experiments/test_remove.py | 1 + 4 files changed, 64 insertions(+), 63 deletions(-) diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index 7dff103b9..3abd986d6 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -86,6 +86,11 @@ class QueueDoneResult(NamedTuple): result: Optional[ExecutorResult] +class ExpRefAndQueueEntry(NamedTuple): + exp_ref_info: Optional["ExpRefInfo"] + queue_entry: Optional["QueueEntry"] + + class BaseStashQueue(ABC): """Naive Git-stash based experiment queue. diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index e9131e1c5..ffec2530d 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -12,6 +12,7 @@ from typing import ( NamedTuple, Optional, Set, + Union, ) from celery.result import AsyncResult @@ -26,7 +27,13 @@ from ..exceptions import UnresolvedQueueExpNamesError from ..executor.base import ExecutorInfo, ExecutorResult from ..refs import CELERY_STASH from ..utils import EXEC_TMP_DIR, get_exp_rwlock -from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult +from .base import ( + BaseStashQueue, + ExpRefAndQueueEntry, + QueueDoneResult, + QueueEntry, + QueueGetResult, +) from .exceptions import CannotKillTasksError from .tasks import run_exp from .utils import fetch_running_exp_from_temp_dir @@ -36,6 +43,8 @@ if TYPE_CHECKING: from dvc_task.proc.manager import ProcessManager from dvc_task.worker import TemporaryWorker + from ..refs import ExpRefInfo + logger = logging.getLogger(__name__) @@ -358,6 +367,7 @@ class LocalCeleryQueue(BaseStashQueue): missing_revs.append(rev) else: to_kill[queue_entry] = rev + if missing_revs: raise UnresolvedQueueExpNamesError(missing_revs) @@ -449,3 +459,31 @@ class LocalCeleryQueue(BaseStashQueue): ) ) return result + + def get_ref_and_entry_by_names( + self, + exp_names: Union[str, List[str]], + git_remote: Optional[str] = None, + ) -> Dict[str, ExpRefAndQueueEntry]: + """Find finished ExpRefInfo or queued or failed QueueEntry by name""" + from ..utils import resolve_name + + if isinstance(exp_names, str): + exp_names = [exp_names] + results: Dict[str, ExpRefAndQueueEntry] = {} + + exp_ref_match: Dict[str, Optional["ExpRefInfo"]] = resolve_name( + self.scm, exp_names, git_remote + ) + if not git_remote: + queue_entry_match: Dict[ + str, Optional["QueueEntry"] + ] = self.match_queue_entry_by_name( + exp_names, self.iter_queued(), self.iter_done() + ) + + for exp_name in exp_names: + exp_ref = exp_ref_match[exp_name] + queue_entry = None if git_remote else queue_entry_match[exp_name] + results[exp_name] = ExpRefAndQueueEntry(exp_ref, queue_entry) + return results diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 9a5509723..b2666fbbf 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -1,19 +1,12 @@ import logging -from typing import ( - TYPE_CHECKING, - Dict, - Iterable, - List, - NamedTuple, - Optional, - Union, -) +from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Union from dvc.repo import locked from dvc.repo.scm_context import scm_context from dvc.scm import iter_revs from .exceptions import UnresolvedExpNamesError +from .queue.base import QueueEntry from .refs import ExpRefInfo from .utils import exp_refs, exp_refs_by_baseline, push_refspec @@ -22,58 +15,12 @@ if TYPE_CHECKING: from dvc.repo.experiments.queue.celery import LocalCeleryQueue from dvc.scm import Git - from .queue.base import QueueEntry + from .queue.base import ExpRefAndQueueEntry logger = logging.getLogger(__name__) -class ExpRefAndQueueEntry(NamedTuple): - exp_ref_list: List["ExpRefInfo"] - queue_entry_list: List["QueueEntry"] - removed: List["str"] - - -def _get_ref_and_entry_by_names( - exp_names: Union[str, List[str]], - scm: "Git", - celery_queue: "LocalCeleryQueue", - git_remote: Optional[str], -) -> ExpRefAndQueueEntry: - from .utils import resolve_name - - exp_ref_list: List["ExpRefInfo"] = [] - queue_entry_list: List["QueueEntry"] = [] - removed: List[str] = [] - if isinstance(exp_names, str): - exp_names = [exp_names] - exp_ref_match: Dict[str, Optional["ExpRefInfo"]] = resolve_name( - scm, exp_names, git_remote - ) - if not git_remote: - queue_entry_match: Dict[ - str, Optional["QueueEntry"] - ] = celery_queue.match_queue_entry_by_name( - exp_names, celery_queue.iter_queued(), celery_queue.iter_done() - ) - - remained = [] - for exp_name in exp_names: - exp_ref = exp_ref_match[exp_name] - queue_entry = None if git_remote else queue_entry_match[exp_name] - if exp_ref or queue_entry: - if exp_ref: - exp_ref_list.append(exp_ref) - if queue_entry: - queue_entry_list.append(queue_entry) - removed.append(exp_name) - else: - remained.append(exp_name) - if remained: - raise UnresolvedExpNamesError(remained) - return ExpRefAndQueueEntry(exp_ref_list, queue_entry_list, removed) - - @locked @scm_context def remove( @@ -103,12 +50,22 @@ def remove( exp_ref_list: List["ExpRefInfo"] = [] queue_entry_list: List["QueueEntry"] = [] if exp_names: - result = _get_ref_and_entry_by_names( - exp_names, repo.scm, celery_queue, git_remote - ) - removed.extend(result.removed) - exp_ref_list.extend(result.exp_ref_list) - queue_entry_list.extend(result.queue_entry_list) + results: Dict[ + str, "ExpRefAndQueueEntry" + ] = celery_queue.get_ref_and_entry_by_names(exp_names, git_remote) + remained: List[str] = [] + for name, result in results.items(): + if not result.exp_ref_info and not result.queue_entry: + remained.append(name) + continue + removed.append(name) + if result.exp_ref_info: + exp_ref_list.append(result.exp_ref_info) + if result.queue_entry: + queue_entry_list.append(result.queue_entry) + + if remained: + raise UnresolvedExpNamesError(remained) elif rev: exp_ref_dict = _resolve_exp_by_baseline(repo, rev, num, git_remote) removed.extend(exp_ref_dict.keys()) diff --git a/tests/unit/repo/experiments/test_remove.py b/tests/unit/repo/experiments/test_remove.py index c9dd3861a..9592484ab 100644 --- a/tests/unit/repo/experiments/test_remove.py +++ b/tests/unit/repo/experiments/test_remove.py @@ -11,6 +11,7 @@ def test_remove_done_tasks(dvc, test_queue, scm, mocker): ref_info_dict = {} for name in success_test: ref_info_dict[name] = mocker.Mock() + ref_info_dict[name].name = name for name in failed_test: ref_info_dict[name] = None