exp remove: rename --all-commits flag and unify the collection of revs (#7155)

fix: #7155
1. rename flags `-A/--all-commits` in `exp remove`
2. add new flag `-n/--num` in `exp remove`
3. unify the revision collection in `exp remove`
4. add unit and func tests for `exp remove`

Co-authored-by: Jorge Orpinel <[email protected]>
This commit is contained in:
karajan1001
2022-03-09 18:54:15 +08:00
committed by Gao
co-authored by Jorge Orpinel
parent 244637e299
commit efc021da73
5 changed files with 231 additions and 103 deletions
+23 -8
View File
@@ -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",
+2 -2
View File
@@ -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:
+138 -74
View File
@@ -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)
+38 -3
View File
@@ -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
+30 -16
View File
@@ -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."
)