contrib/benchmark: Improvements to image handlingsam/benchmark-harness

1 files changed, 17 insertions, 5 deletions

diff --git a/contrib/benchmark/__main__.py b/contrib/benchmark/__main__.py
index cee37bbf5..33ccde00a 100644
--- a/contrib/benchmark/__main__.py
+++ b/contrib/benchmark/__main__.py
@@ -101,7 +101,7 @@ def parse_output_var(line, prefix):
 
 # Creates a Docker image that can be used to test a specific version of
 # BuildStream, as described by 'version_spec'.
-def prepare_version(docker_client, version_spec):
+def prepare_version(docker_client, version_spec, reuse_images=False):
     logging.info("Preparing version '{}'".format(version_spec.name))
 
     # The base image may have a symbolic name, so check that we have the latest
@@ -148,7 +148,7 @@ def prepare_version(docker_client, version_spec):
 
         # FIXME: we never clean up these images; which is perhaps OK but we
         # should at least label them, and if possible reuse them.
-        image = container.commit()
+        image = container.commit('buildstream-benchmarks')
     except:
         container.remove()
         raise
@@ -215,9 +215,15 @@ def run_benchmark(docker_client, version, test_spec):
 
 @click.command()
 @click.option('--config-file', '-c', type=click.Path(),
-              help="YAML description of tests to run, overriding the defaults")
+              help="YAML description of tests to run, overriding the defaults.")
 @click.option('--debug/--no-debug', default=False)
-def run(config_file, debug):
+@click.option('--keep-images/--no-keep-images', default=False,
+              help="Do not delete the Docker images that we create.")
+@click.option('--reuse-images/--no-reuse-images', default=False,
+              help="Reuse test images from previous runs if found. Beware that "
+                   "changes in remote refs will not be picked up if images are "
+                   "reused.")
+def run(config_file, debug, keep_images, reuse_images):
     if debug:
         logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
     else:
@@ -232,7 +238,8 @@ def run(config_file, debug):
 
     versions = []
     for version_spec in version_specs:
-        version = prepare_version(docker_client, version_spec)
+        version = prepare_version(docker_client, version_spec,
+                                  reuse_images=reuse_images)
         versions.append(version)
 
     results = []
@@ -243,6 +250,11 @@ def run(config_file, debug):
             result.record_result_for_version(version, *result_for_version)
         results.append(result)
 
+    if not keep_images:
+        logging.info("Deleting test-specific Docker images")
+        for version in versions:
+            docker_client.images.remove(version.image.id)
+
     end_time = datetime.datetime.now()
     logging.info("BuildStream benchmark runner finished at {}".format(end_time))