From bb53ec2e4eb6937a5d45774d73307d74c74aa82d Mon Sep 17 00:00:00 2001
From: Kareem Khazem <karkhaz@amazon.com>
Date: Mon, 10 Jul 2023 18:41:39 +0100
Subject: [PATCH] Clean up benchmark directories after building (#2583)

This commit makes benchcomp delete fresh copies of benchmark directories after running the benchmark suite, by default. This is to save disk space, especially on CI machines which do not have enough disk space to run the perf suite. The new behavior can be turned off with the new --no-cleanup-run-dirs flag.
---
 tools/benchcomp/benchcomp/cmd_args.py  | 7 +++++++
 tools/benchcomp/benchcomp/entry/run.py | 8 +++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/tools/benchcomp/benchcomp/cmd_args.py b/tools/benchcomp/benchcomp/cmd_args.py
index 6a8e47d422500..d8b7e735824f8 100644
--- a/tools/benchcomp/benchcomp/cmd_args.py
+++ b/tools/benchcomp/benchcomp/cmd_args.py
@@ -138,6 +138,13 @@ def _get_args_dict():
                         "help":
                             "do not make a fresh copy of the benchmark "
                             "directories before running each variant",
+                    }, {
+                        "flags": ["--keep-temps"],
+                        "action": "store_false",
+                        "dest": "cleanup_directory",
+                        "help":
+                            "do not delete fresh copies of benchmark "
+                            "directories after running each variant",
                     }],
                 },
                 "collate": {
diff --git a/tools/benchcomp/benchcomp/entry/run.py b/tools/benchcomp/benchcomp/entry/run.py
index d95ce8ab0d3a5..a870e7e9a1b0d 100644
--- a/tools/benchcomp/benchcomp/entry/run.py
+++ b/tools/benchcomp/benchcomp/entry/run.py
@@ -39,6 +39,7 @@ class _SingleInvocation:
     command_line: str
     directory: pathlib.Path
 
+    cleanup_directory: bool
     env: dict = dataclasses.field(default_factory=dict)
     timeout: int = None
     memout: int = None
@@ -85,6 +86,9 @@ def __call__(self):
                 encoding="utf-8") as handle:
             yaml.dump(suite, handle, default_flow_style=False)
 
+        if self.cleanup_directory and self.copy_benchmarks_dir:
+            shutil.rmtree(self.working_copy)
+
 
 @dataclasses.dataclass
 class _Run:
@@ -95,6 +99,7 @@ class _Run:
     out_dir: str
     out_symlink: str
     copy_benchmarks_dir: bool
+    cleanup_directory: bool
     result: dict = None
 
     def __call__(self):
@@ -110,6 +115,7 @@ def __call__(self):
                     suite_id, variant_id,
                     parse, suite_yaml_out_dir=out_path,
                     copy_benchmarks_dir=self.copy_benchmarks_dir,
+                    cleanup_directory=self.cleanup_directory,
                     **config)
                 invoke()
 
@@ -137,6 +143,6 @@ def get_default_out_prefix():
 def main(args):
     run = _Run(
         args.config, args.out_prefix, args.out_dir, args.out_symlink,
-        args.copy_benchmarks_dir)
+        args.copy_benchmarks_dir, args.cleanup_directory)
     run()
     return run