ax.service¶

Bases: WithDBSettingsBase, BestPointMixin, InstantiationBase

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.

• torch_device – An optional torch.device object, used to choose the device used for generating new points for trials. Works only for torch-based models, such as MBM. Ignored if a generation_strategy is passed in manually. To specify the device for a custom generation_strategy, pass in torch_device as part of model_kwargs. See https://ax.dev/tutorials/generation_strategy.html for a tutorial on generation strategies.

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

• suppress_storage_errors – Whether to suppress SQL storage-related errors if encountered. 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.

• early_stopping_strategy – A BaseEarlyStoppingStrategy that determines whether a trial should be stopped given the current state of the experiment. Used in should_stop_trials_early.

• global_stopping_strategy – A BaseGlobalStoppingStrategy that determines whether the full optimization should be stopped or not.

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

Abandons a trial and adds optional metadata to it.

Parameters:

trial_index – Index of trial within the experiment.

add_tracking_metrics(metric_names: list[str], metric_definitions: dict[str, dict[str, Any]] | None = None, metrics_to_trial_types: dict[str, str] | None = None, canonical_names: dict[str, str] | None = None) → None[source]¶

Add a list of new metrics to the experiment.

If any of the metrics are already defined on the experiment, we raise an error and don’t add any of them to the experiment

Parameters:

• metric_names – Names of metrics to be added.

• metric_definitions – A mapping of metric names to extra kwargs to pass to that metric. Note these are modified in-place. Each Metric must have its is own dictionary (metrics cannot share a single dictionary object).

• metrics_to_trial_types – Only applicable to MultiTypeExperiment. The mapping from metric names to corresponding trial types for each metric. If provided, the metrics will be added with their respective trial types. If not provided, then the default trial type will be used.

• canonical_names – A mapping from metric name (of a particular trial type) to the metric name of the default trial type. Only applicable to MultiTypeExperiment.

attach_trial(parameters: dict[str, None | str | bool | float | int], ttl_seconds: int | None = None, run_metadata: dict[str, Any] | None = None, arm_name: str | None = None) → tuple[dict[str, None | str | bool | float | int], 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: dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]], metadata: dict[str, str | int] | None = None, sample_size: int | None = None) → None[source]¶

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

NOTE: When raw_data does not specify SEM for a given metric, Ax will default to the assumption that the data is noisy (specifically, corrupted by additive zero-mean Gaussian noise) and that the level of noise should be inferred by the optimization model. To indicate that the data is noiseless, set SEM to 0.0, for example:

ax_client.complete_trial(
    trial_index=0,
    raw_data={"my_objective": (objective_mean_value, 0.0)}
)

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 SEM is unknown (then Ax will infer observation noise level). 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, None | str | bool | float | int | Sequence[None | str | bool | float | int] | dict[str, list[str]]]], name: str | None = None, description: str | None = None, owners: list[str] | None = None, objectives: dict[str, ObjectiveProperties] | None = None, parameter_constraints: list[str] | None = None, outcome_constraints: list[str] | None = None, status_quo: dict[str, None | str | bool | float | int] | None = None, overwrite_existing_experiment: bool = False, experiment_type: str | None = None, tracking_metric_names: list[str] | None = None, choose_generation_strategy_kwargs: dict[str, Any] | None = None, support_intermediate_data: bool = False, immutable_search_space_and_opt_config: bool = True, is_test: bool = False, metric_definitions: dict[str, dict[str, Any]] | None = None, default_trial_type: str | None = None, default_runner: Runner | None = 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. 6. “digits” (int) for float-valued range parameters.

• name – Name of the experiment to be created.

• description – Description of the experiment to be created.

• objectives – Mapping from an objective name to object containing: minimize: Whether this experiment represents a minimization problem. threshold: The bound in the objective’s threshold constraint.

• 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 “>=”. Note that parameter constraints may only be placed on range parameters, not choice parameters or fixed parameters.

• 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.

• tracking_metric_names – Names of additional tracking metrics not used for optimization.

• 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.

• support_intermediate_data – Whether trials may report intermediate results for trials that are still running (i.e. have not been completed via ax_client.complete_trial).

• immutable_search_space_and_opt_config – Whether it’s possible to update the search space and optimization config on this experiment after creation. Defaults to True. If set to True, we won’t store or load copies of the search space and optimization config on each generator run, which will improve storage performance.

• is_test – Whether this experiment will be a test experiment (useful for marking test experiments in storage etc). Defaults to False.

• metric_definitions – A mapping of metric names to extra kwargs to pass to that metric. Note these are modified in-place. Each Metric must have its own dictionary (metrics cannot share a single dictionary object).

• default_trial_type – The default trial type if multiple trial types are intended to be used in the experiment. If specified, a MultiTypeExperiment will be created. Otherwise, a single-type Experiment will be created.

• default_runner – The default runner in this experiment. This applies to MultiTypeExperiment (when default_trial_type is specified) and needs to be specified together with default_trial_type. This will be ignored for single-type Experiment (when default_trial_type is not specified).

property early_stopping_strategy: BaseEarlyStoppingStrategy | None¶

The early stopping strategy used on the experiment.

estimate_early_stopping_savings(map_key: str | None = None) → float[source]¶

Estimate early stopping savings using progressions of the MapMetric present on the EarlyStoppingConfig as a proxy for resource usage.

Parameters:

map_key – The name of the map_key by which to estimate early stopping savings, usually steps. If none is specified use some arbitrary map_key in the experiment’s MapData

Returns:

The estimated resource savings as a fraction of total resource usage (i.e. 0.11 estimated savings indicates we would expect the experiment to have used 11% more resources without early stopping present)

property experiment: Experiment¶

Returns the experiment set on this Ax client.

fit_model() → None[source]¶

Fit a model using data collected from the trials so far.

This method will attempt to fit the same model that would be used for generating the next trial. The resulting model may be different from the model that was used to generate the last trial, if the generation node is ready to transition.

This method rarely needs to be called by the user, because model-fitting is usually handled indirectly through AxClient.get_next_trial(). This method instantiates a new model if none is yet available, which may be the case if no trials have been generated using a model-based method.

NOTE: If the current generation node is not model-based, no model may be fit.

classmethod from_json_snapshot(serialized: dict[str, Any], decoder_registry: dict[str, type[T] | Callable[[...], T]] | None = None, class_decoder_registry: None | dict[str, Callable[[dict[str, Any]], Any]] = None, **kwargs) → AxClientSubclass[source]¶

Recreate an AxClient from a JSON snapshot.

property generation_strategy: GenerationStrategy¶

Returns the generation strategy, set on this experiment.

get_best_trial(optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None, use_model_predictions: bool = True) → tuple[int, dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None[source]¶

Identifies the best parameterization tried in the experiment so far.

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

NOTE: TModelPredictArm is of the form:

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

Parameters:

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

• trial_indices – Indices of trials for which to retrieve data. If None will retrieve data from all available trials.

• use_model_predictions – Whether to extract the best point using model predictions or directly observed values. If True, the metric means and covariances in this method’s output will also be based on model predictions and may differ from the observed values.

Returns:

Tuple of trial index, parameterization and model predictions for it.

get_contour_plot(param_x: str | None = None, param_y: str | None = None, metric_name: str | None = None) → 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_current_trial_generation_limit() → tuple[int, bool][source]¶

How many trials this AxClient instance can currently produce via calls to get_next_trial, before more trials are completed, and whether the optimization is complete.

NOTE: If return value of this function is (0, False), no more trials can currently be procuded by this AxClient instance, but optimization is not completed; once more trials are completed with data, more new trials can be generated.

Returns: a two-item tuple of:

• the number of trials that can currently be produced, with -1 meaning unlimited trials,

• whether no more trials can be produced by this AxClient instance at any point (e.g. if the search space is exhausted or generation strategy is completed.

get_feature_importances(relative: bool = True) → 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_hypervolume(optimization_config: MultiObjectiveOptimizationConfig | None = None, trial_indices: Iterable[int] | None = None, use_model_predictions: bool = True) → float[source]¶

Calculate hypervolume of a pareto frontier based on either the posterior means of given observation features or observed data.

Parameters:

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

• trial_indices – Indices of trials for which to retrieve data. If None will retrieve data from all available trials.

• use_model_predictions – Whether to extract the Pareto frontier using model predictions or directly observed values. If True, the metric means and covariances in this method’s output will also be based on model predictions and may differ from the observed values.

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: list[str] | None = None, include_out_of_sample: bool | None = True, parameterizations: dict[int, dict[str, None | str | bool | float | int]] | None = None) → dict[int, dict[str, tuple[float, float]]][source]¶

Retrieve model-estimated means and covariances for all metrics.

NOTE: This method currently only supports one-arm trials.

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.

• include_out_of_sample – Defaults to True. Return predictions for out-of-sample (i.e. not yet completed trials) data in addition to in-sample (i.e. completed trials) data.

• parameterizations – Optional mapping from an int label to Parameterizations. When provided, predictions are performed only on these data points, no predictions from trial data is performed, and include_out_of_sample parameters is ignored.

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_model_predictions_for_parameterizations(parameterizations: list[dict[str, None | str | bool | float | int]], metric_names: list[str] | None = None) → list[dict[str, tuple[float, float]]][source]¶

Retrieve model-estimated means and covariances for all metrics for the provided parameterizations.

Parameters:

• metric_names – Names of the metrics for which to predict. All metrics will be predicted if this argument is not specified.

• parameterizations – List of Parameterizations for which to predict.

Returns:

List[Tuple[float, float]].

Return type:

A list of predicted metric mean and SEM of form

get_next_trial(ttl_seconds: int | None = None, force: bool = False, fixed_features: FixedFeatures | None = None) → tuple[dict[str, None | str | bool | float | int], 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.

• force – If set to True, this function will bypass the global stopping strategy’s decision and generate a new trial anyway.

• fixed_features – A FixedFeatures object containing any features that should be fixed at specified values during generation.

Returns:

Tuple of trial parameterization, trial index

get_next_trials(max_trials: int, ttl_seconds: int | None = None, fixed_features: FixedFeatures | None = None) → tuple[dict[int, dict[str, None | str | bool | float | int]], bool][source]¶

Generate as many trials as currently possible.

NOTE: Useful for running multiple trials in parallel: produces multiple trials, with their number limited by:

• parallelism limit on current generation step,

• number of trials in current generation step,

• number of trials required to complete before moving to next generation step, if applicable,

• and max_trials argument to this method.

Parameters:

• max_trials – Limit on how many trials the call to this method should produce.

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

• fixed_features – A FixedFeatures object containing any features that should be fixed at specified values during generation.

Returns: two-item tuple of:

• mapping from trial indices to parameterizations in those trials,

• boolean indicator of whether optimization is completed and no more trials can be generated going forward.

get_optimization_trace(objective_optimum: float | None = None) → 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.

get_pareto_optimal_parameters(optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None, use_model_predictions: bool = True) → dict[int, tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None]]][source]¶

Identifies the best parameterizations tried in the experiment so far, using model predictions if use_model_predictions is true and using observed values from the experiment otherwise. By default, uses model predictions to account for observation noise.

NOTE: The format of this method’s output is as follows: { trial_index –> (parameterization, (means, covariances) }, where means are a dictionary of form { metric_name –> metric_mean } and covariances are a nested dictionary of form { one_metric_name –> { another_metric_name: covariance } }.

Parameters:

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

• trial_indices – Indices of trials for which to retrieve data. If None will retrieve data from all available trials.

• use_model_predictions – Whether to extract the Pareto frontier using model predictions or directly observed values. If True, the metric means and covariances in this method’s output will also be based on model predictions and may differ from the observed values.

Returns:

• the parameterization of the arm in that trial,

• two-item tuple of metric means dictionary and covariance matrix

  (model-predicted if use_model_predictions=True and observed otherwise).

Raises a NotImplementedError if extracting the Pareto frontier is not possible. Note that the returned dict may be empty.

Return type:

A mapping from trial index to the tuple of

static get_recommended_max_parallelism() → None[source]¶

get_trace(optimization_config: MultiObjectiveOptimizationConfig | None = None) → list[float][source]¶

Get the optimization trace of the given experiment.

The output is equivalent to calling _get_hypervolume or _get_best_trial repeatedly, with an increasing sequence of trial_indices and with use_model_predictions = False, though this does it more efficiently.

Parameters:

• experiment – The experiment to get the trace for.

• optimization_config – An optional optimization config to use for computing the trace. This allows computing the traces under different objectives or constraints without having to modify the experiment.

Returns:

A list of observed hypervolumes or best values.

get_trace_by_progression(optimization_config: OptimizationConfig | None = None, bins: list[float] | None = None, final_progression_only: bool = False) → tuple[list[float], list[float]][source]¶

Get the optimization trace with respect to trial progressions instead of trial_indices (which is the behavior used in get_trace). Note that this method does not take into account the parallelism of trials and essentially assumes that trials are run one after another, in the sense that it considers the total number of progressions “used” at the end of trial k to be the cumulative progressions “used” in trials 0,…,k. This method assumes that the final value of a particular trial is used and does not take the best value of a trial over its progressions.

The best observed value is computed at each value in bins (see below for details). If bins is not supplied, the method defaults to a heuristic of approximately NUM_BINS_PER_TRIAL per trial, where each trial is assumed to run until maximum progression (inferred from the data).

Parameters:

• experiment – The experiment to get the trace for.

• optimization_config – An optional optimization config to use for computing the trace. This allows computing the traces under different objectives or constraints without having to modify the experiment.

• bins – A list progression values at which to calculate the best observed value. The best observed value at bins[i] is defined as the value observed in trials 0,…,j where j = largest trial such that the total progression in trials 0,…,j is less than bins[i].

• final_progression_only – If True, considers the value of the last step to be the value of the trial. If False, considers the best along the curve to be the value of the trial.

Returns:

A tuple containing (1) the list of observed hypervolumes or best values and (2) a list of associated x-values (i.e., progressions) useful for plotting.

get_trial(trial_index: int) → Trial[source]¶

Return a trial on experiment cast as Trial

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

Retrieve the parameterization of the trial by the given index.

get_trials_data_frame() → DataFrame[source]¶

Get a Pandas DataFrame representation of this experiment. The columns will include all the parameters in the search space and all the metrics on this experiment. The rows will each correspond to a trial (if using one-arm trials, which is the case in base AxClient; will correspond to arms in trials in the batch-trial case).

property global_stopping_strategy: BaseGlobalStoppingStrategy | None¶

The global stopping strategy used on the experiment.

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

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

load_experiment_from_database(experiment_name: str, choose_generation_strategy_kwargs: dict[str, Any] | None = 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.

classmethod load_from_json_file(filepath: str = 'ax_client_snapshot.json', **kwargs) → AxClientSubclass[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: dict[str, str] | None = 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 metric_definitions: dict[str, dict[str, Any]]¶

Returns metric definitions for all experiment metrics that can be passed into functions requiring metric_definitions

property metric_names: set[str]¶

Returns the names of all metrics on the attached experiment.

property objective: Objective¶

property objective_name: str¶

Returns the name of the objective in this optimization.

property objective_names: list[str]¶

Returns the name of the objective in this optimization.

remove_tracking_metric(metric_name: str) → None[source]¶

Remove a metric that already exists on the experiment.

Parameters:

metric_name – Unique name of metric to remove.

static save(filepath: str | None = 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.

set_optimization_config(objectives: dict[str, ObjectiveProperties] | None = None, outcome_constraints: list[str] | None = None, metric_definitions: dict[str, dict[str, Any]] | None = None) → None[source]¶

Overwrite experiment’s optimization config

Parameters:

• objectives – Mapping from an objective name to object containing: minimize: Whether this experiment represents a minimization problem. threshold: The bound in the objective’s threshold constraint.

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

• metric_definitions – A mapping of metric names to extra kwargs to pass to that metric

set_search_space(parameters: list[dict[str, None | str | bool | float | int | Sequence[None | str | bool | float | int] | dict[str, list[str]]]], parameter_constraints: list[str] | None = None) → None[source]¶

Sets the search space on the experiment and saves. This is expected to fail on base AxClient as experiment will have immutable search space and optimization config set to True by default

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. 6. “digits” (int) for float-valued range parameters.

• 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 “>=”. Note that parameter constraints may only be placed on range parameters, not choice parameters or fixed parameters.

set_status_quo(params: dict[str, None | str | bool | float | int] | None) → None[source]¶

Set, or unset status quo on the experiment. There may be risk in using this after a trial with the status quo arm has run.

Parameters:

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.

should_stop_trials_early(trial_indices: set[int]) → dict[int, str | None][source]¶

Evaluate whether to early-stop running trials.

Parameters:

trial_indices – Indices of trials to consider for early stopping.

Returns:

A dictionary mapping trial indices that should be early stopped to (optional) messages with the associated reason.

property status_quo: dict[str, None | str | bool | float | int] | None¶

The parameterization of the status quo arm of the experiment.

stop_trial_early(trial_index: int) → None[source]¶

to_json_snapshot(encoder_registry: dict[type, Callable[[Any], dict[str, Any]]] | None = None, class_encoder_registry: dict[type, Callable[[Any], dict[str, Any]]] | None = None) → 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_running_trial_with_intermediate_data(trial_index: int, raw_data: dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]], metadata: dict[str, str | int] | None = None, sample_size: int | None = None) → None[source]¶

Updates the trial with given metric values without completing it. Also adds optional metadata to it. Useful for intermediate results like the metrics of a partially optimized machine learning model. In these cases it should be called instead of complete_trial until it is time to complete the trial.

NOTE: This method will raise an Exception if it is called multiple times with the same raw_data. These cases typically arise when raw_data does not change over time. To avoid this, pass a timestep metric in raw_data, for example:

for ts in range(100):
    raw_data = [({"ts": ts}, {"my_objective": (1.0, 0.0)})]
    ax_client.update_running_trial_with_intermediate_data(
        trial_index=0, raw_data=raw_data
    )

NOTE: When raw_data does not specify SEM for a given metric, Ax will default to the assumption that the data is noisy (specifically, corrupted by additive zero-mean Gaussian noise) and that the level of noise should be inferred by the optimization model. To indicate that the data is noiseless, set SEM to 0.0, for example:

ax_client.update_running_trial_with_intermediate_data(
    trial_index=0,
    raw_data={"my_objective": (objective_mean_value, 0.0)}
)

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 SEM is unknown (then Ax will infer observation noise level). 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.

update_trial_data(trial_index: int, raw_data: dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]], metadata: dict[str, str | int] | None = None, sample_size: int | None = None) → None[source]¶

Attaches additional data or updates the existing data for a trial in a terminal state. For example, if trial was completed with data for only one of the required metrics, this can be used to attach data for the remaining metrics.

NOTE: This does not change the trial status.

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, None | str | bool | float | int]) → bool[source]¶

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

Bases: object

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

full_run() → OptimizationLoop[source]¶

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

get_best_point() → tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None][source]¶

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

get_current_model() → ModelBridge | None[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, None | str | bool | float | int | Sequence[None | str | bool | float | int] | dict[str, list[str]]]], evaluation_function: Callable[[dict[str, None | str | bool | float | int]], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]] | Callable[[dict[str, None | str | bool | float | int], float | None], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]], experiment_name: str | None = None, objective_name: str | None = None, minimize: bool = False, parameter_constraints: list[str] | None = None, outcome_constraints: list[str] | None = None, total_trials: int = 20, arms_per_trial: int = 1, wait_time: int = 0, random_seed: int | None = None, generation_strategy: GenerationStrategy | None = None) → OptimizationLoop[source]¶

Constructs a synchronous OptimizationLoop using an evaluation function.