diff --git a/samples/tools/testbed/Dockerfile b/samples/tools/testbed/Dockerfile new file mode 100644 index 000000000000..6ce06f93a621 --- /dev/null +++ b/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 diff --git a/samples/tools/testbed/README.md b/samples/tools/testbed/README.md index 8c34f3e04af0..76560b1c6bb5 100644 --- a/samples/tools/testbed/README.md +++ b/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. ``` diff --git a/samples/tools/testbed/run_scenarios.py b/samples/tools/testbed/run_scenarios.py index 6020a7591a4b..059d97345d67 100644 --- a/samples/tools/testbed/run_scenarios.py +++ b/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)