experiments: add dvc exp gc (#4803)

* experiments: add `dvc exp gc`

* tests: add exp gc tests

* add exp gc command test

* fix review issue

dvc/command/experiments.py

@@ -8,6 +8,7 @@ from datetime import date, datetime
 from itertools import groupby
 from typing import Iterable, Optional
 
+import dvc.prompt as prompt
 from dvc.command.base import CmdBase, append_doc_link, fix_subparsers
 from dvc.command.metrics import DEFAULT_PRECISION
 from dvc.command.repro import CmdRepro
@@ -477,6 +478,61 @@ class CmdExperimentsRun(CmdRepro):
         return ret
 
 
+class CmdExperimentsGC(CmdRepro):
+    def run(self):
+        from dvc.repo.gc import _raise_error_if_all_disabled
+
+        if not self.repo.experiments:
+            return 0
+
+        _raise_error_if_all_disabled(
+            all_branches=self.args.all_branches,
+            all_tags=self.args.all_tags,
+            all_commits=self.args.all_commits,
+            workspace=self.args.workspace,
+        )
+
+        msg = "This will remove all experiments except those derived from "
+
+        msg += "the workspace"
+        if self.args.all_commits:
+            msg += " and all git commits"
+        elif self.args.all_branches and self.args.all_tags:
+            msg += " and all git branches and tags"
+        elif self.args.all_branches:
+            msg += " and all git branches"
+        elif self.args.all_tags:
+            msg += " and all git tags"
+        msg += " of the current repo."
+        if self.args.queued:
+            msg += " Run queued experiments will be preserved."
+        if self.args.queued:
+            msg += " Run queued experiments will be removed."
+
+        logger.warning(msg)
+
+        msg = "Are you sure you want to proceed?"
+        if not self.args.force and not prompt.confirm(msg):
+            return 1
+
+        removed = self.repo.experiments.gc(
+            all_branches=self.args.all_branches,
+            all_tags=self.args.all_tags,
+            all_commits=self.args.all_commits,
+            workspace=self.args.workspace,
+            queued=self.args.queued,
+        )
+
+        if removed:
+            logger.info(
+                f"Removed {removed} experiments. To remove unused cache files "
+                "use 'dvc gc'."
+            )
+        else:
+            logger.info("No experiments to remove.")
+        return 0
+
+
 def add_parser(subparsers, parent_parser):
     EXPERIMENTS_HELP = "Commands to display and compare experiments."
 
@@ -741,3 +797,62 @@ def add_parser(subparsers, parent_parser):
         ),
     )
     experiments_run_parser.set_defaults(func=CmdExperimentsRun)
+
+    EXPERIMENTS_GC_HELP = "Garbage collect unneeded experiments."
+    EXPERIMENTS_GC_DESCRIPTION = (
+        "Removes all experiments which are not derived from the specified"
+        "Git revisions."
+    )
+    experiments_gc_parser = experiments_subparsers.add_parser(
+        "gc",
+        parents=[parent_parser],
+        description=append_doc_link(
+            EXPERIMENTS_GC_DESCRIPTION, "experiments/gc"
+        ),
+        help=EXPERIMENTS_GC_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    experiments_gc_parser.add_argument(
+        "-w",
+        "--workspace",
+        action="store_true",
+        default=False,
+        help="Keep experiments derived from the current workspace.",
+    )
+    experiments_gc_parser.add_argument(
+        "-a",
+        "--all-branches",
+        action="store_true",
+        default=False,
+        help="Keep experiments derived from the tips of all Git branches.",
+    )
+    experiments_gc_parser.add_argument(
+        "-T",
+        "--all-tags",
+        action="store_true",
+        default=False,
+        help="Keep experiments derived from all Git tags.",
+    )
+    experiments_gc_parser.add_argument(
+        "--all-commits",
+        action="store_true",
+        default=False,
+        help="Keep experiments derived from all Git commits.",
+    )
+    experiments_gc_parser.add_argument(
+        "--queued",
+        action="store_true",
+        default=False,
+        help=(
+            "Keep queued experiments (experiments run queue will be cleared "
+            "by default)."
+        ),
+    )
+    experiments_gc_parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        default=False,
+        help="Force garbage collection - automatically agree to all prompts.",
+    )
+    experiments_gc_parser.set_defaults(func=CmdExperimentsGC)

dvc/repo/experiments/__init__.py

@@ -900,3 +900,8 @@ class Experiments:
         from dvc.repo.experiments.run import run
 
         return run(self.repo, *args, **kwargs)
+
+    def gc(self, *args, **kwargs):
+        from dvc.repo.experiments.gc import gc
+
+        return gc(self.repo, *args, **kwargs)

dvc/repo/experiments/gc.py

@@ -0,0 +1,51 @@
+import logging
+from typing import Optional
+
+from dvc.repo import locked
+
+logger = logging.getLogger(__name__)
+
+
+@locked
+def gc(
+    repo,
+    all_branches: Optional[bool] = False,
+    all_tags: Optional[bool] = False,
+    all_commits: Optional[bool] = False,
+    workspace: Optional[bool] = False,
+    queued: Optional[bool] = False,
+):
+    keep_revs = set(
+        repo.brancher(
+            all_branches=all_branches,
+            all_tags=all_tags,
+            all_commits=all_commits,
+            sha_only=True,
+        )
+    )
+    if workspace:
+        keep_revs.add(repo.scm.get_rev())
+
+    if not keep_revs:
+        return 0
+
+    delete_branches = []
+    for exp_branch in repo.experiments.scm.list_branches():
+        m = repo.experiments.BRANCH_RE.match(exp_branch)
+        if m:
+            rev = repo.scm.resolve_rev(m.group("baseline_rev"))
+            if rev not in keep_revs:
+                delete_branches.append(exp_branch)
+    if delete_branches:
+        repo.experiments.scm.repo.delete_head(*delete_branches, force=True)
+    removed = len(delete_branches)
+
+    delete_stashes = []
+    for _, entry in repo.experiments.stash_revs.items():
+        if not queued or entry.baseline_rev not in keep_revs:
+            delete_stashes.append(entry.index)
+    for index in sorted(delete_stashes, reverse=True):
+        repo.experiments.scm.repo.git.stash("drop", index)
+    removed += len(delete_stashes)
+
+    return removed

tests/func/experiments/test_gc.py

@@ -0,0 +1,120 @@
+import pytest
+from funcy import first
+
+from tests.func.test_repro_multistage import COPY_SCRIPT
+
+
+@pytest.mark.parametrize("queued, expected", [(True, 0), (False, 1)])
+def test_workspace(tmp_dir, scm, dvc, queued, expected):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+
+    stage = dvc.run(
+        cmd="python copy.py params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        name="foo",
+    )
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v1")
+
+    dvc.experiments.run(stage.addressing, params=["foo=2"])
+    dvc.experiments.run(stage.addressing, params=["foo=3"], queue=True)
+
+    removed = dvc.experiments.gc(workspace=True, queued=queued)
+    assert removed == expected
+
+
+@pytest.mark.parametrize("queued, expected", [(True, 0), (False, 2)])
+def test_all_commits(tmp_dir, scm, dvc, queued, expected):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+
+    stage = dvc.run(
+        cmd="python copy.py params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        name="foo",
+    )
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v1")
+
+    results = dvc.experiments.run(stage.addressing, params=["foo=2"])
+    exp_rev = first(results)
+    dvc.experiments.run(stage.addressing, params=["foo=3"], queue=True)
+
+    dvc.experiments.checkout(exp_rev)
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v2")
+
+    dvc.experiments.run(stage.addressing, params=["foo=4"])
+    dvc.experiments.run(stage.addressing, params=["foo=5"], queue=True)
+
+    removed = dvc.experiments.gc(all_commits=True, queued=queued)
+    assert removed == expected
+
+
+@pytest.mark.parametrize("queued, expected", [(True, 2), (False, 3)])
+def test_all_branches(tmp_dir, scm, dvc, queued, expected):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+
+    stage = dvc.run(
+        cmd="python copy.py params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        name="foo",
+    )
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v1")
+
+    results = dvc.experiments.run(stage.addressing, params=["foo=2"])
+    exp_rev = first(results)
+    dvc.experiments.run(stage.addressing, params=["foo=3"], queue=True)
+
+    dvc.experiments.checkout(exp_rev)
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v2")
+
+    scm.branch("branch")
+
+    results = dvc.experiments.run(stage.addressing, params=["foo=4"])
+    dvc.experiments.run(stage.addressing, params=["foo=5"], queue=True)
+    exp_rev = first(results)
+
+    dvc.experiments.checkout(exp_rev)
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v3")
+
+    removed = dvc.experiments.gc(all_branches=True, queued=queued)
+    assert removed == expected
+
+
+@pytest.mark.parametrize("queued, expected", [(True, 2), (False, 3)])
+def test_all_tags(tmp_dir, scm, dvc, queued, expected):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+
+    stage = dvc.run(
+        cmd="python copy.py params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        name="foo",
+    )
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v1")
+
+    results = dvc.experiments.run(stage.addressing, params=["foo=2"])
+    exp_rev = first(results)
+    dvc.experiments.run(stage.addressing, params=["foo=3"], queue=True)
+
+    dvc.experiments.checkout(exp_rev)
+    scm.add(["dvc.yaml", "dvc.lock", "copy.py", "params.yaml", "metrics.yaml"])
+    scm.commit("v2")
+    scm.tag("tag-v2")
+
+    dvc.experiments.run(stage.addressing, params=["foo=4"])
+    dvc.experiments.run(stage.addressing, params=["foo=5"], queue=True)
+
+    removed = dvc.experiments.gc(all_tags=True, queued=queued)
+    assert removed == expected

tests/unit/command/test_experiments.py

@@ -1,10 +1,14 @@
+import pytest
+
 from dvc.cli import parse_args
 from dvc.command.experiments import (
     CmdExperimentsDiff,
+    CmdExperimentsGC,
     CmdExperimentsRun,
     CmdExperimentsShow,
 )
 from dvc.dvcfile import PIPELINE_FILE
+from dvc.exceptions import InvalidArgumentError
 
 from .test_repro import default_arguments as repro_arguments
 
@@ -83,3 +87,38 @@ def test_experiments_run(dvc, mocker):
     cmd.repo.experiments.run.assert_called_with(
         PIPELINE_FILE, **default_arguments
     )
+
+
+def test_experiments_gc(dvc, mocker):
+    cli_args = parse_args(
+        [
+            "exp",
+            "gc",
+            "--workspace",
+            "--all-tags",
+            "--all-branches",
+            "--all-commits",
+            "--queued",
+            "--force",
+        ]
+    )
+    assert cli_args.func == CmdExperimentsGC
+
+    cmd = cli_args.func(cli_args)
+    m = mocker.patch("dvc.repo.experiments.gc.gc", return_value={})
+
+    assert cmd.run() == 0
+
+    m.assert_called_once_with(
+        cmd.repo,
+        workspace=True,
+        all_tags=True,
+        all_branches=True,
+        all_commits=True,
+        queued=True,
+    )
+
+    cli_args = parse_args(["exp", "gc"])
+    cmd = cli_args.func(cli_args)
+    with pytest.raises(InvalidArgumentError):
+        cmd.run()