ax.service

Ax Client

class ax.service.ax_client.AxClient(generation_strategy=None, db_settings=None, enforce_sequential_optimization=True)[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[DBSettings]) – 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.
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)[source]

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

Parameters:
  • trial_index (int) – Index of trial within the experiment.
  • raw_data (Union[Dict[str, Tuple[float, float]], Tuple[float, float], float]) – 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.
  • metadata (Optional[Dict[str, str]]) – Additional metadata to track about this run.
Return type:

None
create_experiment(parameters, name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=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 log_data.
  • 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 >= 2”.
  • outcome_constraints (Optional[List[str]]) – List of string representation of outcome constraints of form “metric_name >= bound”, like “m1 <= 3.”
Return type:

None
experiment

Returns the experiment set on this Ax client

Return type:Experiment
get_best_parameters()[source]

Return the best set of parameters the experiment has knowledge of.

If experiment is in the optimization phase, return the best point determined by the model used in the latest optimization round, otherwise return none.

Custom type TModelPredictArm is defined as Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]], and stands for tuple of two mappings: metric name to its mean value and metric name to a mapping of other mapping name to covariance of the two metrics.

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 (best parameters, model predictions for best parameters). None if no data.
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_report()[source]

Returns HTML of a generated report containing vizualizations.

Return type:str
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
should_stop_early(trial_index, data)[source]

Whether to stop the given parameterization given early data.

Return type:bool

Managed Loop

class ax.service.managed_loop.OptimizationLoop(experiment, total_trials=20, arms_per_trial=1, 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)[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)[source]

Constructs an asynchronous OptimizationLoop using Ax runners and metrics.

Return type:OptimizationLoop
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, wait_time=0)[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)[source]

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

Return type:GenerationStrategy

Instantiation

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)[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

Storage

ax.service.utils.storage.initialize_db(db_settings)[source]

Initialize db connections given settings.

Parameters:db_settings (DBSettings) – Optional[DBSettings] = None
Return type:None
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) – Specifies decoder and xdb tier where experiment is stored.
Returns:

Loaded experiment.
Return type:

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

Save experiment to db.

Parameters:
  • experiment (Experiment) – Experiment object.
  • db_settings (DBSettings) – Specifies decoder and xdb tier where experiment is stored.
Return type:

None