ax.core

Core Classes

AbstractData

Arm

class ax.core.arm.Arm(parameters: Dict[str, Optional[Union[str, bool, float, int]]], name: Optional[str] = None)[source]

Bases: SortableBase

Base class for defining arms.

Randomization in experiments assigns units to a given arm. Thus, the arm encapsulates the parametrization needed by the unit.

clone(clear_name: bool = False) Arm[source]

Create a copy of this arm.

Parameters:

clear_name – whether this cloned copy should set its name to None instead of the name of the arm being cloned. Defaults to False.

property has_name: bool

Return true if arm’s name is not None.

static md5hash(parameters: Dict[str, Optional[Union[str, bool, float, int]]]) str[source]

Return unique identifier for arm’s parameters.

Parameters:

parameters – Parameterization; mapping of param name to value.

Returns:

Hash of arm’s parameters.

property name: str

Get arm name. Throws if name is None.

property name_or_short_signature: str

Returns arm name if exists; else last 4 characters of the hash.

Used for presentation of candidates (e.g. plotting and tables), where the candidates do not yet have names (since names are automatically set upon addition to a trial).

property parameters: Dict[str, Optional[Union[str, bool, float, int]]]

Get mapping from parameter names to values.

property signature: str

Get unique representation of a arm.

BaseTrial

class ax.core.base_trial.BaseTrial(experiment: core.experiment.Experiment, trial_type: Optional[str] = None, ttl_seconds: Optional[int] = None, index: Optional[int] = None)[source]

Bases: ABC, SortableBase

Base class for representing trials.

Trials are containers for arms that are deployed together. There are two kinds of trials: regular Trial, which only contains a single arm, and BatchTrial, which contains an arbitrary number of arms.

Parameters:
  • experiment – Experiment, of which this trial is a part

  • trial_type – Type of this trial, if used in MultiTypeExperiment.

  • ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.

  • index – If specified, the trial’s index will be set accordingly. This should generally not be specified, as in the index will be automatically determined based on the number of existing trials. This is only used for the purpose of loading from storage.

abstract property abandoned_arms: List[Arm]

All abandoned arms, associated with this trial.

property abandoned_reason: Optional[str]
abstract property arms: List[Arm]
abstract property arms_by_name: Dict[str, Arm]
assign_runner() BaseTrial[source]

Assigns default experiment runner if trial doesn’t already have one.

complete(reason: Optional[str] = None) BaseTrial[source]
Stops the trial if functionality is defined on runner

and marks trial completed.

Parameters:

reason – A message containing information why the trial is to be completed.

Returns:

The trial instance.

property completed_successfully: bool

Checks if trial status is COMPLETED.

property deployed_name: Optional[str]

Name of the experiment created in external framework.

This property is derived from the name field in run_metadata.

property did_not_complete: bool

Checks if trial status is terminal, but not COMPLETED.

property experiment: core.experiment.Experiment

The experiment this trial belongs to.

fetch_data(metrics: Optional[List[Metric]] = None, **kwargs: Any) Dict[str, Result[Data, MetricFetchE]][source]

Fetch data for this trial for all metrics on experiment.

Parameters:
  • trial_index – The index of the trial to fetch data for.

  • metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.

  • kwargs – keyword args to pass to underlying metrics’ fetch data functions.

Returns:

Data for this trial.

abstract property generator_runs: List[GeneratorRun]

All generator runs associated with this trial.

property index: int

The index of this trial within the experiment’s trial list.

property is_abandoned: bool

Whether this trial is abandoned.

lookup_data() Data[source]

Lookup cached data on experiment for this trial.

Returns:

If not merging across timestamps, the latest Data object associated with the trial. If merging, all data for trial, merged.

mark_abandoned(reason: Optional[str] = None, unsafe: bool = False) BaseTrial[source]

Mark trial as abandoned.

NOTE: Arms in abandoned trials are considered to be ‘pending points’ in experiment after their abandonment to avoid Ax models suggesting the same arm again as a new candidate. Arms in abandoned trials are also excluded from model training data unless fit_abandoned option is specified to model bridge.

Parameters:
  • abandoned_reason – The reason the trial was abandoned.

  • unsafe – Ignore sanity checks on state transitions.

Returns:

The trial instance.

mark_arm_abandoned(arm_name: str, reason: Optional[str] = None) BaseTrial[source]
mark_as(status: TrialStatus, unsafe: bool = False, **kwargs: Any) BaseTrial[source]

Mark trial with a new TrialStatus.

Parameters:
  • status – The new status of the trial.

  • unsafe – Ignore sanity checks on state transitions.

  • kwargs – Additional keyword args, as can be ued in the respective mark_ methods associated with the trial status.

Returns:

The trial instance.

mark_completed(unsafe: bool = False) BaseTrial[source]

Mark trial as completed.

Parameters:

unsafe – Ignore sanity checks on state transitions.

Returns:

The trial instance.

mark_early_stopped(unsafe: bool = False) BaseTrial[source]

Mark trial as early stopped.

Parameters:

unsafe – Ignore sanity checks on state transitions.

Returns:

The trial instance.

mark_failed(unsafe: bool = False) BaseTrial[source]

Mark trial as failed.

Parameters:

unsafe – Ignore sanity checks on state transitions.

Returns:

The trial instance.

mark_running(no_runner_required: bool = False, unsafe: bool = False) BaseTrial[source]

Mark trial has started running.

Parameters:
  • no_runner_required – Whether to skip the check for presence of a Runner

  • experiment. (on) –

  • unsafe – Ignore sanity checks on state transitions.

Returns:

The trial instance.

mark_staged(unsafe: bool = False) BaseTrial[source]

Mark the trial as being staged for running.

Parameters:

unsafe – Ignore sanity checks on state transitions.

Returns:

The trial instance.

run() BaseTrial[source]

Deploys the trial according to the behavior on the runner.

The runner returns a run_metadata dict containining metadata of the deployment process. It also returns a deployed_name of the trial within the system to which it was deployed. Both these fields are set on the trial.

Returns:

The trial instance.

property run_metadata: Dict[str, Any]

Dict containing metadata from the deployment process.

This is set implicitly during trial.run().

property runner: Optional[Runner]

The runner object defining how to deploy the trial.

property status: TrialStatus

The status of the trial in the experimentation lifecycle.

stop(new_status: TrialStatus, reason: Optional[str] = None) BaseTrial[source]

Stops the trial according to the behavior on the runner.

The runner returns a stop_metadata dict containining metadata of the stopping process.

Parameters:
  • new_status – The new TrialStatus. Must be one of {TrialStatus.COMPLETED, TrialStatus.ABANDONED, TrialStatus.EARLY_STOPPED}

  • reason – A message containing information why the trial is to be stopped.

Returns:

The trial instance.

property stop_metadata: Dict[str, Any]

Dict containing metadata from the stopping process.

This is set implicitly during trial.stop().

property time_completed: Optional[datetime]

Completion time of the trial.

property time_created: datetime

Creation time of the trial.

property time_run_started: Optional[datetime]

Time the trial was started running (i.e. collecting data).

property time_staged: Optional[datetime]

Staged time of the trial.

property trial_type: Optional[str]

The type of the trial.

Relevant for experiments containing different kinds of trials (e.g. different deployment types).

property ttl_seconds: Optional[int]

This trial’s time-to-live once ran, in seconds. If not set, trial will never be automatically considered failed (i.e. infinite TTL). Reflects after how many seconds since the time the trial was run it will be considered failed unless completed.

update_run_metadata(metadata: Dict[str, Any]) Dict[str, Any][source]

Updates the run metadata dict stored on this trial and returns the updated dict.

update_stop_metadata(metadata: Dict[str, Any]) Dict[str, Any][source]

Updates the stop metadata dict stored on this trial and returns the updated dict.

class ax.core.base_trial.TrialStatus(value)[source]

Bases: int, Enum

Enum of trial status.

General lifecycle of a trial is::

CANDIDATE --> STAGED --> RUNNING --> COMPLETED
          ------------->         --> FAILED (retryable)
                                 --> EARLY_STOPPED (deemed unpromising)
          -------------------------> ABANDONED (non-retryable)

Trial is marked as a CANDIDATE immediately upon its creation.

Trials may be abandoned at any time prior to completion or failure. The difference between abandonment and failure is that the FAILED state is meant to express a possibly transient or retryable error, so trials in that state may be re-run and arm(s) in them may be resuggested by Ax models to be added to new trials.

ABANDONED trials on the other end, indicate that the trial (and arms(s) in it) should not be rerun or added to new trials. A trial might be marked ABANDONED as a result of human-initiated action (if some trial in experiment is poorly-performing, deterministically failing etc., and should not be run again in the experiment). It might also be marked ABANDONED in an automated way if the trial’s execution encounters an error that indicates that the arm(s) in the trial should bot be evaluated in the experiment again (e.g. the parameterization in a given arm deterministically causes trial evaluation to fail). Note that it’s also possible to abandon a single arm in a BatchTrial via batch.mark_arm_abandoned.

Early-stopped refers to trials that were deemed unpromising by an early-stopping strategy and therefore terminated.

Additionally, when trials are deployed, they may be in an intermediate staged state (e.g. scheduled but waiting for resources) or immediately transition to running. Note that STAGED trial status is not always applicable and depends on the Runner trials are deployed with (and whether a Runner is present at all; for example, in Ax Service API, trials are marked as RUNNING immediately when generated from get_next_trial, skipping the STAGED status).

NOTE: Data for abandoned trials (or abandoned arms in batch trials) is not passed to the model as part of training data, unless fit_abandoned option is specified to model bridge.

ABANDONED = 5
CANDIDATE = 0
COMPLETED = 3
DISPATCHED = 6
EARLY_STOPPED = 7
FAILED = 2
RUNNING = 4
STAGED = 1
property expecting_data: bool

True if trial is expecting data.

property is_abandoned: bool

True if this trial is an abandoned one.

property is_candidate: bool

True if this trial is a candidate.

property is_completed: bool

True if this trial is a successfully completed one.

property is_deployed: bool

True if trial has been deployed but not completed.

property is_early_stopped: bool

True if this trial is an early stopped one.

property is_failed: bool

True if this trial is a failed one.

property is_running: bool

True if this trial is a running one.

property is_terminal: bool

True if trial is completed.

ax.core.base_trial.immutable_once_run(func: Callable) Callable[source]

Decorator for methods that should throw Error when trial is running or has ever run and immutable.

BatchTrial

class ax.core.batch_trial.AbandonedArm(name: str, time: datetime, reason: Optional[str] = None)[source]

Bases: SortableBase

Class storing metadata of arm that has been abandoned within a BatchTrial.

name: str
reason: Optional[str] = None
time: datetime
class ax.core.batch_trial.BatchTrial(experiment: core.experiment.Experiment, generator_run: Optional[GeneratorRun] = None, trial_type: Optional[str] = None, optimize_for_power: Optional[bool] = False, ttl_seconds: Optional[int] = None, index: Optional[int] = None, lifecycle_stage: Optional[LifecycleStage] = None)[source]

Bases: BaseTrial

Batched trial that has multiple attached arms, meant to be deployed and evaluated together, and possibly arm weights, which are a measure of how much of the total resources allocated to evaluating a batch should go towards evaluating the specific arm. For instance, for field experiments the weights could describe the fraction of the total experiment population assigned to the different treatment arms. Interpretation of the weights is defined in Runner.

NOTE: A BatchTrial is not just a trial with many arms; it is a trial, for which it is important that the arms are evaluated simultaneously, e.g. in an A/B test where the evaluation results are subject to nonstationarity. For cases where multiple arms are evaluated separately and independently of each other, use multiple Trial objects with a single arm each.

Parameters:
  • experiment – Experiment, to which this trial is attached

  • generator_run – GeneratorRun, associated with this trial. This can a also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set.

  • trial_type – Type of this trial, if used in MultiTypeExperiment.

  • optimize_for_power – Whether to optimize the weights of arms in this trial such that the experiment’s power to detect effects of certain size is as high as possible. Refer to documentation of BatchTrial.set_status_quo_and_optimize_power for more detail.

  • ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.

  • index – If specified, the trial’s index will be set accordingly. This should generally not be specified, as in the index will be automatically determined based on the number of existing trials. This is only used for the purpose of loading from storage.

  • lifecycle_stage – The stage of the experiment lifecycle that this trial represents

property abandoned_arm_names: Set[str]

Set of names of arms that have been abandoned within this trial.

property abandoned_arms: List[Arm]

List of arms that have been abandoned within this trial.

property abandoned_arms_metadata: List[AbandonedArm]
add_arm(*args, **kwargs)
add_arms_and_weights(*args, **kwargs)
add_generator_run(*args, **kwargs)
property arm_weights: MutableMapping[Arm, float]

The set of arms and associated weights for the trial.

These are constructed by merging the arms and weights from each generator run that is attached to the trial.

property arms: List[Arm]

All arms contained in the trial.

property arms_by_name: Dict[str, Arm]

Map from arm name to object for all arms in trial.

clone() BatchTrial[source]

Clone the trial.

Returns:

A new instance of the trial.

property experiment: core.experiment.Experiment

The experiment this batch belongs to.

property generator_run_structs: List[GeneratorRunStruct]

List of generator run structs attached to this trial.

Struct holds generator_run object and the weight with which it was added.

property generator_runs: List[GeneratorRun]

All generator runs associated with this trial.

property in_design_arms: List[Arm]
property index: int

The index of this batch within the experiment’s batch list.

property is_factorial: bool

Return true if the trial’s arms are a factorial design with no linked factors.

property lifecycle_stage: Optional[LifecycleStage]
mark_arm_abandoned(arm_name: str, reason: Optional[str] = None) BatchTrial[source]

Mark a arm abandoned.

Usually done after deployment when one arm causes issues but user wants to continue running other arms in the batch.

NOTE: Abandoned arms are considered to be ‘pending points’ in experiment after their abandonment to avoid Ax models suggesting the same arm again as a new candidate. Abandoned arms are also excluded from model training data unless fit_abandoned option is specified to model bridge.

Parameters:
  • arm_name – The name of the arm to abandon.

  • reason – The reason for abandoning the arm.

Returns:

The batch instance.

normalized_arm_weights(total: float = 1, trunc_digits: Optional[int] = None) MutableMapping[Arm, float][source]

Returns arms with a new set of weights normalized to the given total.

This method is useful for many runners where we need to normalize weights to a certain total without mutating the weights attached to a trial.

Parameters:
  • total – The total weight to which to normalize. Default is 1, in which case arm weights can be interpreted as probabilities.

  • trunc_digits – The number of digits to keep. If the resulting total weight is not equal to total, re-allocate weight in such a way to maintain relative weights as best as possible.

Returns:

Mapping from arms to the new set of weights.

run() BatchTrial[source]

Deploys the trial according to the behavior on the runner.

The runner returns a run_metadata dict containining metadata of the deployment process. It also returns a deployed_name of the trial within the system to which it was deployed. Both these fields are set on the trial.

Returns:

The trial instance.

set_status_quo_and_optimize_power(*args, **kwargs)
set_status_quo_with_weight(*args, **kwargs)
property status_quo: Optional[Arm]

The control arm for this batch.

unset_status_quo() None[source]

Set the status quo to None.

property weights: List[float]

Weights corresponding to arms contained in the trial.

class ax.core.batch_trial.GeneratorRunStruct(generator_run: GeneratorRun, weight: float)[source]

Bases: SortableBase

Stores GeneratorRun object as well as the weight with which it was added.

generator_run: GeneratorRun
weight: float
class ax.core.batch_trial.LifecycleStage(value)[source]

Bases: int, Enum

An enumeration.

BAKEOFF = 2
EXPLORATION = 0
EXPLORATION_CONCURRENT = 4
ITERATION = 1
OFFLINE_OPTIMIZED = 3

Data

class ax.core.data.Data(df: Optional[pandas.DataFrame] = None, description: Optional[str] = None)[source]

Bases: Base, SerializationMixin

Class storing data for an experiment.

The dataframe is retrieved via the df property. The data can be stored to an external store for future use by attaching it to an experiment using experiment.attach_data() (this requires a description to be set.)

df

DataFrame with underlying data, and required columns.

description

Human-readable description of data.

COLUMN_DATA_TYPES = {'arm_name': <class 'str'>, 'end_time': pandas.Timestamp, 'fidelities': <class 'str'>, 'frac_nonnull': <class 'numpy.float64'>, 'mean': <class 'numpy.float64'>, 'metric_name': <class 'str'>, 'n': <class 'numpy.int64'>, 'random_split': <class 'numpy.int64'>, 'sem': <class 'numpy.float64'>, 'start_time': pandas.Timestamp, 'trial_index': <class 'numpy.int64'>}
REQUIRED_COLUMNS = {'arm_name', 'mean', 'metric_name', 'sem'}
classmethod column_data_types(extra_column_types: Optional[Dict[str, Type]] = None, excluded_columns: Optional[Iterable[str]] = None) Dict[str, Type][source]

Type specification for all supported columns.

copy_structure_with_df(df: pandas.DataFrame) Data[source]

Serialize the structural properties needed to initialize this Data. Used for storage and to help construct new similar Data. All kwargs other than df and description are considered structural.

classmethod deserialize_init_args(args: Dict[str, Any]) Dict[str, Any][source]

Given a dictionary, extract the properties needed to initialize the metric. Used for storage.

property df: pandas.DataFrame
property df_hash: str

Compute hash of pandas DataFrame.

This first serializes the DataFrame and computes the md5 hash on the resulting string. Note that this may cause performance issue for very large DataFrames.

Parameters:

df – The DataFrame for which to compute the hash.

Returns

str: The hash of the DataFrame.

static from_evaluations(evaluations: Dict[str, Dict[str, Union[float, floating, integer, Tuple[Union[float, floating, integer], Optional[Union[float, floating, integer]]]]]], trial_index: int, sample_sizes: Optional[Dict[str, int]] = None, start_time: Optional[int] = None, end_time: Optional[int] = None) Data[source]

Convert dict of evaluations to Ax data object.

Parameters:
  • evaluations – Map from arm name to metric outcomes, which itself is a mapping of metric names to means or tuples of mean and an SEM). If SEM is not specified, it will be set to None and inferred from data.

  • trial_index – Trial index to which this data belongs.

  • sample_sizes – Number of samples collected for each arm.

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

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

Returns:

Ax Data object.

static from_fidelity_evaluations(evaluations: Dict[str, List[Tuple[Dict[str, Optional[Union[str, bool, float, int]]], Dict[str, Union[float, floating, integer, Tuple[Union[float, floating, integer], Optional[Union[float, floating, integer]]]]]]]], trial_index: int, sample_sizes: Optional[Dict[str, int]] = None, start_time: Optional[int] = None, end_time: Optional[int] = None) Data[source]

Convert dict of fidelity evaluations to Ax data object.

Parameters:
  • evaluations – Map from arm name to list of (fidelity, metric outcomes) where metric outcomes is itself a mapping of metric names to means or tuples of mean and SEM. If SEM is not specified, it will be set to None and inferred from data.

  • trial_index – Trial index to which this data belongs.

  • sample_sizes – Number of samples collected for each arm.

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

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

Returns:

Ax Data object.

static from_multiple_data(data: Iterable[Data], subset_metrics: Optional[Iterable[str]] = None) Data[source]

Combines multiple data objects into one (with the concatenated underlying dataframe).

Parameters:
  • data – Iterable of Ax Data objects to combine.

  • subset_metrics – If specified, combined Data will only contain metrics, names of which appear in this iterable, in the underlying dataframe.

get_filtered_results(**filters: Dict[str, Any]) pandas.DataFrame[source]

Return filtered subset of data.

Parameters:

filter – Column names and values they must match.

Returns

df: The filtered DataFrame.

property metric_names: Set[str]

Set of metric names that appear in the underlying dataframe of this object.

classmethod required_columns() Set[str][source]

Names of columns that must be present in the underlying DataFrame.

classmethod serialize_init_args(obj: Any) Dict[str, Any][source]

Serialize the class-dependent properties needed to initialize this Data. Used for storage and to help construct new similar Data.

classmethod supported_columns(extra_column_names: Optional[Iterable[str]] = None) Set[str][source]

Names of columns supported (but not necessarily required) by this class.

property true_df: pandas.DataFrame

Return the DataFrame being used as the source of truth (avoid using except for caching).

ax.core.data.clone_without_metrics(data: Data, excluded_metric_names: Iterable[str]) Data[source]

Returns a new data object where rows containing the metrics specified by metric_names are filtered out. Used to sanitize data before using it as training data for a model that requires data rectangularity.

Parameters:
  • data – Original data to clone.

  • excluded_metric_names – Metrics to avoid copying

Returns:

new version of the original data without specified metrics.

ax.core.data.custom_data_class(column_data_types: Optional[Dict[str, Type]] = None, required_columns: Optional[Set[str]] = None, time_columns: Optional[Set[str]] = None) Type[Data][source]

Creates a custom data class with additional columns.

All columns and their designations on the base data class are preserved, the inputs here are appended to the definitions on the base class.

Parameters:
  • column_data_types – Dict from column name to column type.

  • required_columns – Set of additional columns required for this data object.

  • time_columns – Set of additional columns to cast to timestamp.

Returns:

New data subclass with amended column definitions.

ax.core.data.set_single_trial(data: Data) Data[source]

Returns a new Data object where we set all rows to have the same trial index (i.e. 0). This is meant to be used with our IVW transform, which will combine multiple observations of the same metric.

Experiment

class ax.core.experiment.DataType(value)[source]

Bases: Enum

An enumeration.

DATA = 1
MAP_DATA = 3
class ax.core.experiment.Experiment(search_space: SearchSpace, name: Optional[str] = None, optimization_config: Optional[OptimizationConfig] = None, tracking_metrics: Optional[List[Metric]] = None, runner: Optional[Runner] = None, status_quo: Optional[Arm] = None, description: Optional[str] = None, is_test: bool = False, experiment_type: Optional[str] = None, properties: Optional[Dict[str, Any]] = None, default_data_type: Optional[DataType] = None)[source]

Bases: Base

Base class for defining an experiment.

add_tracking_metric(metric: Metric) Experiment[source]

Add a new metric to the experiment.

Parameters:

metric – Metric to be added.

add_tracking_metrics(metrics: List[Metric]) Experiment[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:

metrics – Metrics to be added.

property arms_by_name: Dict[str, Arm]

The arms belonging to this experiment, by their name.

property arms_by_signature: Dict[str, Arm]

The arms belonging to this experiment, by their signature.

attach_data(data: Data, combine_with_last_data: bool = False, overwrite_existing_data: bool = False) int[source]

Attach data to experiment. Stores data in experiment._data_by_trial, to be looked up via experiment.lookup_data_for_trial.

Parameters:
  • data – Data object to store.

  • combine_with_last_data

    By default, when attaching data, it’s identified by its timestamp, and experiment.lookup_data_for_trial returns data by most recent timestamp. Sometimes, however, we want to combine the data from multiple calls to attach_data into one dataframe. This might be because:

    • We attached data for some metrics at one point and data for

    the rest of the metrics later on. - We attached data for some fidelity at one point and data for another fidelity later one.

    To achieve that goal, set combine_with_last_data to True. In this case, we will take the most recent previously attached data, append the newly attached data to it, attach a new Data object with the merged result, and delete the old one. Afterwards, calls to lookup_data_for_trial will return this new combined data object. This operation will also validate that the newly added data does not contain observations for metrics that already have observations at the same fidelity in the most recent data.

  • overwrite_existing_data – By default, we keep around all data that has ever been attached to the experiment. However, if we know that the incoming data contains all the information we need for a given trial, we can replace the existing data for that trial, thereby reducing the amount we need to store in the database.

Returns:

Timestamp of storage in millis.

attach_fetch_results(results: Mapping[int, Mapping[str, Result[Data, MetricFetchE]]], combine_with_last_data: bool = False, overwrite_existing_data: bool = False) Optional[int][source]

UNSAFE: Prefer to use attach_data directly instead.

Attach fetched data results to the Experiment so they will not have to be fetched again. Returns the timestamp from attachment, which is used as a dict key for _data_by_trial.

NOTE: Any Errs in the results passed in will silently be dropped! This will cause the Experiment to fail to find them in the _data_by_trial cache and attempt to refetch at fetch time. If this is not your intended behavior you MUST resolve your results first and use attach_data directly instead.

property data_by_trial: Dict[int, OrderedDict]

Data stored on the experiment, indexed by trial index and storage time.

First key is trial index and second key is storage time in milliseconds. For a given trial, data is ordered by storage time, so first added data will appear first in the list.

property default_data_constructor: Type
property default_data_type: DataType
property default_trial_type: Optional[str]

Default trial type assigned to trials in this experiment.

In the base experiment class this is always None. For experiments with multiple trial types, use the MultiTypeExperiment class.

property experiment_type: Optional[str]

The type of the experiment.

fetch_data(metrics: Optional[List[Metric]] = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) Data[source]

Fetches data for all trials on this experiment and for either the specified metrics or all metrics currently on the experiment, if metrics argument is not specified.

NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whether a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.

Parameters:
  • metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.

  • kwargs – keyword args to pass to underlying metrics’ fetch data functions.

Returns:

Data for the experiment.

fetch_data_results(metrics: Optional[List[Metric]] = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) Dict[int, Dict[str, Result[Data, MetricFetchE]]][source]

Fetches data for all trials on this experiment and for either the specified metrics or all metrics currently on the experiment, if metrics argument is not specified.

If a metric fetch fails, the Exception will be captured in the MetricFetchResult along with a message.

NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whether a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.

Parameters:
  • metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.

  • kwargs – keyword args to pass to underlying metrics’ fetch data functions.

Returns:

A nested Dictionary from trial_index => metric_name => result

fetch_trials_data(trial_indices: Iterable[int], metrics: Optional[List[Metric]] = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) Data[source]

Fetches data for specific trials on the experiment.

NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whetner a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.

Parameters:
  • trial_indices – Indices of trials, for which to fetch data.

  • metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.

  • kwargs – Keyword args to pass to underlying metrics’ fetch data functions.

Returns:

Data for the specific trials on the experiment.

fetch_trials_data_results(trial_indices: Iterable[int], metrics: Optional[List[Metric]] = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) Dict[int, Dict[str, Result[Data, MetricFetchE]]][source]

Fetches data for specific trials on the experiment.

If a metric fetch fails, the Exception will be captured in the MetricFetchResult along with a message.

NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whether a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.

Parameters:
  • trial_indices – Indices of trials, for which to fetch data.

  • metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.

  • kwargs – keyword args to pass to underlying metrics’ fetch data functions.

Returns:

A nested Dictionary from trial_index => metric_name => result

get_trials_by_indices(trial_indices: Iterable[int]) List[BaseTrial][source]

Grabs trials on this experiment by their indices.

property has_name: bool

Return true if experiment’s name is not None.

property immutable_search_space_and_opt_config: bool

Boolean representing whether search space and metrics on this experiment are mutable (by default they are).

NOTE: For experiments with immutable search spaces and metrics, generator runs will not store copies of search space and metrics, which improves storage layer performance. Not keeping copies of those on generator runs also disables keeping track of changes to search space and metrics, thereby necessitating that those attributes be immutable on experiment.

property is_moo_problem: bool

Whether the experiment’s optimization config contains multiple objectives.

property is_test: bool

Get whether the experiment is a test.

lookup_data(trial_indices: Optional[Iterable[int]] = None) Data[source]

Lookup data for all trials on this experiment and for either the specified metrics or all metrics currently on the experiment, if metrics argument is not specified.

Parameters:

trial_indices – Indices of trials, for which to fetch data.

Returns:

Data for the experiment.

lookup_data_for_trial(trial_index: int) Tuple[Data, int][source]

Lookup stored data for a specific trial.

Returns latest data object, and its storage timestamp, present for this trial. Returns empty data and -1 if no data present.

Parameters:

trial_index – The index of the trial to lookup data for.

Returns:

The requested data object, and its storage timestamp in milliseconds.

lookup_data_for_ts(timestamp: int) Data[source]

Collect data for all trials stored at this timestamp.

Useful when many trials’ data was fetched and stored simultaneously and user wants to retrieve same collection of data later.

Can also be used to lookup specific data for a single trial when storage time is known.

Parameters:

timestamp – Timestamp in millis at which data was stored.

Returns:

Data object with all data stored at the timestamp.

property metrics: Dict[str, Metric]

The metrics attached to the experiment.

property name: str

Get experiment name. Throws if name is None.

new_batch_trial(generator_run: Optional[GeneratorRun] = None, trial_type: Optional[str] = None, optimize_for_power: Optional[bool] = False, ttl_seconds: Optional[int] = None, lifecycle_stage: Optional[LifecycleStage] = None) BatchTrial[source]

Create a new batch trial associated with this experiment.

Parameters:
  • generator_run – GeneratorRun, associated with this trial. This can a also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set.

  • trial_type – Type of this trial, if used in MultiTypeExperiment.

  • optimize_for_power – Whether to optimize the weights of arms in this trial such that the experiment’s power to detect effects of certain size is as high as possible. Refer to documentation of BatchTrial.set_status_quo_and_optimize_power for more detail.

  • ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.

  • lifecycle_stage – The stage of the experiment lifecycle that this trial represents

new_trial(generator_run: Optional[GeneratorRun] = None, trial_type: Optional[str] = None, ttl_seconds: Optional[int] = None) Trial[source]

Create a new trial associated with this experiment.

Parameters:
  • generator_run – GeneratorRun, associated with this trial. Trial has only one arm attached to it and this generator_run must therefore contain one arm. This arm can also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set.

  • trial_type – Type of this trial, if used in MultiTypeExperiment.

  • ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.

property num_abandoned_arms: int

How many arms attached to this experiment are abandoned.

property num_trials: int

How many trials are associated with this experiment.

property optimization_config: Optional[OptimizationConfig]

The experiment’s optimization config.

property parameters: Dict[str, Parameter]

The parameters in the experiment’s search space.

remove_tracking_metric(metric_name: str) Experiment[source]

Remove a metric that already exists on the experiment.

Parameters:

metric_name – Unique name of metric to remove.

reset_runners(runner: Runner) None[source]

Replace all candidate trials runners.

Parameters:

runner – New runner to replace with.

runner_for_trial(trial: BaseTrial) Optional[Runner][source]

The default runner to use for a given trial.

In the base experiment class, this is always the default experiment runner. For experiments with multiple trial types, use the MultiTypeExperiment class.

property running_trial_indices: Set[int]

Indices of running trials, associated with the experiment.

property search_space: SearchSpace

The search space for this experiment.

When setting a new search space, all parameter names and types must be preserved. However, if no trials have been created, all modifications are allowed.

property status_quo: Optional[Arm]

The existing arm that new arms will be compared against.

property sum_trial_sizes: int

Sum of numbers of arms attached to each trial in this experiment.

supports_trial_type(trial_type: Optional[str]) bool[source]

Whether this experiment allows trials of the given type.

The base experiment class only supports None. For experiments with multiple trial types, use the MultiTypeExperiment class.

property time_created: datetime

Creation time of the experiment.

property tracking_metrics: List[Metric]
property trial_indices_by_status: Dict[TrialStatus, Set[int]]

Indices of trials associated with the experiment, grouped by trial status.

property trials: Dict[int, BaseTrial]

The trials associated with the experiment.

NOTE: If some trials on this experiment specify their TTL, RUNNING trials will be checked for whether their TTL elapsed during this call. Found past- TTL trials will be marked as FAILED.

property trials_by_status: Dict[TrialStatus, List[BaseTrial]]

Trials associated with the experiment, grouped by trial status.

property trials_expecting_data: List[BaseTrial]

the list of all trials for which data has arrived or is expected to arrive.

Type:

List[BaseTrial]

update_tracking_metric(metric: Metric) Experiment[source]

Redefine a metric that already exists on the experiment.

Parameters:

metric – New metric definition.

warm_start_from_old_experiment(old_experiment: Experiment, copy_run_metadata_keys: Optional[List[str]] = None, trial_statuses_to_copy: Optional[List[TrialStatus]] = None) List[Trial][source]

Copy all completed trials with data from an old Ax expeirment to this one. This function checks that the parameters of each trial are members of the current experiment’s search_space.

NOTE: Currently only handles experiments with 1-arm Trial-s, not BatchTrial-s as there has not yet been need for support of the latter.

Parameters:
  • old_experiment – The experiment from which to transfer trials and data

  • copy_run_metadata_keys – A list of keys denoting which items to copy over from each trial’s run_metadata.

  • trial_statuses_to_copy – All trials with a status in this list will be copied. By default, copies all COMPLETED, ABANDONED, and EARLY_STOPPED trials.

Returns:

List of trials successfully copied from old_experiment to this one

ax.core.experiment.add_arm_and_maybe_clear_name(new_trial: Trial, old_trial: Trial) None[source]

GeneratorRun

class ax.core.generator_run.ArmWeight(arm: Arm, weight: float)[source]

Bases: Base

NamedTuple for tying together arms and weights.

arm: Arm
weight: float
class ax.core.generator_run.GeneratorRun(arms: List[Arm], weights: Optional[List[float]] = None, optimization_config: Optional[OptimizationConfig] = None, search_space: Optional[SearchSpace] = None, model_predictions: Optional[Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]]] = None, best_arm_predictions: Optional[Tuple[Arm, Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]] = None, type: Optional[str] = None, fit_time: Optional[float] = None, gen_time: Optional[float] = None, model_key: Optional[str] = None, model_kwargs: Optional[Dict[str, Any]] = None, bridge_kwargs: Optional[Dict[str, Any]] = None, gen_metadata: Optional[Dict[str, Any]] = None, model_state_after_gen: Optional[Dict[str, Any]] = None, generation_step_index: Optional[int] = None, candidate_metadata_by_arm_signature: Optional[Dict[str, Optional[Dict[str, Any]]]] = None)[source]

Bases: SortableBase

An object that represents a single run of a generator.

This object is created each time the gen method of a generator is called. It stores the arms and (optionally) weights that were generated by the run. When we add a generator run to a trial, its arms and weights will be merged with those from previous generator runs that were already attached to the trial.

property arm_signatures: Set[str]

Returns signatures of arms generated by this run.

property arm_weights: MutableMapping[Arm, float]

Mapping from arms to weights (order matches order in arms property).

property arms: List[Arm]

Returns arms generated by this run.

property best_arm_predictions: Optional[Tuple[Arm, Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]]
property candidate_metadata_by_arm_signature: Optional[Dict[str, Optional[Dict[str, Any]]]]

Retrieves model-produced candidate metadata as a mapping from arm name (for the arm the candidate became when added to experiment) to the metadata dict.

clone() GeneratorRun[source]

Return a deep copy of a GeneratorRun.

property fit_time: Optional[float]
property gen_metadata: Optional[Dict[str, Any]]

Returns metadata generated by this run.

property gen_time: Optional[float]
property generator_run_type: Optional[str]

The type of the generator run.

property index: Optional[int]

The index of this generator run within a trial’s list of generator run structs. This field is set when the generator run is added to a trial.

property model_predictions: Optional[Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]]]
property model_predictions_by_arm: Optional[Dict[str, Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]
property optimization_config: Optional[OptimizationConfig]

The optimization config used during generation of this run.

property param_df: pandas.DataFrame

Constructs a Pandas dataframe with the parameter values for each arm.

Useful for inspecting the contents of a generator run.

Returns:

a dataframe with the generator run’s arms.

Return type:

pd.DataFrame

property search_space: Optional[SearchSpace]

The search used during generation of this run.

property time_created: datetime

Creation time of the batch.

property weights: List[float]

Returns weights associated with arms generated by this run.

class ax.core.generator_run.GeneratorRunType(value)[source]

Bases: Enum

Class for enumerating generator run types.

MANUAL = 1
STATUS_QUO = 0
ax.core.generator_run.extract_arm_predictions(model_predictions: Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]], arm_idx: int) Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]][source]

Extract a particular arm from model_predictions.

Parameters:
  • model_predictions – Mean and Cov for all arms.

  • arm_idx – Index of arm in prediction list.

Returns:

(mean, cov) for specified arm.

MapData

class ax.core.map_data.MapData(df: Optional[pandas.DataFrame] = None, map_key_infos: Optional[Iterable[MapKeyInfo]] = None, description: Optional[str] = None)[source]

Bases: Data

Class storing mapping-like results for an experiment.

Data is stored in a dataframe, and axilary information ((key name, default value) pairs) are stored in a collection of MapKeyInfo objects.

Mapping-like results occur whenever a metric is reported as a collection of results, each element corresponding to a tuple of values.

The simplest case is a sequence. For instance a time series is a mapping from the 1-tuple (timestamp) to (mean, sem) results.

Another example: MultiFidelity results. This is a mapping from (fidelity_feature_1, …, fidelity_feature_n) to (mean, sem) results.

The dataframe is retrieved via the map_df property. The data can be stored to an external store for future use by attaching it to an experiment using experiment.attach_data() (this requires a description to be set.)

DEDUPLICATE_BY_COLUMNS = ['arm_name', 'metric_name']
classmethod deserialize_init_args(args: Dict[str, Any]) Dict[str, Any][source]

Given a dictionary, extract the properties needed to initialize the metric. Used for storage.

property df: pandas.DataFrame

Returns a Data shaped DataFrame

static from_map_evaluations(evaluations: Dict[str, List[Tuple[Dict[str, Hashable], Dict[str, Union[float, floating, integer, Tuple[Union[float, floating, integer], Optional[Union[float, floating, integer]]]]]]]], trial_index: int, map_key_infos: Optional[Iterable[MapKeyInfo]] = None) MapData[source]
static from_multiple_data(data: Iterable[Data], subset_metrics: Optional[Iterable[str]] = None) MapData[source]

Downcast instances of Data into instances of MapData with empty map_key_infos if necessary then combine as usual (filling in empty cells with default values).

static from_multiple_map_data(data: Iterable[MapData], subset_metrics: Optional[Iterable[str]] = None) MapData[source]
property map_df: pandas.DataFrame
property map_key_infos: Iterable[MapKeyInfo]
property map_key_to_type: Dict[str, Type]
property map_keys: List[str]
classmethod serialize_init_args(obj: Any) Dict[str, Any][source]

Serialize the class-dependent properties needed to initialize this Data. Used for storage and to help construct new similar Data.

subsample(map_key: Optional[str] = None, keep_every: Optional[int] = None, limit_rows_per_group: Optional[int] = None, limit_total_rows: Optional[int] = None, include_first_last: bool = True) MapData[source]

Subsample the map_key column in an equally-spaced manner (if there is a self.map_keys is length one, then map_key can be set to None). The values of the map_key column are not taken into account, so this function is most reasonable when those values are equally-spaced. There are three ways that this can be done:

  1. If keep_every = k is set, then every kth row of the DataFrame in the

    map_key column is kept after grouping by DEDUPLICATE_BY_COLUMNS. In other words, every kth step of each (arm, metric) will be kept.

  2. If limit_rows_per_group = n, the method will find the (arm, metric)

    pair with the largest number of rows in the map_key column and select an approprioate keep_every such that each (arm, metric) has at most n rows in the map_key column.

  3. If limit_total_rows = n, the method will select an appropriate

    keep_every such that the total number of rows per metric is less than n.

If multiple of keep_every, limit_rows_per_group, limit_total_rows, then the priority is in the order above: 1. keep_every, 2. limit_rows_per_group, and 3. limit_total_rows.

Note that we want all curves to be subsampled with nearly the same spacing. Internally, the method converts limit_rows_per_group and limit_total_rows to a keep_every quantity that will satisfy the original request.

When include_first_last is True, then the method will use the keep_every as a guideline and for each group, produce (nearly) evenly spaced points that include the first and last points.

property true_df: pandas.DataFrame

Return the DataFrame being used as the source of truth (avoid using except for caching).

class ax.core.map_data.MapKeyInfo(key: str, default_value: T)[source]

Bases: Generic[T], SortableBase

Helper class storing map keys and auxilary info for use in MapData

property default_value: T
property key: str
property value_type: Type

MapMetric

class ax.core.map_metric.MapMetric(name: str, lower_is_better: Optional[bool] = None, properties: Optional[Dict[str, Any]] = None)[source]

Bases: Metric

Base class for representing metrics that return MapData.

The fetch_trial_data method is the essential method to override when subclassing, which specifies how to retrieve a Metric, for a given trial.

A MapMetric must return a MapData object, which requires (at minimum) the following:

https://ax.dev/api/_modules/ax/core/data.html#Data.required_columns

lower_is_better

Flag for metrics which should be minimized.

properties

Properties specific to a particular metric.

data_constructor

alias of MapData

Metric

class ax.core.metric.Metric(name: str, lower_is_better: Optional[bool] = None, properties: Optional[Dict[str, Any]] = None)[source]

Bases: SortableBase, SerializationMixin

Base class for representing metrics.

The fetch_trial_data method is the essential method to override when subclassing, which specifies how to retrieve a Metric, for a given trial.

A Metric must return a Data object, which requires (at minimum) the following:

https://ax.dev/api/_modules/ax/core/data.html#Data.required_columns

lower_is_better

Flag for metrics which should be minimized.

properties

Properties specific to a particular metric.

clone() Metric[source]

Create a copy of this Metric.

data_constructor

alias of Data

fetch_experiment_data(experiment: core.experiment.Experiment, **kwargs: Any) Dict[int, MetricFetchResult][source]

Fetch this metric’s data for an experiment.

Returns Dict of trial_index => Result

classmethod fetch_experiment_data_multi(experiment: core.experiment.Experiment, metrics: Iterable[Metric], trials: Optional[Iterable[core.base_trial.BaseTrial]] = None, **kwargs: Any) Dict[int, Dict[str, MetricFetchResult]][source]

Fetch multiple metrics data for an experiment.

Returns Dict of trial_index => (metric_name => Result) Default behavior calls fetch_trial_data_multi for each trial. Subclasses should override to batch data computation across trials + metrics.

property fetch_multi_group_by_metric: Type[Metric]

Metric class, with which to group this metric in Experiment._metrics_by_class, which is used to combine metrics on experiment into groups and then fetch their data via Metric.fetch_trial_data_multi for each group.

NOTE: By default, this property will just return the class on which it is defined; however, in some cases it is useful to group metrics by their superclass, in which case this property should return that superclass.

fetch_trial_data(trial: core.base_trial.BaseTrial, **kwargs: Any) MetricFetchResult[source]

Fetch data for one trial.

classmethod fetch_trial_data_multi(trial: core.base_trial.BaseTrial, metrics: Iterable[Metric], **kwargs: Any) Dict[str, MetricFetchResult][source]

Fetch multiple metrics data for one trial.

Returns Dict of metric_name => Result Default behavior calls fetch_trial_data for each metric. Subclasses should override this to trial data computation for multiple metrics.

classmethod is_available_while_running() bool[source]

Whether metrics of this class are available while the trial is running. Metrics that are not available while the trial is running are assumed to be available only upon trial completion. For such metrics, data is assumed to never change once the trial is completed.

NOTE: If this method returns False, data-fetching via experiment.fetch_data will return the data cached on the experiment (for the metrics of the given class) whenever its available. Data is cached on experiment when attached via experiment.attach_data.

classmethod lookup_or_fetch_experiment_data_multi(experiment: core.experiment.Experiment, metrics: Iterable[Metric], trials: Optional[Iterable[core.base_trial.BaseTrial]] = None, **kwargs: Any) Tuple[Dict[int, Dict[str, MetricFetchResult]], bool][source]

Fetch or lookup (with fallback to fetching) data for given metrics, depending on whether they are available while running. Return a tuple containing the data, along with a boolean that will be True if new data was fetched, and False if all data was looked up from cache.

If metric is available while running, its data can change (and therefore we should always re-fetch it). If metric is available only upon trial completion, its data does not change, so we can look up that data on the experiment and only fetch the data that is not already attached to the experiment.

NOTE: If fetching data for a metrics class that is only available upon trial completion, data fetched in this function (data that was not yet available on experiment) will be attached to experiment.

property name: str

Get name of metric.

class ax.core.metric.MetricFetchE(message: 'str', exception: 'Optional[Exception]')[source]

Bases: object

exception: Optional[Exception]
message: str

MultiTypeExperiment

class ax.core.multi_type_experiment.MultiTypeExperiment(name: str, search_space: SearchSpace, default_trial_type: str, default_runner: Runner, optimization_config: Optional[OptimizationConfig] = None, status_quo: Optional[Arm] = None, description: Optional[str] = None, is_test: bool = False, experiment_type: Optional[str] = None, properties: Optional[Dict[str, Any]] = None, default_data_type: Optional[DataType] = None)[source]

Bases: Experiment

Class for experiment with multiple trial types.

A canonical use case for this is tuning a large production system with limited evaluation budget and a simulator which approximates evaluations on the main system. Trial deployment and data fetching is separate for the two systems, but the final data is combined and fed into multi-task models.

See the Multi-Task Modeling tutorial for more details.

name

Name of the experiment.

description

Description of the experiment.

add_tracking_metric(metric: Metric, trial_type: str, canonical_name: Optional[str] = None) MultiTypeExperiment[source]

Add a new metric to the experiment.

Parameters:
  • metric – The metric to add.

  • trial_type – The trial type for which this metric is used.

  • canonical_name – The default metric for which this metric is a proxy.

add_trial_type(trial_type: str, runner: Runner) MultiTypeExperiment[source]

Add a new trial_type to be supported by this experiment.

Parameters:
  • trial_type – The new trial_type to be added.

  • runner – The default runner for trials of this type.

property default_trial_type: Optional[str]

Default trial type assigned to trials in this experiment.

property default_trials: Set[int]

Return the indicies for trials of the default type.

fetch_data(metrics: Optional[List[Metric]] = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) Data[source]

Fetches data for all trials on this experiment and for either the specified metrics or all metrics currently on the experiment, if metrics argument is not specified.

NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whether a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.

Parameters:
  • metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.

  • kwargs – keyword args to pass to underlying metrics’ fetch data functions.

Returns:

Data for the experiment.

property metric_to_trial_type: Dict[str, str]

Map metrics to trial types.

Adds in default trial type for OC metrics to custom defined trial types..

remove_tracking_metric(metric_name: str) MultiTypeExperiment[source]

Remove a metric that already exists on the experiment.

Parameters:

metric_name – Unique name of metric to remove.

reset_runners(runner: Runner) None[source]

Replace all candidate trials runners.

Parameters:

runner – New runner to replace with.

runner_for_trial(trial: BaseTrial) Optional[Runner][source]

The default runner to use for a given trial.

Looks up the appropriate runner for this trial type in the trial_type_to_runner.

supports_trial_type(trial_type: Optional[str]) bool[source]

Whether this experiment allows trials of the given type.

Only trial types defined in the trial_type_to_runner are allowed.

update_runner(trial_type: str, runner: Runner) MultiTypeExperiment[source]

Update the default runner for an existing trial_type.

Parameters:
  • trial_type – The new trial_type to be added.

  • runner – The new runner for trials of this type.

update_tracking_metric(metric: Metric, trial_type: str, canonical_name: Optional[str] = None) MultiTypeExperiment[source]

Update an existing metric on the experiment.

Parameters:
  • metric – The metric to add.

  • trial_type – The trial type for which this metric is used.

  • canonical_name – The default metric for which this metric is a proxy.

Objective

class ax.core.objective.MultiObjective(objectives: Optional[List[Objective]] = None, **extra_kwargs: Any)[source]

Bases: Objective

Class for an objective composed of a multiple component objectives.

The Acquisition function determines how the objectives are weighted.

objectives

List of objectives.

clone() Objective[source]

Create a copy of the objective.

property metric: Metric

Override base method to error.

property metrics: List[Metric]

Get the objective metrics.

property objective_weights: Iterable[Tuple[Objective, float]]

Get the objectives and weights.

property objectives: List[Objective]

Get the objectives.

weights: List[float]
class ax.core.objective.Objective(metric: Metric, minimize: Optional[bool] = None)[source]

Bases: SortableBase

Base class for representing an objective.

minimize

If True, minimize metric.

clone() Objective[source]

Create a copy of the objective.

get_unconstrainable_metrics() List[Metric][source]

Return a list of metrics that are incompatible with OutcomeConstraints.

property metric: Metric

Get the objective metric.

property metric_names: List[str]

Get a list of objective metric names.

property metrics: List[Metric]

Get a list of objective metrics.

class ax.core.objective.ScalarizedObjective(metrics: List[Metric], weights: Optional[List[float]] = None, minimize: bool = False)[source]

Bases: Objective

Class for an objective composed of a linear scalarization of metrics.

metrics

List of metrics.

weights

Weights for scalarization; default to 1.

Type:

List[float]

clone() Objective[source]

Create a copy of the objective.

property metric: Metric

Override base method to error.

property metric_weights: Iterable[Tuple[Metric, float]]

Get the metrics and weights.

property metrics: List[Metric]

Get the metrics.

weights: List[float]

Observation

class ax.core.observation.Observation(features: ObservationFeatures, data: ObservationData, arm_name: Optional[str] = None)[source]

Bases: Base

Represents an observation.

A set of features (ObservationFeatures) and corresponding measurements (ObservationData). Optionally, an arm name associated with the features.

features
Type:

ObservationFeatures

data
Type:

ObservationData

arm_name
Type:

Optional[str]

class ax.core.observation.ObservationData(metric_names: List[str], means: ndarray, covariance: ndarray)[source]

Bases: Base

Outcomes observed at a point.

The “point” corresponding to this ObservationData would be an ObservationFeatures object.

metric_names

A list of k metric names that were observed

means

a k-array of observed means

covariance

a (k x k) array of observed covariances

property covariance_matrix: Dict[str, Dict[str, float]]

Extract covariance matric from this observation data as mapping from metric name (m1) to mapping of another metric name (m2) to the covariance of the two metrics (m1 and m2).

property means_dict: Dict[str, float]

Extract means from this observation data as mapping from metric name to mean.

class ax.core.observation.ObservationFeatures(parameters: Dict[str, Optional[Union[str, bool, float, int]]], trial_index: Optional[int64] = None, start_time: Optional[pandas.Timestamp] = None, end_time: Optional[pandas.Timestamp] = None, random_split: Optional[int64] = None, metadata: Optional[Dict[str, Any]] = None)[source]

Bases: Base

The features of an observation.

These include both the arm parameters and the features of the observation found in the Data object: trial index, times, and random split. This object is meant to contain everything needed to represent this observation in a model feature space. It is essentially a row of Data joined with the arm parameters.

An ObservationFeatures object would typically have a corresponding ObservationData object that provides the observed outcomes.

parameters

arm parameters

trial_index

trial index

start_time

batch start time

end_time

batch end time

random_split

random split

clone(replace_parameters: Optional[Dict[str, Optional[Union[str, bool, float, int]]]] = None) ObservationFeatures[source]

Make a copy of these ObservationFeatures.

Parameters:

replace_parameters – An optimal parameterization, to which to set the parameters of the cloned ObservationFeatures. Useful when transforming observation features in a way that requires a change to parameterization –– for example, while casting it to a hierarchical search space.

static from_arm(arm: Arm, trial_index: Optional[int64] = None, start_time: Optional[pandas.Timestamp] = None, end_time: Optional[pandas.Timestamp] = None, random_split: Optional[int64] = None, metadata: Optional[Dict[str, Any]] = None) ObservationFeatures[source]

Convert a Arm to an ObservationFeatures, including additional data as specified.

update_features(new_features: ObservationFeatures) ObservationFeatures[source]

Updates the existing ObservationFeatures with the fields of the the input.

Adds all of the new parameters to the existing parameters and overwrites any other fields that are not None on the new input features.

ax.core.observation.get_feature_cols(data: Data) List[str][source]
ax.core.observation.get_feature_cols_from_map_data(map_data: MapData) List[str][source]
ax.core.observation.observations_from_data(experiment: Experiment, data: Data, include_abandoned: bool = False) List[Observation][source]

Convert Data to observations.

Converts a Data object to a list of Observation objects. Pulls arm parameters from from experiment. Overrides fidelity parameters in the arm with those found in the Data object.

Uses a diagonal covariance matrix across metric_names.

Parameters:
  • experiment – Experiment with arm parameters.

  • data – Data of observations.

  • include_abandoned – Whether data for abandoned trials and arms should be included in the observations, returned from this function.

Returns:

List of Observation objects.

ax.core.observation.observations_from_map_data(experiment: Experiment, map_data: MapData, include_abandoned: bool = False, map_keys_as_parameters: bool = False, limit_total_rows: Optional[int] = None, limit_rows_per_group: Optional[int] = None) List[Observation][source]

Convert MapData to observations.

Converts a MapData object to a list of Observation objects. Pulls arm parameters from experiment. Overrides fidelity parameters in the arm with those found in the Data object.

Uses a diagonal covariance matrix across metric_names.

Parameters:
  • experiment – Experiment with arm parameters.

  • map_data – MapData of observations.

  • include_abandoned – Whether data for abandoned trials and arms should be included in the observations, returned from this function.

  • map_keys_as_parameters – Whether map_keys should be returned as part of the parameters of the Observation objects.

  • limit_total_rows – If specified, uses MapData.subsample() with limit_total_rows equal to the specified value on the first map_key (map_data.map_keys[0]) to subsample the MapData. This is useful in, e.g., cases where learning curves are frequently updated, leading to an intractable number of Observation objects created.

  • limit_rows_per_group – If specified, uses MapData.subsample() with limit_rows_per_group equal to the specified value on the first map_key (map_data.map_keys[0]) to subsample the MapData.

Returns:

List of Observation objects.

ax.core.observation.recombine_observations(observation_features: List[ObservationFeatures], observation_data: List[ObservationData]) List[Observation][source]
ax.core.observation.separate_observations(observations: List[Observation], copy: bool = False) Tuple[List[ObservationFeatures], List[ObservationData]][source]

Split out observations into features+data.

Parameters:

observations – input observations

Returns:

ObservationFeatures observation_data: ObservationData

Return type:

observation_features

OptimizationConfig

class ax.core.optimization_config.MultiObjectiveOptimizationConfig(objective: Objective, outcome_constraints: Optional[List[OutcomeConstraint]] = None, objective_thresholds: Optional[List[ObjectiveThreshold]] = None, risk_measure: Optional[RiskMeasure] = None)[source]

Bases: OptimizationConfig

An optimization configuration for multi-objective optimization, which comprises multiple objective, outcome constraints, objective thresholds, and an optional risk measure.

There is no minimum or maximum number of outcome constraints, but an individual metric can have at most two constraints–which is how we represent metrics with both upper and lower bounds.

ObjectiveThresholds should be present for every objective. A good rule of thumb is to set them 10% below the minimum acceptable value for each metric.

property all_constraints: List[OutcomeConstraint]

Get all constraints and thresholds.

clone_with_args(objective: ~typing.Optional[~ax.core.objective.Objective] = None, outcome_constraints: ~typing.Optional[~typing.List[~ax.core.outcome_constraint.OutcomeConstraint]] = [OutcomeConstraint( >= 0%)], objective_thresholds: ~typing.Optional[~typing.List[~ax.core.outcome_constraint.ObjectiveThreshold]] = [ObjectiveThreshold( <= 0%)], risk_measure: ~typing.Optional[~ax.core.risk_measures.RiskMeasure] = RiskMeasure(risk_measure=, options={})) MultiObjectiveOptimizationConfig[source]

Make a copy of this optimization config.

property metrics: Dict[str, Metric]
property objective: Objective

Get objective.

property objective_thresholds: List[ObjectiveThreshold]

Get objective thresholds.

property objective_thresholds_dict: Dict[str, ObjectiveThreshold]

Get a mapping from objective metric name to the corresponding threshold.

class ax.core.optimization_config.OptimizationConfig(objective: Objective, outcome_constraints: Optional[List[OutcomeConstraint]] = None, risk_measure: Optional[RiskMeasure] = None)[source]

Bases: Base

An optimization configuration, which comprises an objective, outcome constraints and an optional risk measure.

There is no minimum or maximum number of outcome constraints, but an individual metric can have at most two constraints–which is how we represent metrics with both upper and lower bounds.

property all_constraints: List[OutcomeConstraint]

Get outcome constraints.

clone() OptimizationConfig[source]

Make a copy of this optimization config.

clone_with_args(objective: ~typing.Optional[~ax.core.objective.Objective] = None, outcome_constraints: ~typing.Optional[~typing.List[~ax.core.outcome_constraint.OutcomeConstraint]] = [OutcomeConstraint( >= 0%)], risk_measure: ~typing.Optional[~ax.core.risk_measures.RiskMeasure] = RiskMeasure(risk_measure=, options={})) OptimizationConfig[source]

Make a copy of this optimization config.

property is_moo_problem: bool
property metrics: Dict[str, Metric]
property objective: Objective

Get objective.

property outcome_constraints: List[OutcomeConstraint]

Get outcome constraints.

ax.core.optimization_config.check_objective_thresholds_match_objectives(objectives_by_name: Dict[str, Objective], objective_thresholds: List[ObjectiveThreshold]) None[source]

Error if thresholds on objective_metrics bound from the wrong direction or if there is a mismatch between objective thresholds and objectives.

OutcomeConstraint

class ax.core.outcome_constraint.ObjectiveThreshold(metric: Metric, bound: float, relative: bool = True, op: Optional[ComparisonOp] = None)[source]

Bases: OutcomeConstraint

Class for representing Objective Thresholds.

An objective threshold represents the threshold for an objective metric to contribute to hypervolume calculations. A list containing the objective threshold for each metric collectively form a reference point.

Objective thresholds may bound the metric from above or from below. The bound can be expressed as an absolute measurement or relative to the status quo (if applicable).

The direction of the bound is inferred from the Metric’s lower_is_better attribute.

metric

Metric to constrain.

bound

The bound in the constraint.

relative

Whether you want to bound on an absolute or relative scale. If relative, bound is the acceptable percent change.

op

automatically inferred, but manually overwritable. specifies whether metric should be greater or equal to, or less than or equal to, some bound.

clone() ObjectiveThreshold[source]

Create a copy of this ObjectiveThreshold.

class ax.core.outcome_constraint.OutcomeConstraint(metric: Metric, op: ComparisonOp, bound: float, relative: bool = True)[source]

Bases: SortableBase

Base class for representing outcome constraints.

Outcome constraints may of the form metric >= bound or metric <= bound, where the bound can be expressed as an absolute measurement or relative to the status quo (if applicable).

metric

Metric to constrain.

op

Specifies whether metric should be greater or equal to, or less than or equal to, some bound.

bound

The bound in the constraint.

relative

Whether you want to bound on an absolute or relative scale. If relative, bound is the acceptable percent change.

clone() OutcomeConstraint[source]

Create a copy of this OutcomeConstraint.

property metric: Metric
property op: ComparisonOp
class ax.core.outcome_constraint.ScalarizedOutcomeConstraint(metrics: List[Metric], op: ComparisonOp, bound: float, relative: bool = True, weights: Optional[List[float]] = None)[source]

Bases: OutcomeConstraint

Class for presenting outcome constraints composed of a linear scalarization of metrics.

metrics

List of metrics.

weights

Weights for scalarization; default to 1.0 / len(metrics).

Type:

List[float]

op

Specifies whether metric should be greater or equal to, or less than or equal to, some bound.

bound

The bound in the constraint.

relative

Whether you want to bound on an absolute or relative scale. If relative, bound is the acceptable percent change.

clone() ScalarizedOutcomeConstraint[source]

Create a copy of this ScalarizedOutcomeConstraint.

property metric: Metric

Override base method to error.

property metric_weights: Iterable[Tuple[Metric, float]]

Get the objective metrics and weights.

property metrics: List[Metric]
property op: ComparisonOp
weights: List[float]

Parameter

class ax.core.parameter.ChoiceParameter(name: str, parameter_type: ParameterType, values: List[Optional[Union[str, bool, float, int]]], is_ordered: Optional[bool] = None, is_task: bool = False, is_fidelity: bool = False, target_value: Optional[Union[str, bool, float, int]] = None, sort_values: Optional[bool] = None, dependents: Optional[Dict[Optional[Union[str, bool, float, int]], List[str]]] = None)[source]

Bases: Parameter

Parameter object that specifies a discrete set of values.

Parameters:
  • name – Name of the parameter.

  • parameter_type – Enum indicating the type of parameter value (e.g. string, int).

  • values – List of allowed values for the parameter.

  • is_ordered – If False, the parameter is a categorical variable. Defaults to False if parameter_type is STRING and values is longer than 2, else True.

  • is_task – Treat the parameter as a task parameter for modeling.

  • is_fidelity – Whether this parameter is a fidelity parameter.

  • target_value – Target value of this parameter if it’s fidelity.

  • sort_values – Whether to sort values before encoding. Defaults to False if parameter_type is STRING, else True.

  • dependents – Optional mapping for parameters in hierarchical search spaces; format is { value -> list of dependent parameter names }.

add_values(values: List[Optional[Union[str, bool, float, int]]]) ChoiceParameter[source]

Add input list to the set of allowed values for parameter.

Cast all input values to the parameter type.

Parameters:

values – Values being added to the allowed list.

clone() ChoiceParameter[source]
property dependents: Dict[Optional[Union[str, bool, float, int]], List[str]]
property is_ordered: bool
property is_task: bool
set_values(values: List[Optional[Union[str, bool, float, int]]]) ChoiceParameter[source]

Set the list of allowed values for parameter.

Cast all input values to the parameter type.

Parameters:

values – New list of allowed values.

property sort_values: bool
validate(value: Optional[Union[str, bool, float, int]]) bool[source]

Checks that the input is in the list of allowed values.

Parameters:

value – Value being checked.

Returns:

True if valid, False otherwise.

property values: List[Optional[Union[str, bool, float, int]]]
class ax.core.parameter.FixedParameter(name: str, parameter_type: ParameterType, value: Optional[Union[str, bool, float, int]], is_fidelity: bool = False, target_value: Optional[Union[str, bool, float, int]] = None, dependents: Optional[Dict[Optional[Union[str, bool, float, int]], List[str]]] = None)[source]

Bases: Parameter

Parameter object that specifies a single fixed value.

clone() FixedParameter[source]
property dependents: Dict[Optional[Union[str, bool, float, int]], List[str]]
set_value(value: Optional[Union[str, bool, float, int]]) FixedParameter[source]
validate(value: Optional[Union[str, bool, float, int]]) bool[source]

Checks that the input is equal to the fixed value.

Parameters:

value – Value being checked.

Returns:

True if valid, False otherwise.

property value: Optional[Union[str, bool, float, int]]
class ax.core.parameter.Parameter[source]

Bases: SortableBase

cast(value: Optional[Union[str, bool, float, int]]) Optional[Union[str, bool, float, int]][source]
clone() Parameter[source]
property dependents: Dict[Optional[Union[str, bool, float, int]], List[str]]
property is_fidelity: bool
property is_hierarchical: bool
property is_numeric: bool
is_valid_type(value: Optional[Union[str, bool, float, int]]) bool[source]

Whether a given value’s type is allowed by this parameter.

property name: str
property parameter_type: ParameterType
property python_type: Union[Type[int], Type[float], Type[str], Type[bool]]

The python type for the corresponding ParameterType enum.

Used primarily for casting values of unknown type to conform to that of the parameter.

property target_value: Optional[Union[str, bool, float, int]]
abstract validate(value: Optional[Union[str, bool, float, int]]) bool[source]
class ax.core.parameter.ParameterType(value)[source]

Bases: Enum

An enumeration.

BOOL: int = 0
FLOAT: int = 2
INT: int = 1
STRING: int = 3
property is_numeric: bool
class ax.core.parameter.RangeParameter(name: str, parameter_type: ParameterType, lower: float, upper: float, log_scale: bool = False, logit_scale: bool = False, digits: Optional[int] = None, is_fidelity: bool = False, target_value: Optional[Union[str, bool, float, int]] = None)[source]

Bases: Parameter

Parameter object that specifies a range of values.

cast(value: Optional[Union[str, bool, float, int]]) Optional[Union[str, bool, float, int]][source]
clone() RangeParameter[source]
property digits: Optional[int]

Number of digits to round values to for float type.

Upper and lower bound are re-cast after this property is changed.

is_valid_type(value: Optional[Union[str, bool, float, int]]) bool[source]

Same as default except allows floats whose value is an int for Int parameters.

property log_scale: bool

Whether the parameter’s random values should be sampled from log space.

property logit_scale: bool

Whether the parameter’s random values should be sampled from logit space.

property lower: float

Lower bound of the parameter range.

Value is cast to parameter type upon set and also validated to ensure the bound is strictly less than upper bound.

set_digits(digits: int) RangeParameter[source]
set_log_scale(log_scale: bool) RangeParameter[source]
set_logit_scale(logit_scale: bool) RangeParameter[source]
update_range(lower: Optional[float] = None, upper: Optional[float] = None) RangeParameter[source]

Set the range to the given values.

If lower or upper is not provided, it will be left at its current value.

Parameters:
  • lower – New value for the lower bound.

  • upper – New value for the upper bound.

property upper: float

Upper bound of the parameter range.

Value is cast to parameter type upon set and also validated to ensure the bound is strictly greater than lower bound.

validate(value: Optional[Union[str, bool, float, int]]) bool[source]

Returns True if input is a valid value for the parameter.

Checks that value is of the right type and within the valid range for the parameter. Returns False if value is None.

Parameters:

value – Value being checked.

Returns:

True if valid, False otherwise.

ParameterConstraint

class ax.core.parameter_constraint.OrderConstraint(lower_parameter: Parameter, upper_parameter: Parameter)[source]

Bases: ParameterConstraint

Constraint object for specifying one parameter to be smaller than another.

clone() OrderConstraint[source]

Clone.

clone_with_transformed_parameters(transformed_parameters: Dict[str, Parameter]) OrderConstraint[source]

Clone, but replace parameters with transformed versions.

property constraint_dict: Dict[str, float]

Weights on parameters for linear constraint representation.

property lower_parameter: Parameter

Parameter with lower value.

property parameters: List[Parameter]

Parameters.

property upper_parameter: Parameter

Parameter with higher value.

class ax.core.parameter_constraint.ParameterConstraint(constraint_dict: Dict[str, float], bound: float)[source]

Bases: SortableBase

Base class for linear parameter constraints.

Constraints are expressed using a map from parameter name to weight followed by a bound.

The constraint is satisfied if w * v <= b where:

w is the vector of parameter weights. v is a vector of parameter values. b is the specified bound. * is the dot product operator.

property bound: float

Get bound of the inequality of the constraint.

check(parameter_dict: Dict[str, Union[int, float]]) bool[source]

Whether or not the set of parameter values satisfies the constraint.

Does a weighted sum of the parameter values based on the constraint_dict and checks that the sum is less than the bound.

Parameters:

parameter_dict – Map from parameter name to parameter value.

Returns:

Whether the constraint is satisfied.

clone() ParameterConstraint[source]

Clone.

clone_with_transformed_parameters(transformed_parameters: Dict[str, Parameter]) ParameterConstraint[source]

Clone, but replaced parameters with transformed versions.

property constraint_dict: Dict[str, float]

Get mapping from parameter names to weights.

class ax.core.parameter_constraint.SumConstraint(parameters: List[Parameter], is_upper_bound: bool, bound: float)[source]

Bases: ParameterConstraint

Constraint on the sum of parameters being greater or less than a bound.

clone() SumConstraint[source]

Clone.

To use the same constraint, we need to reconstruct the original bound. We do this by re-applying the original bound weighting.

clone_with_transformed_parameters(transformed_parameters: Dict[str, Parameter]) SumConstraint[source]

Clone, but replace parameters with transformed versions.

property constraint_dict: Dict[str, float]

Weights on parameters for linear constraint representation.

property op: ComparisonOp

Whether the sum is constrained by a <= or >= inequality.

property parameters: List[Parameter]

Parameters.

ax.core.parameter_constraint.validate_constraint_parameters(parameters: List[Parameter]) None[source]

Basic validation of parameters used in a constraint.

Parameters:

parameters – Parameters used in constraint.

Raises:

ValueError if the parameters are not valid for use.

ParameterDistribution

class ax.core.parameter_distribution.ParameterDistribution(parameters: List[str], distribution_class: str, distribution_parameters: Optional[Dict[str, Any]], multiplicative: bool = False)[source]

Bases: SortableBase

A class for defining parameter distributions.

Intended for robust optimization use cases. This could be used to specify the distribution of an environmental variable or the distribution of the input noise.

clone() ParameterDistribution[source]

Clone.

property distribution: rv_frozen

Get the distribution object.

property distribution_class: str

The name of the scipy distribution class.

property distribution_parameters: Dict[str, Any]

The parameters of the distribution.

is_environmental(search_space: RobustSearchSpace) bool[source]

Check if the parameters are environmental variables of the given search space.

Parameters:

search_space – The search space to check.

Returns:

A boolean denoting whether the parameters are environmental variables.

RiskMeasure

class ax.core.risk_measures.RiskMeasure(risk_measure: str, options: Dict[str, Union[int, float, bool, List[float]]])[source]

Bases: SortableBase

A class for defining risk measures.

This can be used with a RobustSearchSpace, to convert the predictions over `ParameterDistribution`s to robust metrics, which then get used in candidate generation to recommend robust candidates.

See ax/modelbridge/modelbridge_utils.py for RISK_MEASURE_NAME_TO_CLASS, which lists the supported risk measures, and for extract_risk_measure helper, which extracts the BoTorch risk measure.

clone() RiskMeasure[source]

Clone.

Runner

class ax.core.runner.Runner[source]

Bases: Base, SerializationMixin, ABC

Abstract base class for custom runner classes

clone() Runner[source]

Create a copy of this Runner.

poll_available_capacity() int[source]

Checks how much available capacity there is to schedule trial evaluations. Required for runners used with Ax Scheduler.

NOTE: This method might be difficult to implement in some systems. Returns -1 if capacity of the system is “unlimited” or “unknown” (meaning that the Scheduler should be trying to schedule as many trials as is possible without violating scheduler settings). There is no need to artificially force this method to limit capacity; Scheduler has other limitations in place to limit number of trials running at once, like the SchedulerOptions.max_pending_trials setting, or more granular control in the form of the max_parallelism setting in each of the GenerationStep`s of a `GenerationStrategy).

Returns:

An integer, representing how many trials there is available capacity for; -1 if capacity is “unlimited” or not possible to know in advance.

poll_trial_status(trials: Iterable[core.base_trial.BaseTrial]) Dict[core.base_trial.TrialStatus, Set[int]][source]

Checks the status of any non-terminal trials and returns their indices as a mapping from TrialStatus to a list of indices. Required for runners used with Ax Scheduler.

NOTE: Does not need to handle waiting between polling calls while trials are running; this function should just perform a single poll.

Parameters:

trials – Trials to poll.

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

abstract run(trial: core.base_trial.BaseTrial) Dict[str, Any][source]

Deploys a trial based on custom runner subclass implementation.

Parameters:

trial – The trial to deploy.

Returns:

Dict of run metadata from the deployment process.

run_multiple(trials: Iterable[core.base_trial.BaseTrial]) Dict[int, Dict[str, Any]][source]

Runs a single evaluation for each of the given trials. Useful when deploying multiple trials at once is more efficient than deploying them one-by-one. Used in Ax Scheduler.

NOTE: By default simply loops over run_trial. Should be overwritten if deploying multiple trials in batch is preferable.

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 a BatchTrial if contains multiple arms.

Returns:

Dict of trial index to the run metadata of that trial from the deployment process.

property staging_required: bool

Whether the trial goes to staged or running state once deployed.

stop(trial: core.base_trial.BaseTrial, reason: Optional[str] = None) Dict[str, Any][source]

Stop a trial based on custom runner subclass implementation.

Optional method.

Parameters:
  • trial – The trial to stop.

  • reason – A message containing information why the trial is to be stopped.

Returns:

A dictionary of run metadata from the stopping process.

SearchSpace

class ax.core.search_space.HierarchicalSearchSpace(parameters: List[Parameter], parameter_constraints: Optional[List[ParameterConstraint]] = None)[source]

Bases: SearchSpace

cast_observation_features(observation_features: ObservationFeatures) ObservationFeatures[source]

Cast parameterization of given observation features to the hierarchical structure of the given search space; return the newly cast observation features with the full parameterization stored in metadata under Keys.FULL_PARAMETERIZATION.

For each parameter in given parameterization, cast it to the proper type specified in this search space and remove it from the parameterization if that parameter should not be in the arm within the search space due to its hierarchical structure.

check_membership(parameterization: Dict[str, Optional[Union[str, bool, float, int]]], raise_error: bool = False, check_all_parameters_present: bool = True) bool[source]

Whether the given parameterization belongs in the search space.

Checks that the given parameter values have the same name/type as search space parameters, are contained in the search space domain, and satisfy the parameter constraints.

Parameters:
  • parameterization – Dict from parameter name to value to validate.

  • raise_error – If true parameterization does not belong, raises an error with detailed explanation of why.

  • check_all_parameters_present – Ensure that parameterization specifies values for all parameters as expected by the search space and its hierarchical structure.

Returns:

Whether the parameterization is contained in the search space.

flatten() SearchSpace[source]

Returns a flattened SearchSpace with all the parameters in the given HierarchicalSearchSpace; ignores their hierarchical structure.

flatten_observation_features(observation_features: ObservationFeatures, inject_dummy_values_to_complete_flat_parameterization: bool = False) ObservationFeatures[source]

Flatten observation features that were previously cast to the hierarchical structure of the given search space; return the newly flattened observation features. This method re-injects parameter values that were removed from observation features during casting (as they are saved in observation features metadata).

Parameters:
  • observation_features – Observation features corresponding to one point to flatten.

  • inject_dummy_values_to_complete_flat_parameterization – Whether to inject values for parameters that are not in the parameterization if they are not recorded in the observation features’ metadata (this can happen if e.g. the point wasn’t generated by Ax but attached manually).

hierarchical_structure_str(parameter_names_only: bool = False) str[source]

String representation of the hierarchical structure.

Parameters:

parameter_names_only – Whether parameter should show up just as names (instead of full parameter strings), useful for a more concise representation.

property root: Parameter

Root of the hierarchical search space tree, as identified during HierarchicalSearchSpace construction.

class ax.core.search_space.RobustSearchSpace(parameters: List[Parameter], parameter_distributions: List[ParameterDistribution], num_samples: int, environmental_variables: Optional[List[Parameter]] = None, parameter_constraints: Optional[List[ParameterConstraint]] = None)[source]

Bases: SearchSpace

Search space for robust optimization that supports environmental variables and input noise.

In addition to the usual search space properties, this allows specifying environmental variables (parameters) and input noise distributions.

clone() RobustSearchSpace[source]
is_environmental_variable(parameter_name: str) bool[source]

Check if a given parameter is an environmental variable.

Parameters:

parameter – A string denoting the name of the parameter.

Returns:

A boolean denoting whether the given parameter_name corresponds to an environmental variable of this search space.

property parameters: Dict[str, Parameter]

Get all parameters and environmental variables.

We include environmental variables here to support transform_search_space and other similar functionality. It also helps avoid having to overwrite a bunch of parent methods.

update_parameter(parameter: Parameter) None[source]
class ax.core.search_space.RobustSearchSpaceDigest(sample_param_perturbations: ~typing.Optional[~typing.Callable[[], ~numpy.ndarray]] = None, sample_environmental: ~typing.Optional[~typing.Callable[[], ~numpy.ndarray]] = None, environmental_variables: ~typing.List[str] = <factory>, multiplicative: bool = False)[source]

Bases: object

Container for lightweight representation of properties that are unique to the RobustSearchSpace. This is used to append the SearchSpaceDigest.

NOTE: Both sample_param_perturbations and sample_environmental should require no inputs and return a num_samples x d-dim array of samples from the corresponding parameter distributions, where d is the number of non-environmental parameters for distribution_sampler and the number of environmental variables for environmental_sampler.

sample_param_perturbations

An optional callable for sampling from the parameter distributions representing input perturbations.

Type:

Optional[Callable[[], numpy.ndarray]]

sample_environmental

An optional callable for sampling from the distributions of the environmental variables.

Type:

Optional[Callable[[], numpy.ndarray]]

environmental_variables

A list of environmental variable names.

Type:

List[str]

multiplicative

Denotes whether the distribution is multiplicative. Only relevant if paired with a distribution_sampler.

Type:

bool

environmental_variables: List[str]
multiplicative: bool = False
sample_environmental: Optional[Callable[[], ndarray]] = None
sample_param_perturbations: Optional[Callable[[], ndarray]] = None
class ax.core.search_space.SearchSpace(parameters: List[Parameter], parameter_constraints: Optional[List[ParameterConstraint]] = None)[source]

Bases: Base

Base object for SearchSpace object.

Contains a set of Parameter objects, each of which have a name, type, and set of valid values. The search space also contains a set of ParameterConstraint objects, which can be used to define restrictions across parameters (e.g. p_a < p_b).

add_parameter(parameter: Parameter) None[source]
add_parameter_constraints(parameter_constraints: List[ParameterConstraint]) None[source]
cast_arm(arm: Arm) Arm[source]

Cast parameterization of given arm to the types in this SearchSpace.

For each parameter in given arm, cast it to the proper type specified in this search space. Throws if there is a mismatch in parameter names. This is mostly useful for int/float, which user can be sloppy with when hand written.

Parameters:

arm – Arm to cast.

Returns:

New casted arm.

check_all_parameters_present(parameterization: Dict[str, Optional[Union[str, bool, float, int]]], raise_error: bool = False) bool[source]

Whether a given parameterization contains all the parameters in the search space.

Parameters:
  • parameterization – Dict from parameter name to value to validate.

  • raise_error – If true parameterization does not belong, raises an error with detailed explanation of why.

Returns:

Whether the parameterization is contained in the search space.

check_membership(parameterization: Dict[str, Optional[Union[str, bool, float, int]]], raise_error: bool = False, check_all_parameters_present: bool = True) bool[source]

Whether the given parameterization belongs in the search space.

Checks that the given parameter values have the same name/type as search space parameters, are contained in the search space domain, and satisfy the parameter constraints.

Parameters:
  • parameterization – Dict from parameter name to value to validate.

  • raise_error – If true parameterization does not belong, raises an error with detailed explanation of why.

  • check_all_parameters_present – Ensure that parameterization specifies values for all parameters as expected by the search space.

Returns:

Whether the parameterization is contained in the search space.

check_types(parameterization: Dict[str, Optional[Union[str, bool, float, int]]], allow_none: bool = True, raise_error: bool = False) bool[source]

Checks that the given parameterization’s types match the search space.

Parameters:
  • parameterization – Dict from parameter name to value to validate.

  • allow_none – Whether None is a valid parameter value.

  • raise_error – If true and parameterization does not belong, raises an error with detailed explanation of why.

Returns:

Whether the parameterization has valid types.

clone() SearchSpace[source]
construct_arm(parameters: Optional[Dict[str, Optional[Union[str, bool, float, int]]]] = None, name: Optional[str] = None) Arm[source]

Construct new arm using given parameters and name. Any missing parameters fallback to the experiment defaults, represented as None.

property is_hierarchical: bool
property is_robust: bool
out_of_design_arm() Arm[source]

Create a default out-of-design arm.

An out of design arm contains values for some parameters which are outside of the search space. In the modeling conversion, these parameters are all stripped down to an empty dictionary, since the point is already outside of the modeled space.

Returns:

New arm w/ null parameter values.

property parameter_constraints: List[ParameterConstraint]
property parameters: Dict[str, Parameter]
property range_parameters: Dict[str, Parameter]
set_parameter_constraints(parameter_constraints: List[ParameterConstraint]) None[source]
property tunable_parameters: Dict[str, Parameter]
update_parameter(parameter: Parameter) None[source]
class ax.core.search_space.SearchSpaceDigest(feature_names: ~typing.List[str], bounds: ~typing.List[~typing.Tuple[~typing.Union[int, float], ~typing.Union[int, float]]], ordinal_features: ~typing.List[int] = <factory>, categorical_features: ~typing.List[int] = <factory>, discrete_choices: ~typing.Dict[int, ~typing.List[~typing.Union[int, float]]] = <factory>, task_features: ~typing.List[int] = <factory>, fidelity_features: ~typing.List[int] = <factory>, target_fidelities: ~typing.Dict[int, ~typing.Union[int, float]] = <factory>, robust_digest: ~typing.Optional[~ax.core.search_space.RobustSearchSpaceDigest] = None)[source]

Bases: object

Container for lightweight representation of search space properties.

This is used for communicating between modelbridge and models. This is an ephemeral object and not meant to be stored / serialized.

feature_names

A list of parameter names.

Type:

List[str]

bounds

A list [(l_0, u_0), …, (l_d, u_d)] of tuples representing the lower and upper bounds on the respective parameter (both inclusive).

Type:

List[Tuple[Union[int, float], Union[int, float]]]

ordinal_features

A list of indices corresponding to the parameters to be considered as ordinal discrete parameters. The corresponding bounds are assumed to be integers, and parameter i is assumed to take on values l_i, l_i+1, …, u_i.

Type:

List[int]

categorical_features

A list of indices corresponding to the parameters to be considered as categorical discrete parameters. The corresponding bounds are assumed to be integers, and parameter i is assumed to take on values l_i, l_i+1, …, u_i.

Type:

List[int]

discrete_choices

A dictionary mapping indices of discrete (ordinal or categorical) parameters to their respective sets of values provided as a list.

Type:

Dict[int, List[Union[int, float]]]

task_features

A list of parameter indices to be considered as task parameters.

Type:

List[int]

fidelity_features

A list of parameter indices to be considered as fidelity parameters.

Type:

List[int]

target_fidelities

A dictionary mapping parameter indices (of fidelity parameters) to their respective target fidelity value. Only used when generating candidates.

Type:

Dict[int, Union[int, float]]

robust_digest

An optional RobustSearchSpaceDigest that carries the additional attributes if using a RobustSearchSpace.

Type:

Optional[ax.core.search_space.RobustSearchSpaceDigest]

bounds: List[Tuple[Union[int, float], Union[int, float]]]
categorical_features: List[int]
discrete_choices: Dict[int, List[Union[int, float]]]
feature_names: List[str]
fidelity_features: List[int]
ordinal_features: List[int]
robust_digest: Optional[RobustSearchSpaceDigest] = None
target_fidelities: Dict[int, Union[int, float]]
task_features: List[int]

Trial

class ax.core.trial.Trial(experiment: core.experiment.Experiment, generator_run: Optional[GeneratorRun] = None, trial_type: Optional[str] = None, ttl_seconds: Optional[int] = None, index: Optional[int] = None)[source]

Bases: BaseTrial

Trial that only has one attached arm and no arm weights.

Parameters:
  • experiment – Experiment, to which this trial is attached.

  • generator_run – GeneratorRun, associated with this trial. Trial has only one generator run (of just one arm) attached to it. This can also be set later through add_arm or add_generator_run, but a trial’s associated genetor run is immutable once set.

  • trial_type – Type of this trial, if used in MultiTypeExperiment.

  • ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.

  • index – If specified, the trial’s index will be set accordingly. This should generally not be specified, as in the index will be automatically determined based on the number of existing trials. This is only used for the purpose of loading from storage.

property abandoned_arms: List[Arm]

Abandoned arms attached to this trial.

add_arm(*args, **kwargs)
add_generator_run(*args, **kwargs)
property arm: Optional[Arm]

The arm associated with this batch.

property arms: List[Arm]

All arms attached to this trial.

Returns:

list of a single arm

attached to this trial if there is one, else None.

Return type:

arms

property arms_by_name: Dict[str, Arm]

Dictionary of all arms attached to this trial with their names as keys.

Returns:

dictionary of a single

arm name to arm if one is attached to this trial, else None.

Return type:

arms

property generator_run: Optional[GeneratorRun]

Generator run attached to this trial.

property generator_runs: List[GeneratorRun]

All generator runs associated with this trial.

get_metric_mean(metric_name: str) float[source]

Metric mean for the arm attached to this trial, retrieved from the latest data available for the metric for the trial.

property objective_mean: float

Objective mean for the arm attached to this trial, retrieved from the latest data available for the objective for the trial.

Note: the retrieved objective is the experiment-level objective at the time of the call to objective_mean, which is not necessarily the objective that was set at the time the trial was created or ran.

Core Types

class ax.core.types.ComparisonOp(value)[source]

Bases: Enum

Class for enumerating comparison operations.

GEQ: int = 0
LEQ: int = 1
ax.core.types.merge_model_predict(predict: Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]], predict_append: Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]]) Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]][source]

Append model predictions to an existing set of model predictions.

TModelPredict is of the form:

{metric_name: [mean1, mean2, …], {metric_name: {metric_name: [var1, var2, …]}})

This will append the predictions

Parameters:
  • predict – Initial set of predictions.

  • other_predict – Predictions to be appended.

Returns:

TModelPredict with the new predictions appended.

Core Utils

class ax.core.utils.MissingMetrics(objective, outcome_constraints, tracking_metrics)[source]

Bases: tuple

objective: Dict[str, Set[Tuple[str, int]]]

Alias for field number 0

outcome_constraints: Dict[str, Set[Tuple[str, int]]]

Alias for field number 1

tracking_metrics: Dict[str, Set[Tuple[str, int]]]

Alias for field number 2

ax.core.utils.best_feasible_objective(optimization_config: OptimizationConfig, values: Dict[str, ndarray]) ndarray[source]

Compute the best feasible objective value found by each iteration.

Parameters:
  • optimization_config – Optimization config.

  • values – Dictionary from metric name to array of value at each iteration. If optimization config contains outcome constraints, values for them must be present in values.

Returns: Array of cumulative best feasible value.

ax.core.utils.get_missing_metrics(data: Data, optimization_config: OptimizationConfig) MissingMetrics[source]

Return all arm_name, trial_index pairs, for which some of the observatins of optimization config metrics are missing.

Parameters:
  • data – Data to search.

  • optimization_config – provides metric_names to search for.

Returns:

A NamedTuple(missing_objective, Dict[str, missing_outcome_constraint])

ax.core.utils.get_missing_metrics_by_name(data: Data, metric_names: Iterable[str]) Dict[str, Set[Tuple[str, int]]][source]

Return all arm_name, trial_index pairs missing some observations of specified metrics.

Parameters:
  • data – Data to search.

  • metric_names – list of metrics to search for.

Returns:

A Dict[str, missing_metrics], one entry for each metric_name.

ax.core.utils.get_model_times(experiment: Experiment) Tuple[float, float][source]

Get total times spent fitting the model and generating candidates in the course of the experiment.