Open in GitHubRun in Google Colab

Automating Orchestration

Previously, we've demonstrated using Ax for ask-tell optimization, a paradigm in which we "ask" Ax for candidate configurations and "tell" Ax our observations. This can be effective in many scenerios, and it can be automated through use of flow control statements like for and while loops. However there are some situations where it would be beneficial to allow Ax to orchestrate the entire optimization: deploying trials to external systems, polling their status, and reading reading their results. This can be common in a number of real world engineering tasks, including:

• Large scale machine learning experiments running workloads on high-performance computing clusters
• A/B tests conducted using an external experimentation platform
• Materials science optimizations utilizing a self-driving laboratory

Ax's Client can orchestrate automated adaptive experiments like this using its method run_trials. Users create custom classes which implement Ax's IMetric and IRunner protocols to handle data fetching and trial deployment respectively. Then, users simply configure their Client as they would normally and call run_trials; Ax will deploy trials, fetch data, generate candidates, and repeat as necessary. Ax can manage complex orchestration tasks including launching multiple trials in parallel while still respecting a user-defined concurrency limit, and gracefully handling trial failure by allowing the experiment to continue even if some trials do not complete successfully or data fetching fails.

In this tutorial we will optimize the Hartmann6 function as before, but we will configure custom Runners and Metrics to mimic an external execution system. The Runner will calculate Hartmann6 with the appropriate parameters, write the result to a file, and tell Ax the trial is ready after 5 seconds. The Metric will find the appropriate file and report the results back to Ax.

Learning Objectives

• Learn when it can be appropriate and/or advantageous to run Ax in a closed-loop
• Configure custom Runners and Metrics, allowing Ax to deploy trials and fetch data automatically
• Understand tradeoffs between parallelism and optimization performance

Prerequisites

• Understanding of adaptive experimentation and Bayesian optimization
• Familiarity with configuring and conducting experiments in Ax

Step 1: Import Necessary Modules

First, ensure you have all the necessary imports:

import os
import time
from typing import Any, Mapping

import numpy as np
from ax.api.client import Client
from ax.api.configs import RangeParameterConfig
from ax.api.protocols.metric import IMetric
from ax.api.protocols.runner import IRunner, TrialStatus
from ax.api.types import TParameterization

Step 2: Defining our custom Runner and Metric

As stated before, we will be creating custom Runner and Metric classes to mimic an external system. Let's start by defining our Hartmann6 function as before.

# Hartmann6 function
def hartmann6(x1, x2, x3, x4, x5, x6):
    alpha = np.array([1.0, 1.2, 3.0, 3.2])
    A = np.array([
        [10, 3, 17, 3.5, 1.7, 8],
        [0.05, 10, 17, 0.1, 8, 14],
        [3, 3.5, 1.7, 10, 17, 8],
        [17, 8, 0.05, 10, 0.1, 14]
    ])
    P = 10**-4 * np.array([
        [1312, 1696, 5569, 124, 8283, 5886],
        [2329, 4135, 8307, 3736, 1004, 9991],
        [2348, 1451, 3522, 2883, 3047, 6650],
        [4047, 8828, 8732, 5743, 1091, 381]
    ])

    outer = 0.0
    for i in range(4):
        inner = 0.0
        for j, x in enumerate([x1, x2, x3, x4, x5, x6]):
            inner += A[i, j] * (x - P[i, j])**2
        outer += alpha[i] * np.exp(-inner)
    return -outer

hartmann6(0.1, 0.45, 0.8, 0.25, 0.552, 1.0)

Next, we will define the MockRunner. The MockRunner requires two methods: run_trial and poll_trial.

run_trial deploys a trial to the external system with the given parameters. In this case, we will simply save a file containing the result of a call to the Hartmann6 function.

poll_trial queries the external system to see if the trial has completed, failed, or if it's still running. In this mock example, we will check to see how many seconds have elapsed since the run_trial was called and only report a trial as completed once 5 seconds have elapsed.

Runner's may also optionally implement a stop_trial method to terminate a trial's execution before it has completed. This is necessary for using early stopping in closed-loop experimentation, but we will skip this for now.

class MockRunner(IRunner):
    def run_trial(
        self, trial_index: int, parameterization: TParameterization
    ) -> dict[str, Any]:
        file_name = f"{int(time.time())}.txt"

        x1 = parameterization["x1"]
        x2 = parameterization["x2"]
        x3 = parameterization["x3"]
        x4 = parameterization["x4"]
        x5 = parameterization["x5"]
        x6 = parameterization["x6"]

        result = hartmann6(x1, x2, x3, x4, x5, x6)

        with open(file_name, "w") as f:
            f.write(f"{result}")

        return {"file_name": file_name}

    def poll_trial(
        self, trial_index: int, trial_metadata: Mapping[str, Any]
    ) -> TrialStatus:
        file_name = trial_metadata["file_name"]
        time_elapsed = time.time() - int(file_name[:4])

        if time_elapsed < 5:
            return TrialStatus.RUNNING

        return TrialStatus.COMPLETED

It's worthwhile to instantiate your Runner and test it is behaving as expected. Let's deploy a mock trial by manually calling run_trial and ensuring it creates a file.

runner = MockRunner()

trial_metadata = runner.run_trial(
    trial_index=-1,
    parameterization={
        "x1": 0.1,
        "x2": 0.45,
        "x3": 0.8,
        "x4": 0.25,
        "x5": 0.552,
        "x6": 1.0,
    },
)

os.path.exists(trial_metadata["file_name"])

Now, we will implement the Metric. Metrics only need to implement a fetch method, which returns a progression value (i.e. a step in a timeseries) and an observation value. Note that the observation can either be a simple float or a (mean, SEM) pair if the external system can report observed noise.

In this case, we have neither a relevant progression value nor observed noise so we will simply read the file and report (0, value).

class MockMetric(IMetric):
    def fetch(
        self,
        trial_index: int,
        trial_metadata: Mapping[str, Any],
    ) -> tuple[int, float | tuple[float, float]]:
        file_name = trial_metadata["file_name"]

        with open(file_name, 'r') as file:
            value = float(file.readline())
            return (0, value)

Again, let's validate the Metric created above by instantiating it and reporting the value from the file generated during testing of the Runner.

# Note: all Metrics must have a name. This will become relevant when attaching metrics to the Client
hartmann6_metric = MockMetric(name="hartmann6")

hartmann6_metric.fetch(trial_index=-1, trial_metadata=trial_metadata)

Step 3: Initialize the Client and Configure the Experiment

Finally, we can initialize the Client and configure the experiment as before. This will be familiar to readers of the Getting Started with Ax tutorial -- the only difference is we will attach the previously defined Runner and Metric by calling configure_runner and configure_metrics respectively.

Note that when initializing hartmann6_metric we set name=hartmann6, matching the objective we now set in configure_optimization. The configure_metrics method uses this name to ensure that data fetched by this Metric is used correctly during the experiment. Be careful to correctly set the name of the Metric to reflect its use as an objective or outcome constraint.

client = Client()
# Define six float parameters for the Hartmann6 function
parameters = [
    RangeParameterConfig(name=f"x{i + 1}", parameter_type="float", bounds=(0, 1))
    for i in range(6)
]

client.configure_experiment(
    parameters=parameters,
    # The following arguments are only necessary when saving to the DB
    name="hartmann6_experiment",
    description="Optimization of the Hartmann6 function",
    owner="developer",
)
client.configure_optimization(objective="-hartmann6")

client.configure_runner(runner=runner)
client.configure_metrics(metrics=[hartmann6_metric])

Step 5: Run trials

Once the Client has been configured, we can begin running trials.

Internally, Ax uses a class named Scheduler to orchestrate the trial deployment, polling, data fetching, and candidate generation.

The run_trials method provides users with control over various orchestration settings as well as the total maximum number of trials to evaluate:

• parallelism defines the maximum number of trials that may be run at once. If your external system supports multiple evaluations in parallel, increasing this number can significantly decrease experimentation time. However, it is important to note that as parallelism increases, optimiztion performance often decreases. This is because adaptive experimentation methods rely on previously observed data for candidate generation -- the more tirals that have been observed prior to generation of a new candidate, the more accurate Ax's model will be for generation of that candidate.
• tolerated_trial_failure_rate sets the proportion of trials are allowed to fail before Ax raises an Exception. Depending on how expensive a single trial is to evaluate or how unreliable trials are expected to be, the experimenter may want to be notified as soon as a single trial fails or they may not care until more than half the trials are failing. Set this value as is appropriate for your context.
• initial_seconds_between_polls sets the frequency at which the status of a trial is checked and the results are attempted to be fetched. Set this to be low for trials that are expected to complete quickly or high for trials the are expected to take a long time.

client.run_trials(
    max_trials=30,
    parallelism=3,
    tolerated_trial_failure_rate=0.1,
    initial_seconds_between_polls=1,
)

Step 6: Analyze Results

As before, Ax can compute the best parameterization observed and produce a number of analyses to help interpret the results of the experiment.

It is also worth noting that the experiment can be resumed at any time using Ax's storage functionality. When configured to use a SQL databse, the Client saves a snapshot of itself at various points throughout the call to run_trials, making it incredibly easy to continue optimization after an unexpected failure. You can learn more about storage in Ax here.

best_parameters, prediction, index, name = client.get_best_parameterization()
print("Best Parameters:", best_parameters)
print("Prediction (mean, variance):", prediction)

# display=True instructs Ax to sort then render the resulting analyses
cards = client.compute_analyses(display=True)

Conclusion

This tutorial demonstrates how to use Ax's Client for closed-loop optimization using the Hartmann6 function as an example. This style of optimization is useful in scenarios where trials are evaluated on some external system or when experimenters wish to take advantage of parallel evaluation, trial failure handling, or simply to manage long-running optimization tasks without human intervention. You can define your own Runner and Metric classes to communicate with whatever external systems you wish to interface with, and control optimization using the OrchestrationConfig.