diff --git a/dvc/commands/experiments/remove.py b/dvc/commands/experiments/remove.py index 4226c4845..d2687aa1b 100644 --- a/dvc/commands/experiments/remove.py +++ b/dvc/commands/experiments/remove.py @@ -3,24 +3,44 @@ import logging from dvc.cli.command import CmdBase from dvc.cli.utils import append_doc_link +from dvc.exceptions import InvalidArgumentError logger = logging.getLogger(__name__) class CmdExperimentsRemove(CmdBase): + def raise_error_if_all_disabled(self): + if not any( + [ + self.args.experiment, + self.args.all_commits, + self.args.rev, + self.args.queue, + ] + ): + raise InvalidArgumentError( + "Either provide an `experiment` argument, or use the " + "`--rev` or `--all-commits` flag." + ) + def run(self): + self.raise_error_if_all_disabled() + self.repo.experiments.remove( exp_names=self.args.experiment, + all_commits=self.args.all_commits, + rev=self.args.rev, + num=self.args.num, queue=self.args.queue, - clear_all=self.args.all, - remote=self.args.git_remote, + git_remote=self.args.git_remote, ) return 0 def add_parser(experiments_subparsers, parent_parser): + from . import add_rev_selection_flags EXPERIMENTS_REMOVE_HELP = "Remove experiments." experiments_remove_parser = experiments_subparsers.add_parser( @@ -31,15 +51,10 @@ def add_parser(experiments_subparsers, parent_parser): formatter_class=argparse.RawDescriptionHelpFormatter, ) remove_group = experiments_remove_parser.add_mutually_exclusive_group() + add_rev_selection_flags(experiments_remove_parser, "Remove", False) remove_group.add_argument( "--queue", action="store_true", help="Remove all queued experiments." ) - remove_group.add_argument( - "-A", - "--all", - action="store_true", - help="Remove all committed experiments.", - ) remove_group.add_argument( "-g", "--git-remote", diff --git a/dvc/repo/experiments/exceptions.py b/dvc/repo/experiments/exceptions.py index d6a7a304e..527200fdc 100644 --- a/dvc/repo/experiments/exceptions.py +++ b/dvc/repo/experiments/exceptions.py @@ -1,4 +1,4 @@ -from typing import Iterable, List +from typing import Collection, Iterable from dvc.exceptions import InvalidArgumentError @@ -24,7 +24,7 @@ class AmbiguousExpRefInfo(InvalidArgumentError): class UnresolvedExpNamesError(InvalidArgumentError): def __init__( - self, unresolved_list: List[str], *args, git_remote: str = None + self, unresolved_list: Collection[str], *args, git_remote: str = None ): unresolved_names = ";".join(unresolved_list) if not git_remote: diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py index 65c1ac56c..6246f9489 100644 --- a/dvc/repo/experiments/remove.py +++ b/dvc/repo/experiments/remove.py @@ -1,12 +1,31 @@ import logging -from typing import List, Optional +from typing import ( + TYPE_CHECKING, + Collection, + List, + Mapping, + Optional, + Set, + Union, +) -from dvc.exceptions import InvalidArgumentError from dvc.repo import locked from dvc.repo.scm_context import scm_context -from dvc.scm import RevError +from dvc.scm import iter_revs + +from .base import ExpRefInfo +from .exceptions import UnresolvedExpNamesError +from .utils import ( + exp_refs, + exp_refs_by_baseline, + push_refspec, + remove_exp_refs, + resolve_name, +) + +if TYPE_CHECKING: + from dvc.scm import Git -from .utils import exp_refs, push_refspec, remove_exp_refs, resolve_name logger = logging.getLogger(__name__) @@ -15,101 +34,146 @@ logger = logging.getLogger(__name__) @scm_context def remove( repo, - exp_names=None, - queue=False, - clear_all=False, - remote=None, - **kwargs, -): - if not any([exp_names, queue, clear_all]): + 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: + if not any([exp_names, queue, all_commits, rev]): return 0 removed = 0 if queue: removed += _clear_stash(repo) - if clear_all: - removed += _clear_all(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 exp_names: - removed += _remove_exp_by_names(repo, remote, exp_names) + _resolve_exp_by_name( + repo, exp_names, commit_ref_set, queued_ref_set, git_remote + ) + + if rev: + _resolve_exp_by_baseline(repo, rev, num, commit_ref_set, git_remote) + + if commit_ref_set: + removed += _remove_commited_exps(repo.scm, commit_ref_set, git_remote) + + if queued_ref_set: + removed += _remove_queued_exps(repo, queued_ref_set) + return removed +def _resolve_exp_by_name( + repo, + exp_names: Union[str, List[str]], + commit_ref_set: Set["ExpRefInfo"], + queued_ref_set: Set[int], + git_remote: Optional[str], +): + remained = set() + if isinstance(exp_names, str): + exp_names = [exp_names] + + exp_ref_dict = resolve_name(repo.scm, exp_names, git_remote) + for exp_name, exp_ref in exp_ref_dict.items(): + if exp_ref is None: + remained.add(exp_name) + else: + commit_ref_set.add(exp_ref) + + 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) + remained.remove(exp_name) + + if remained: + raise UnresolvedExpNamesError(remained) + + +def _resolve_exp_by_baseline( + repo, + rev: str, + num: int, + commit_ref_set: Set["ExpRefInfo"], + git_remote: Optional[str] = None, +): + rev_dict = iter_revs(repo.scm, [rev], num) + rev_set = set(rev_dict.keys()) + ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote) + + for _, ref_info_list in ref_info_dict.items(): + for ref_info in ref_info_list: + commit_ref_set.add(ref_info) + + def _clear_stash(repo): removed = len(repo.experiments.stash) repo.experiments.stash.clear() return removed -def _clear_all(repo): - ref_infos = list(exp_refs(repo.scm)) - remove_exp_refs(repo.scm, ref_infos) +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 _get_exp_stash_index(repo, ref_or_rev: str) -> Optional[int]: +def _get_queued_index_by_names( + repo, + exp_name_set: Set[str], +) -> Mapping[str, Optional[int]]: + from scmrepo.exceptions import RevError as InternalRevError + + result = {} stash_revs = repo.experiments.stash_revs for _, entry in stash_revs.items(): - if entry.name == ref_or_rev: - return entry.stash_index + if entry.name in exp_name_set: + result[entry.name] = entry.stash_index - from dvc.scm import resolve_rev - - try: - rev = resolve_rev(repo.scm, ref_or_rev) - if rev in stash_revs: - return stash_revs.get(rev).stash_index - except RevError: - pass - return None + for exp_name in exp_name_set: + if exp_name in result: + continue + try: + rev = repo.scm.resolve_rev(exp_name) + if rev in stash_revs: + result[exp_name] = stash_revs.get(rev).stash_index + except InternalRevError: + result[exp_name] = None + return result def _remove_commited_exps( - repo, remote: Optional[str], exp_names: List[str] -) -> List[str]: - remain_list = [] - remove_list = [] - ref_info_dict = resolve_name(repo.scm, exp_names, remote) - for exp_name, ref_info in ref_info_dict.items(): - if ref_info: - remove_list.append(ref_info) - else: - remain_list.append(exp_name) - if remove_list: - if not remote: - remove_exp_refs(repo.scm, remove_list) - else: - from dvc.scm import TqdmGit + scm: "Git", exp_ref_list: Collection["ExpRefInfo"], remote: Optional[str] +) -> int: + if remote: + from dvc.scm import TqdmGit - for ref_info in remove_list: - with TqdmGit(desc="Pushing git refs") as pbar: - push_refspec( - repo.scm, - remote, - None, - str(ref_info), - progress=pbar.update_git, - ) - return remain_list + for ref_info in exp_ref_list: + with TqdmGit(desc="Pushing git refs") as pbar: + push_refspec( + scm, + remote, + None, + str(ref_info), + progress=pbar.update_git, + ) + else: + remove_exp_refs(scm, exp_ref_list) + return len(exp_ref_list) -def _remove_queued_exps(repo, refs_or_revs: List[str]) -> List[str]: - remain_list = [] - for ref_or_rev in refs_or_revs: - stash_index = _get_exp_stash_index(repo, ref_or_rev) - if stash_index is None: - remain_list.append(ref_or_rev) - else: - repo.experiments.stash.drop(stash_index) - return remain_list - - -def _remove_exp_by_names(repo, remote, exp_names: List[str]) -> int: - remained = _remove_commited_exps(repo, remote, exp_names) - if not remote: - remained = _remove_queued_exps(repo, remained) - if remained: - raise InvalidArgumentError( - "'{}' is not a valid experiment".format(";".join(remained)) - ) - return len(exp_names) - len(remained) +def _remove_queued_exps(repo, indexes: Collection[int]) -> int: + index_list = list(indexes) + index_list.sort(reverse=True) + for index in index_list: + repo.experiments.stash.drop(index) + return len(index_list) diff --git a/tests/func/experiments/test_remove.py b/tests/func/experiments/test_remove.py index d51cbcca4..2a0eb2635 100644 --- a/tests/func/experiments/test_remove.py +++ b/tests/func/experiments/test_remove.py @@ -17,7 +17,12 @@ def test_remove_experiments_by_ref(tmp_dir, scm, dvc, exp_stage, caplog): ref_list.append(str(ref_info)) with pytest.raises(InvalidArgumentError): - assert dvc.experiments.remove(ref_list[:2] + ["non-exist"]) + 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 + + 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 @@ -85,7 +90,7 @@ 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(clear_all=True) + removed = dvc.experiments.remove(all_commits=True) assert removed == 2 assert len(dvc.experiments.stash) == 2 assert scm.get_ref(str(ref_info2)) is None @@ -110,7 +115,8 @@ def test_remove_remote(tmp_dir, scm, dvc, exp_stage, git_upstream, use_url): assert git_upstream.tmp_dir.scm.get_ref(str(ref_info)) == exp dvc.experiments.remove( - remote=remote, exp_names=[str(ref_info_list[0]), ref_info_list[1].name] + git_remote=remote, + exp_names=[str(ref_info_list[0]), ref_info_list[1].name], ) assert git_upstream.tmp_dir.scm.get_ref(str(ref_info_list[0])) is None @@ -118,3 +124,32 @@ def test_remove_remote(tmp_dir, scm, dvc, exp_stage, git_upstream, use_url): assert ( git_upstream.tmp_dir.scm.get_ref(str(ref_info_list[2])) == exp_list[2] ) + + +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) + + results = dvc.experiments.run( + exp_stage.addressing, params=["foo=2"], queue=True, name="queue2" + ) + baseline_queue_ref = first(results) + + scm.commit("new_baseline") + + results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"]) + ref_info = first(exp_refs_by_rev(scm, first(results))) + new_exp_ref = str(ref_info) + + results = dvc.experiments.run( + 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 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 diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 7cc5c1fe4..7ff828936 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -338,18 +338,22 @@ def test_experiments_pull(dvc, scm, mocker): ) -@pytest.mark.parametrize( - "queue,clear_all,remote", - [(True, False, None), (False, True, None), (False, False, True)], -) -def test_experiments_remove(dvc, scm, mocker, queue, clear_all, remote): - if queue: - args = ["--queue"] - if clear_all: - args = ["--all"] - if remote: - args = ["--git-remote", "myremote", "exp-123", "exp-234"] - cli_args = parse_args(["experiments", "remove"] + args) +def test_experiments_remove(dvc, scm, mocker, capsys, caplog): + cli_args = parse_args( + [ + "experiments", + "remove", + "--all-commits", + "--rev", + "foo", + "--num", + "2", + "--git-remote", + "myremote", + "exp-123", + "exp-234", + ] + ) assert cli_args.func == CmdExperimentsRemove cmd = cli_args.func(cli_args) @@ -358,10 +362,20 @@ def test_experiments_remove(dvc, scm, mocker, queue, clear_all, remote): assert cmd.run() == 0 m.assert_called_once_with( cmd.repo, - exp_names=["exp-123", "exp-234"] if remote else [], - queue=queue, - clear_all=clear_all, - remote="myremote" if remote else None, + exp_names=["exp-123", "exp-234"], + all_commits=True, + rev="foo", + num=2, + queue=False, + git_remote="myremote", + ) + + cmd = cli_args.func(parse_args(["exp", "remove"])) + with pytest.raises(InvalidArgumentError) as excinfo: + cmd.run() + assert ( + str(excinfo.value) == "Either provide an `experiment` argument" + ", or use the `--rev` or `--all-commits` flag." )