ax.storage

JSON

ax.storage.json_store.decoder module

ax.storage.json_store.decoder.ax_class_from_json_dict(_class: Type, object_json: Dict[str, Any]) → Any[source]

Reinstantiates an Ax class registered in DECODER_REGISTRY from a JSON dict.

ax.storage.json_store.decoder.data_from_json(data_by_trial_json: Dict[str, Any]) → Dict[int, OrderedDict[int, Data]][source]

Load Ax Data from JSON.

ax.storage.json_store.decoder.experiment_from_json(object_json: Dict[str, Any]) → ax.core.experiment.Experiment[source]

Load Ax Experiment from JSON.

ax.storage.json_store.decoder.generation_step_from_json(generation_step_json: Dict[str, Any]) → ax.modelbridge.generation_strategy.GenerationStep[source]

Load generation step from JSON.

ax.storage.json_store.decoder.generation_strategy_from_json(generation_strategy_json: Dict[str, Any]) → ax.modelbridge.generation_strategy.GenerationStrategy[source]

Load generation strategy from JSON.

ax.storage.json_store.decoder.generator_run_from_json(object_json: Dict[str, Any]) → ax.core.generator_run.GeneratorRun[source]

Load Ax GeneratorRun from JSON.

ax.storage.json_store.decoder.multi_type_experiment_from_json(object_json: Dict[str, Any]) → ax.core.multi_type_experiment.MultiTypeExperiment[source]

Load AE MultiTypeExperiment from JSON.

ax.storage.json_store.decoder.object_from_json(object_json: Any) → Any[source]

Recursively load objects from a JSON-serializable dictionary.

ax.storage.json_store.decoder.parameter_constraints_from_json(parameter_constraint_json: List[Dict[str, Any]], parameters: List[ax.core.parameter.Parameter]) → List[ax.core.parameter_constraint.ParameterConstraint][source]

Load ParameterConstraints from JSON.

Order and SumConstraint are tied to a search space, and require that SearchSpace’s parameters to be passed in for decoding.

Parameters

  • parameter_constraint_json – JSON representation of parameter constraints.

  • parameters – Parameter definitions for decoding via parameter names.

Returns

Python classes for parameter constraints.

Return type

parameter_constraints

ax.storage.json_store.decoder.search_space_from_json(search_space_json: Dict[str, Any]) → ax.core.search_space.SearchSpace[source]

Load a SearchSpace from JSON.

This function is necessary due to the coupled loading of SearchSpace and parameter constraints.

ax.storage.json_store.decoder.simple_benchmark_problem_from_json(object_json: Dict[str, Any]) → ax.benchmark.benchmark_problem.SimpleBenchmarkProblem[source]

Load a benchmark problem from JSON.

ax.storage.json_store.decoder.simple_experiment_from_json(object_json: Dict[str, Any]) → ax.core.simple_experiment.SimpleExperiment[source]

Load AE SimpleExperiment from JSON.

ax.storage.json_store.decoder.trials_from_json(experiment: ax.core.experiment.Experiment, trials_json: Dict[str, Any]) → Dict[int, ax.core.base_trial.BaseTrial][source]

Load Ax Trials from JSON.

ax.storage.json_store.decoders module

ax.storage.json_store.decoders.batch_trial_from_json(experiment: core.experiment.Experiment, index: int, trial_type: Optional[str], status: ax.core.base_trial.TrialStatus, time_created: datetime.datetime, time_completed: Optional[datetime.datetime], time_staged: Optional[datetime.datetime], time_run_started: Optional[datetime.datetime], abandoned_reason: Optional[str], run_metadata: Optional[Dict[str, Any]], generator_run_structs: List[ax.core.batch_trial.GeneratorRunStruct], runner: Optional[ax.core.runner.Runner], abandoned_arms_metadata: Dict[str, ax.core.batch_trial.AbandonedArm], num_arms_created: int, status_quo: Optional[ax.core.arm.Arm], status_quo_weight_override: float, optimize_for_power: Optional[bool], ttl_seconds: Optional[int] = None, generation_step_index: Optional[int] = None, properties: Optional[Dict[str, Any]] = None) → ax.core.batch_trial.BatchTrial[source]

Load Ax BatchTrial from JSON.

Other classes don’t need explicit deserializers, because we can just use their constructors (see decoder.py). However, the constructor for Batch does not allow us to exactly recreate an existing object.

ax.storage.json_store.decoders.class_from_json(json: Dict[str, Any]) → Type[Any][source]

Load any class registered in CLASS_DECODER_REGISTRY from JSON.

ax.storage.json_store.decoders.transform_type_from_json(object_json: Dict[str, Any]) → Type[ax.modelbridge.transforms.base.Transform][source]

Load the transform type from JSON.

ax.storage.json_store.decoders.trial_from_json(experiment: core.experiment.Experiment, index: int, trial_type: Optional[str], status: ax.core.base_trial.TrialStatus, time_created: datetime.datetime, time_completed: Optional[datetime.datetime], time_staged: Optional[datetime.datetime], time_run_started: Optional[datetime.datetime], abandoned_reason: Optional[str], run_metadata: Optional[Dict[str, Any]], generator_run: ax.core.generator_run.GeneratorRun, runner: Optional[ax.core.runner.Runner], num_arms_created: int, ttl_seconds: Optional[int] = None, generation_step_index: Optional[int] = None, properties: Optional[Dict[str, Any]] = None) → ax.core.trial.Trial[source]

Load Ax trial from JSON.

Other classes don’t need explicit deserializers, because we can just use their constructors (see decoder.py). However, the constructor for Trial does not allow us to exactly recreate an existing object.

ax.storage.json_store.encoder module

ax.storage.json_store.encoder.object_to_json(obj: Any) → Any[source]

Convert an Ax object to a JSON-serializable dictionary.

The root node passed to this function should always be an instance of a core Ax class or a JSON-compatible python builtin. The sub-fields of the input will then be recursively passed to this function.

e.g. if we pass an instance of Experiment, we will first fall through to the line object_dict = ENCODER_REGISTRY[_type](object), which will convert the Experiment to a (shallow) dictionary, where search subfield remains “unconverted”, i.e.: {“name”: <name: string>, “search_space”: <search space: SearchSpace>}. We then pass each item of the dictionary back into this function to recursively convert the entire object.

ax.storage.json_store.encoders module

ax.storage.json_store.encoders.arm_to_dict(arm: ax.core.arm.Arm) → Dict[str, Any][source]

Convert Ax arm to a dictionary.

ax.storage.json_store.encoders.batch_to_dict(batch: ax.core.batch_trial.BatchTrial) → Dict[str, Any][source]

Convert Ax batch to a dictionary.

ax.storage.json_store.encoders.benchmark_problem_to_dict(benchmark_problem: ax.benchmark.benchmark_problem.BenchmarkProblem) → Dict[str, Any][source]

Converts an Ax benchmark problem to a serializable dictionary.

ax.storage.json_store.encoders.botorch_model_to_dict(model: ax.models.torch.botorch_modular.model.BoTorchModel) → Dict[str, Any][source]

Convert Ax model to a dictionary.

ax.storage.json_store.encoders.botorch_modular_to_dict(class_type: Type[Any]) → Dict[str, Any][source]

Convert any class to a dictionary.

ax.storage.json_store.encoders.choice_parameter_to_dict(parameter: ax.core.parameter.ChoiceParameter) → Dict[str, Any][source]

Convert Ax choice parameter to a dictionary.

ax.storage.json_store.encoders.data_to_dict(data: ax.core.data.Data) → Dict[str, Any][source]

Convert Ax data to a dictionary.

ax.storage.json_store.encoders.experiment_to_dict(experiment: ax.core.experiment.Experiment) → Dict[str, Any][source]

Convert Ax experiment to a dictionary.

ax.storage.json_store.encoders.fixed_parameter_to_dict(parameter: ax.core.parameter.FixedParameter) → Dict[str, Any][source]

Convert Ax fixed parameter to a dictionary.

ax.storage.json_store.encoders.generation_step_to_dict(generation_step: ax.modelbridge.generation_strategy.GenerationStep) → Dict[str, Any][source]

Converts Ax generation step to a dictionary.

ax.storage.json_store.encoders.generation_strategy_to_dict(generation_strategy: ax.modelbridge.generation_strategy.GenerationStrategy) → Dict[str, Any][source]

Converts Ax generation strategy to a dictionary.

ax.storage.json_store.encoders.generator_run_to_dict(generator_run: ax.core.generator_run.GeneratorRun) → Dict[str, Any][source]

Convert Ax generator run to a dictionary.

ax.storage.json_store.encoders.metric_to_dict(metric: ax.core.metric.Metric) → Dict[str, Any][source]

Convert Ax metric to a dictionary.

ax.storage.json_store.encoders.multi_objective_optimization_config_to_dict(multi_objective_optimization_config: ax.core.optimization_config.MultiObjectiveOptimizationConfig) → Dict[str, Any][source]

Convert Ax optimization config to a dictionary.

ax.storage.json_store.encoders.multi_objective_to_dict(objective: ax.core.objective.MultiObjective) → Dict[str, Any][source]

Convert Ax objective to a dictionary.

ax.storage.json_store.encoders.multi_type_experiment_to_dict(experiment: ax.core.multi_type_experiment.MultiTypeExperiment) → Dict[str, Any][source]

Convert AE multitype experiment to a dictionary.

ax.storage.json_store.encoders.objective_to_dict(objective: ax.core.objective.Objective) → Dict[str, Any][source]

Convert Ax objective to a dictionary.

ax.storage.json_store.encoders.observation_features_to_dict(obs_features: ax.core.observation.ObservationFeatures) → Dict[str, Any][source]

Converts Ax observation features to a dictionary

ax.storage.json_store.encoders.optimization_config_to_dict(optimization_config: ax.core.optimization_config.OptimizationConfig) → Dict[str, Any][source]

Convert Ax optimization config to a dictionary.

ax.storage.json_store.encoders.order_parameter_constraint_to_dict(parameter_constraint: ax.core.parameter_constraint.OrderConstraint) → Dict[str, Any][source]

Convert Ax order parameter constraint to a dictionary.

ax.storage.json_store.encoders.outcome_constraint_to_dict(outcome_constraint: ax.core.outcome_constraint.OutcomeConstraint) → Dict[str, Any][source]

Convert Ax outcome constraint to a dictionary.

ax.storage.json_store.encoders.parameter_constraint_to_dict(parameter_constraint: ax.core.parameter_constraint.ParameterConstraint) → Dict[str, Any][source]

Convert Ax sum parameter constraint to a dictionary.

ax.storage.json_store.encoders.range_parameter_to_dict(parameter: ax.core.parameter.RangeParameter) → Dict[str, Any][source]

Convert Ax range parameter to a dictionary.

ax.storage.json_store.encoders.runner_to_dict(runner: ax.runners.synthetic.SyntheticRunner) → Dict[str, Any][source]

Convert Ax synthetic runner to a dictionary.

ax.storage.json_store.encoders.scalarized_objective_to_dict(objective: ax.core.objective.ScalarizedObjective) → Dict[str, Any][source]

Convert Ax objective to a dictionary.

ax.storage.json_store.encoders.search_space_to_dict(search_space: ax.core.search_space.SearchSpace) → Dict[str, Any][source]

Convert Ax search space to a dictionary.

ax.storage.json_store.encoders.simple_experiment_to_dict(experiment: ax.core.simple_experiment.SimpleExperiment) → Dict[str, Any][source]

Convert AE simple experiment to a dictionary.

ax.storage.json_store.encoders.sum_parameter_constraint_to_dict(parameter_constraint: ax.core.parameter_constraint.SumConstraint) → Dict[str, Any][source]

Convert Ax sum parameter constraint to a dictionary.

ax.storage.json_store.encoders.surrogate_to_dict(surrogate: ax.models.torch.botorch_modular.surrogate.Surrogate) → Dict[str, Any][source]

Convert Ax surrogate to a dictionary.

ax.storage.json_store.encoders.transform_type_to_dict(transform_type: Type[ax.modelbridge.transforms.base.Transform]) → Dict[str, Any][source]

Convert a transform class to a dictionary.

ax.storage.json_store.encoders.trial_to_dict(trial: ax.core.trial.Trial) → Dict[str, Any][source]

Convert Ax trial to a dictionary.

ax.storage.json_store.load module

ax.storage.json_store.load.load_experiment(filepath: str) → ax.core.experiment.Experiment[source]

Load experiment from file.

  1. Read file.

  2. Convert dictionary to Ax experiment instance.

ax.storage.json_store.registry module

ax.storage.json_store.save module

ax.storage.json_store.save.save_experiment(experiment: ax.core.experiment.Experiment, filepath: str) → None[source]

Save experiment to file.

  1. Convert Ax experiment to JSON-serializable dictionary.

  2. Write to file.

SQLAlchemy (MySQL / SQLite)

ax.storage.sqa_store.base_decoder module

class ax.storage.sqa_store.decoder.Decoder(config: ax.storage.sqa_store.sqa_config.SQAConfig)[source]

Bases: object

Class that contains methods for loading an Ax experiment from SQLAlchemy.

Instantiate with an instance of Config to customize the functionality. For even more flexibility, create a subclass.

config

Metadata needed to save and load an experiment to SQLAlchemy.

abandoned_arm_from_sqa(abandoned_arm_sqa: ax.storage.sqa_store.sqa_classes.SQAAbandonedArm) → ax.core.batch_trial.AbandonedArm[source]

Convert SQLAlchemy AbandonedArm to Ax AbandonedArm.

arm_from_sqa(arm_sqa: ax.storage.sqa_store.sqa_classes.SQAArm) → ax.core.arm.Arm[source]

Convert SQLAlchemy Arm to Ax Arm.

data_from_sqa(data_sqa: ax.storage.sqa_store.sqa_classes.SQAData) → ax.core.data.Data[source]

Convert SQLAlchemy Data to AE Data.

experiment_from_sqa(experiment_sqa: ax.storage.sqa_store.sqa_classes.SQAExperiment, reduced_state: bool = False) → ax.core.experiment.Experiment[source]

Convert SQLAlchemy Experiment to Ax Experiment.

Parameters

  • experiment_sqa – SQAExperiment to decode.

  • reduced_state – Whether to load experiment with a slightly reduced state (without abandoned arms on experiment and without model state, search space, and optimization config on generator runs).

generation_strategy_from_sqa(gs_sqa: ax.storage.sqa_store.sqa_classes.SQAGenerationStrategy, experiment: Optional[ax.core.experiment.Experiment] = None, reduced_state: bool = False) → ax.modelbridge.generation_strategy.GenerationStrategy[source]

Convert SQALchemy generation strategy to Ax GenerationStrategy.

generator_run_from_sqa(generator_run_sqa: ax.storage.sqa_store.sqa_classes.SQAGeneratorRun, reduced_state: bool = False) → ax.core.generator_run.GeneratorRun[source]

Convert SQLAlchemy GeneratorRun to Ax GeneratorRun.

Parameters

  • generator_run_sqa – SQAGeneratorRun to decode.

  • reduced_state – Whether to load generator runs with a slightly reduced state

  • model state, search space, and optimization config) ((without) –

get_enum_name(value: Optional[int], enum: Optional[enum.Enum]) → Optional[str][source]

Given an enum value (int) and an enum (of ints), return the corresponding enum name. If the value is not present in the enum, throw an error.

metric_from_sqa(metric_sqa: ax.storage.sqa_store.sqa_classes.SQAMetric) → Union[ax.core.metric.Metric, ax.core.objective.Objective, ax.core.outcome_constraint.OutcomeConstraint][source]

Convert SQLAlchemy Metric to Ax Metric, Objective, or OutcomeConstraint.

metric_from_sqa_util(metric_sqa: ax.storage.sqa_store.sqa_classes.SQAMetric) → ax.core.metric.Metric[source]

Convert SQLAlchemy Metric to Ax Metric

opt_config_and_tracking_metrics_from_sqa(metrics_sqa: List[ax.storage.sqa_store.sqa_classes.SQAMetric]) → Tuple[Optional[ax.core.optimization_config.OptimizationConfig], List[ax.core.metric.Metric]][source]

Convert a list of SQLAlchemy Metrics to a a tuple of Ax OptimizationConfig and tracking metrics.

parameter_constraint_from_sqa(parameter_constraint_sqa: ax.storage.sqa_store.sqa_classes.SQAParameterConstraint, parameters: List[ax.core.parameter.Parameter]) → ax.core.parameter_constraint.ParameterConstraint[source]

Convert SQLAlchemy ParameterConstraint to Ax ParameterConstraint.

parameter_from_sqa(parameter_sqa: ax.storage.sqa_store.sqa_classes.SQAParameter) → ax.core.parameter.Parameter[source]

Convert SQLAlchemy Parameter to Ax Parameter.

runner_from_sqa(runner_sqa: ax.storage.sqa_store.sqa_classes.SQARunner) → ax.core.runner.Runner[source]

Convert SQLAlchemy Runner to Ax Runner.

search_space_from_sqa(parameters_sqa: List[ax.storage.sqa_store.sqa_classes.SQAParameter], parameter_constraints_sqa: List[ax.storage.sqa_store.sqa_classes.SQAParameterConstraint]) → Optional[ax.core.search_space.SearchSpace][source]

Convert a list of SQLAlchemy Parameters and ParameterConstraints to an Ax SearchSpace.

trial_from_sqa(trial_sqa: ax.storage.sqa_store.sqa_classes.SQATrial, experiment: ax.core.experiment.Experiment, reduced_state: bool = False) → ax.core.base_trial.BaseTrial[source]

Convert SQLAlchemy Trial to Ax Trial.

Parameters

  • trial_sqa – SQATrial to decode.

  • reduced_state – Whether to load trial’s generator run(s) with a slightly

  • state (reduced) –

ax.storage.sqa_store.base_encoder module

class ax.storage.sqa_store.encoder.Encoder(config: ax.storage.sqa_store.sqa_config.SQAConfig)[source]

Bases: object

Class that contains methods for storing an Ax experiment to SQLAlchemy.

Instantiate with an instance of Config to customize the functionality. For even more flexibility, create a subclass.

config

Metadata needed to save and load an experiment to SQLAlchemy.

abandoned_arm_to_sqa(abandoned_arm: ax.core.batch_trial.AbandonedArm) → ax.storage.sqa_store.sqa_classes.SQAAbandonedArm[source]

Convert Ax AbandonedArm to SQLAlchemy.

arm_to_sqa(arm: ax.core.arm.Arm, weight: Optional[float] = 1.0) → ax.storage.sqa_store.sqa_classes.SQAArm[source]

Convert Ax Arm to SQLAlchemy.

data_to_sqa(data: ax.core.data.Data, trial_index: Optional[int], timestamp: int) → ax.storage.sqa_store.sqa_classes.SQAData[source]

Convert AE data to SQLAlchemy.

experiment_to_sqa(experiment: ax.core.experiment.Experiment) → Tuple[ax.storage.sqa_store.sqa_classes.SQAExperiment, List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax Experiment to SQLAlchemy and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

In addition to creating and storing a new Experiment object, we need to create and store copies of the Trials, Metrics, Parameters, ParameterConstraints, and Runner owned by this Experiment.

generation_strategy_to_sqa(generation_strategy: ax.modelbridge.generation_strategy.GenerationStrategy, experiment_id: Optional[int]) → Tuple[ax.storage.sqa_store.sqa_classes.SQAGenerationStrategy, List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert an Ax GenerationStrategy to SQLAlchemy, preserving its state, so that the restored generation strategy can be resumed from the point at which it was interrupted and stored.

generator_run_to_sqa(generator_run: ax.core.generator_run.GeneratorRun, weight: Optional[float] = None) → Tuple[ax.storage.sqa_store.sqa_classes.SQAGeneratorRun, List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax GeneratorRun to SQLAlchemy and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

In addition to creating and storing a new GeneratorRun object, we need to create and store copies of the Arms, Metrics, Parameters, and ParameterConstraints owned by this GeneratorRun.

get_enum_value(value: Optional[str], enum: Optional[enum.Enum]) → Optional[int][source]

Given an enum name (string) and an enum (of ints), return the corresponding enum value. If the name is not present in the enum, throw an error.

get_metric_type_and_properties(metric: ax.core.metric.Metric) → Tuple[int, Dict[str, Any]][source]

Given an Ax Metric, convert its type into a member of MetricType enum, and construct a dictionary to be stored in the database properties json blob.

metric_to_sqa(metric: ax.core.metric.Metric) → ax.storage.sqa_store.sqa_classes.SQAMetric[source]

Convert Ax Metric to SQLAlchemy.

multi_objective_to_sqa(objective: ax.core.objective.MultiObjective) → Tuple[ax.storage.sqa_store.sqa_classes.SQAMetric, List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax Multi Objective to SQLAlchemy and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

Returns: A parent SQAMetric, whose children are the SQAMetric-s

corresponding to metrics attribute of MultiObjective. NOTE: The parent is used as a placeholder for storage purposes.

objective_threshold_to_sqa(objective_threshold: ax.core.outcome_constraint.ObjectiveThreshold) → ax.storage.sqa_store.sqa_classes.SQAMetric[source]

Convert Ax OutcomeConstraint to SQLAlchemy.

objective_to_sqa(objective: ax.core.objective.Objective) → Tuple[ax.storage.sqa_store.sqa_classes.SQAMetric, List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax Objective to SQLAlchemy and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

optimization_config_to_sqa(optimization_config: Optional[ax.core.optimization_config.OptimizationConfig]) → Tuple[List[ax.storage.sqa_store.sqa_classes.SQAMetric], List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax OptimizationConfig to a list of SQLAlchemy Metrics and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

outcome_constraint_to_sqa(outcome_constraint: ax.core.outcome_constraint.OutcomeConstraint) → ax.storage.sqa_store.sqa_classes.SQAMetric[source]

Convert Ax OutcomeConstraint to SQLAlchemy.

parameter_constraint_to_sqa(parameter_constraint: ax.core.parameter_constraint.ParameterConstraint) → ax.storage.sqa_store.sqa_classes.SQAParameterConstraint[source]

Convert Ax ParameterConstraint to SQLAlchemy.

parameter_to_sqa(parameter: ax.core.parameter.Parameter) → ax.storage.sqa_store.sqa_classes.SQAParameter[source]

Convert Ax Parameter to SQLAlchemy.

runner_to_sqa(runner: ax.core.runner.Runner, trial_type: Optional[str] = None) → ax.storage.sqa_store.sqa_classes.SQARunner[source]

Convert Ax Runner to SQLAlchemy.

scalarized_objective_to_sqa(objective: ax.core.objective.ScalarizedObjective) → Tuple[ax.storage.sqa_store.sqa_classes.SQAMetric, List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax Scalarized Objective to SQLAlchemy and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

Returns: A parent SQAMetric, whose children are the SQAMetric-s

corresponding to metrics attribute of ScalarizedObjective. NOTE: The parent is used as a placeholder for storage purposes.

search_space_to_sqa(search_space: Optional[ax.core.search_space.SearchSpace]) → Tuple[List[ax.storage.sqa_store.sqa_classes.SQAParameter], List[ax.storage.sqa_store.sqa_classes.SQAParameterConstraint], List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax SearchSpace to a list of SQLAlchemy Parameters and ParameterConstraints and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

trial_to_sqa(trial: ax.core.base_trial.BaseTrial) → Tuple[ax.storage.sqa_store.sqa_classes.SQATrial, List[Tuple[ax.utils.common.base.Base, ax.storage.sqa_store.db.SQABase]]][source]

Convert Ax Trial to SQLAlchemy and compile a list of (object, sqa_counterpart) tuples to set db_id on user-facing classes after the conversion is complete and the SQL session is flushed (SQLAlchemy classes receive their id attributes during session.flush()).

In addition to creating and storing a new Trial object, we need to create and store the GeneratorRuns and Runner that it owns.

classmethod validate_experiment_metadata(experiment: ax.core.experiment.Experiment, existing_sqa_experiment: Optional[ax.storage.sqa_store.sqa_classes.SQAExperiment], owners: Optional[List[str]] = None) → None[source]

Validates required experiment metadata.

Does not expect owners kwarg, present for use in subclasses.

ax.storage.sqa_store.db module

class ax.storage.sqa_store.db.SQABase[source]

Bases: object

Metaclass for SQLAlchemy classes corresponding to core Ax classes.

property attributes

Return a list of the column attributes and relationship fields on this SQABase instance. Used for iterating over the fields to determine equality, perform updates, etc.

equals(other)
fields_equal(other: ax.storage.sqa_store.db.SQABase, field: str) → bool[source]

Check if field on self is equal to field on other.

static list_equals(l1: List[T], l2: List[T]) → bool[source]

Compare equality of two lists.

Assumptions:

– The lists do not contain duplicates

Checking equality is then the same as checking that the lists are the same length, and that one is a subset of the other.

static list_update(l1: List[T], l2: List[T]) → List[T][source]
Given an existing list (l1) and an new version (l2):

– update the existing items in l1 that have matching items in l2 – delete existing items in l1 that don’t have matching items in l2 – add items in l2 that don’t exist in l1

e.g. list_update([1,2,3], [1,5]) => [1,5]
list_update([Arm(name=”0_0”)], [Arm(name=”0_0”), Arm(name=”0_1”)]) =>

[Arm(name=”0_0”), Arm(name=”0_1”)] where Arm(name=”0_0”) has been updated, not replaced, so that we don’t delete/recreate the DB row

static list_update_without_unique_id(l1: List[T], l2: List[T]) → List[T][source]

Merge a new list (l2) into an existing list (l1) This method works for lists whose items do not have a unique_id field. If the lists are equal, return the old one. Else, return the new one.

update(other: ax.storage.sqa_store.db.SQABase) → None[source]

Merge other into self.

update_field(other: ax.storage.sqa_store.db.SQABase, field: str) → None[source]

Update field on self to be equal to field on other.

ax.storage.sqa_store.db.create_all_tables(engine: sqlalchemy.engine.base.Engine) → None[source]

Create all tables that inherit from Base.

Parameters

engine – a SQLAlchemy engine with a connection to a MySQL or SQLite DB.

Note

In order for all tables to be correctly created, all modules that define a mapped class that inherits from Base must be imported.

ax.storage.sqa_store.db.create_mysql_engine_from_creator(creator: Callable, echo: bool = False, pool_recycle: int = 10, **kwargs: Any) → sqlalchemy.engine.base.Engine[source]

Create a SQLAlchemy engine with the MySQL dialect given a creator function.

Parameters

  • creator – a callable which returns a DBAPI connection.

  • echo – if True, set engine to be verbose.

  • pool_recycle – number of seconds after which to recycle connections. -1 means no timeout. Default is 10 seconds.

  • **kwargs – keyword args passed to create_engine

Returns

SQLAlchemy engine with connection to MySQL DB.

Return type

Engine

ax.storage.sqa_store.db.create_mysql_engine_from_url(url: str, echo: bool = False, pool_recycle: int = 10, **kwargs: Any) → sqlalchemy.engine.base.Engine[source]

Create a SQLAlchemy engine with the MySQL dialect given a database url.

Parameters

  • url – a database url that can include username, password, hostname, database name as well as optional keyword arguments for additional configuration. e.g. dialect+driver://username:password@host:port/database.

  • echo – if True, set engine to be verbose.

  • pool_recycle – number of seconds after which to recycle connections. -1 means no timeout. Default is 10 seconds.

  • **kwargs – keyword args passed to create_engine

Returns

SQLAlchemy engine with connection to MySQL DB.

Return type

Engine

ax.storage.sqa_store.db.create_test_engine(path: Optional[str] = None, echo: bool = True) → sqlalchemy.engine.base.Engine[source]

Creates a SQLAlchemy engine object for use in unit tests.

Parameters

  • path – if None, use in-memory SQLite; else attempt to create a SQLite DB in the path provided.

  • echo – if True, set engine to be verbose.

Returns

an instance of SQLAlchemy engine.

Return type

Engine

ax.storage.sqa_store.db.get_engine() → sqlalchemy.engine.base.Engine[source]

Fetch a SQLAlchemy engine, if already initialized.

If not initialized, need to either call init_engine_and_session_factory or get_session explicitly.

Returns

an instance of a SQLAlchemy engine with a connection to a DB.

Return type

Engine

ax.storage.sqa_store.db.get_session() → sqlalchemy.orm.session.Session[source]

Fetch a SQLAlchemy session with a connection to a DB.

Unless init_engine_and_session_factory is called first with custom args, this will automatically initialize a connection to xdb.adaptive_experiment.

Returns

an instance of a SQLAlchemy session.

Return type

Session

ax.storage.sqa_store.db.init_engine_and_session_factory(url: Optional[str] = None, creator: Optional[Callable] = None, echo: bool = False, force_init: bool = False, **kwargs: Any) → None[source]

Initialize the global engine and SESSION_FACTORY for SQLAlchemy.

The initialization needs to only happen once. Note that it is possible to re-initialize the engine by setting the force_init flag to True, but this should only be used if you are absolutely certain that you know what you are doing.

Parameters

  • url – a database url that can include username, password, hostname, database name as well as optional keyword arguments for additional configuration. e.g. dialect+driver://username:password@host:port/database. Either this argument or creator argument must be specified.

  • creator – a callable which returns a DBAPI connection. Either this argument or url argument must be specified.

  • echo – if True, logging for engine is enabled.

  • force_init – if True, allows re-initializing engine and session factory.

  • **kwargs – keyword arguments passed to create_mysql_engine_from_creator

ax.storage.sqa_store.db.init_test_engine_and_session_factory(tier_or_path: Optional[str] = None, echo: bool = False, force_init: bool = False, **kwargs: Any) → None[source]

Initialize the global engine and SESSION_FACTORY for SQLAlchemy, using an in-memory SQLite database.

The initialization needs to only happen once. Note that it is possible to re-initialize the engine by setting the force_init flag to True, but this should only be used if you are absolutely certain that you know what you are doing.

Parameters

  • tier_or_path – the name of the DB tier.

  • echo – if True, logging for engine is enabled.

  • force_init – if True, allows re-initializing engine and session factory.

  • **kwargs – keyword arguments passed to create_mysql_engine_from_creator

ax.storage.sqa_store.db.optional_session_scope(session: Optional[sqlalchemy.orm.session.Session] = None) → AbstractContextManager[sqlalchemy.orm.session.Session][source]
ax.storage.sqa_store.db.session_scope() → Generator[Session, None, None][source]

Provide a transactional scope around a series of operations.

ax.storage.sqa_store.json module

class ax.storage.sqa_store.json.JSONEncodedMediumText(object_pairs_hook: Any = None, *args: List[Any], **kwargs: Dict[Any, Any])[source]

Bases: ax.storage.sqa_store.json.JSONEncodedObject

Class for JSON-encoding objects in SQLAlchemy, backed by MEDIUMTEXT (MySQL).

See description in JSONEncodedObject.

impl: sqlalchemy.sql.sqltypes.VARCHAR = Text(length=16777215)
class ax.storage.sqa_store.json.JSONEncodedObject(object_pairs_hook: Any = None, *args: List[Any], **kwargs: Dict[Any, Any])[source]

Bases: sqlalchemy.sql.type_api.TypeDecorator

Class for JSON-encoding objects in SQLAlchemy.

Represents an object that is automatically marshalled and unmarshalled to/from the corresponding JSON string. By itself, this data structure does not track any changes.

impl: sqlalchemy.sql.sqltypes.VARCHAR = VARCHAR(length=4096)
process_bind_param(value: Any, dialect: Any) → Optional[str][source]

Receive a bound parameter value to be converted.

Subclasses override this method to return the value that should be passed along to the underlying TypeEngine object, and from there to the DBAPI execute() method.

The operation could be anything desired to perform custom behavior, such as transforming or serializing data. This could also be used as a hook for validating logic.

This operation should be designed with the reverse operation in mind, which would be the process_result_value method of this class.

Parameters

  • value – Data to operate upon, of any type expected by this method in the subclass. Can be None.

  • dialect – the Dialect in use.

process_result_value(value: Any, dialect: Any) → Any[source]

Receive a result-row column value to be converted.

Subclasses should implement this method to operate on data fetched from the database.

Subclasses override this method to return the value that should be passed back to the application, given a value that is already processed by the underlying TypeEngine object, originally from the DBAPI cursor method fetchone() or similar.

The operation could be anything desired to perform custom behavior, such as transforming or serializing data. This could also be used as a hook for validating logic.

Parameters

  • value – Data to operate upon, of any type expected by this method in the subclass. Can be None.

  • dialect – the Dialect in use.

This operation should be designed to be reversible by the “process_bind_param” method of this class.

class ax.storage.sqa_store.json.JSONEncodedText(object_pairs_hook: Any = None, *args: List[Any], **kwargs: Dict[Any, Any])[source]

Bases: ax.storage.sqa_store.json.JSONEncodedObject

Class for JSON-encoding objects in SQLAlchemy, backed by TEXT (MySQL).

See description in JSONEncodedObject.

impl

alias of sqlalchemy.sql.sqltypes.Text

ax.storage.sqa_store.load module

ax.storage.sqa_store.load.load_experiment(experiment_name: str, config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None, reduced_state: bool = False) → ax.core.experiment.Experiment[source]

Load experiment by name.

Parameters

  • experiment_name – Name of the expeirment to load.

  • config – SQAConfig, from which to retrieve the decoder. Optional, defaults to base SQAConfig.

  • reduced_state – Whether to load experiment with a slightly reduced state (without abandoned arms on experiment and withoug model state, search space, and optimization config on generator runs).

ax.storage.sqa_store.load.load_generation_strategy_by_experiment_name(experiment_name: str, config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None, experiment: Optional[ax.core.experiment.Experiment] = None, reduced_state: bool = False) → ax.modelbridge.generation_strategy.GenerationStrategy[source]

Finds a generation strategy attached to an experiment specified by a name and restores it from its corresponding SQA object.

ax.storage.sqa_store.load.load_generation_strategy_by_id(gs_id: int, config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None, experiment: Optional[ax.core.experiment.Experiment] = None, reduced_state: bool = False) → ax.modelbridge.generation_strategy.GenerationStrategy[source]

Finds a generation strategy stored by a given ID and restores it.

ax.storage.sqa_store.save module

ax.storage.sqa_store.save.save_experiment(experiment: ax.core.experiment.Experiment, config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None) → None[source]

Save experiment (using default SQAConfig).

ax.storage.sqa_store.save.save_generation_strategy(generation_strategy: ax.modelbridge.generation_strategy.GenerationStrategy, config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None) → int[source]

Save generation strategy (using default SQAConfig if no config is specified). If the generation strategy has an experiment set, the experiment will be saved first.

Returns

The ID of the saved generation strategy.

ax.storage.sqa_store.save.save_new_trial(experiment: ax.core.experiment.Experiment, trial: ax.core.base_trial.BaseTrial, config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None) → None[source]

Add new trial to the experiment (using default SQAConfig).

ax.storage.sqa_store.save.update_generation_strategy(generation_strategy: ax.modelbridge.generation_strategy.GenerationStrategy, generator_runs: List[ax.core.generator_run.GeneratorRun], config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None) → None[source]

Update generation strategy’s current step and attach generator runs (using default SQAConfig).

ax.storage.sqa_store.save.update_trial(experiment: ax.core.experiment.Experiment, trial: ax.core.base_trial.BaseTrial, config: Optional[ax.storage.sqa_store.sqa_config.SQAConfig] = None) → None[source]

Update trial and attach data (using default SQAConfig).

ax.storage.sqa_store.structs module

class ax.storage.sqa_store.structs.DBSettings[source]

Bases: tuple

Defines behavior for loading/saving experiment to/from db. Either creator or url must be specified as a way to connect to the SQL db.

property creator

Alias for field number 0

property decoder

Alias for field number 1

property encoder

Alias for field number 2

property url

Alias for field number 3

ax.storage.sqa_store.sqa_classes module

class ax.storage.sqa_store.sqa_classes.SQAAbandonedArm(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

abandoned_reason: Optional[str] = Column(None, String(length=255), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
immutable_fields = ['name']
name: str = Column(None, String(length=100), table=None, nullable=False)
time_abandoned: datetime.datetime = Column(None, IntTimestamp(), table=None, nullable=False, default=ColumnDefault(<function datetime.now>))
trial_id: int = Column(None, Integer(), ForeignKey('trial_v2.id'), table=None)
unique_id = 'name'
class ax.storage.sqa_store.sqa_classes.SQAArm(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

generator_run_id: int = Column(None, Integer(), ForeignKey('generator_run_v2.id'), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
immutable_fields = ['parameters']
name: Optional[str] = Column(None, String(length=100), table=None)
parameters: Dict[str, Union[str, bool, float, int, None]] = Column(None, JSONEncodedText(), table=None, nullable=False)
unique_id = 'name'
weight: float = Column(None, Float(), table=None, nullable=False, default=ColumnDefault(1.0))
class ax.storage.sqa_store.sqa_classes.SQAData(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

data_json: str = Column(None, Text(length=4294967295), table=None, nullable=False)
description: Optional[str] = Column(None, String(length=255), table=None)
experiment_id: int = Column(None, Integer(), ForeignKey('experiment_v2.id'), table=None)
generation_strategy_id: Optional[int] = Column(None, Integer(), ForeignKey('generation_strategy.id'), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
time_created: int = Column(None, BigInteger(), table=None, nullable=False)
trial_index: Optional[int] = Column(None, Integer(), table=None)
unique_id = 'time_created'
class ax.storage.sqa_store.sqa_classes.SQAExperiment(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

data: List[ax.storage.sqa_store.sqa_classes.SQAData] = <RelationshipProperty at 0x7fe9eb6fdac8; no key>
default_trial_type: Optional[str] = Column(None, String(length=100), table=None)
description: Optional[str] = Column(None, String(length=255), table=None)
experiment_type: Optional[int] = Column(None, Integer(), table=None)
generation_strategy: Optional[ax.storage.sqa_store.sqa_classes.SQAGenerationStrategy] = <RelationshipProperty at 0x7fe9eb6fddc8; no key>
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
ignore_during_update_fields = ['time_created']
immutable_fields = ['name']
is_test: bool = Column(None, Boolean(), table=None, nullable=False, default=ColumnDefault(False))
metrics: List[ax.storage.sqa_store.sqa_classes.SQAMetric] = <RelationshipProperty at 0x7fe9eb6fdb48; no key>
name: str = Column(None, String(length=100), table=None, nullable=False)
parameter_constraints: List[ax.storage.sqa_store.sqa_classes.SQAParameterConstraint] = <RelationshipProperty at 0x7fe9eb6fdc48; no key>
parameters: List[ax.storage.sqa_store.sqa_classes.SQAParameter] = <RelationshipProperty at 0x7fe9eb6fdbc8; no key>
properties: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None, default=ColumnDefault({}))
runners: List[ax.storage.sqa_store.sqa_classes.SQARunner] = <RelationshipProperty at 0x7fe9eb6fdcc8; no key>
status_quo_name: Optional[str] = Column(None, String(length=100), table=None)
status_quo_parameters: Optional[Dict[str, Union[str, bool, float, int, None]]] = Column(None, JSONEncodedText(), table=None)
time_created: datetime.datetime = Column(None, IntTimestamp(), table=None, nullable=False)
trials: List[ax.storage.sqa_store.sqa_classes.SQATrial] = <RelationshipProperty at 0x7fe9eb6fdd48; no key>
class ax.storage.sqa_store.sqa_classes.SQAGenerationStrategy(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

curr_index: int = Column(None, Integer(), table=None, nullable=False)
experiment_id: Optional[int] = Column(None, Integer(), ForeignKey('experiment_v2.id'), table=None)
generator_runs: List[ax.storage.sqa_store.sqa_classes.SQAGeneratorRun] = <RelationshipProperty at 0x7fe9eb6e8dc8; no key>
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
name: str = Column(None, String(length=100), table=None, nullable=False)
steps: List[Dict[str, Any]] = Column(None, JSONEncodedObject(length=4096), table=None, nullable=False)
class ax.storage.sqa_store.sqa_classes.SQAGeneratorRun(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

arms: List[ax.storage.sqa_store.sqa_classes.SQAArm] = <RelationshipProperty at 0x7fe9eb6e8648; no key>
best_arm_name: Optional[str] = Column(None, String(length=100), table=None)
best_arm_parameters: Optional[Dict[str, Union[str, bool, float, int, None]]] = Column(None, JSONEncodedText(), table=None)
best_arm_predictions: Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]] = Column(None, JSONEncodedObject(length=4096), table=None)
bridge_kwargs: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None)
candidate_metadata_by_arm_signature: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None)
fit_time: Optional[float] = Column(None, Float(), table=None)
gen_metadata: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None)
gen_time: Optional[float] = Column(None, Float(), table=None)
generation_step_index: Optional[int] = Column(None, Integer(), table=None)
generation_strategy_id: Optional[int] = Column(None, Integer(), ForeignKey('generation_strategy.id'), table=None)
generator_run_type: Optional[int] = Column(None, Integer(), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
ignore_during_update_fields = ['time_created']
index: Optional[int] = Column(None, Integer(), table=None)
metrics: List[ax.storage.sqa_store.sqa_classes.SQAMetric] = <RelationshipProperty at 0x7fe9eb6e86c8; no key>
model_key: Optional[str] = Column(None, String(length=100), table=None)
model_kwargs: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None)
model_predictions: Optional[Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]]] = Column(None, JSONEncodedObject(length=4096), table=None)
model_state_after_gen: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None)
parameter_constraints: List[ax.storage.sqa_store.sqa_classes.SQAParameterConstraint] = <RelationshipProperty at 0x7fe9eb6e87c8; no key>
parameters: List[ax.storage.sqa_store.sqa_classes.SQAParameter] = <RelationshipProperty at 0x7fe9eb6e8748; no key>
time_created: datetime.datetime = Column(None, IntTimestamp(), table=None, nullable=False, default=ColumnDefault(<function datetime.now>))
trial_id: Optional[int] = Column(None, Integer(), ForeignKey('trial_v2.id'), table=None)
unique_id = 'index'
weight: Optional[float] = Column(None, Float(), table=None)
class ax.storage.sqa_store.sqa_classes.SQAMetric(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

bound: Optional[float] = Column(None, Float(), table=None)
canonical_name: Optional[str] = Column(None, String(length=100), table=None)
experiment_id: Optional[int] = Column(None, Integer(), ForeignKey('experiment_v2.id'), table=None)
generator_run_id: Optional[int] = Column(None, Integer(), ForeignKey('generator_run_v2.id'), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
immutable_fields = ['name']
intent: ax.storage.utils.MetricIntent = Column(None, StringEnum(length=100), table=None, nullable=False)
lower_is_better: Optional[bool] = Column(None, Boolean(), table=None)
metric_type: int = Column(None, Integer(), table=None, nullable=False)
minimize: Optional[bool] = Column(None, Boolean(), table=None)
name: str = Column(None, String(length=255), table=None, nullable=False)
op: Optional[ax.core.types.ComparisonOp] = Column(None, IntEnum(), table=None)
properties: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None, default=ColumnDefault({}))
relative: Optional[bool] = Column(None, Boolean(), table=None)
scalarized_objective_children_metrics = <RelationshipProperty at 0x7fe9eb7d7ec8; no key>
scalarized_objective_id = Column(None, Integer(), ForeignKey('metric_v2.id'), table=None)
scalarized_objective_weight: Optional[float] = Column(None, Float(), table=None)
trial_type: Optional[str] = Column(None, String(length=100), table=None)
unique_id = 'name'
class ax.storage.sqa_store.sqa_classes.SQAParameter(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

choice_values: Optional[List[Union[str, bool, float, int, None]]] = Column(None, JSONEncodedObject(length=4096), table=None)
digits: Optional[int] = Column(None, Integer(), table=None)
domain_type: ax.storage.utils.DomainType = Column(None, IntEnum(), table=None, nullable=False)
experiment_id: Optional[int] = Column(None, Integer(), ForeignKey('experiment_v2.id'), table=None)
fixed_value: Union[str, bool, float, int, None] = Column(None, JSONEncodedObject(length=4096), table=None)
generator_run_id: Optional[int] = Column(None, Integer(), ForeignKey('generator_run_v2.id'), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
immutable_fields = ['name']
is_fidelity: Optional[bool] = Column(None, Boolean(), table=None)
is_ordered: Optional[bool] = Column(None, Boolean(), table=None)
is_task: Optional[bool] = Column(None, Boolean(), table=None)
log_scale: Optional[bool] = Column(None, Boolean(), table=None)
lower: Optional[float] = Column(None, Float(), table=None)
name: str = Column(None, String(length=100), table=None, nullable=False)
parameter_type: ax.core.parameter.ParameterType = Column(None, IntEnum(), table=None, nullable=False)
target_value: Union[str, bool, float, int, None] = Column(None, JSONEncodedObject(length=4096), table=None)
unique_id = 'name'
upper: Optional[float] = Column(None, Float(), table=None)
class ax.storage.sqa_store.sqa_classes.SQAParameterConstraint(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

bound: float = Column(None, Float(), table=None, nullable=False)
constraint_dict: Dict[str, float] = Column(None, JSONEncodedObject(length=4096), table=None, nullable=False)
experiment_id: Optional[int] = Column(None, Integer(), ForeignKey('experiment_v2.id'), table=None)
generator_run_id: Optional[int] = Column(None, Integer(), ForeignKey('generator_run_v2.id'), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
immutable_fields = ['type', 'constraint_dict', 'bound']
type: ax.storage.sqa_store.sqa_enum.IntEnum = Column(None, IntEnum(), table=None, nullable=False)
class ax.storage.sqa_store.sqa_classes.SQARunner(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

experiment_id: Optional[int] = Column(None, Integer(), ForeignKey('experiment_v2.id'), table=None)
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
properties: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None, default=ColumnDefault({}))
runner_type: int = Column(None, Integer(), table=None, nullable=False)
trial_id: Optional[int] = Column(None, Integer(), ForeignKey('trial_v2.id'), table=None)
trial_type: Optional[str] = Column(None, String(length=100), table=None)
class ax.storage.sqa_store.sqa_classes.SQATrial(*args: Any, **kwargs: Any)[source]

Bases: sqlalchemy.ext.declarative.declarative_base

abandoned_arms: List[ax.storage.sqa_store.sqa_classes.SQAAbandonedArm] = <RelationshipProperty at 0x7fe9eb6fd6c8; no key>
abandoned_reason: Optional[str] = Column(None, String(length=100), table=None)
deployed_name: Optional[str] = Column(None, String(length=100), table=None)
experiment_id: int = Column(None, Integer(), ForeignKey('experiment_v2.id'), table=None)
generation_step_index: Optional[int] = Column(None, Integer(), table=None)
generator_runs: List[ax.storage.sqa_store.sqa_classes.SQAGeneratorRun] = <RelationshipProperty at 0x7fe9eb6fd748; no key>
id: int = Column(None, Integer(), table=None, primary_key=True, nullable=False)
ignore_during_update_fields = ['time_created']
immutable_fields = ['is_batch']
index: int = Column(None, Integer(), table=None, nullable=False)
is_batch: bool = Column('is_batched', Boolean(), table=None, nullable=False, default=ColumnDefault(True))
num_arms_created: int = Column(None, Integer(), table=None, nullable=False, default=ColumnDefault(0))
optimize_for_power: Optional[bool] = Column(None, Boolean(), table=None)
properties: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None, default=ColumnDefault({}))
run_metadata: Optional[Dict[str, Any]] = Column(None, JSONEncodedText(), table=None)
runner: ax.storage.sqa_store.sqa_classes.SQARunner = <RelationshipProperty at 0x7fe9eb6fd7c8; no key>
status: ax.core.base_trial.TrialStatus = Column(None, IntEnum(), table=None, nullable=False, default=ColumnDefault(<TrialStatus.CANDIDATE: 0>))
status_quo_name: Optional[str] = Column(None, String(length=100), table=None)
time_completed: Optional[datetime.datetime] = Column(None, IntTimestamp(), table=None)
time_created: datetime.datetime = Column(None, IntTimestamp(), table=None, nullable=False)
time_run_started: Optional[datetime.datetime] = Column(None, IntTimestamp(), table=None)
time_staged: Optional[datetime.datetime] = Column(None, IntTimestamp(), table=None)
trial_type: Optional[str] = Column(None, String(length=100), table=None)
ttl_seconds: Optional[int] = Column(None, Integer(), table=None)
unique_id = 'index'

ax.storage.sqa_store.sqa_config module

class ax.storage.sqa_store.sqa_config.SQAConfig[source]

Bases: tuple

Metadata needed to save and load an experiment to SQLAlchemy.

class_to_sqa_class

Mapping of user-facing class to SQLAlchemy class that it will be encoded to. This allows overwriting of the default classes to provide custom save functionality.

experiment_type_enum

Enum containing valid Experiment types.

generator_run_type_enum

Enum containing valid Generator Run types.

property class_to_sqa_class

Alias for field number 0

property experiment_type_enum

Alias for field number 1

property generator_run_type_enum

Alias for field number 2

ax.storage.sqa_store.sqa_enum module

class ax.storage.sqa_store.sqa_enum.BaseNullableEnum(enum: Any, *arg: List[Any], **kw: Dict[Any, Any])[source]

Bases: sqlalchemy.sql.type_api.TypeDecorator

process_bind_param(value: Any, dialect: Any) → Any[source]

Receive a bound parameter value to be converted.

Subclasses override this method to return the value that should be passed along to the underlying TypeEngine object, and from there to the DBAPI execute() method.

The operation could be anything desired to perform custom behavior, such as transforming or serializing data. This could also be used as a hook for validating logic.

This operation should be designed with the reverse operation in mind, which would be the process_result_value method of this class.

Parameters

  • value – Data to operate upon, of any type expected by this method in the subclass. Can be None.

  • dialect – the Dialect in use.

process_result_value(value: Any, dialect: Any) → Any[source]

Receive a result-row column value to be converted.

Subclasses should implement this method to operate on data fetched from the database.

Subclasses override this method to return the value that should be passed back to the application, given a value that is already processed by the underlying TypeEngine object, originally from the DBAPI cursor method fetchone() or similar.

The operation could be anything desired to perform custom behavior, such as transforming or serializing data. This could also be used as a hook for validating logic.

Parameters

  • value – Data to operate upon, of any type expected by this method in the subclass. Can be None.

  • dialect – the Dialect in use.

This operation should be designed to be reversible by the “process_bind_param” method of this class.

class ax.storage.sqa_store.sqa_enum.IntEnum(enum: Any, *arg: List[Any], **kw: Dict[Any, Any])[source]

Bases: ax.storage.sqa_store.sqa_enum.BaseNullableEnum

impl

alias of sqlalchemy.sql.sqltypes.SmallInteger

class ax.storage.sqa_store.sqa_enum.StringEnum(enum: Any, *arg: List[Any], **kw: Dict[Any, Any])[source]

Bases: ax.storage.sqa_store.sqa_enum.BaseNullableEnum

impl = VARCHAR(length=100)

ax.storage.sqa_store.timestamp module

class ax.storage.sqa_store.timestamp.IntTimestamp(*args, **kwargs)[source]

Bases: sqlalchemy.sql.type_api.TypeDecorator

impl

alias of sqlalchemy.sql.sqltypes.Integer

process_bind_param(value: Optional[datetime.datetime], dialect: sqlalchemy.engine.interfaces.Dialect) → Optional[int][source]

Receive a bound parameter value to be converted.

Subclasses override this method to return the value that should be passed along to the underlying TypeEngine object, and from there to the DBAPI execute() method.

The operation could be anything desired to perform custom behavior, such as transforming or serializing data. This could also be used as a hook for validating logic.

This operation should be designed with the reverse operation in mind, which would be the process_result_value method of this class.

Parameters

  • value – Data to operate upon, of any type expected by this method in the subclass. Can be None.

  • dialect – the Dialect in use.

process_result_value(value: Optional[int], dialect: sqlalchemy.engine.interfaces.Dialect) → Optional[datetime.datetime][source]

Receive a result-row column value to be converted.

Subclasses should implement this method to operate on data fetched from the database.

Subclasses override this method to return the value that should be passed back to the application, given a value that is already processed by the underlying TypeEngine object, originally from the DBAPI cursor method fetchone() or similar.

The operation could be anything desired to perform custom behavior, such as transforming or serializing data. This could also be used as a hook for validating logic.

Parameters

  • value – Data to operate upon, of any type expected by this method in the subclass. Can be None.

  • dialect – the Dialect in use.

This operation should be designed to be reversible by the “process_bind_param” method of this class.

ax.storage.sqa_store.utils module

ax.storage.sqa_store.utils.is_foreign_key_field(field: str) → bool[source]

Return true if field name is a foreign key field, i.e. ends in _id.

ax.storage.sqa_store.validation module

ax.storage.sqa_store.validation.consistency_exactly_one(instance: ax.storage.sqa_store.db.SQABase, exactly_one_fields: List[str]) → Any[source]

Ensure that exactly one of exactly_one_fields has a value set.

ax.storage.sqa_store.validation.validate_metric(mapper: sqlalchemy.orm.mapper.Mapper, connection: sqlalchemy.engine.base.Connection, target: ax.storage.sqa_store.db.SQABase) → None[source]
ax.storage.sqa_store.validation.validate_parameter(mapper: sqlalchemy.orm.mapper.Mapper, connection: sqlalchemy.engine.base.Connection, target: ax.storage.sqa_store.db.SQABase) → None[source]
ax.storage.sqa_store.validation.validate_parameter_constraint(mapper: sqlalchemy.orm.mapper.Mapper, connection: sqlalchemy.engine.base.Connection, target: ax.storage.sqa_store.db.SQABase) → None[source]
ax.storage.sqa_store.validation.validate_runner(mapper: sqlalchemy.orm.mapper.Mapper, connection: sqlalchemy.engine.base.Connection, target: ax.storage.sqa_store.db.SQABase) → None[source]

Registries

ax.storage.botorch_modular_registry.ACQUISITION_FUNCTION_REGISTRY: Dict[Type[botorch.acquisition.acquisition.AcquisitionFunction], str] = {<class 'botorch.acquisition.monte_carlo.qExpectedImprovement'>: 'qExpectedImprovement', <class 'botorch.acquisition.knowledge_gradient.qKnowledgeGradient'>: 'qKnowledgeGradient', <class 'botorch.acquisition.max_value_entropy_search.qMaxValueEntropy'>: 'qMaxValueEntropy', <class 'botorch.acquisition.knowledge_gradient.qMultiFidelityKnowledgeGradient'>: 'qMultiFidelityKnowledgeGradient', <class 'botorch.acquisition.max_value_entropy_search.qMultiFidelityMaxValueEntropy'>: 'qMultiFidelityMaxValueEntropy', <class 'botorch.acquisition.monte_carlo.qNoisyExpectedImprovement'>: 'qNoisyExpectedImprovement'}

Mapping of BoTorch MarginalLogLikelihood classes to class name strings.

ax.storage.botorch_modular_registry.ACQUISITION_REGISTRY: Dict[Type[ax.models.torch.botorch_modular.acquisition.Acquisition], str] = {<class 'ax.models.torch.botorch_modular.acquisition.Acquisition'>: 'Acquisition', <class 'ax.models.torch.botorch_modular.kg.KnowledgeGradient'>: 'KnowledgeGradient', <class 'ax.models.torch.botorch_modular.mes.MaxValueEntropySearch'>: 'MaxValueEntropySearch', <class 'ax.models.torch.botorch_modular.kg.MultiFidelityKnowledgeGradient'>: 'MultiFidelityKnowledgeGradient', <class 'ax.models.torch.botorch_modular.mes.MultiFidelityMaxValueEntropySearch'>: 'MultiFidelityMaxValueEntropySearch'}

Mapping of BoTorch Model classes to class name strings.

ax.storage.botorch_modular_registry.CLASS_TO_REGISTRY: Dict[Any, Dict[Type[Any], str]] = {<class 'ax.models.torch.botorch_modular.acquisition.Acquisition'>: {<class 'ax.models.torch.botorch_modular.acquisition.Acquisition'>: 'Acquisition', <class 'ax.models.torch.botorch_modular.kg.KnowledgeGradient'>: 'KnowledgeGradient', <class 'ax.models.torch.botorch_modular.mes.MaxValueEntropySearch'>: 'MaxValueEntropySearch', <class 'ax.models.torch.botorch_modular.kg.MultiFidelityKnowledgeGradient'>: 'MultiFidelityKnowledgeGradient', <class 'ax.models.torch.botorch_modular.mes.MultiFidelityMaxValueEntropySearch'>: 'MultiFidelityMaxValueEntropySearch'}, <class 'botorch.acquisition.acquisition.AcquisitionFunction'>: {<class 'botorch.acquisition.monte_carlo.qExpectedImprovement'>: 'qExpectedImprovement', <class 'botorch.acquisition.knowledge_gradient.qKnowledgeGradient'>: 'qKnowledgeGradient', <class 'botorch.acquisition.max_value_entropy_search.qMaxValueEntropy'>: 'qMaxValueEntropy', <class 'botorch.acquisition.knowledge_gradient.qMultiFidelityKnowledgeGradient'>: 'qMultiFidelityKnowledgeGradient', <class 'botorch.acquisition.max_value_entropy_search.qMultiFidelityMaxValueEntropy'>: 'qMultiFidelityMaxValueEntropy', <class 'botorch.acquisition.monte_carlo.qNoisyExpectedImprovement'>: 'qNoisyExpectedImprovement'}, <class 'gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood'>: {<class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>: 'ExactMarginalLogLikelihood', <class 'gpytorch.mlls.sum_marginal_log_likelihood.SumMarginalLogLikelihood'>: 'SumMarginalLogLikelihood'}, <class 'botorch.models.model.Model'>: {<class 'botorch.models.gp_regression.FixedNoiseGP'>: 'FixedNoiseGP', <class 'botorch.models.gp_regression_fidelity.FixedNoiseMultiFidelityGP'>: 'FixedNoiseMultiFidelityGP', <class 'botorch.models.multitask.FixedNoiseMultiTaskGP'>: 'FixedNoiseMultiTaskGP', <class 'botorch.models.model_list_gp_regression.ModelListGP'>: 'ModelListGP', <class 'botorch.models.multitask.MultiTaskGP'>: 'MultiTaskGP', <class 'botorch.models.gp_regression.SingleTaskGP'>: 'SingleTaskGP', <class 'botorch.models.gp_regression_fidelity.SingleTaskMultiFidelityGP'>: 'SingleTaskMultiFidelityGP'}}

Reverse registries for decoding.

ax.storage.botorch_modular_registry.MLL_REGISTRY: Dict[Type[gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood], str] = {<class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>: 'ExactMarginalLogLikelihood', <class 'gpytorch.mlls.sum_marginal_log_likelihood.SumMarginalLogLikelihood'>: 'SumMarginalLogLikelihood'}

Overarching mapping from encoded classes to registry map.

ax.storage.botorch_modular_registry.MODEL_REGISTRY: Dict[Type[botorch.models.model.Model], str] = {<class 'botorch.models.gp_regression.FixedNoiseGP'>: 'FixedNoiseGP', <class 'botorch.models.gp_regression_fidelity.FixedNoiseMultiFidelityGP'>: 'FixedNoiseMultiFidelityGP', <class 'botorch.models.multitask.FixedNoiseMultiTaskGP'>: 'FixedNoiseMultiTaskGP', <class 'botorch.models.model_list_gp_regression.ModelListGP'>: 'ModelListGP', <class 'botorch.models.multitask.MultiTaskGP'>: 'MultiTaskGP', <class 'botorch.models.gp_regression.SingleTaskGP'>: 'SingleTaskGP', <class 'botorch.models.gp_regression_fidelity.SingleTaskMultiFidelityGP'>: 'SingleTaskMultiFidelityGP'}

Mapping of Botorch AcquisitionFunction classes to class name strings.

ax.storage.botorch_modular_registry.REVERSE_MLL_REGISTRY: Dict[str, Type[gpytorch.mlls.marginal_log_likelihood.MarginalLogLikelihood]] = {'ExactMarginalLogLikelihood': <class 'gpytorch.mlls.exact_marginal_log_likelihood.ExactMarginalLogLikelihood'>, 'SumMarginalLogLikelihood': <class 'gpytorch.mlls.sum_marginal_log_likelihood.SumMarginalLogLikelihood'>}

Overarching mapping from encoded classes to reverse registry map.

ax.storage.botorch_modular_registry.register_acquisition(acq_class: Type[ax.models.torch.botorch_modular.acquisition.Acquisition]) → None[source]

Add a custom acquisition class to the SQA and JSON registries.

ax.storage.metric_registry.register_metric(metric_cls: Type[ax.core.metric.Metric], val: Optional[int] = None) → None[source]

Add a custom metric class to the SQA and JSON registries. For the SQA registry, if no int is specified, use a hash of the class name.

ax.storage.runner_registry.register_runner(runner_cls: Type[ax.core.runner.Runner], val: Optional[int] = None) → None[source]

Add a custom runner class to the SQA and JSON registries. For the SQA registry, if no int is specified, use a hash of the class name.

Utilities

class ax.storage.utils.DomainType[source]

Bases: enum.Enum

Class for enumerating domain types.

CHOICE: int = 2
FIXED: int = 0
RANGE: int = 1
class ax.storage.utils.EncodeDecodeFieldsMap(python_only, encoded_only, python_to_encoded)[source]

Bases: tuple

property encoded_only

Alias for field number 1

property python_only

Alias for field number 0

property python_to_encoded

Alias for field number 2

class ax.storage.utils.MetricIntent[source]

Bases: enum.Enum

Class for enumerating metric use types.

ADDITIONAL_OBJECTIVE: str = 'additional_objective'
MULTI_OBJECTIVE: str = 'multi_objective'
OBJECTIVE: str = 'objective'
OBJECTIVE_THRESHOLD: str = 'objective_threshold'
OUTCOME_CONSTRAINT: str = 'outcome_constraint'
SCALARIZED_OBJECTIVE: str = 'scalarized_objective'
TRACKING: str = 'tracking'
class ax.storage.utils.ParameterConstraintType[source]

Bases: enum.Enum

Class for enumerating parameter constraint types.

Linear constraint is base type whereas other constraint types are special types of linear constraints.

LINEAR: int = 0
ORDER: int = 1
SUM: int = 2
ax.storage.utils.remove_prefix(text, prefix)[source]