exp: add exp remove command (#5443)

* exp: add remove function for removing specific experiment(s)

* add `dvc exp remove` command

* update tests

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="<experiment>",
+    )
+    experiments_remove_parser.set_defaults(func=CmdExperimentsRemove)
+
 
 def _add_run_common(parser):
     """Add common args for 'exp run' and 'exp resume'."""

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)

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():

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]

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))

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

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,
+    )