diff --git a/dvc/commands/experiments/remove.py b/dvc/commands/experiments/remove.py index d2687aa1b..92ebf9bdf 100644 --- a/dvc/commands/experiments/remove.py +++ b/dvc/commands/experiments/remove.py @@ -4,6 +4,7 @@ import logging from dvc.cli.command import CmdBase from dvc.cli.utils import append_doc_link from dvc.exceptions import InvalidArgumentError +from dvc.ui import ui logger = logging.getLogger(__name__) @@ -27,7 +28,7 @@ class CmdExperimentsRemove(CmdBase): self.raise_error_if_all_disabled() - self.repo.experiments.remove( + removed_list = self.repo.experiments.remove( exp_names=self.args.experiment, all_commits=self.args.all_commits, rev=self.args.rev, @@ -35,6 +36,8 @@ class CmdExperimentsRemove(CmdBase): queue=self.args.queue, git_remote=self.args.git_remote, ) + removed = ",".join(removed_list) + ui.write(f"Removed experiments: {removed}") return 0 diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 6246f9489..01bf07b7d 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -1,19 +1,10 @@ import logging -from typing import ( - TYPE_CHECKING, - Collection, - List, - Mapping, - Optional, - Set, - Union, -) +from typing import TYPE_CHECKING, Dict, List, Mapping, Optional, Set, Union from dvc.repo import locked from dvc.repo.scm_context import scm_context from dvc.scm import iter_revs -from .base import ExpRefInfo from .exceptions import UnresolvedExpNamesError from .utils import ( exp_refs, @@ -24,8 +15,11 @@ from .utils import ( ) if TYPE_CHECKING: + from dvc.repo import Repo from dvc.scm import Git + from .base import ExpRefInfo + logger = logging.getLogger(__name__) @@ -33,48 +27,50 @@ logger = logging.getLogger(__name__) @locked @scm_context def remove( - repo, + repo: "Repo", exp_names: Union[None, str, List[str]] = None, rev: Optional[str] = None, all_commits: bool = False, num: int = 1, queue: bool = False, git_remote: Optional[str] = None, -) -> int: +) -> List[str]: + removed: List[str] = [] if not any([exp_names, queue, all_commits, rev]): - return 0 - - removed = 0 - if queue: - removed += _clear_stash(repo) - if all_commits: - removed += _clear_all_commits(repo.scm, git_remote) return removed - commit_ref_set: Set[ExpRefInfo] = set() - queued_ref_set: Set[int] = set() + if queue: + removed.extend(_clear_stash(repo)) + if all_commits: + removed.extend(_clear_all_commits(repo, git_remote)) + return removed + + commit_ref_dict: Dict["ExpRefInfo", str] = {} + queued_ref_dict: Dict[int, str] = {} if exp_names: _resolve_exp_by_name( - repo, exp_names, commit_ref_set, queued_ref_set, git_remote + repo, exp_names, commit_ref_dict, queued_ref_dict, git_remote ) if rev: - _resolve_exp_by_baseline(repo, rev, num, commit_ref_set, git_remote) + _resolve_exp_by_baseline(repo, rev, num, commit_ref_dict, git_remote) - if commit_ref_set: - removed += _remove_commited_exps(repo.scm, commit_ref_set, git_remote) + if commit_ref_dict: + removed.extend( + _remove_commited_exps(repo.scm, commit_ref_dict, git_remote) + ) - if queued_ref_set: - removed += _remove_queued_exps(repo, queued_ref_set) + if queued_ref_dict: + removed.extend(_remove_queued_exps(repo, queued_ref_dict)) return removed def _resolve_exp_by_name( - repo, + repo: "Repo", exp_names: Union[str, List[str]], - commit_ref_set: Set["ExpRefInfo"], - queued_ref_set: Set[int], + commit_ref_dict: Dict["ExpRefInfo", str], + queued_ref_dict: Dict[int, str], git_remote: Optional[str], ): remained = set() @@ -86,13 +82,13 @@ def _resolve_exp_by_name( if exp_ref is None: remained.add(exp_name) else: - commit_ref_set.add(exp_ref) + commit_ref_dict[exp_ref] = exp_name if not git_remote: stash_index_dict = _get_queued_index_by_names(repo, remained) for exp_name, stash_index in stash_index_dict.items(): if stash_index is not None: - queued_ref_set.add(stash_index) + queued_ref_dict[stash_index] = exp_name remained.remove(exp_name) if remained: @@ -103,7 +99,7 @@ def _resolve_exp_by_baseline( repo, rev: str, num: int, - commit_ref_set: Set["ExpRefInfo"], + commit_ref_dict: Dict["ExpRefInfo", str], git_remote: Optional[str] = None, ): rev_dict = iter_revs(repo.scm, [rev], num) @@ -112,19 +108,27 @@ def _resolve_exp_by_baseline( for _, ref_info_list in ref_info_dict.items(): for ref_info in ref_info_list: - commit_ref_set.add(ref_info) + if ref_info not in commit_ref_dict: + commit_ref_dict[ref_info] = ref_info.name -def _clear_stash(repo): - removed = len(repo.experiments.stash) +def _clear_stash(repo: "Repo") -> List[str]: + removed_name_list = [] + for entry in repo.experiments.stash.list(): + new_sha = entry.new_sha.decode("utf8") + name = repo.experiments.get_exact_name(new_sha) + if not name: + name = new_sha[:7] + removed_name_list.append(name) repo.experiments.stash.clear() - return removed + return removed_name_list -def _clear_all_commits(scm, git_remote): - ref_infos = list(exp_refs(scm, git_remote)) - _remove_commited_exps(scm, ref_infos, git_remote) - return len(ref_infos) +def _clear_all_commits(repo, git_remote) -> List: + ref_infos = { + ref_info: ref_info.name for ref_info in exp_refs(repo.scm, git_remote) + } + return _remove_commited_exps(repo.scm, ref_infos, git_remote) def _get_queued_index_by_names( @@ -152,12 +156,12 @@ def _get_queued_index_by_names( def _remove_commited_exps( - scm: "Git", exp_ref_list: Collection["ExpRefInfo"], remote: Optional[str] -) -> int: + scm: "Git", exp_ref_dict: Mapping["ExpRefInfo", str], remote: Optional[str] +) -> List[str]: if remote: from dvc.scm import TqdmGit - for ref_info in exp_ref_list: + for ref_info in exp_ref_dict: with TqdmGit(desc="Pushing git refs") as pbar: push_refspec( scm, @@ -167,13 +171,13 @@ def _remove_commited_exps( progress=pbar.update_git, ) else: - remove_exp_refs(scm, exp_ref_list) - return len(exp_ref_list) + remove_exp_refs(scm, exp_ref_dict) + return list(exp_ref_dict.values()) -def _remove_queued_exps(repo, indexes: Collection[int]) -> int: - index_list = list(indexes) +def _remove_queued_exps(repo, indexes: Mapping[int, str]) -> List[str]: + index_list = list(indexes.keys()) index_list.sort(reverse=True) for index in index_list: repo.experiments.stash.drop(index) - return len(index_list) + return list(indexes.values()) diff --git a/tests/func/experiments/test_remove.py b/tests/func/experiments/test_remove.py index 2a0eb2635..a519f3edd 100644 --- a/tests/func/experiments/test_remove.py +++ b/tests/func/experiments/test_remove.py @@ -7,34 +7,39 @@ from dvc.repo.experiments.utils import exp_refs_by_rev def test_remove_experiments_by_ref(tmp_dir, scm, dvc, exp_stage, caplog): queue_length = 3 - ref_list = [] + ref_info_list = [] + ref_name_list = [] for i in range(queue_length): results = dvc.experiments.run( exp_stage.addressing, params=[f"foo={i}"] ) ref_info = first(exp_refs_by_rev(scm, first(results))) - ref_list.append(str(ref_info)) + ref_info_list.append(ref_info) + ref_name_list.append(str(ref_info)) with pytest.raises(InvalidArgumentError): - dvc.experiments.remove(ref_list[:2] + ["non-exist"]) - assert scm.get_ref(str(ref_list[0])) is not None - assert scm.get_ref(str(ref_list[1])) is not None - assert scm.get_ref(str(ref_list[2])) is not None + dvc.experiments.remove(ref_name_list[:2] + ["non-exist"]) + assert scm.get_ref(ref_name_list[0]) is not None + assert scm.get_ref(ref_name_list[1]) is not None + assert scm.get_ref(ref_name_list[2]) is not None - assert dvc.experiments.remove(ref_list[:2]) - assert scm.get_ref(str(ref_list[0])) is None - assert scm.get_ref(str(ref_list[1])) is None - assert scm.get_ref(str(ref_list[2])) is not None + assert set(dvc.experiments.remove(ref_name_list[:2])) == set( + ref_name_list[:2] + ) + assert scm.get_ref(ref_name_list[0]) is None + assert scm.get_ref(ref_name_list[1]) is None + assert scm.get_ref(ref_name_list[2]) is not None def test_remove_all_queued_experiments(tmp_dir, scm, dvc, exp_stage): queue_length = 3 - + ref_list = [] for i in range(queue_length): - dvc.experiments.run( + result = dvc.experiments.run( exp_stage.addressing, params=[f"foo={i}"], queue=True ) + ref_list.append(first(result)[:7]) results = dvc.experiments.run( exp_stage.addressing, params=[f"foo={queue_length}"] @@ -42,7 +47,7 @@ def test_remove_all_queued_experiments(tmp_dir, scm, dvc, exp_stage): ref_info = first(exp_refs_by_rev(scm, first(results))) assert len(dvc.experiments.stash) == queue_length - assert dvc.experiments.remove(queue=True) == queue_length + assert set(dvc.experiments.remove(queue=True)) == set(ref_list) assert len(dvc.experiments.stash) == 0 assert scm.get_ref(str(ref_info)) is not None @@ -71,7 +76,9 @@ def test_remove_special_queued_experiments(tmp_dir, scm, dvc, exp_stage): assert scm.get_ref(str(ref_info1)) is not None assert scm.get_ref(str(ref_info2)) is not None - assert dvc.experiments.remove(["queue1", rev2[:5], str(ref_info1)]) == 3 + assert set( + dvc.experiments.remove(["queue1", rev2[:5], str(ref_info1)]) + ) == {"queue1", rev2[:5], str(ref_info1)} assert rev1 not in dvc.experiments.stash_revs assert rev2 not in dvc.experiments.stash_revs assert rev3 in dvc.experiments.stash_revs @@ -90,8 +97,10 @@ def test_remove_all(tmp_dir, scm, dvc, exp_stage): ref_info2 = first(exp_refs_by_rev(scm, first(results))) dvc.experiments.run(exp_stage.addressing, params=["foo=4"], queue=True) - removed = dvc.experiments.remove(all_commits=True) - assert removed == 2 + assert set(dvc.experiments.remove(all_commits=True)) == { + ref_info1.name, + ref_info2.name, + } assert len(dvc.experiments.stash) == 2 assert scm.get_ref(str(ref_info2)) is None assert scm.get_ref(str(ref_info1)) is None @@ -130,8 +139,7 @@ def test_remove_experiments_by_rev(tmp_dir, scm, dvc, exp_stage): baseline = scm.get_rev() results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"]) - ref_info = first(exp_refs_by_rev(scm, first(results))) - baseline_exp_ref = str(ref_info) + baseline_exp_ref = first(exp_refs_by_rev(scm, first(results))) results = dvc.experiments.run( exp_stage.addressing, params=["foo=2"], queue=True, name="queue2" @@ -148,8 +156,8 @@ def test_remove_experiments_by_rev(tmp_dir, scm, dvc, exp_stage): exp_stage.addressing, params=["foo=4"], queue=True, name="queue4" ) new_queue_ref = first(results) - assert dvc.experiments.remove(rev=baseline) == 1 - assert scm.get_ref(baseline_exp_ref) is None + assert dvc.experiments.remove(rev=baseline) == [baseline_exp_ref.name] + assert scm.get_ref(str(baseline_exp_ref)) is None assert baseline_queue_ref in dvc.experiments.stash_revs assert scm.get_ref(new_exp_ref) is not None assert new_queue_ref in dvc.experiments.stash_revs