Hyperparameter sweeps#

This example shows how you can run a hyperparameter sweep in Beaker using beaker-py.


You can find the source code for this example on GitHub.


Add the following files to a directory of your choosing:

# This Dockerfile defines the image that we'll use for all of the sweep experiments
# that we submit to Beaker.

FROM python:3.9

COPY entrypoint.py .

ENTRYPOINT ["python", "entrypoint.py"]
This is the script that will run on Beaker as the Docker image's "entrypoint".

All it does is write out a simple JSON file with a random number in it to
the experiment's result directory. This is just meant to simulate the results
of a training/evaluation pipeline.

import json
import random
import sys

# NOTE: it's important that this file is called 'metrics.json'. That tells Beaker
# to collect metrics for the task from this file.
OUTPUT_PATH = "/output/metrics.json"

def main(x: int):
    with open(OUTPUT_PATH, "w") as out_file:
        json.dump({"result": random.random()}, out_file)

if __name__ == "__main__":
This script will upload an image to Beaker and then submit a bunch
of experiments with different inputs. It will wait for all experiments to finish
and then collect the results.

See the output of 'python run.py --help' for usage.

import argparse
import uuid

import petname
from rich import print, progress, table, traceback

from beaker import *

def unique_name() -> str:
    """Helper function to generate a unique name for the image, group, and each experiment."""
    return petname.generate() + "-" + str(uuid.uuid4())[:8]  # type: ignore

def main(image: str, workspace: str):
    beaker = Beaker.from_env(default_workspace=workspace)
    sweep_name = unique_name()
    print(f"Starting sweep '{sweep_name}'...\n")

    # Using the `beaker.session()` context manager is not necessary, but it does
    # speed things up since it allows the Beaker client to reuse the same TCP connection
    # for all requests made within-context.
    with beaker.session():
        # Upload image to Beaker.
        print(f"Uploading image '{image}' to Beaker...")
        beaker_image = beaker.image.create(unique_name(), image)
            f"Image uploaded as '{beaker_image.full_name}', view at {beaker.image.url(beaker_image)}\n"

        # Launch experiments.
        experiments = []
        for x in progress.track(range(5), description="Launching experiments..."):
            spec = ExperimentSpec.new(
                description=f"Run {x+1} of sweep {sweep_name}",
            experiment = beaker.experiment.create(f"{sweep_name}-{x+1}", spec)

        # Create group.
        print("Creating group for sweep...")
        group = beaker.group.create(
            sweep_name, *experiments, description="Group for sweep {sweep_name}"
        print(f"Group '{group.full_name}' created, view at {beaker.group.url(group)}\n")

        # Wait for experiments to finish.
        print("Waiting for experiments to finalize...\n")
        experiments = beaker.experiment.wait_for(*experiments)

        # Display results as a table.
        results_table = table.Table(title="Results for sweep")
        for x, experiment in enumerate(
            progress.track(experiments, description="Gathering results...")
            metrics = beaker.experiment.metrics(experiment)
            assert metrics is not None
            results_table.add_row(f"x = {x}", f"{metrics['result']:.4f}")

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Run a hyperparameter sweep in Beaker")
        "image", type=str, help="""The tag of the local Docker image built from the Dockerfile."""
    parser.add_argument("workspace", type=str, help="""The Beaker workspace to use.""")
    opts = parser.parse_args()

    main(image=opts.image, workspace=opts.workspace)

Running it#

To run it, first build the Docker image:

docker build -t $image .

Then launch the sweep with:

workspace=ai2/my-sweep  # change this to the workspace of your choosing
cluster=ai2/petew-cpu  # change this to the cluster of your choosing
python run.py $image $workspace $cluster