diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index 6521d78d0..6cd4bd02a 100644 --- a/dvc/command/experiments.py +++ b/dvc/command/experiments.py @@ -678,6 +678,16 @@ class CmdExperimentsPull(CmdBase): return 0 +class CmdExperimentsRemove(CmdBase): + def run(self): + + self.repo.experiments.remove( + exp_names=self.args.experiment, queue=self.args.queue, + ) + + return 0 + + def add_parser(subparsers, parent_parser): EXPERIMENTS_HELP = "Commands to run and compare experiments." @@ -1163,6 +1173,25 @@ def add_parser(subparsers, parent_parser): ) experiments_pull_parser.set_defaults(func=CmdExperimentsPull) + EXPERIMENTS_REMOVE_HELP = "Remove local experiments." + experiments_remove_parser = experiments_subparsers.add_parser( + "remove", + parents=[parent_parser], + description=append_doc_link(EXPERIMENTS_REMOVE_HELP, "exp/remove"), + help=EXPERIMENTS_REMOVE_HELP, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + experiments_remove_parser.add_argument( + "--queue", action="store_true", help="Remove all queued experiments.", + ) + experiments_remove_parser.add_argument( + "experiment", + nargs="*", + help="Experiments to remove.", + metavar="", + ) + experiments_remove_parser.set_defaults(func=CmdExperimentsRemove) + def _add_run_common(parser): """Add common args for 'exp run' and 'exp resume'.""" diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index abd6965c0..2438d2dd8 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -851,3 +851,8 @@ class Experiments: from dvc.repo.experiments.ls import ls return ls(self.repo, *args, **kwargs) + + def remove(self, *args, **kwargs): + from dvc.repo.experiments.remove import remove + + return remove(self.repo, *args, **kwargs) diff --git a/dvc/repo/experiments/gc.py b/dvc/repo/experiments/gc.py index 3dfeade22..f12e607ba 100644 --- a/dvc/repo/experiments/gc.py +++ b/dvc/repo/experiments/gc.py @@ -3,8 +3,7 @@ from typing import Optional from dvc.repo import locked -from .base import EXEC_APPLY, EXEC_BRANCH, EXEC_CHECKPOINT -from .utils import exp_refs +from .utils import exp_refs, remove_exp_refs logger = logging.getLogger(__name__) @@ -32,22 +31,13 @@ def gc( if not keep_revs: return 0 - exec_branch = repo.scm.get_ref(EXEC_BRANCH, follow=False) - exec_apply = repo.scm.get_ref(EXEC_APPLY) - exec_checkpoint = repo.scm.get_ref(EXEC_CHECKPOINT) - - removed = 0 - for ref_info in exp_refs(repo.scm): - if ref_info.baseline_sha not in keep_revs: - ref = repo.scm.get_ref(str(ref_info)) - if exec_branch and str(ref_info): - repo.scm.remove_ref(EXEC_BRANCH) - if exec_apply and exec_apply == ref: - repo.scm.remove_ref(EXEC_APPLY) - if exec_checkpoint and exec_checkpoint == ref: - repo.scm.remove_ref(EXEC_CHECKPOINT) - repo.scm.remove_ref(str(ref_info)) - removed += 1 + to_remove = [ + ref_info + for ref_info in exp_refs(repo.scm) + if ref_info.baseline_sha not in keep_revs + ] + remove_exp_refs(repo.scm, to_remove) + removed = len(to_remove) delete_stashes = [] for _, entry in repo.experiments.stash_revs.items(): diff --git a/dvc/repo/experiments/remove.py b/dvc/repo/experiments/remove.py new file mode 100644 index 000000000..279733b01 --- /dev/null +++ b/dvc/repo/experiments/remove.py @@ -0,0 +1,63 @@ +import logging + +from dvc.exceptions import InvalidArgumentError +from dvc.repo import locked +from dvc.repo.scm_context import scm_context + +from .base import EXPS_NAMESPACE, ExpRefInfo +from .utils import exp_refs_by_name, remove_exp_refs + +logger = logging.getLogger(__name__) + + +@locked +@scm_context +def remove(repo, exp_names=None, queue=False, **kwargs): + if not exp_names and not queue: + return 0 + + removed = 0 + if queue: + removed += len(repo.experiments.stash) + repo.experiments.stash.clear() + if exp_names: + ref_infos = list(_get_exp_refs(repo, exp_names)) + remove_exp_refs(repo.scm, ref_infos) + removed += len(ref_infos) + return removed + + +def _get_exp_refs(repo, exp_names): + cur_rev = repo.scm.get_rev() + for name in exp_names: + if name.startswith(EXPS_NAMESPACE): + if not repo.scm.get_ref(name): + raise InvalidArgumentError( + f"'{name}' is not a valid experiment name" + ) + yield ExpRefInfo.from_ref(name) + else: + + exp_refs = list(exp_refs_by_name(repo.scm, name)) + if not exp_refs: + raise InvalidArgumentError( + f"'{name}' is not a valid experiment name" + ) + yield _get_ref(exp_refs, name, cur_rev) + + +def _get_ref(ref_infos, name, cur_rev): + if len(ref_infos) > 1: + for info in ref_infos: + if info.baseline_sha == cur_rev: + return info + msg = [ + ( + f"Ambiguous name '{name}' refers to multiple " + "experiments. Use full refname to remove one of " + "the following:" + ), + ] + msg.extend([f"\t{info}" for info in ref_infos]) + raise InvalidArgumentError("\n".join(msg)) + return ref_infos[0] diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index ebe8fcb9c..97b7d8138 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -83,3 +83,21 @@ def exp_commits( shas.update(scm.branch_revs(str(ref_info), ref_info.baseline_sha)) shas.add(ref_info.baseline_sha) yield from shas + + +def remove_exp_refs(scm: "Git", ref_infos: Iterable["ExpRefInfo"]): + from .base import EXEC_APPLY, EXEC_BRANCH, EXEC_CHECKPOINT + + exec_branch = scm.get_ref(EXEC_BRANCH, follow=False) + exec_apply = scm.get_ref(EXEC_APPLY) + exec_checkpoint = scm.get_ref(EXEC_CHECKPOINT) + + for ref_info in ref_infos: + ref = scm.get_ref(str(ref_info)) + if exec_branch and str(ref_info): + scm.remove_ref(EXEC_BRANCH) + if exec_apply and exec_apply == ref: + scm.remove_ref(EXEC_APPLY) + if exec_checkpoint and exec_checkpoint == ref: + scm.remove_ref(EXEC_CHECKPOINT) + scm.remove_ref(str(ref_info)) diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index dcbaf30f1..7c81f14f5 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -593,3 +593,18 @@ def test_run_metrics(tmp_dir, scm, dvc, exp_stage, mocker): main(["exp", "run", "-m"]) assert show_mock.called_once() + + +def test_remove(tmp_dir, scm, dvc, exp_stage): + results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"]) + exp = first(results) + ref_info = first(exp_refs_by_rev(scm, exp)) + dvc.experiments.run(exp_stage.addressing, params=["foo=3"], queue=True) + + removed = dvc.experiments.remove([str(ref_info)]) + assert removed == 1 + assert scm.get_ref(str(ref_info)) is None + + removed = dvc.experiments.remove(queue=True) + assert removed == 1 + assert len(dvc.experiments.stash) == 0 diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 2745036ce..78a80b7e2 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -9,6 +9,7 @@ from dvc.command.experiments import ( CmdExperimentsList, CmdExperimentsPull, CmdExperimentsPush, + CmdExperimentsRemove, CmdExperimentsRun, CmdExperimentsShow, ) @@ -240,3 +241,17 @@ def test_experiments_pull(dvc, scm, mocker): jobs=1, run_cache=True, ) + + +def test_experiments_remove(dvc, scm, mocker): + cli_args = parse_args(["experiments", "remove", "--queue"]) + assert cli_args.func == CmdExperimentsRemove + + cmd = cli_args.func(cli_args) + m = mocker.patch("dvc.repo.experiments.remove.remove", return_value={}) + + assert cmd.run() == 0 + + m.assert_called_once_with( + cmd.repo, exp_names=[], queue=True, + )