ax.core¶

Arm¶

class ax.core.arm.Arm(parameters: dict[str, None | str | bool | float | int], name: str | None = 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, None | 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 8 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, None | 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: str | None = None, ttl_seconds: int | None = None, index: int | None = 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: str | None¶

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: str | None = 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: str | None¶

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.

property failed_reason: str | None¶

fetch_data(metrics: list[Metric] | None = None, **kwargs: Any) → Data[source]¶

Fetch data for this trial for all metrics on experiment.

# NOTE: This can be lossy (ex. a MapData could get implicitly cast to a Data and # lose rows)if some if Experiment.default_data_type is misconfigured!

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.

fetch_data_results(metrics: list[Metric] | None = None, **kwargs: Any) → dict[str, Result[Data, MetricFetchE]][source]¶

Fetch data results 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:

MetricFetchResults 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: str | None = 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: str | None = 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(reason: str | None = None, 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 on the experiment.

• 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: Runner | None¶

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: str | None = 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: datetime | None¶

Completion time of the trial.

property time_created: datetime¶

Creation time of the trial.

property time_run_started: datetime | None¶

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

property time_staged: datetime | None¶

Staged time of the trial.

property trial_type: str | None¶

The type of the trial.

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

property ttl_seconds: int | None¶

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. Additionally, data from MapMetrics is typically excluded unless the corresponding trial is completed.

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: str | None = None)[source]¶

Bases: SortableBase

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

name: str¶

reason: str | None = None¶

time: datetime¶

class ax.core.batch_trial.BatchTrial(experiment: core.experiment.Experiment, generator_run: GeneratorRun | None = None, generator_runs: list[GeneratorRun] | None = None, trial_type: str | None = None, optimize_for_power: bool | None = False, ttl_seconds: int | None = None, index: int | None = None, lifecycle_stage: LifecycleStage | None = 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.

• generator_runs – GeneratorRuns, 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. This cannot be combined with the generator_run argument.

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

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

Attaches data to the trial

Parameters:

• raw_data – Map from arm name to metric outcomes.

• sample_sizes – Dict from arm name to sample size.

• metadata – Additional metadata to track about this run. importantly the start_date and end_date

• complete_trial – Whether to mark trial as complete after attaching data. Defaults to False.

clone() → BatchTrial[source]¶

Clone the trial and attach it to the current experiment.

clone_to(experiment: core.experiment.Experiment | None = None, include_sq: bool = True) → BatchTrial[source]¶

Clone the trial and attach it to a specified experiment. If None provided, attach it to the current experiment.

Parameters:

• experiment – The experiment to which the cloned trial will belong. If unspecified, uses the current experiment.

• include_sq – Whether to include status quo in the cloned 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: LifecycleStage | None¶

mark_arm_abandoned(arm_name: str, reason: str | None = 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: int | None = 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: Arm | None¶

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.BaseData(df: DataFrame | None = None, description: str | None = 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. For BaseData, the one required column is “arm_name”.

description

Human-readable description of data.

COLUMN_DATA_TYPES: dict[str, Any] = {'arm_name': <class 'str'>, 'end_time': <class 'pandas._libs.tslibs.timestamps.Timestamp'>, 'fidelities': <class 'str'>, 'frac_nonnull': <class 'numpy.float64'>, 'mean': <class 'numpy.float64'>, 'metric_name': <class 'str'>, 'n': <class 'int'>, 'random_split': <class 'int'>, 'sem': <class 'numpy.float64'>, 'start_time': <class 'pandas._libs.tslibs.timestamps.Timestamp'>, 'trial_index': <class 'int'>}

REQUIRED_COLUMNS = {'arm_name'}

classmethod column_data_types(extra_column_types: dict[str, type] | None = None, excluded_columns: Iterable[str] | None = None) → dict[str, type][source]

Type specification for all supported columns.

copy_structure_with_df(df: DataFrame) → TBaseData[source]

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

classmethod deserialize_init_args(args: dict[str, Any], decoder_registry: dict[str, type[T] | Callable[[...], T]] | None = None, class_decoder_registry: dict[str, Callable[[dict[str, Any]], Any]] | None = None) → dict[str, Any][source]

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

property df: 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.

classmethod from_evaluations(evaluations: dict[str, dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]], trial_index: int, sample_sizes: dict[str, int] | None = None, start_time: int | str | None = None, end_time: int | str | None = None) → TBaseData[source]

Convert dict of evaluations to Ax data object.

Parameters:

• evaluations – Map from arm name to outcomes, which itself is a mapping of outcome names to values, 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 or iso format. Milliseconds will be automatically converted to iso format because iso format automatically works with the pandas column type Timestamp.

• end_time – Optional end time of run of the trial that produced this data, in milliseconds or iso format. Milliseconds will be automatically converted to iso format because iso format automatically works with the pandas column type Timestamp.

Returns:

Ax object of the enclosing class.

classmethod from_fidelity_evaluations(evaluations: dict[str, list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]], trial_index: int, sample_sizes: dict[str, int] | None = None, start_time: int | None = None, end_time: int | None = None) → TBaseData[source]

Convert dict of fidelity evaluations to Ax data object.

Parameters:

• evaluations – Map from arm name to list of (fidelity, outcomes) where outcomes is itself a mapping of outcome names to values, 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 object of type cls.

classmethod from_multiple(data: Iterable[TBaseData]) → TBaseData[source]

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

Parameters:

data – Iterable of Ax objects of this class to combine.

get_filtered_results(**filters: dict[str, Any]) → DataFrame[source]

Return filtered subset of data.

Parameters:

filter – Column names and values they must match.

Returns

df: The filtered DataFrame.

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: Iterable[str] | None = None) → set[str][source]

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

property true_df: DataFrame

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

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

Bases: BaseData

Class storing numerical 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. For BaseData, the required columns are “arm_name”, “metric_name”, “mean”, and “sem”, the latter two of which must be numeric.

description

Human-readable description of data.

REQUIRED_COLUMNS: set[str] = {'arm_name', 'mean', 'metric_name', 'sem'}

clone() → Data[source]

Returns a new Data object with the same underlying dataframe.

filter(trial_indices: Iterable[int] | None = None, metric_names: Iterable[str] | None = None) → Data[source]

Construct a new object with the subset of rows corresponding to the provided trial indices AND metric names. If either trial_indices or metric_names are not provided, that dimension will not be filtered.

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

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

Parameters:

• data – Iterable of Ax objects of this class to combine.

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

property metric_names: set[str]

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

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

Returns a new data object where rows containing the outcomes 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: dict[str, type] | None = None, required_columns: set[str] | None = None, time_columns: set[str] | None = 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 outcome.

Experiment¶

class ax.core.experiment.Experiment(search_space: SearchSpace, name: str | None = None, optimization_config: OptimizationConfig | None = None, tracking_metrics: list[Metric] | None = None, runner: Runner | None = None, status_quo: Arm | None = None, description: str | None = None, is_test: bool = False, experiment_type: str | None = None, properties: dict[str, Any] | None = None, default_data_type: DataType | None = None, auxiliary_experiments_by_purpose: None | dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]] = 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.

property arms_by_signature_for_deduplication: dict[str, Arm]¶

The arms belonging to this experiment that should be used for deduplication in GenerationStrategy, by their signature.

In its current form, this includes all arms except for those that are associated with a FAILED trial. - The CANDIDATE, STAGED, RUNNING, and ABANDONED arms are included as pending points during generation, so they should be less likely

to get suggested by the model again.

• The EARLY_STOPPED and COMPLETED trials were already evaluated, so

the model will have data for these and is unlikely to suggest them again.

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) → int | None[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.

attach_trial(parameterizations: list[dict[str, None | str | bool | float | int]], arm_names: list[str] | None = None, ttl_seconds: int | None = None, run_metadata: dict[str, Any] | None = None, optimize_for_power: bool = False) → tuple[dict[str, dict[str, None | str | bool | float | int]], int][source]¶

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

Parameters:

• parameterizations – List of parameterization for the new trial. If only one is provided a single-arm Trial is created. If multiple arms are provided a BatchTrial is created.

• arm_names – Names of arm(s) in the new trial.

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

• run_metadata – Metadata to attach to the trial.

• optimize_for_power – For BatchTrial only. 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.

Returns:

Tuple of arm name to parameterization dict, and trial index from newly created trial.

clone_with(search_space: SearchSpace | None = None, name: str | None = None, optimization_config: OptimizationConfig | None = None, tracking_metrics: list[Metric] | None = None, runner: Runner | None = None, status_quo: Arm | None = None, description: str | None = None, is_test: bool | None = None, properties: dict[str, Any] | None = None, trial_indices: list[int] | None = None, data: Data | None = None) → Experiment[source]¶

Return a copy of this experiment with some attributes replaced.

NOTE: This method only retains the latest data attached to the experiment. This is the same data that would be accessed using common APIs such as Experiment.lookup_data().

Parameters:

• search_space – New search space. If None, it uses the cloned search space of the original experiment.

• name – New experiment name. If None, it adds cloned_experiment_ prefix to the original experiment name.

• optimization_config – New optimization config. If None, it clones the same optimization_config from the orignal experiment.

• tracking_metrics – New list of metrics to track. If None, it clones the tracking metrics already attached to the main experiment.

• runner – New runner. If None, it clones the existing runner.

• status_quo – New status quo arm. If None, it clones the existing status quo.

• description – New description. If None, it uses the same description.

• is_test – Whether the cloned experiment should be considered a test. If None, it uses the same value.

• properties – New properties dictionary. If None, it uses a copy of the same properties.

• trial_indices – If specified, only clones the specified trials. If None, clones all trials.

• data – If specified, attach this data to the cloned experiment. If None, clones the latest data attached to the original experiment if the experiment has any data.

property completed_trials: list[BaseTrial]¶

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

Type:

list[BaseTrial]

property data_by_trial: dict[int, OrderedDict[int, Data]]¶

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: str | None¶

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: str | None¶

The type of the experiment.

fetch_data(metrics: list[Metric] | None = 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.

NOTE: This can be lossy (ex. a MapData could get implicitly cast to a Data and lose rows) if Experiment.default_data_type is misconfigured!

Parameters:

• metrics – If provided, fetch data for these metrics; otherwise, fetch data for all metrics defined on the experiment.

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

Returns:

Data for the experiment.

fetch_data_results(metrics: list[Metric] | None = 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: list[Metric] | None = 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.

NOTE: This can be lossy (ex. a MapData could get implicitly cast to a Data and lose rows) if Experiment.default_data_type is misconfigured!

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: list[Metric] | None = 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: Iterable[int] | None = None) → Data[source]¶

Lookup stored data for trials on this experiment.

For each trial, returns latest data object present for this trial. Returns empty data if no data is present. In particular, this method will not fetch data from metrics - to do that, use fetch_data() instead.

Parameters:

trial_indices – Indices of trials for which to fetch data. If omitted, lookup data for all trials on the experiment.

Returns:

Data for the trials on 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 is present. In particular, this method will not fetch data from metrics - to do that, use fetch_data() instead.

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 metric_config_summary_df: DataFrame¶

Creates a dataframe with information about each metric in the experiment. The resulting dataframe has one row per metric, and the following columns:

• Name: the name of the metric.

• Type: the metric subclass (e.g., Metric, BraninMetric).

• Goal: the goal for this for this metric, based on the optimization config (minimize, maximize, constraint or track).

• Bound: the bound of this metric (e.g., “<=10.0”) if it is being used as part of an ObjectiveThreshold or OutcomeConstraint.

• Lower is Better: whether the user prefers this metric to be lower, if provided.

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: GeneratorRun | None = None, generator_runs: list[GeneratorRun] | None = None, trial_type: str | None = None, optimize_for_power: bool | None = False, ttl_seconds: int | None = None, lifecycle_stage: LifecycleStage | None = 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.

• generator_runs – GeneratorRuns, 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. This cannot be combined with the generator_run argument.

• 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: GeneratorRun | None = None, trial_type: str | None = None, ttl_seconds: int | None = 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: OptimizationConfig | None¶

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) → Runner | None[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: Arm | None¶

The existing arm that new arms will be compared against.

stop_trial_runs(trials: list[BaseTrial], reasons: list[str | None] | None = None) → None[source]¶

Stops the jobs that execute given trials.

Used if, for example, TTL for a trial was specified and expired, or poor early results suggest the trial is not worth running to completion.

Override default implementation on the Runner if its desirable to stop trials in bulk.

Parameters:

• trials – Trials to be stopped.

• reasons – A list of strings describing the reasons for why the trials are to be stopped (in the same order).

property sum_trial_sizes: int¶

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

supports_trial_type(trial_type: str | None) → 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 trial_indices_expecting_data: set[int]¶

Set of indices of trials, statuses of which indicate that we expect these trials to have data, either already or in the future.

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.

validate_trials(trials: Iterable[BaseTrial]) → None[source]¶

Raise ValueError if any of the trials in the input are not from this experiment.

warm_start_from_old_experiment(old_experiment: Experiment, copy_run_metadata_keys: list[str] | None = None, trial_statuses_to_copy: list[TrialStatus] | None = None, search_space_check_membership_raise_error: bool = True) → 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. Defaults to old_experiment.runner.run_metadata_report_keys.

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

• search_space_check_membership_raise_error – Whether to raise an exception if the warm started trials being imported fall outside of the defined search space.

Returns:

List of trials successfully copied from old_experiment to this one