ax.service

Ax Client

class ax.service.ax_client.AxClient(generation_strategy=None, db_settings=None, enforce_sequential_optimization=True, random_seed=None)[source]

Bases: object

Convenience handler for management of experimentation cycle through a service-like API. External system manages scheduling of the cycle and makes calls to this client to get next suggestion in the experiment and log back data from the evaluation of that suggestion.

Note: AxClient expects to only propose 1 arm (suggestion) per trial; support for use cases that require use of batches is coming soon.

Two custom types used in this class for convenience are TParamValue and TParameterization. Those are shortcuts for Union[str, bool, float, int] and Dict[str, Union[str, bool, float, int]], respectively.

Parameters

  • generation_strategy (Optional[GenerationStrategy]) – Optional generation strategy. If not set, one is intelligently chosen based on properties of search space.

  • db_settings (Optional[Any]) – Settings for saving and reloading the underlying experiment to a database.

  • enforce_sequential_optimization (bool) – Whether to enforce that when it is reasonable to switch models during the optimization (as prescribed by num_arms in generation strategy), Ax will wait for enough trials to be completed with data to proceed. Defaults to True. If set to False, Ax will keep generating new trials from the previous model until enough data is gathered. Use this only if necessary; otherwise, it is more resource-efficient to optimize sequentially, by waiting until enough data is available to use the next model.

  • random_seed (Optional[int]) –

    Optional integer random seed, set to fix the optimization random seed for reproducibility. Works only for Sobol quasi-random generator and for BoTorch-powered models. For the latter models, the trials generated from the same optimization setup with the same seed, will be mostly similar, but the exact parameter values may still vary and trials latter in the optimizations will diverge more and more. This is because a degree of randomness is essential for high performance of the Bayesian optimization models and is not controlled by the seed.

    Note: In multi-threaded environments, the random seed is thread-safe, but does not actually guarantee reproducibility. Whether the outcomes will be exactly the same for two same operations that use the random seed, depends on whether the threads modify the random state in the same order across the two operations.

attach_trial(parameters)[source]

Attach a new trial with the given parameterization to the experiment.

Parameters

parameters (Dict[str, Union[str, bool, float, int, None]]) – Parameterization of the new trial.

Return type

Tuple[Dict[str, Union[str, bool, float, int, None]], int]

Returns

Tuple of parameterization and trial index from newly created trial.

complete_trial(trial_index, raw_data, metadata=None, sample_size=None)[source]

Completes the trial with given metric values and adds optional metadata to it.

Parameters
Return type

None

create_experiment(parameters, name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, status_quo=None)[source]

Create a new experiment and save it if DBSettings available.

Parameters

  • parameters (List[Dict[str, Union[str, bool, float, int, None, List[Union[str, bool, float, int, None]]]]]) – List of dictionaries representing parameters in the experiment search space. Required elements in the dictionaries are: “name” (name of this parameter, string), “type” (type of the parameter: “range”, “fixed”, or “choice”, string), and “bounds” for range parameters (list of two values, lower bound first), “values” for choice parameters (list of values), and “value” for fixed parameters (single value).

  • objective – Name of the metric used as objective in this experiment. This metric must be present in raw_data argument to complete_trial.

  • name (Optional[str]) – Name of the experiment to be created.

  • minimize (bool) – Whether this experiment represents a minimization problem.

  • parameter_constraints (Optional[List[str]]) – List of string representation of parameter constraints, such as “x3 >= x4” or “x3 + x4 + x5 >= 2”. For sum constraints, any number of arguments is accepted, and acceptable operators are “<=” and “>=”.

  • outcome_constraints (Optional[List[str]]) – List of string representation of outcome constraints of form “metric_name >= bound”, like “m1 <= 3.”

  • status_quo (Optional[Dict[str, Union[str, bool, float, int, None]]]) – Parameterization of the current state of the system. If set, this will be added to each trial to be evaluated alongside test configurations.

Return type

None

property experiment

Returns the experiment set on this Ax client

Return type

Experiment

static from_json_snapshot(serialized)[source]

Recreate an AxClient from a JSON snapshot.

Return type

AxClient

get_best_parameters()[source]

Given an experiment, identifies the best arm.

First attempts according to do so with models used in optimization and its corresponding predictions if available. Falls back to the best raw objective based on the data fetched from the experiment.

TModelPredictArm is of the form:

({metric_name: mean}, {metric_name_1: {metric_name_2: cov_1_2}})

Parameters

experiment – Experiment, on which to identify best raw objective arm.

Return type

Optional[Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]]

Returns

Tuple of parameterization and model predictions for it.

get_contour_plot(param_x=None, param_y=None, metric_name=None)[source]

Retrieves a plot configuration for a contour plot of the response surface. For response surfaces with more than two parameters, selected two parameters will appear on the axes, and remaining parameters will be affixed to the middle of their range. If contour params arguments are not provided, the first two parameters in the search space will be used. If contour metrics are not provided, objective will be used.

Parameters

  • param_x (Optional[str]) – name of parameters to use on x-axis for the contour response surface plots.

  • param_y (Optional[str]) – name of parameters to use on y-axis for the contour response surface plots.

  • metric_name (Optional[str]) – Name of the metric, for which to plot the response surface.

Return type

AxPlotConfig

get_next_trial()[source]

Generate trial with the next set of parameters to try in the iteration process.

Note: Service API currently supports only 1-arm trials.

Return type

Tuple[Dict[str, Union[str, bool, float, int, None]], int]

Returns

Tuple of trial parameterization, trial index

get_optimization_trace(objective_optimum=None)[source]

Retrieves the plot configuration for optimization trace, which shows the evolution of the objective mean over iterations.

Parameters

objective_optimum (Optional[float]) – Optimal objective, if known, for display in the visualization.

Return type

AxPlotConfig

Recommends maximum number of trials that can be scheduled in parallel at different stages of optimization.

Some optimization algorithms profit significantly from sequential optimization (e.g. suggest a few points, get updated with data for them, repeat). This setting indicates how many trials should be in flight (generated, but not yet completed with data).

The output of this method is mapping of form {num_trials -> max_parallelism_setting}, where the max_parallelism_setting is used for num_trials trials. If max_parallelism_setting is -1, as many of the trials can be ran in parallel, as necessary. If num_trials in a tuple is -1, then the corresponding max_parallelism_setting should be used for all subsequent trials.

For example, if the returned list is [(5, -1), (12, 6), (-1, 3)], the schedule could be: run 5 trials in parallel, run 6 trials in parallel twice, run 3 trials in parallel for as long as needed. Here, ‘running’ a trial means obtaining a next trial from AxClient through get_next_trials and completing it with data when available.

Return type

List[Tuple[int, int]]

Returns

Mapping of form {num_trials -> max_parallelism_setting}.

get_report()[source]

Returns HTML of a generated report containing vizualizations.

Return type

str

get_trial_parameters(trial_index)[source]

Retrieve the parameterization of the trial by the given index.

Return type

Dict[str, Union[str, bool, float, int, None]]

static load(filepath='ax_client_snapshot.json')[source]

Restore an AxClient and its state from a JSON-serialized snapshot, residing in a .json file by the given path.

Return type

AxClient

load_experiment(experiment_name)[source]

[Work in progress] Load an existing experiment.

Parameters

experiment_name (str) – Name of the experiment.

Return type

None

Returns

Experiment object.

log_trial_failure(trial_index, metadata=None)[source]

Mark that the given trial has failed while running.

Parameters

  • trial_index (int) – Index of trial within the experiment.

  • metadata (Optional[Dict[str, str]]) – Additional metadata to track about this run.

Return type

None

save(filepath='ax_client_snapshot.json')[source]

Save a JSON-serialized snapshot of this AxClient’s settings and state to a .json file by the given path.

Return type

None

should_stop_early(trial_index, data)[source]

Whether to stop the given parameterization given early data.

Return type

bool

to_json_snapshot()[source]

Serialize this AxClient to JSON to be able to interrupt and restart optimization and save it to file by the provided path.

Return type

Dict[str, Any]

class ax.service.ax_client.DBSettings[source]

Bases: tuple

Defines behavior for loading/saving experiment to/from db. Either creator or url must be specified as a way to connect to the SQL db.

property creator

Alias for field number 0

property decoder

Alias for field number 1

property encoder

Alias for field number 2

property url

Alias for field number 3

ax.service.ax_client.logger = <Logger ax.service.ax_client (INFO)>

Managed Loop

class ax.service.managed_loop.OptimizationLoop(experiment, total_trials=20, arms_per_trial=1, random_seed=None, wait_time=0, run_async=False)[source]

Bases: object

Managed optimization loop, in which Ax oversees deployment of trials and gathering data.

full_run()[source]

Runs full optimization loop as defined in the provided optimization plan.

Return type

OptimizationLoop

get_best_point()[source]

Obtains the best point encountered in the course of this optimization.

Return type

Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]

get_current_model()[source]

Obtain the most recently used model in optimization.

Return type

Optional[ModelBridge]

run_trial()[source]

Run a single step of the optimization plan.

Return type

None

static with_evaluation_function(parameters, evaluation_function, experiment_name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, total_trials=20, arms_per_trial=1, wait_time=0, random_seed=None)[source]

Constructs a synchronous OptimizationLoop using an evaluation function.

Return type

OptimizationLoop

classmethod with_runners_and_metrics(parameters, path_to_runner, paths_to_metrics, experiment_name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, total_trials=20, arms_per_trial=1, wait_time=0, random_seed=None)[source]

Constructs an asynchronous OptimizationLoop using Ax runners and metrics.

Return type

OptimizationLoop

ax.service.managed_loop.logger = <Logger ax.service.managed_loop (INFO)>
ax.service.managed_loop.optimize(parameters, evaluation_function, experiment_name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, total_trials=20, arms_per_trial=1, random_seed=None)[source]

Construct and run a full optimization loop.

Return type

Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]], Experiment, Optional[ModelBridge]]

Utils

Dispatch

ax.service.utils.dispatch.choose_generation_strategy(search_space, arms_per_trial=1, enforce_sequential_optimization=True, random_seed=None)[source]

Select an appropriate generation strategy based on the properties of the search space.

Return type

GenerationStrategy

ax.service.utils.dispatch.logger = <Logger ax.service.utils.dispatch (INFO)>

Instantiation

ax.service.utils.instantiation.COMPARISON_OPS = {'<=': <ComparisonOp.LEQ: 1>, '>=': <ComparisonOp.GEQ: 0>}
ax.service.utils.instantiation.PARAM_CLASSES = ['range', 'choice', 'fixed']
ax.service.utils.instantiation.PARAM_TYPES = {'bool': <class 'bool'>, 'float': <class 'float'>, 'int': <class 'int'>, 'str': <class 'str'>}
ax.service.utils.instantiation.TParameterRepresentation = typing.Dict[str, typing.Union[str, bool, float, int, NoneType, typing.List[typing.Union[str, bool, float, int, NoneType]]]]
ax.service.utils.instantiation.constraint_from_str(representation, parameters)[source]

Parse string representation of a parameter constraint.

Return type

ParameterConstraint

ax.service.utils.instantiation.make_experiment(parameters, name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, status_quo=None)[source]

Instantiation wrapper that allows for creation of SimpleExperiment without importing or instantiating any Ax classes.

Return type

Experiment

ax.service.utils.instantiation.outcome_constraint_from_str(representation)[source]

Parse string representation of an outcome constraint.

Return type

OutcomeConstraint

ax.service.utils.instantiation.parameter_from_json(representation)[source]

Instantiate a parameter from JSON representation.

Return type

Parameter

ax.service.utils.instantiation.raw_data_to_evaluation(raw_data, objective_name)[source]

Format the trial evaluation data to a standard TTrialEvaluation (mapping from metric names to a tuple of mean and SEM) representation, or to a TFidelityTrialEvaluation.

Note: this function expects raw_data to be data for a Trial, not a BatchedTrial.

Return type

Union[Dict[str, Tuple[float, Optional[float]]], List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]]

Storage

ax.service.utils.storage.load_experiment(name, db_settings)[source]

Load experiment from the db. Service API only supports Experiment.

Parameters

  • name (str) – Experiment name.

  • db_settings (DBSettings) – Defines behavior for loading/saving experiment to/from db.

Returns

Loaded experiment.

Return type

ax.core.Experiment

ax.service.utils.storage.load_experiment_and_generation_strategy(experiment_name, db_settings)[source]

Load experiment and the corresponding generation strategy from the DB.

Parameters

  • name – Experiment name.

  • db_settings (DBSettings) – Defines behavior for loading/saving experiment to/from db.

Return type

Tuple[Experiment, GenerationStrategy]

Returns

A tuple of the loaded experiment and generation strategy.

ax.service.utils.storage.save_experiment(experiment, db_settings)[source]

Save experiment to db.

Parameters

  • experiment (Experiment) – Experiment object.

  • db_settings (DBSettings) – Defines behavior for loading/saving experiment to/from db.

Return type

None

ax.service.utils.storage.save_experiment_and_generation_strategy(experiment, generation_strategy, db_settings)[source]

Save experiment and generation strategy to DB.

Parameters

  • experiment (Experiment) – Experiment object.

  • generation_strategy (GenerationStrategy) – Corresponding generation strategy.

  • db_settings (DBSettings) – Defines behavior for loading/saving experiment to/from db.

Return type

None