Add ability to configure test docker container (#8329)

Co-authored-by: Stephanie Crater <[email protected]>

tfb

@@ -103,4 +103,4 @@ fi
 
 test -t 1 && USE_TTY="-t"
 docker build -t techempower/tfb --build-arg USER_ID=$(id -u) --build-arg GROUP_ID=$(id -g) - < ${SCRIPT_ROOT}/Dockerfile
-exec docker run -i ${USE_TTY} ${EXTRA_DOCKER_ARGS} --rm --network tfb -e USER_ID=$(id -u) -v /var/run/docker.sock:/var/run/docker.sock -v ${SCRIPT_ROOT}:/FrameworkBenchmarks techempower/tfb "${@}"
+exec docker run -i ${USE_TTY} --rm --network tfb -e USER_ID=$(id -u) -v /var/run/docker.sock:/var/run/docker.sock -v ${SCRIPT_ROOT}:/FrameworkBenchmarks techempower/tfb "${@}"

toolset/run-tests.py

@@ -187,6 +187,15 @@ def main(argv=None):
         nargs='+',
         default=[1, 10, 20, 50, 100],
         help='List of cached query levels to benchmark')
+    parser.add_argument(
+        '--test-container-memory',
+        default=None,
+        help='Amount of memory to be given to the test container')
+    parser.add_argument(
+        '--extra-docker-runtime-args',
+        nargs='*',
+        default=None,
+        help='Extra docker arguments to be passed to the test container')
 
     # Network options
     parser.add_argument(

toolset/utils/benchmark_config.py

@@ -51,6 +51,8 @@ class BenchmarkConfig:
         self.database_docker_host = None
         self.client_docker_host = None
         self.network = None
+        self.test_container_memory = args.test_container_memory
+        self.extra_docker_runtime_args = args.extra_docker_runtime_args
 
         if self.network_mode is None:
             self.network = 'tfb'

toolset/utils/docker_helper.py

@@ -13,8 +13,6 @@ from toolset.databases import databases
 
 from psutil import virtual_memory
 
-# total memory limit allocated for the test container
-mem_limit = int(round(virtual_memory().total * .95))
 
 class DockerHelper:
     def __init__(self, benchmarker=None):
@@ -207,6 +205,17 @@ class DockerHelper:
                 # to the webserver from IDE
                 if hasattr(test, 'debug_port'):
                     ports[test.debug_port] = test.debug_port
+                    
+            # Total memory limit allocated for the test container
+            if self.benchmarker.config.test_container_memory is not None:
+                mem_limit = self.benchmarker.config.test_container_memory
+            else:
+                mem_limit = int(round(virtual_memory().total * .95))
+
+            # Convert extra docker runtime args to a dictionary
+            extra_docker_args = {}
+            if self.benchmarker.config.extra_docker_runtime_args is not None:
+                extra_docker_args = {key: int(value) if value.isdigit() else value for key, value in (pair.split(":") for pair in self.benchmarker.config.extra_docker_runtime_args)}
 
             container = self.server.containers.run(
                 "techempower/tfb.test.%s" % test.name,
@@ -225,7 +234,9 @@ class DockerHelper:
                 mem_limit=mem_limit,
                 sysctls=sysctl,
                 remove=True,
-                log_config={'type': None})
+                log_config={'type': None},
+                **extra_docker_args
+                )
 
             watch_thread = Thread(
                 target=watch_container,