ax.service¶

Ax Client¶

class ax.service.ax_client.AxClient(generation_strategy: Optional[ax.modelbridge.generation_strategy.GenerationStrategy] = None, db_settings: Optional[ax.storage.sqa_store.structs.DBSettings] = None, enforce_sequential_optimization: bool = True, random_seed: Optional[int] = None, verbose_logging: bool = True, suppress_storage_errors: bool = False)[source]¶

Bases: ax.service.utils.with_db_settings_base.WithDBSettingsBase

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 generation strategy. If not set, one is intelligently chosen based on properties of search space.

• db_settings – Settings for saving and reloading the underlying experiment to a database. Expected to be of type ax.storage.sqa_store.structs.DBSettings and require SQLAlchemy.

• enforce_sequential_optimization – Whether to enforce that when it is reasonable to switch models during the optimization (as prescribed by num_trials 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 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.

• verbose_logging – Whether Ax should log significant optimization events, defaults to True.

• suppress_storage_errors – Whether to suppress SQL storage-related errors if encounted. Only use if SQL storage is not important for the given use case, since this will only log, but not raise, an exception if its encountered while saving to DB or loading from it.

abandon_trial(trial_index: int, reason: Optional[str] = None) → None[source]¶

Abandons a trial and adds optional metadata to it.

Parameters

trial_index – Index of trial within the experiment.

attach_trial(parameters: Dict[str, Union[str, bool, float, int, None]], ttl_seconds: Optional[int] = None) → Tuple[Dict[str, Union[str, bool, float, int, None]], int][source]¶

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

Parameters

• parameters – Parameterization of the new trial.

• ttl_seconds – If specified, will consider the trial failed after this many seconds. Used to detect dead trials that were not marked failed properly.

Returns

Tuple of parameterization and trial index from newly created trial.

complete_trial(trial_index: int, raw_data: Union[Dict[str, Tuple[float, Optional[float]]], Tuple[float, Optional[float]], float, List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]], metadata: Optional[Dict[str, Union[str, int]]] = None, sample_size: Optional[int] = None) → None[source]¶

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

Parameters

• trial_index – Index of trial within the experiment.

• raw_data – Evaluation data for the trial. Can be a mapping from metric name to a tuple of mean and SEM, just a tuple of mean and SEM if only one metric in optimization, or just the mean if there is no SEM. Can also be a list of (fidelities, mapping from metric name to a tuple of mean and SEM).

• metadata – Additional metadata to track about this run.

• sample_size – Number of samples collected for the underlying arm, optional.

create_experiment(parameters: List[Dict[str, Union[str, bool, float, int, None, List[Union[str, bool, float, int, None]]]]], name: Optional[str] = None, objective_name: Optional[str] = None, minimize: bool = False, parameter_constraints: Optional[List[str]] = None, outcome_constraints: Optional[List[str]] = None, status_quo: Optional[Dict[str, Union[str, bool, float, int, None]]] = None, overwrite_existing_experiment: bool = False, experiment_type: Optional[str] = None, choose_generation_strategy_kwargs: Optional[Dict[str, Any]] = None) → None[source]¶

Create a new experiment and save it if DBSettings available.

Parameters

• parameters – List of dictionaries representing parameters in the experiment search space. Required elements in the dictionaries are: 1. “name” (name of parameter, string), 2. “type” (type of parameter: “range”, “fixed”, or “choice”, string), and one of the following: 3a. “bounds” for range parameters (list of two values, lower bound first), 3b. “values” for choice parameters (list of values), or 3c. “value” for fixed parameters (single value). Optional elements are: 1. “log_scale” (for float-valued range parameters, bool), 2. “value_type” (to specify type that values of this parameter should take; expects “float”, “int”, “bool” or “str”), 3. “is_fidelity” (bool) and “target_value” (float) for fidelity parameters, 4. “is_ordered” (bool) for choice parameters, and 5. “is_task” (bool) for task parameters.

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

• name – Name of the experiment to be created.

• minimize – Whether this experiment represents a minimization problem.

• parameter_constraints – List of string representation of parameter constraints, such as “x3 >= x4” or “-x3 + 2*x4 - 3.5*x5 >= 2”. For the latter constraints, any number of arguments is accepted, and acceptable operators are “<=” and “>=”.

• outcome_constraints – List of string representation of outcome constraints of form “metric_name >= bound”, like “m1 <= 3.”

• status_quo – Parameterization of the current state of the system. If set, this will be added to each trial to be evaluated alongside test configurations.

• overwrite_existing_experiment – If an experiment has already been set on this AxClient instance, whether to reset it to the new one. If overwriting the experiment, generation strategy will be re-selected for the new experiment and restarted. To protect experiments in production, one cannot overwrite existing experiments if the experiment is already stored in the database, regardless of the value of overwrite_existing_experiment.

• choose_generation_strategy_kwargs – Keyword arguments to pass to choose_generation_strategy function which determines what generation strategy should be used when none was specified on init.

property experiment¶

Returns the experiment set on this Ax client.

static from_json_snapshot(serialized: Dict[str, Any], **kwargs) → ax.service.ax_client.AxClient[source]¶

Recreate an AxClient from a JSON snapshot.

property generation_strategy¶

Returns the generation strategy, set on this experiment.

get_best_parameters() → Optional[Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]][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.

Returns

Tuple of parameterization and model predictions for it.

get_contour_plot(param_x: Optional[str] = None, param_y: Optional[str] = None, metric_name: Optional[str] = None) → ax.plot.base.AxPlotConfig[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 – name of parameters to use on x-axis for the contour response surface plots.

• param_y – name of parameters to use on y-axis for the contour response surface plots.

• metric_name – Name of the metric, for which to plot the response surface.

get_feature_importances(relative: bool = True) → ax.plot.base.AxPlotConfig[source]¶

Get a bar chart showing feature_importances for a metric.

A drop-down controls the metric for which the importances are displayed.

Parameters

relative – Whether the values are displayed as percentiles or as raw importance metrics.

get_max_parallelism() → List[Tuple[int, int]][source]¶

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

Some optimization algorithms profit significantly from sequential optimization (i.e. suggest a few points, get updated with data for them, repeat, see https://ax.dev/docs/bayesopt.html). Parallelism setting indicates how many trials should be running simulteneously (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 with any parallelism, 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.

Returns

Mapping of form {num_trials -> max_parallelism_setting}.

get_model_predictions(metric_names: Optional[List[str]] = None) → Dict[int, Dict[str, Tuple[float, float]]][source]¶

Retrieve model-estimated means and covariances for all metrics. Note: this function retrieves the predictions for the ‘in-sample’ arms, which means that the return mapping on this function will only contain predictions for trials that have been completed with data.

Parameters

metric_names – Names of the metrics, for which to retrieve predictions. All metrics on experiment will be retrieved if this argument was not specified.

Returns

A mapping from trial index to a mapping of metric names to tuples of predicted metric mean and SEM, of form: { trial_index -> { metric_name: ( mean, SEM ) } }.

get_next_trial(ttl_seconds: Optional[int] = None) → Tuple[Dict[str, Union[str, bool, float, int, None]], int][source]¶

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

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

Parameters

ttl_seconds – If specified, will consider the trial failed after this many seconds. Used to detect dead trials that were not marked failed properly.

Returns

Tuple of trial parameterization, trial index

get_optimization_trace(objective_optimum: Optional[float] = None) → ax.plot.base.AxPlotConfig[source]¶

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

Parameters

objective_optimum – Optimal objective, if known, for display in the visualization.

static get_recommended_max_parallelism() → None[source]¶

get_trial_parameters(trial_index: int) → Dict[str, Union[str, bool, float, int, None]][source]¶

Retrieve the parameterization of the trial by the given index.

get_trials_data_frame() → pandas.DataFrame[source]¶

static load(filepath: Optional[str] = None) → None[source]¶

static load_experiment(experiment_name: str) → None[source]¶

load_experiment_from_database(experiment_name: str, choose_generation_strategy_kwargs: Optional[Dict[str, Any]] = None) → None[source]¶

Load an existing experiment from database using the DBSettings passed to this AxClient on instantiation.

Parameters

experiment_name – Name of the experiment.

Returns

Experiment object.

static load_from_json_file(filepath: str = 'ax_client_snapshot.json', **kwargs) → ax.service.ax_client.AxClient[source]¶

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

log_trial_failure(trial_index: int, metadata: Optional[Dict[str, str]] = None) → None[source]¶

Mark that the given trial has failed while running.

Parameters

• trial_index – Index of trial within the experiment.

• metadata – Additional metadata to track about this run.

property objective_name¶

Returns the name of the objective in this optimization.

static save(filepath: Optional[str] = None) → None[source]¶

save_to_json_file(filepath: str = 'ax_client_snapshot.json') → None[source]¶

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

to_json_snapshot() → Dict[str, Any][source]¶

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

Returns

A JSON-safe dict representation of this AxClient.

update_trial_data(trial_index: int, raw_data: Union[Dict[str, Tuple[float, Optional[float]]], Tuple[float, Optional[float]], float, List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]], metadata: Optional[Dict[str, Union[str, int]]] = None, sample_size: Optional[int] = None) → None[source]¶

Attaches additional data for completed trial (for example, if trial was completed with data for only one of the required metrics and more data needs to be attached).

Parameters

• trial_index – Index of trial within the experiment.

• raw_data – Evaluation data for the trial. Can be a mapping from metric name to a tuple of mean and SEM, just a tuple of mean and SEM if only one metric in optimization, or just the mean if there is no SEM. Can also be a list of (fidelities, mapping from metric name to a tuple of mean and SEM).

• metadata – Additional metadata to track about this run.

• sample_size – Number of samples collected for the underlying arm, optional.

verify_trial_parameterization(trial_index: int, parameterization: Dict[str, Union[str, bool, float, int, None]]) → bool[source]¶

Whether the given parameterization matches that of the arm in the trial specified in the trial index.

Managed Loop¶

class ax.service.managed_loop.OptimizationLoop(experiment: ax.core.experiment.Experiment, total_trials: int = 20, arms_per_trial: int = 1, random_seed: Optional[int] = None, wait_time: int = 0, run_async: bool = False, generation_strategy: Optional[ax.modelbridge.generation_strategy.GenerationStrategy] = None)[source]¶

Bases: object

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

full_run() → ax.service.managed_loop.OptimizationLoop[source]¶

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

get_best_point() → Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]][source]¶

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

get_current_model() → Optional[ax.modelbridge.base.ModelBridge][source]¶

Obtain the most recently used model in optimization.

run_trial() → None[source]¶

Run a single step of the optimization plan.

static with_evaluation_function(parameters: List[Dict[str, Union[str, bool, float, int, None, List[Union[str, bool, float, int, None]]]]], evaluation_function: Callable[[Dict[str, Union[str, bool, float, int, None]], Optional[float]], Union[Dict[str, Tuple[float, Optional[float]]], Tuple[float, Optional[float]], float, List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]]], experiment_name: Optional[str] = None, objective_name: Optional[str] = None, minimize: bool = False, parameter_constraints: Optional[List[str]] = None, outcome_constraints: Optional[List[str]] = None, total_trials: int = 20, arms_per_trial: int = 1, wait_time: int = 0, random_seed: Optional[int] = None, generation_strategy: Optional[ax.modelbridge.generation_strategy.GenerationStrategy] = None) → OptimizationLoop[source]¶

Constructs a synchronous OptimizationLoop using an evaluation function.

classmethod with_runners_and_metrics(parameters: List[Dict[str, Union[str, bool, float, int, None, List[Union[str, bool, float, int, None]]]]], path_to_runner: str, paths_to_metrics: List[str], experiment_name: Optional[str] = None, objective_name: Optional[str] = None, minimize: bool = False, parameter_constraints: Optional[List[str]] = None, outcome_constraints: Optional[List[str]] = None, total_trials: int = 20, arms_per_trial: int = 1, wait_time: int = 0, random_seed: Optional[int] = None) → OptimizationLoop[source]¶

Constructs an asynchronous OptimizationLoop using Ax runners and metrics.

ax.service.managed_loop.optimize(parameters: List[Dict[str, Union[str, bool, float, int, None, List[Union[str, bool, float, int, None]]]]], evaluation_function: Callable[[Dict[str, Union[str, bool, float, int, None]], Optional[float]], Union[Dict[str, Tuple[float, Optional[float]]], Tuple[float, Optional[float]], float, List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]]], experiment_name: Optional[str] = None, objective_name: Optional[str] = None, minimize: bool = False, parameter_constraints: Optional[List[str]] = None, outcome_constraints: Optional[List[str]] = None, total_trials: int = 20, arms_per_trial: int = 1, random_seed: Optional[int] = None, generation_strategy: Optional[ax.modelbridge.generation_strategy.GenerationStrategy] = None) → Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]], ax.core.experiment.Experiment, Optional[ax.modelbridge.base.ModelBridge]][source]¶

Construct and run a full optimization loop.

Best Point Identification¶

ax.service.utils.best_point.get_best_from_model_predictions(experiment: ax.core.experiment.Experiment) → Optional[Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]][source]¶

Given an experiment, returns the best predicted parameterization and corresponding prediction based on the most recent Trial with predictions. If no trials have predictions returns None.

Only some models return predictions. For instance GPEI does while Sobol does not.

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.

Returns

Tuple of parameterization and model predictions for it.

ax.service.utils.best_point.get_best_parameters(experiment: ax.core.experiment.Experiment) → Optional[Tuple[Dict[str, Union[str, bool, float, int, None]], Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]][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.

Returns

Tuple of parameterization and model predictions for it.

ax.service.utils.best_point.get_best_raw_objective_point(experiment: ax.core.experiment.Experiment, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None) → Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, float]]][source]¶

Given an experiment, identifies the arm that had the best raw objective, based on the data fetched from the experiment.

Parameters

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

• optimization_config – Optimization config to use in absence or in place of the one stored on the experiment.

Returns

Tuple of parameterization and a mapping from metric name to a tuple of

the corresponding objective mean and SEM.

Instantiation¶

ax.service.utils.instantiation.constraint_from_str(representation: str, parameters: Dict[str, ax.core.parameter.Parameter]) → ax.core.parameter_constraint.ParameterConstraint[source]¶

Parse string representation of a parameter constraint.

ax.service.utils.instantiation.data_from_evaluations(evaluations: Dict[str, Union[Dict[str, Tuple[float, Optional[float]]], Tuple[float, Optional[float]], float, List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]]], trial_index: int, sample_sizes: Dict[str, int], start_time: Optional[int] = None, end_time: Optional[int] = None) → ax.core.data.Data[source]¶

Transforms evaluations into Ax Data.

Each evaluation is either a trial evaluation: {metric_name -> (mean, SEM)} or a fidelity trial evaluation for multi-fidelity optimizations: [(fidelities, {metric_name -> (mean, SEM)})].

Parameters

• evalutions – Mapping from arm name to evaluation.

• trial_index – Index of the trial, for which the evaluations are.

• sample_sizes – Number of samples collected for each arm, may be empty if unavailable.

• start_time – Optional start time of run of the trial that produced this data, in milliseconds.

• end_time – Optional end time of run of the trial that produced this data, in milliseconds.

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

Utilities for RESTful-like instantiation of Ax classes needed in AxClient.

ax.service.utils.instantiation.make_experiment(parameters: List[Dict[str, Union[str, bool, float, int, None, List[Union[str, bool, float, int, None]]]]], name: Optional[str] = None, objective_name: Optional[str] = None, minimize: bool = False, parameter_constraints: Optional[List[str]] = None, outcome_constraints: Optional[List[str]] = None, status_quo: Optional[Dict[str, Union[str, bool, float, int, None]]] = None, experiment_type: Optional[str] = None) → ax.core.experiment.Experiment[source]¶

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

ax.service.utils.instantiation.outcome_constraint_from_str(representation: str) → ax.core.outcome_constraint.OutcomeConstraint[source]¶

Parse string representation of an outcome constraint.

ax.service.utils.instantiation.parameter_from_json(representation: Dict[str, Union[str, bool, float, int, None, List[Union[str, bool, float, int, None]]]]) → ax.core.parameter.Parameter[source]¶

Instantiate a parameter from JSON representation.

ax.service.utils.instantiation.raw_data_to_evaluation(raw_data: Union[Dict[str, Tuple[float, Optional[float]]], Tuple[float, Optional[float]], float, List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]], objective_name: str, start_time: Optional[int] = None, end_time: Optional[int] = None) → Union[Dict[str, Tuple[float, Optional[float]]], Tuple[float, Optional[float]], float, List[Tuple[Dict[str, Union[str, bool, float, int, None]], Dict[str, Tuple[float, Optional[float]]]]]][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.

WithDBSettingsBase¶

class ax.service.utils.with_db_settings_base.WithDBSettingsBase(db_settings: Optional[ax.storage.sqa_store.structs.DBSettings] = None, logging_level: int = 20)[source]¶

Bases: object

Helper class providing methods for saving changes made to an experiment if db_settings property is set to a non-None value on the instance.

property db_settings¶

DB settings set on this instance; guaranteed to be non-None.

property db_settings_set¶

Whether non-None DB settings are set on this instance.