ax.service¶
Ax Client¶
- class ax.service.ax_client.AxClient(generation_strategy: GenerationStrategy | None = None, db_settings: None = None, enforce_sequential_optimization: bool = True, random_seed: int | None = None, torch_device: device | None = None, verbose_logging: bool = True, suppress_storage_errors: bool = False, early_stopping_strategy: BaseEarlyStoppingStrategy | None = None, global_stopping_strategy: BaseGlobalStoppingStrategy | None = None)[source]¶
Bases:
AnalysisBase
,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 inshould_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.
- NOTE:
- 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 toget_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 thisAxClient
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
- 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_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.
- 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 objective: Objective¶
- 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.
- 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.
- 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 whenraw_data
does not change over time. To avoid this, pass a timestep metric inraw_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.
Managed Loop¶
- class ax.service.managed_loop.OptimizationLoop(experiment: Experiment, 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]]]]], total_trials: int = 20, arms_per_trial: int = 1, random_seed: int | None = None, wait_time: int = 0, run_async: bool = False, generation_strategy: GenerationStrategy | None = None)[source]¶
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.
- 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.
- classmethod with_runners_and_metrics(parameters: list[dict[str, None | str | bool | float | int | Sequence[None | str | bool | float | int] | dict[str, list[str]]]], path_to_runner: str, paths_to_metrics: list[str], 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) OptimizationLoop [source]¶
Constructs an asynchronous OptimizationLoop using Ax runners and metrics.
- ax.service.managed_loop.optimize(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, random_seed: int | None = None, generation_strategy: GenerationStrategy | None = None) tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None, Experiment, ModelBridge | None] [source]¶
Construct and run a full optimization loop.
Interactive Loop¶
- ax.service.interactive_loop.ax_client_candidate_generator(queue: Queue[tuple[dict[str, None | str | bool | float | int], int]], stop_event: Event, num_trials: int, ax_client: AxClient, lock: allocate_lock) None [source]¶
Thread-safe method for generating the next trial from the AxClient and enqueueing it to the candidate queue. The number of candidates pre-generated is controlled by the maximum size of the queue. Generation stops when num_trials trials are attached to the AxClient’s experiment.
- ax.service.interactive_loop.ax_client_data_attacher(queue: Queue[tuple[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]]]]]], stop_event: Event, ax_client: AxClient, lock: allocate_lock) None [source]¶
Thread-safe method for attaching evaluation outcomes to the AxClient from the outcome queue. If the AxClient’s lock is acquired all data in the outcome queue is attached at once, then the lock released. Stops when the event is set.
- ax.service.interactive_loop.interactive_optimize(num_trials: int, candidate_queue_maxsize: int, candidate_generator_function: Callable[[...], None], data_attacher_function: Callable[[...], None], elicitation_function: Callable[[...], Any], candidate_generator_kwargs: dict[str, Any] | None = None, data_attacher_kwargs: dict[str, Any] | None = None, elicitation_function_kwargs: dict[str, Any] | None = None) bool [source]¶
Function to facilitate running Ax experiments with candidate pregeneration (the generation of candidate points while waiting for trial evaluation). This can be useful in many contexts, especially in interactive experiments in which trial evaluation entails eliciting a response from a human. Candidate pregeneration uses the time waiting for trail evaluation to generate new candidates from the data available. Note that this is a tradeoff – a larger pregeneration queue will result in more trials being generated with less data compared to a smaller pregeneration queue (or no pregeneration as in conventional Ax usage) and should only be used in contexts where it is necessary for the user to not experience any “lag” while candidates are being generated.
- Parameters:
num_trials – The total number of trials to be run.
candidate_queue_maxsize – The maximum number of candidates to pregenerate.
candidate_generator_function – A function taking in a queue and event that enqueues candidates (generated by any means). See ax_client_candidate_generator for an example.
data_attacher_function – A function taking in a queue and event that attaches observations to Ax. See ax_client_data_attacher for an example.
elicitation_function – Function from parameterization (as returned by AxClient.get_next_trial) to outcome (as expected by AxClient.complete_trial). If None, elicitation is aborted by the user.
candidate_generator_kwargs – kwargs to be passed into candidate_generator_function when it is spawned as a thread.
data_attacher_kwargs – kwargs to be passed into data_attacher_function when it is spawned as a thread.
elicitation_function_kwargs – kwargs to be passed into elicitation_function
- Returns:
True if optimization was completed and False if aborted.
- ax.service.interactive_loop.interactive_optimize_with_client(ax_client: AxClient, num_trials: int, candidate_queue_maxsize: int, elicitation_function: Callable[[tuple[dict[str, None | str | bool | float | int], 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]]]]]) bool [source]¶
Implementation of interactive_loop using the AxClient. Extract results of the experiment from the AxClient passed in.
The basic structure is as follows: One thread tries for a lock on the AxClient, generates a candidate, and enqueues it to a candidate queue. Another thread tries for the same lock, takes all the trial outcomes in the outcome queue, and attaches them to the AxClient. The main thread pops a candidate off the candidate queue, elicits response from the user, and puts the response onto the outcome queue.
Scheduler¶
- exception ax.service.scheduler.FailureRateExceededError(message: str, hint: str = '')[source]¶
Bases:
AxError
Error that indicates the optimization was aborted due to excessive failure rate.
- class ax.service.scheduler.MessageOutput(text: str, priority: OutputPriority | int)[source]¶
Bases:
object
Message to be shown in the output of the scheduler.
- priority: OutputPriority | int¶
- class ax.service.scheduler.OptimizationResult[source]¶
Bases:
NamedTuple
- class ax.service.scheduler.OutputPriority(value)[source]¶
Bases:
IntEnum
Priority of a message. Messages with higher priority will be shown first, and messages with the same priority will be sorted alphabetically.
- DEBUG = 10¶
- ERROR = 50¶
- INFO = 20¶
- NOTSET = 0¶
- TOPLINE = 30¶
- WARNING = 40¶
- class ax.service.scheduler.Scheduler(experiment: Experiment, generation_strategy: GenerationStrategyInterface, options: SchedulerOptions, db_settings: None = None, _skip_experiment_save: bool = False)[source]¶
Bases:
AnalysisBase
,BestPointMixin
Closed-loop manager class for Ax optimization.
- experiment¶
Experiment, in which results of the optimization will be recorded.
- generation_strategy¶
Generation strategy for the optimization, describes models that will be used in optimization.
- options¶
SchedulerOptions for this scheduler instance.
- 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.
- _skip_experiment_save¶
If True, scheduler will not re-save the experiment passed to it. Use only if the experiment had just been saved, as otherwise experiment state could get corrupted.
- DEFAULT_FETCH_KWARGS = {'overwrite_existing_data': True}¶
- property candidate_trials: list[BaseTrial]¶
Candidate trials on the experiment this scheduler is running.
Note: if the experiment is a MultiTypeExperiment, then this will only fetch trials of type Scheduler.trial_type.
- Returns:
List of trials that are currently candidates.
- error_if_failure_rate_exceeded(force_check: bool = False) None [source]¶
Raises an exception if the failure rate (set in scheduler options) has been exceeded at any point during the optimization.
NOTE: Both FAILED and ABANDONED trial statuses count towards the failure rate.
- Parameters:
force_check – Indicates whether to force a failure-rate check regardless of the number of trials that have been executed. If False (default), the check will be skipped if the optimization has fewer than five failed trials. If True, the check will be performed unless there are 0 failures.
- experiment: Experiment¶
- property failed_abandoned_trial_indices: set[int]¶
Failed or abandoned trials.
Note: if the experiment is a MultiTypeExperiment, then this will only fetch trials of type Scheduler.trial_type.
- Returns:
List of trials that are currently running.
- classmethod from_stored_experiment(experiment_name: str, options: SchedulerOptions, db_settings: None = None, generation_strategy: GenerationStrategy | None = None, reduced_state: bool = True, **kwargs: Any) Scheduler [source]¶
Create a
Scheduler
with a previously stored experiment, which the scheduler should resume.- Parameters:
experiment_name – Experiment to load and resume.
options –
SchedulerOptions
, with which to set up the new scheduler.db_settings – Optional
DBSettings
, which to use for reloading the experiment; also passed asdb_settings
argument to the scheduler constructor.generation_strategy – Generation strategy to use to provide candidates for the resumed optimization. Provide this argument only if the experiment does not already have a generation strategy associated with it.
kwargs – Kwargs to pass through to the
Scheduler
constructor.
- generate_candidates(num_trials: int = 1, reduce_state_generator_runs: bool = False, remove_stale_candidates: bool = False) tuple[list[BaseTrial], Exception | None] [source]¶
Fetch the latest data and generate new candidate trials.
- Parameters:
num_trials – Number of candidate trials to generate.
reduce_state_generator_runs – Flag to determine whether to save model state for every generator run (default) or to only save model state on the final generator run of each batch.
remove_stale_candidates – If true, mark any existing candidate trials failed before trial generation because: - they should not be treated as pending points - they will no longer be relevant
- Returns:
List of trials, empty if generation is not possible.
- generation_strategy: GenerationStrategyInterface¶
- 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.
- NOTE:
- 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_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
- get_trace(optimization_config: OptimizationConfig | 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.
- logger: LoggerAdapter¶
- markdown_messages: dict[str, MessageOutput]¶
- property options: SchedulerOptions¶
Scheduler options.
- property pending_trials: list[BaseTrial]¶
Running or staged trials on the experiment this scheduler is running.
Note: if the experiment is a MultiTypeExperiment, then this will only fetch trials of type Scheduler.trial_type.
- Returns:
List of trials that are currently running or staged.
- poll_and_process_results(poll_all_trial_statuses: bool = False) bool [source]¶
- Takes the following actions:
Poll trial runs for their statuses
Find trials to fetch data for
Apply new trial statuses
Fetch data
Early-stop trials where possible
Save modified trials, having either new statuses or new data
- Returns:
A boolean representing whether any trial evaluations completed or have been marked as failed or abandoned, changing the number of currently running trials.
- poll_trial_status(poll_all_trial_statuses: bool = False) dict[TrialStatus, set[int]] [source]¶
Polling function, checks the status of any non-terminal trials and returns their indices as a mapping from TrialStatus to a list of indices.
NOTE: Does not need to handle waiting between polling while trials are running; that logic is handled in
Scheduler.poll
, which calls this function.- Returns:
A dictionary mapping TrialStatus to a list of trial indices that have the respective status at the time of the polling. This does not need to include trials that at the time of polling already have a terminal (ABANDONED, FAILED, COMPLETED) status (but it may).
- report_results(force_refit: bool = False) dict[str, Any] [source]¶
Optional user-defined function for reporting intermediate and final optimization results (e.g. make some API call, write to some other db). This function is called whenever new results are available during the optimization.
- Parameters:
force_refit – Whether to force the implementation of this method to refit the model on generation strategy before using it to produce results to report (e.g. if using model to visualize data).
- Returns:
An optional dictionary with any relevant data about optimization.
- run(max_new_trials: int, timeout_hours: float | None = None) bool [source]¶
Schedules trial evaluation(s) if stopping criterion is not triggered, maximum parallelism is not currently reached, and capacity allows. Logs any failures / issues.
- Parameters:
max_new_trials – Maximum number of new trials this function should generate and run (useful when generating and running trials in batches). Note that this function might also re-deploy existing
CANDIDATE
trials that failed to deploy before, which will not count against this number.timeout_hours – Maximum number of hours, for which to run the optimization. This function will abort after running for timeout_hours even if stopping criterion has not been reached. If set to None, no optimization timeout will be applied.
- Returns:
Boolean representing success status.
- run_all_trials(timeout_hours: float | None = None, idle_callback: Callable[[Scheduler], Any] | None = None) OptimizationResult [source]¶
Run all trials until
should_consider_optimization_complete
yields true (by default,should_consider_optimization_complete
will yield true when reaching thenum_trials
setting, passed to scheduler on instantiation as part ofSchedulerOptions
).NOTE: This function is available only when
SchedulerOptions.num_trials
is specified.- Parameters:
timeout_hours – Limit on length of ths optimization; if reached, the optimization will abort even if completon criterion is not yet reached.
idle_callback – Callable that takes a Scheduler instance as an argument to deliver information while the trials are still running. Any output of idle_callback will not be returned, so idle_callback must expose information in some other way. For example, it could print something about the state of the scheduler or underlying experiment to STDOUT, write something to a database, or modify a Plotly figure or other object in place. ax.service.utils.report_utils.get_figure_and_callback is a helper function for generating a callback that will update a Plotly figure.
Example
>>> trials_info = {"n_completed": None} >>> >>> def write_n_trials(scheduler: Scheduler) -> None: ... trials_info["n_completed"] = len(scheduler.experiment.trials) >>> >>> scheduler.run_all_trials( ... timeout_hours=0.1, idle_callback=write_n_trials ... ) >>> print(trials_info["n_completed"])
- run_n_trials(max_trials: int, ignore_global_stopping_strategy: bool = False, timeout_hours: float | None = None, idle_callback: Callable[[Scheduler], Any] | None = None) OptimizationResult [source]¶
Run up to
max_trials
trials; will run allmax_trials
unless completion criterion is reached. For baseScheduler
, completion criterion is reaching total number of trials set inSchedulerOptions
, so if that option is not specified, this function will run exactlymax_trials
trials always.- Parameters:
max_trials – Maximum number of trials to run.
ignore_global_stopping_strategy – If set, Scheduler will skip the global stopping strategy in
should_consider_optimization_complete
.timeout_hours – Limit on length of ths optimization; if reached, the optimization will abort even if completon criterion is not yet reached.
idle_callback – Callable that takes a Scheduler instance as an argument to deliver information while the trials are still running. Any output of idle_callback will not be returned, so idle_callback must expose information in some other way. For example, it could print something about the state of the scheduler or underlying experiment to STDOUT, write something to a database, or modify a Plotly figure or other object in place. ax.service.utils.report_utils.get_figure_and_callback is a helper function for generating a callback that will update a Plotly figure.
Example
>>> trials_info = {"n_completed": None} >>> >>> def write_n_trials(scheduler: Scheduler) -> None: ... trials_info["n_completed"] = len(scheduler.experiment.trials) >>> >>> scheduler.run_n_trials( ... max_trials=3, idle_callback=write_n_trials ... ) >>> print(trials_info["n_completed"]) 3
- run_trials(trials: Iterable[BaseTrial]) dict[int, dict[str, Any]] [source]¶
Deployment function, runs a single evaluation for each of the given trials.
Override default implementation on the
Runner
if its desirable to deploy trials in bulk.NOTE: the retry_on_exception decorator applied to this function should also be applied to its subclassing override if one is provided and retry behavior is desired.
- Parameters:
trials – Iterable of trials to be deployed, each containing arms with parameterizations to be evaluated. Can be a
Trial
if contains only one arm or aBatchTrial
if contains multiple arms.- Returns:
Dict of trial index to the run metadata of that trial from the deployment process.
- run_trials_and_yield_results(max_trials: int, ignore_global_stopping_strategy: bool = False, timeout_hours: float | None = None, idle_callback: Callable[[Scheduler], None] | None = None) Generator[dict[str, Any], None, None] [source]¶
Make continuous calls to run and process_results to run up to
max_trials
trials, until completion criterion is reached. This is the ‘main’ method of aScheduler
.- Parameters:
max_trials – Maximum number of trials to run in this generator. The generator will run trials until a completion criterion is reached, a completion signal is received from the generation strategy, or
max_trials
trials have been run (whichever happens first).ignore_global_stopping_strategy – If set, Scheduler will skip the global stopping strategy in
should_consider_optimization_complete
.timeout_hours – Maximum number of hours, for which to run the optimization. This function will abort after running for timeout_hours even if stopping criterion has not been reached. If set to None, no optimization timeout will be applied.
idle_callback – Callable that takes a Scheduler instance as an argument to deliver information while the trials are still running. Any output of idle_callback will not be returned, so idle_callback must expose information in some other way. For example, it could print something about the state of the scheduler or underlying experiment to STDOUT, write something to a database, or modify a Plotly figure or other object in place. ax.service.utils.report_utils.get_figure_and_callback is a helper function for generating a callback that will update a Plotly figure.
- property runner: Runner¶
Runner
specified on the experiment associated with thisScheduler
instance.
- property running_trial_indices: set[int]¶
Currently running trials.
- Returns:
List of trials that are currently running.
- property running_trials: list[BaseTrial]¶
Currently running trials.
Note: if the experiment is a MultiTypeExperiment, then this will only fetch trials of type Scheduler.trial_type.
- Returns:
List of trials that are currently running.
- should_abort_optimization(timeout_hours: float | None = None) bool [source]¶
Checks whether this scheduler has reached some intertuption / abort criterion, such as an overall optimization timeout, tolerated failure rate, etc.
- should_consider_optimization_complete() tuple[bool, str] [source]¶
Whether this scheduler should consider this optimization complete and not run more trials (and conclude the optimization via
_complete_optimization
).NOTE: An optimization is considered complete when a generation strategy signaled completion or when the
should_consider_optimization_complete
method on this scheduler evaluates toTrue
. Theshould_consider_optimization_complete
method is also responsible for checking global_stopping_strategy’s decision as well. Alongside the stop decision, this function returns a string describing the reason for stopping the optimization.
- property standard_generation_strategy: GenerationStrategy¶
Used for operations in the scheduler that can only be done with and instance of
GenerationStrategy
.
- summarize_final_result() OptimizationResult [source]¶
Get some summary of result: which trial did best, what were the metric values, what were encountered failures, etc.
- property trial_type: str | None¶
Trial type for the experiment this scheduler is running.
This returns None if the experiment is not a MultitypeExperiment
- Returns:
Trial type for the experiment this scheduler is running if the experiment is a MultiTypeExperiment and None otherwise.
- property trials: list[BaseTrial]¶
All trials.
Note: if the experiment is a MultiTypeExperiment, then this will only fetch trials of type Scheduler.trial_type.
- Returns:
List of trials that are currently running.
- property trials_expecting_data: list[BaseTrial]¶
Trials expecting data.
Note: if the experiment is a MultiTypeExperiment, then this will only fetch trials of type Scheduler.trial_type.
- wait_for_completed_trials_and_report_results(idle_callback: Callable[[Scheduler], None] | None = None, force_refit: bool = False) dict[str, Any] [source]¶
Continuously poll for successful trials, with limited exponential backoff, and process the results. Stop once at least one successful trial has been found. This function can be overridden to a different waiting function as needed; it must call poll_and_process_results to ensure that trials that completed their evaluation are appropriately marked as ‘COMPLETED’ in Ax.
- Parameters:
idle_callback – Callable that takes a Scheduler instance as an argument to deliver information while the trials are still running. Any output of idle_callback will not be returned, so idle_callback must expose information in some other way. For example, it could print something about the state of the scheduler or underlying experiment to STDOUT, write something to a database, or modify a Plotly figure or other object in place. ax.service.utils.report_utils.get_figure_and_callback is a helper function for generating a callback that will update a Plotly figure.
force_refit – Whether to force a refit of the model during report_results.
- Returns:
Results of the optimization so far, represented as a dict. The contents of the dict depend on the implementation of report_results in the given Scheduler subclass.
- exception ax.service.scheduler.SchedulerInternalError(message: str, hint: str = '')[source]¶
Bases:
AxError
Error that indicates an error within the Scheduler logic.
- ax.service.scheduler.get_fitted_model_bridge(scheduler: Scheduler, force_refit: bool = False) ModelBridge [source]¶
Returns a fitted ModelBridge object. If the model is fit already, directly returns the already fitted model. Otherwise, fits and returns a new one.
- Parameters:
scheduler – The scheduler object from which to get the fitted model.
force_refit – If True, will force a data lookup and a refit of the model.
- Returns:
A ModelBridge object fitted to the observations of the scheduler’s experiment.
- class ax.service.utils.scheduler_options.SchedulerOptions(max_pending_trials: int = 10, trial_type: ~ax.service.utils.scheduler_options.TrialType = TrialType.TRIAL, batch_size: int | None = None, total_trials: int | None = None, tolerated_trial_failure_rate: float = 0.5, min_failed_trials_for_failure_rate_check: int = 5, log_filepath: str | None = None, logging_level: int = 20, ttl_seconds_for_trials: int | None = None, init_seconds_between_polls: int | None = 1, min_seconds_before_poll: float = 1.0, seconds_between_polls_backoff_factor: float = 1.5, run_trials_in_batches: bool = False, debug_log_run_metadata: bool = False, early_stopping_strategy: ~ax.early_stopping.strategies.base.BaseEarlyStoppingStrategy | None = None, global_stopping_strategy: ~ax.global_stopping.strategies.base.BaseGlobalStoppingStrategy | None = None, suppress_storage_errors_after_retries: bool = False, wait_for_running_trials: bool = True, fetch_kwargs: dict[str, ~typing.Any] = <factory>, validate_metrics: bool = True, status_quo_weight: float = 0.0, enforce_immutable_search_space_and_opt_config: bool = True, mt_experiment_trial_type: str | None = None, force_candidate_generation: bool = False)[source]¶
Bases:
object
Settings for a scheduler instance.
- max_pending_trials¶
Maximum number of pending trials the scheduler can have
STAGED
orRUNNING
at once, required. If looking to useRunner.poll_available_capacity
as a primary guide for how many trials should be pending at a given time, set this limit to a high number, as an upper bound on number of trials that should not be exceeded.- Type:
- trial_type¶
Type of trials (1-arm
Trial
or multi-armBatch Trial
) that will be deployed using the scheduler. Defaults to 1-arm Trial. NOTE: useBatchTrial
only if need to evaluate multiple arms together, e.g. in an A/B-test influenced by data nonstationarity. For cases where just deploying multiple arms at once is beneficial but the trials are evaluated independently, implementrun_trials
method in scheduler subclass, to deploy multiple 1-arm trials at the same time.
- batch_size¶
If using BatchTrial the number of arms to be generated and deployed per trial.
- Type:
int | None
- total_trials¶
Limit on number of trials a given
Scheduler
should run. If no stopping criteria are implemented on a given scheduler, exhaustion of this number of trials will be used as default stopping criterion inScheduler.run_all_trials
. Required to be non-null if usingScheduler.run_all_trials
(not required forScheduler.run_n_trials
).- Type:
int | None
- tolerated_trial_failure_rate¶
Fraction of trials in this optimization that are allowed to fail without the whole optimization ending. Expects value between 0 and 1. NOTE: Failure rate checks begin once min_failed_trials_for_failure_rate_check trials have failed; after that point if the ratio of failed trials to total trials ran so far exceeds the failure rate, the optimization will halt.
- Type:
- min_failed_trials_for_failure_rate_check¶
The minimum number of trials that must fail in Scheduler in order to start checking failure rate.
- Type:
- ttl_seconds_for_trials¶
Optional TTL for all trials created within this
Scheduler
, in seconds. Trials that remainRUNNING
for more than their TTL seconds will be markedFAILED
once the TTL elapses and may be re-suggested by the Ax optimization models.- Type:
int | None
- init_seconds_between_polls¶
Initial wait between rounds of polling, in seconds. Relevant if using the default wait- for-completed-runs functionality of the base
Scheduler
(ifwait_for_completed_trials_and_report_results
is not overridden). With the default waiting, every time a poll returns that no trial evaluations completed, wait time will increase; once some completed trial evaluations are found, it will reset back to this value. Specify 0 to not introduce any wait between polls.- Type:
int | None
- min_seconds_before_poll¶
Minimum number of seconds between beginning to run a trial and the first poll to check trial status.
- Type:
- timeout_hours¶
Number of hours after which the optimization will abort.
- run_trials_in_batches¶
If True and
poll_available_capacity
is implemented to return non-null results, trials will be dispatched in groups via run_trials instead of one-by-one viarun_trial
. This allows to save time, IO calls or computation in cases where dispatching trials in groups is more efficient then sequential deployment. The size of the groups will be determined as the minimum ofself.poll_available_capacity()
and the number of generator runs that the generation strategy is able to produce without more data or reaching its allowed max paralellism limit.- Type:
- early_stopping_strategy¶
A
BaseEarlyStoppingStrategy
that determines whether a trial should be stopped given the current state of the experiment. Used inshould_stop_trials_early
.
- global_stopping_strategy¶
A
BaseGlobalStoppingStrategy
that determines whether the full optimization should be stopped or not.
- suppress_storage_errors_after_retries¶
Whether to fully suppress SQL storage-related errors if encountered, after retrying the call multiple times. Only use if SQL storage is not important for the given use case, since this will only log, but not raise, an exception if it’s encountered while saving to DB or loading from it.
- Type:
- validate_metrics¶
Whether to raise an error if there is a problem with the metrics attached to the experiment.
- Type:
- status_quo_weight¶
The weight of the status quo arm. This is only used if the scheduler is using a BatchTrial. This requires that the status_quo be set on the experiment.
- Type:
- enforce_immutable_search_space_and_opt_config¶
Whether to enforce that the search space and optimization config are immutable. If true, will add “immutable_search_space_and_opt_config”: True to experiment properties
- Type:
- mt_experiment_trial_type¶
Type of trial to run for MultiTypeExperiments. This is currently required for MultiTypeExperiments. This is ignored for “regular” or single type experiments. If you don’t know what a single type experiment is, you don’t need this.
- Type:
str | None
- force_candidate_generation¶
Whether to force candidate generation even if the generation strategy is not ready to generate candidates, meaning one of the transition criteria with block_gen_if_met is met. This is not yet implemented.
- Type:
- early_stopping_strategy: BaseEarlyStoppingStrategy | None = None¶
- global_stopping_strategy: BaseGlobalStoppingStrategy | None = None¶
- class ax.service.utils.scheduler_options.TrialType(value)[source]¶
Bases:
Enum
An enumeration.
- BATCH_TRIAL = 1¶
- TRIAL = 0¶
Utils¶
Analysis¶
- class ax.service.utils.analysis_base.AnalysisBase(db_settings: None = None, logging_level: int = 20, suppress_all_errors: bool = False)[source]¶
Bases:
WithDBSettingsBase
Base class for analysis functionality shared between AxClient and Scheduler.
- compute_analyses(analyses: Iterable[Analysis] | None = None) list[AnalysisCard] [source]¶
Compute Analyses for the Experiment and GenerationStrategy associated with this Scheduler instance and save them to the DB if possible. If an Analysis fails to compute (e.g. due to a missing metric), it will be skipped and a warning will be logged.
- Parameters:
analyses – Analyses to compute. If None, the Scheduler will choose a set of Analyses to compute based on the Experiment and GenerationStrategy.
- experiment: Experiment¶
- generation_strategy: GenerationStrategyInterface¶
Best Point Identification¶
- class ax.service.utils.best_point_mixin.BestPointMixin[source]¶
Bases:
object
- get_best_parameters(optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None, use_model_predictions: bool = True) tuple[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 parameterization and model predictions for it.
- NOTE:
- abstract 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.
- NOTE:
- abstract 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_improvement_over_baseline(experiment: Experiment, generation_strategy: GenerationStrategy, baseline_arm_name: str | None = None) float [source]¶
Returns the scalarized improvement over baseline, if applicable.
- Returns:
For Single Objective cases, returns % improvement of objective. Positive indicates improvement over baseline. Negative indicates regression. For Multi Objective cases, throws NotImplementedError
- abstract 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
- abstract get_trace() 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.
- abstract get_trace_by_progression(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.
- ax.service.utils.best_point.extract_Y_from_data(experiment: Experiment, metric_names: list[str], data: Data | None = None) tuple[Tensor, Tensor] [source]¶
Converts the experiment observation data into a tensor.
NOTE: This requires block design for observations. It will error out if any trial is missing data for any of the given metrics or if the data is missing the trial_index.
- Parameters:
experiment – The experiment to extract the data from.
metric_names – List of metric names to extract data for.
data – An optional Data object to use instead of the experiment data. Note that the experiment must have a corresponding COMPLETED or EARLY_STOPPED trial for each trial_index in the data.
- Returns:
A two-element Tuple containing a tensor of observed metrics and a tensor of trial_indices.
- ax.service.utils.best_point.fill_missing_thresholds_from_nadir(experiment: Experiment, optimization_config: OptimizationConfig) list[ObjectiveThreshold] [source]¶
Get the objective thresholds from the optimization config and fill the missing thresholds based on the nadir point.
- Parameters:
experiment – The experiment, whose data is used to calculate the nadir point.
optimization_config – Optimization config to get the objective thresholds and the objective directions from.
- Returns:
A list of objective thresholds, one for each objective in optimization config.
- ax.service.utils.best_point.get_best_by_raw_objective(experiment: Experiment, optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None) tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None [source]¶
Given an experiment, identifies the arm that had 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.
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.
- Returns:
Tuple of parameterization, and model predictions for it.
- ax.service.utils.best_point.get_best_by_raw_objective_with_trial_index(experiment: Experiment, optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None) tuple[int, dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None [source]¶
Given an experiment, identifies the arm that had 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.
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.
- Returns:
Tuple of trial index, parameterization, and model predictions for it.
- ax.service.utils.best_point.get_best_parameters(experiment: Experiment, models_enum: type[ModelRegistryBase], optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None) tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None [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.
models_enum – Registry of all models that may be in the experiment’s generation strategy.
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.
- Returns:
Tuple of parameterization and model predictions for it.
- ax.service.utils.best_point.get_best_parameters_from_model_predictions(experiment: Experiment, models_enum: type[ModelRegistryBase], trial_indices: Iterable[int] | None = None) tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None [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.
models_enum – Registry of all models that may be in the experiment’s generation strategy.
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.
- Returns:
Tuple of parameterization and model predictions for it.
- ax.service.utils.best_point.get_best_parameters_from_model_predictions_with_trial_index(experiment: Experiment, models_enum: type[ModelRegistryBase], optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None) tuple[int, dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None [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.
models_enum – Registry of all models that may be in the experiment’s generation strategy.
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.
- Returns:
Tuple of trial index, parameterization, and model predictions for it.
- ax.service.utils.best_point.get_best_parameters_with_trial_index(experiment: Experiment, models_enum: type[ModelRegistryBase], optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None) tuple[int, dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None [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.
models_enum – Registry of all models that may be in the experiment’s generation strategy.
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.
- Returns:
Tuple of trial index, parameterization, and model predictions for it.
- ax.service.utils.best_point.get_best_raw_objective_point(experiment: Experiment, optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None) tuple[dict[str, None | str | bool | float | int], dict[str, tuple[float, float]]] [source]¶
- ax.service.utils.best_point.get_best_raw_objective_point_with_trial_index(experiment: Experiment, optimization_config: OptimizationConfig | None = None, trial_indices: Iterable[int] | None = None) tuple[int, dict[str, None | str | bool | float | int], 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 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.
- Returns:
- Tuple of parameterization and a mapping from metric name to a tuple of
the corresponding objective mean and SEM.
- ax.service.utils.best_point.get_pareto_optimal_parameters(experiment: Experiment, generation_strategy: GenerationStrategy, 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:
experiment – Experiment, from which to find Pareto-optimal arms.
generation_strategy – Generation strategy containing the modelbridge.
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).
- Return type:
A mapping from trial index to the tuple of
- ax.service.utils.best_point_utils.select_baseline_name_default_first_trial(experiment: Experiment, baseline_arm_name: str | None) tuple[str, bool] [source]¶
Choose a baseline arm from arms on the experiment. Logic: 1. If
baseline_arm_name
provided, validate that arm existsand return that arm name.
If
experiment.status_quo
is set, return its arm name.If there is at least one trial on the experiment, use the first trial’s first arm as the baseline.
Error if 1-3 all don’t apply.
- Returns:
- baseline arm name (str)
true when baseline selected from first arm of experiment (bool)
raise ValueError if no valid baseline found
- Return type:
Tuple
Instantiation¶
- class ax.service.utils.instantiation.FixedFeatures(parameters: dict[str, None | str | bool | float | int], trial_index: int | None = None)[source]¶
Bases:
object
Class for representing fixed features via the Service API.
- class ax.service.utils.instantiation.InstantiationBase[source]¶
Bases:
object
This is a lightweight stateless class that bundles together instantiation utils. It is used both on its own and as a mixin to AxClient, with the intent that these methods can be overridden by its subclasses for specific use cases.
- static build_objective_threshold(objective: str, objective_properties: ObjectiveProperties) str [source]¶
Constructs constraint string for an objective threshold interpretable by make_experiment()
- Parameters:
objective – Name of the objective
objective_properties – Object containing: minimize: Whether this experiment represents a minimization problem. threshold: The bound in the objective’s threshold constraint.
- classmethod build_objective_thresholds(objectives: dict[str, ObjectiveProperties]) list[str] [source]¶
Construct a list of constraint string for an objective thresholds interpretable by make_experiment()
- Parameters:
objectives – Mapping of name of the objective to Object containing: minimize: Whether this experiment represents a minimization problem. threshold: The bound in the objective’s threshold constraint.
- static constraint_from_str(representation: str, parameters: dict[str, Parameter]) ParameterConstraint [source]¶
Parse string representation of a parameter constraint.
- classmethod make_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, parameter_constraints: list[str] | None = None, outcome_constraints: list[str] | None = None, status_quo: dict[str, None | str | bool | float | int] | None = None, experiment_type: str | None = None, tracking_metric_names: list[str] | None = None, metric_definitions: dict[str, dict[str, Any]] | None = None, objectives: dict[str, str] | None = None, objective_thresholds: list[str] | None = None, support_intermediate_data: bool = False, immutable_search_space_and_opt_config: bool = True, auxiliary_experiments_by_purpose: None | dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]] = None, default_trial_type: str | None = None, default_runner: Runner | None = None, is_test: bool = False) Experiment [source]¶
Instantiation wrapper that allows for Ax Experiment creation without importing or instantiating any Ax classes.
- 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, 5. “is_task” (bool) for task parameters, and 6. “digits” (int) for float-valued range parameters.
name – Name of the experiment to be created.
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.
experiment_type – String indicating type of the experiment (e.g. name of a product in which it is used), if any.
tracking_metric_names – Names of additional tracking metrics not used for optimization.
metric_definitions – A mapping of metric names to extra kwargs to pass to that metric
objectives – Mapping from an objective name to “minimize” or “maximize” representing the direction for that objective.
objective_thresholds – A list of objective threshold constraints for multi- objective optimization, in the same string format as outcome_constraints argument.
support_intermediate_data – Whether trials may report metrics results for incomplete runs.
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.
auxiliary_experiments_by_purpose – Dictionary of auxiliary experiments for different use cases (e.g., transfer learning).
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 only applies to MultiTypeExperiment (when default_trial_type is specified).
is_test – Whether this experiment will be a test experiment (useful for marking test experiments in storage etc). Defaults to False.
- static make_fixed_observation_features(fixed_features: FixedFeatures) ObservationFeatures [source]¶
Construct ObservationFeatures from FixedFeatures.
- Parameters:
fixed_features – The fixed features for generation.
- Returns:
The new ObservationFeatures object.
- classmethod make_objective_thresholds(objective_thresholds: list[str], status_quo_defined: bool, metric_definitions: dict[str, dict[str, Any]] | None = None) list[ObjectiveThreshold] [source]¶
- classmethod make_objectives(objectives: dict[str, str], metric_definitions: dict[str, dict[str, Any]] | None = None) list[Objective] [source]¶
- classmethod make_optimization_config(objectives: dict[str, str], objective_thresholds: list[str], outcome_constraints: list[str], status_quo_defined: bool, metric_definitions: dict[str, dict[str, Any]] | None = None) OptimizationConfig [source]¶
- classmethod make_optimization_config_from_properties(objectives: dict[str, ObjectiveProperties] | None = None, outcome_constraints: list[str] | None = None, metric_definitions: dict[str, dict[str, Any]] | None = None, status_quo_defined: bool = False) OptimizationConfig | None [source]¶
Makes optimization config based on ObjectiveProperties objects
- 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.”
status_quo_defined – bool for whether the experiment has a status quo
metric_definitions – A mapping of metric names to extra kwargs to pass to that metric
- classmethod make_outcome_constraints(outcome_constraints: list[str], status_quo_defined: bool, metric_definitions: dict[str, dict[str, Any]] | None = None) list[OutcomeConstraint] [source]¶
- classmethod make_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) SearchSpace [source]¶
- classmethod objective_threshold_constraint_from_str(representation: str, metric_definitions: dict[str, dict[str, Any]] | None = None) ObjectiveThreshold [source]¶
- static optimization_config_from_objectives(objectives: list[Objective], objective_thresholds: list[ObjectiveThreshold], outcome_constraints: list[OutcomeConstraint]) OptimizationConfig [source]¶
Parse objectives and constraints to define optimization config.
The resulting optimization config will be regular single-objective config if objectives is a list of one element and a multi-objective config otherwise.
NOTE: If passing in multiple objectives, objective_thresholds must be a non-empty list defining constraints for each objective.
- class ax.service.utils.instantiation.MetricObjective(value)[source]¶
Bases:
Enum
An enumeration.
- MAXIMIZE = 2¶
- MINIMIZE = 1¶
- class ax.service.utils.instantiation.ObjectiveProperties(minimize: bool, threshold: float | None = None)[source]¶
Bases:
object
Class that holds properties of objective functions. Can be used to define an the objectives argument of ax_client.create_experiment, e.g.:
- ax_client.create_experiment(
name=”moo_experiment”, parameters=[…], objectives={
# threshold arguments are optional “a”: ObjectiveProperties(minimize=False, threshold=ref_point[0]), “b”: ObjectiveProperties(minimize=False, threshold=ref_point[1]),
},
)
- Parameters:
minimize (-) – Boolean indicating whether the objective is to be minimized or maximized.
threshold (-) – Optional float representing the smallest objective value (resp. largest if minimize=True) that is considered valuable in the context of multi-objective optimization. In BoTorch and in the literature, this is also known as an element of the reference point vector that defines the hyper-volume of the Pareto front.
Reporting¶
- ax.service.utils.report_utils.compare_to_baseline(experiment: Experiment, optimization_config: OptimizationConfig | None, comparison_arm_names: list[str] | None, baseline_arm_name: str | None = None) str | None [source]¶
Calculate metric improvement of the experiment against baseline. Returns the message(s) added to markdown_messages.
- ax.service.utils.report_utils.compare_to_baseline_impl(comparison_list: list[tuple[str, bool, str, float, str, float]]) str | None [source]¶
Implementation of compare_to_baseline, taking in a list of arm comparisons. Can be used directly with the output of ‘maybe_extract_baseline_comparison_values’
- ax.service.utils.report_utils.compute_maximum_map_values(experiment: Experiment, map_key: str | None = None) dict[int, float] [source]¶
A function that returns a map from trial_index to the maximum map value reached. If map_key is not specified, it uses the first map_key.
- ax.service.utils.report_utils.exp_to_df(exp: Experiment, metrics: list[Metric] | None = None, run_metadata_fields: list[str] | None = None, trial_properties_fields: list[str] | None = None, trial_attribute_fields: list[str] | None = None, additional_fields_callables: None | dict[str, Callable[[Experiment], dict[int, str | float]]] = None, always_include_field_columns: bool = False, show_relative_metrics: bool = False, **kwargs: Any) DataFrame [source]¶
Transforms an experiment to a DataFrame with rows keyed by trial_index and arm_name, metrics pivoted into one row. If the pivot results in more than one row per arm (or one row per
arm * map_keys
combination ifmap_keys
are present), results are omitted and warning is produced. Only supportsExperiment
.Transforms an
Experiment
into apd.DataFrame
.- Parameters:
exp – An
Experiment
that may have pending trials.metrics – Override list of metrics to return. Return all metrics if
None
.run_metadata_fields – Fields to extract from
trial.run_metadata
for trial inexperiment.trials
. If there are multiple arms per trial, these fields will be replicated across the arms of a trial.trial_properties_fields – Fields to extract from
trial._properties
for trial inexperiment.trials
. If there are multiple arms per trial, these fields will be replicated across the arms of a trial. Output columns names will be prepended with"trial_properties_"
.trial_attribute_fields – Fields to extract from trial attributes for each trial in
experiment.trials
. If there are multiple arms per trial, these fields will be replicated across the arms of a trial.additional_fields_callables – A dictionary of field names to callables, with each being a function from experiment to a trials_dict of the form {trial_index: value}. An example of a custom callable like this is the function compute_maximum_map_values.
always_include_field_columns – If True, even if all trials have missing values, include field columns anyway. Such columns are by default omitted (False).
show_relative_metrics – If True, show % metric changes relative to the provided status quo arm. If no status quo arm is provided, raise a warning and show raw metric values. If False, show raw metric values (default).
- Returns:
A dataframe of inputs, metadata and metrics by trial and arm (and
map_keys
, if present). If no trials are available, returns an empty dataframe. If no metric ouputs are available, returns a dataframe of inputs and metadata. Columns include:trial_index
arm_name
trial_status
generation_method
any elements of exp.runner.run_metadata_report_keys that are present in the trial.run_metadata of each trial
one column per metric (named after the metric.name)
one column per parameter (named after the parameter.name)
- Return type:
DataFrame
- ax.service.utils.report_utils.get_figure_and_callback(plot_fn: Callable[[Scheduler], plotly.graph_objects.Figure]) tuple[plotly.graph_objects.Figure, Callable[[Scheduler], None]] [source]¶
Produce a figure and a callback for updating the figure in place.
A likely use case is that plot_fn takes a Scheduler instance and returns a plotly Figure. Then get_figure_and_callback will produce a figure and callback that updates that figure according to plot_fn when the callback is passed to Scheduler.run_n_trials or Scheduler.run_all_trials.
- Parameters:
plot_fn – A function for producing a Plotly figure from a scheduler. If plot_fn raises a RuntimeError, the update wil be skipped and optimization will proceed.
Example
>>> def _plot(scheduler: Scheduler): >>> standard_plots = get_standard_plots(scheduler.experiment) >>> return standard_plots[0] >>> >>> fig, callback = get_figure_and_callback(_plot)
- ax.service.utils.report_utils.get_standard_plots(experiment: Experiment, model: ModelBridge | None, data: Data | None = None, true_objective_metric_name: str | None = None, early_stopping_strategy: BaseEarlyStoppingStrategy | None = None, limit_points_per_plot: int | None = None, global_sensitivity_analysis: bool = True) list[plotly.graph_objects.Figure] [source]¶
Extract standard plots for single-objective optimization.
Extracts a list of plots from an
Experiment
andModelBridge
of general interest to an Ax user. Currently not supported are - TODO: multi-objective optimization - TODO: ChoiceParameter plots- Parameters:
experiment (-) – The
Experiment
from which to obtain standard plots.model (-) – The
ModelBridge
used to suggest trial parameters.true_objective_metric_name (-) – Name of the metric to use as the true objective.
early_stopping_strategy (-) – Early stopping strategy used throughout the experiment; used for visualizing when curves are stopped.
limit_points_per_plot (-) – Limit the number of points used per metric in each curve plot. Passed to _get_curve_plot_dropdown.
global_sensitivity_analysis (-) – If True, plot total Variance-based sensitivity analysis for the model parameters. If False, plot sensitivities based on GP kernel lengthscales. Defaults to True.
- Returns:
a plot of objective value vs. trial index, to show experiment progression
a plot of objective value vs. range parameter values, only included if the model associated with generation_strategy can create predictions. This consists of:
a plot_slice plot if the search space contains one range parameter
an interact_contour plot if the search space contains multiple range parameters
- ax.service.utils.report_utils.maybe_extract_baseline_comparison_values(experiment: Experiment, optimization_config: OptimizationConfig | None, comparison_arm_names: list[str] | None, baseline_arm_name: str | None) list[tuple[str, bool, str, float, str, float]] | None [source]¶
Extracts the baseline values from the experiment, for use in comparing the baseline arm to the optimal results. Requires the user specifies the names of the arms to compare to.
- Returns:
(metric_name, minimize, baseline_arm_name, baseline_value, comparison_arm_name, comparison_arm_value, )
- Return type:
List of tuples containing
- ax.service.utils.report_utils.pareto_frontier_scatter_2d_plotly(experiment: Experiment, metric_names: tuple[str, str], reference_point: tuple[float, float] | None = None, minimize: bool | tuple[bool, bool] | None = None) plotly.graph_objects.Figure [source]¶
- ax.service.utils.report_utils.warn_if_unpredictable_metrics(experiment: Experiment, generation_strategy: GenerationStrategy, model_fit_threshold: float, metric_names: list[str] | None = None, model_fit_metric_name: str = 'coefficient_of_determination') str | None [source]¶
Warn if any optimization config metrics are considered unpredictable, i.e., their coefficient of determination is less than model_fit_threshold. :param experiment: The experiment containing the data and optimization_config.
If there is no optimization config, this function checks all metrics attached to the experiment.
- Parameters:
generation_strategy – The generation strategy containing the model.
model_fit_threshold – If a model’s coefficient of determination is below this threshold, that metric is considered unpredictable.
metric_names – If specified, only check these metrics.
model_fit_metric_name – Name of the metric to apply the model fit threshold to.
- Returns:
A string warning the user about unpredictable metrics, if applicable.
WithDBSettingsBase¶
- class ax.service.utils.with_db_settings_base.WithDBSettingsBase(db_settings: None = None, logging_level: int = 20, suppress_all_errors: bool = False)[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.
- ax.service.utils.with_db_settings_base.try_load_generation_strategy(experiment_name: str, decoder: None, experiment: Experiment | None = None, reduced_state: bool = False) GenerationStrategy | None [source]¶
Load generation strategy by experiment name, if it exists.
EarlyStopping¶
- ax.service.utils.early_stopping.get_early_stopping_metrics(experiment: Experiment, early_stopping_strategy: BaseEarlyStoppingStrategy | None) list[str] [source]¶
A helper function that returns a list of metric names on which a given early_stopping_strategy is operating.
- ax.service.utils.early_stopping.should_stop_trials_early(early_stopping_strategy: BaseEarlyStoppingStrategy | None, trial_indices: set[int], experiment: Experiment) dict[int, str | None] [source]¶
Evaluate whether to early-stop running trials.
- Parameters:
early_stopping_strategy – A
BaseEarlyStoppingStrategy
that determines whether a trial should be stopped given the state of an experiment.trial_indices – Indices of trials to consider for early stopping.
experiment – The experiment containing the trials.
- Returns:
A dictionary mapping trial indices that should be early stopped to (optional) messages with the associated reason.