From 78fb5275e5d79d197aef978a67d84b25978443ca Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Tue, 28 Feb 2023 18:05:02 +0900
Subject: [PATCH] exp: add `dvc exp clean` command

---
 dvc/commands/experiments/__init__.py       |  2 ++
 dvc/commands/experiments/clean.py          | 25 ++++++++++++++++++++++
 dvc/repo/experiments/__init__.py           |  5 +++++
 dvc/repo/experiments/clean.py              | 12 +++++++++++
 tests/func/experiments/test_experiments.py |  6 ++++++
 tests/unit/command/test_experiments.py     | 13 +++++++++++
 6 files changed, 63 insertions(+)
 create mode 100644 dvc/commands/experiments/clean.py
 create mode 100644 dvc/repo/experiments/clean.py

diff --git a/dvc/commands/experiments/__init__.py b/dvc/commands/experiments/__init__.py
index 7546b0ae7..d7e77d81e 100644
--- a/dvc/commands/experiments/__init__.py
+++ b/dvc/commands/experiments/__init__.py
@@ -4,6 +4,7 @@ from dvc.cli.utils import append_doc_link, fix_subparsers, hide_subparsers_from_
 from dvc.commands.experiments import (
     apply,
     branch,
+    clean,
     diff,
     exec_run,
     gc,
@@ -21,6 +22,7 @@ from dvc.commands.experiments import (
 SUB_COMMANDS = [
     apply,
     branch,
+    clean,
     diff,
     exec_run,
     gc,
diff --git a/dvc/commands/experiments/clean.py b/dvc/commands/experiments/clean.py
new file mode 100644
index 000000000..890eb34f3
--- /dev/null
+++ b/dvc/commands/experiments/clean.py
@@ -0,0 +1,25 @@
+import argparse
+import logging
+
+from dvc.cli.command import CmdBase
+from dvc.cli.utils import append_doc_link
+
+logger = logging.getLogger(__name__)
+
+
+class CmdExperimentsClean(CmdBase):
+    def run(self):
+        self.repo.experiments.clean()
+        return 0
+
+
+def add_parser(experiments_subparsers, parent_parser):
+    EXPERIMENTS_CLEAN_HELP = "Cleanup dvc exp internal tempfiles."
+    experiments_clean_parser = experiments_subparsers.add_parser(
+        "clean",
+        parents=[parent_parser],
+        description=append_doc_link(EXPERIMENTS_CLEAN_HELP, "exp/clean"),
+        help=EXPERIMENTS_CLEAN_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    experiments_clean_parser.set_defaults(func=CmdExperimentsClean)
diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index c4c74d95a..d731086da 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -494,3 +494,8 @@ class Experiments:
         from dvc.repo.experiments.remove import remove
 
         return remove(self.repo, *args, **kwargs)
+
+    def clean(self, *args, **kwargs):
+        from dvc.repo.experiments.clean import clean
+
+        return clean(self.repo, *args, **kwargs)
diff --git a/dvc/repo/experiments/clean.py b/dvc/repo/experiments/clean.py
new file mode 100644
index 000000000..653f2e2f3
--- /dev/null
+++ b/dvc/repo/experiments/clean.py
@@ -0,0 +1,12 @@
+from typing import TYPE_CHECKING
+
+from dvc.ui import ui
+
+if TYPE_CHECKING:
+    from dvc.repo import Repo
+
+
+def clean(repo: "Repo"):
+    ui.write("Cleaning up dvc-task messages...")
+    repo.experiments.celery_queue.celery.clean()
+    ui.write("Done!")
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index fc340e005..e9d5f83c9 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -633,3 +633,9 @@ def test_experiment_run_dry(tmp_dir, scm, dvc, exp_stage):
     dvc.experiments.run(exp_stage.addressing, dry=True)
 
     assert len(dvc.experiments.ls()["master"]) == 0
+
+
+def test_clean(tmp_dir, scm, dvc, mocker):
+    clean = mocker.spy(dvc.experiments.celery_queue.celery, "clean")
+    dvc.experiments.clean()
+    clean.assert_called_once_with()
diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py
index be667c5fd..8fdbb3109 100644
--- a/tests/unit/command/test_experiments.py
+++ b/tests/unit/command/test_experiments.py
@@ -8,6 +8,7 @@ import pytest
 from dvc.cli import DvcParserError, parse_args
 from dvc.commands.experiments.apply import CmdExperimentsApply
 from dvc.commands.experiments.branch import CmdExperimentsBranch
+from dvc.commands.experiments.clean import CmdExperimentsClean
 from dvc.commands.experiments.diff import CmdExperimentsDiff
 from dvc.commands.experiments.gc import CmdExperimentsGC
 from dvc.commands.experiments.init import CmdExperimentsInit
@@ -926,3 +927,15 @@ def test_experiments_save(dvc, scm, mocker):
     m.assert_called_once_with(
         cmd.repo, name="exp-name", force=True, include_untracked=[]
     )
+
+
+def test_experiments_clean(dvc, scm, mocker):
+    cli_args = parse_args(["experiments", "clean"])
+    assert cli_args.func == CmdExperimentsClean
+
+    cmd = cli_args.func(cli_args)
+    m = mocker.patch("dvc.repo.experiments.clean.clean", return_value={})
+
+    assert cmd.run() == 0
+
+    m.assert_called_once_with(cmd.repo)