Allow users to specify the Docker image to use with Testbed

samples/tools/testbed/Dockerfile

@@ -0,0 +1,16 @@
+# Host a jsPsych experiment in Azure
+FROM python:3.11
+MAINTAINER AutoGen
+
+# Upgrade pip
+RUN pip install --upgrade pip
+
+# Set the image to the Pacific Timezone
+RUN ln -snf /usr/share/zoneinfo/US/Pacific /etc/localtime && echo "US/Pacific" > /etc/timezone
+
+# Pre-load autogen dependencies, but not autogen itself since we'll often want to install the latest from source
+RUN pip install pyautogen[teachable,lmm,graphs]
+RUN pip uninstall --yes pyautogen
+
+# Pre-load popular packages as per https://learnpython.com/blog/most-popular-python-packages/
+RUN pip install numpy pandas matplotlib seaborn scikit-learn requests urllib3 nltk pillow pytest

samples/tools/testbed/README.md

@@ -46,6 +46,10 @@ options:
                 The requirements file to pip install before running the scenario. This file must be found in
                 the 'includes' directory. (default: requirements.txt)
 
+  -d DOCKER_IMAGE, --docker-image DOCKER_IMAGE
+                The Docker image to use when running scenarios. Can not be used together with --native.
+                (default: 'autogen/testbed:default', which will be created if not present)
+
   --native      Run the scenarios natively rather than in docker.
                 NOTE: This is not advisable, and should be done with great caution.
 ```

samples/tools/testbed/run_scenarios.py

@@ -15,8 +15,12 @@
 # Location of the global includes dir. The contents of this directory will be copied to the Docker environment.
 GLOBAL_INCLUDES_DIR = "includes"
 
+# This is the tag given to the image that is *built* when no other image is provided.
+# Do not use this field to specify the name of an existing image (e.g., on Dockerhub)
+DEFAULT_DOCKER_IMAGE_TAG = "autogen/testbed:default"
 
-def run_scenarios(scenario, n_repeats, is_native, config_list, requirements, results_dir="results"):
+
+def run_scenarios(scenario, n_repeats, is_native, config_list, requirements, docker_image=None, results_dir="results"):
     """
     Run a set testbed scenarios a given number of times.
 
@@ -103,7 +107,7 @@ def run_scenarios(scenario, n_repeats, is_native, config_list, requirements, res
                     if is_native:
                         run_scenario_natively(results_repetition)
                     else:
-                        run_scenario_in_docker(results_repetition, requirements)
+                        run_scenario_in_docker(results_repetition, requirements, docker_image=docker_image)
 
 
 def expand_scenario(scenario_dir, scenario, output_dir):
@@ -244,7 +248,7 @@ def run_scenario_natively(work_dir):
     return
 
 
-def run_scenario_in_docker(work_dir, requirements, timeout=600):
+def run_scenario_in_docker(work_dir, requirements, timeout=600, docker_image=None):
     """
     Run a scenario in a Docker environment.
 
@@ -253,20 +257,34 @@ def run_scenario_in_docker(work_dir, requirements, timeout=600):
         timeout (Optional, int): the number of seconds to allow a Docker container to run before timing out
     """
 
-    # Create a docker client
     client = docker.from_env()
-    image_name = "python:3.11"
-
-    # Pull a suitable image
-    try:
-        image = client.images.get(image_name)
-    except docker.errors.ImageNotFound:
-        # pull the image
-        print("Pulling image", image_name)
+    image = None
+
+    # If the docker_image is None, then we will fetch DEFAULT_DOCKER_IMAGE_TAG, if present,
+    # or build it if missing.
+    if docker_image is None:
+        # Pull a suitable image
+        try:
+            image = client.images.get(DEFAULT_DOCKER_IMAGE_TAG)
+        except docker.errors.ImageNotFound:
+            print(f"Building default Docker image '{DEFAULT_DOCKER_IMAGE_TAG}'. This may take a few minutes...")
+            try:
+                build_default_docker_image(client, DEFAULT_DOCKER_IMAGE_TAG)
+                image = client.images.get(DEFAULT_DOCKER_IMAGE_TAG)
+            except docker.errors.DockerException:
+                print(f"Failed to build image '{DEFAULT_DOCKER_IMAGE_TAG}'")
+
+    # Otherwise get the requested image
+    else:
         try:
-            image = client.images.pull(image_name)
-        except docker.errors.DockerException:
-            print("Failed to pull image", image_name)
+            image = client.images.get(docker_image)
+        except docker.errors.ImageNotFound:
+            # pull the image
+            print(f"Pulling image '{docker_image}'")
+            try:
+                image = client.images.pull(docker_image)
+            except docker.errors.DockerException:
+                print(f"Failed to pull image '{docker_image}'")
 
     # Prepare the run script
     with open(os.path.join(work_dir, "run.sh"), "wt", newline="\n") as f:
@@ -351,6 +369,12 @@ def run_scenario_in_docker(work_dir, requirements, timeout=600):
         f.write(logs)
 
 
+def build_default_docker_image(docker_client, image_tag):
+    for segment in docker_client.api.build(path=".", dockerfile="Dockerfile", rm=True, tag=image_tag, decode=True):
+        if "stream" in segment:
+            sys.stdout.write(segment["stream"])
+
+
 ###############################################################################
 if __name__ == "__main__":
     script_name = os.path.basename(__file__)
@@ -382,6 +406,15 @@ def run_scenario_in_docker(work_dir, requirements, timeout=600):
         + "' directory. (default: requirements.txt)",
         default=None,
     )
+    parser.add_argument(
+        "-d",
+        "--docker-image",
+        type=str,
+        help="The Docker image to use when running scenarios. Can not be used together with --native. (default: '"
+        + DEFAULT_DOCKER_IMAGE_TAG
+        + "', which will be created if not present)",
+        default=None,
+    )
     parser.add_argument(
         "--native",
         action="store_true",
@@ -395,6 +428,10 @@ def run_scenario_in_docker(work_dir, requirements, timeout=600):
     if len(config_list) == 0:
         raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), args.config)
 
+    # Don't allow both --docker-image and --native on the same command
+    if args.docker_image is not None and args.native:
+        sys.exit("The options --native and --docker-image can not be used together. Exiting.")
+
     # Warn if running natively
     if args.native:
         if IS_WIN32:
@@ -434,4 +471,4 @@ def run_scenario_in_docker(work_dir, requirements, timeout=600):
             f"The environment file '{env_file}' does not exist (perhaps this is your first time setting up the testbed). A default environment file has been provided, but you may want to edit it to include your API keys and configurations.\n"
         )
 
-    run_scenarios(args.scenario, args.repeat, is_native, config_list, requirements)
+    run_scenarios(args.scenario, args.repeat, is_native, config_list, requirements, docker_image=args.docker_image)