From 80e181f2fc7e70dff46396903c8c1ac717c204b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Fri, 6 Nov 2020 19:05:41 +0900 Subject: [PATCH] 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 | 115 ++++++++++++++++++++++++ dvc/repo/experiments/__init__.py | 5 ++ dvc/repo/experiments/gc.py | 51 +++++++++++ tests/func/experiments/test_gc.py | 120 +++++++++++++++++++++++++ tests/unit/command/test_experiments.py | 39 ++++++++ 5 files changed, 330 insertions(+) create mode 100644 dvc/repo/experiments/gc.py create mode 100644 tests/func/experiments/test_gc.py diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index fc1a1b69b..1530f1648 100644 --- a/dvc/command/experiments.py +++ b/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) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index b6ab1d43c..bd3002aa3 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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) diff --git a/dvc/repo/experiments/gc.py b/dvc/repo/experiments/gc.py new file mode 100644 index 000000000..74d3bdbd2 --- /dev/null +++ b/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 diff --git a/tests/func/experiments/test_gc.py b/tests/func/experiments/test_gc.py new file mode 100644 index 000000000..20d1d870f --- /dev/null +++ b/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 diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 4e799bc61..22fb6a0ef 100644 --- a/tests/unit/command/test_experiments.py +++ b/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()