def _load_generation_strategy_by_id(
gs_id: int,
decoder: Decoder,
experiment: Optional[Experiment] = None,
reduced_state: bool = False,
) -> GenerationStrategy:
"""Finds a generation strategy stored by a given ID and restores it."""
if reduced_state:
return decoder.generation_strategy_from_sqa(
gs_sqa=_get_generation_strategy_sqa_reduced_state(gs_id=gs_id,
decoder=decoder),
experiment=experiment,
reduced_state=reduced_state,
)
exp_sqa_class = decoder.config.class_to_sqa_class[Experiment]
immutable_opt_config_and_search_space = ((
_get_experiment_immutable_opt_config_and_search_space(
experiment_name=experiment.name,
# pyre-ignore Incompatible parameter type [6]: Expected
# `Type[SQAExperiment]` for 2nd parameter `exp_sqa_class`
# to call `_get_experiment_immutable_opt_config_and_search_space`
# but got `Type[ax.storage.sqa_store.db.SQABase]`.
exp_sqa_class=exp_sqa_class,
)) if experiment is not None else False)
if immutable_opt_config_and_search_space:
gs_sqa = _get_generation_strategy_sqa_immutable_opt_config_and_search_space(
gs_id=gs_id, decoder=decoder)
else:
gs_sqa = _get_generation_strategy_sqa(gs_id=gs_id, decoder=decoder)
return decoder.generation_strategy_from_sqa(gs_sqa=gs_sqa,
experiment=experiment,
reduced_state=reduced_state)
def _load_experiment(experiment_name: str,
decoder: Decoder,
reduced_state: bool = False) -> Experiment:
"""Load experiment by name, using given Decoder instance.
1) Get SQLAlchemy object from DB.
2) Convert to corresponding Ax object.
"""
# pyre-ignore Incompatible variable type [9]: exp_sqa_class is declared to have type
# `Type[SQAExperiment]` but is used as type `Type[ax.storage.sqa_store.db.SQABase]`
exp_sqa_class: Type[SQAExperiment] = decoder.config.class_to_sqa_class[
Experiment]
if reduced_state:
return decoder.experiment_from_sqa(
experiment_sqa=_get_experiment_sqa_reduced_state(
experiment_name=experiment_name, exp_sqa_class=exp_sqa_class),
reduced_state=reduced_state,
)
immutable_opt_config_and_search_space = (
_get_experiment_immutable_opt_config_and_search_space(
experiment_name=experiment_name, exp_sqa_class=exp_sqa_class))
if immutable_opt_config_and_search_space:
experiment_sqa = _get_experiment_sqa_immutable_opt_config_and_search_space(
experiment_name=experiment_name, exp_sqa_class=exp_sqa_class)
else:
experiment_sqa = _get_experiment_sqa(experiment_name=experiment_name,
exp_sqa_class=exp_sqa_class)
return decoder.experiment_from_sqa(experiment_sqa=experiment_sqa)
def setUp(self):
init_test_engine_and_session_factory(force_init=True)
self.config = SQAConfig()
self.encoder = Encoder(config=self.config)
self.decoder = Decoder(config=self.config)
self.experiment = get_experiment_with_batch_trial()
self.dummy_parameters = [
get_range_parameter(), # w
get_range_parameter2(), # x
]
def test_storage_error_handling(self, mock_save_fails):
"""Check that if `suppress_storage_errors` is True, AxClient won't
visibly fail if encountered storage errors.
"""
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
ax_client = AxClient(db_settings=db_settings,
suppress_storage_errors=True)
ax_client.create_experiment(
name="test_experiment",
parameters=[
{
"name": "x",
"type": "range",
"bounds": [-5.0, 10.0]
},
{
"name": "y",
"type": "range",
"bounds": [0.0, 15.0]
},
],
minimize=True,
)
for _ in range(3):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(trial_index=trial_index,
raw_data=branin(*parameters.values()))
def save_or_update_trials(
experiment: Experiment,
trials: List[BaseTrial],
config: Optional[SQAConfig] = None,
batch_size: Optional[int] = None,
reduce_state_generator_runs: bool = False,
) -> None:
"""Add new trials to the experiment, or update if already exists
(using default SQAConfig).
Note that new data objects (whether attached to existing or new trials)
will also be added to the experiment, but existing data objects in the
database will *not* be updated or removed.
"""
config = config or SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
_save_or_update_trials(
experiment=experiment,
trials=trials,
encoder=encoder,
decoder=decoder,
batch_size=batch_size,
reduce_state_generator_runs=reduce_state_generator_runs,
)
def load_experiment(
experiment_name: str, config: Optional[SQAConfig] = None
) -> Experiment:
"""Load experiment by name (uses default SQAConfig)."""
config = config or SQAConfig()
decoder = Decoder(config=config)
return _load_experiment(experiment_name=experiment_name, decoder=decoder)
def load_generation_strategy_by_id(
gs_id: int, config: Optional[SQAConfig] = None
) -> GenerationStrategy:
"""Finds a generation strategy stored by a given ID and restores it."""
config = config or SQAConfig()
decoder = Decoder(config=config)
return _load_generation_strategy_by_id(gs_id=gs_id, decoder=decoder)
def _load_generation_strategy_by_experiment_name(
experiment_name: str, decoder: Decoder
) -> GenerationStrategy:
"""Load a generation strategy attached to an experiment specified by a name,
using given Decoder instance.
1) Get SQLAlchemy object from DB.
2) Convert to corresponding Ax object.
"""
exp_sqa_class = decoder.config.class_to_sqa_class[Experiment]
gs_sqa_class = decoder.config.class_to_sqa_class[GenerationStrategy]
with session_scope() as session:
gs_sqa = (
session.query(gs_sqa_class)
.join(exp_sqa_class.generation_strategy) # pyre-ignore[16]
# pyre-fixme[16]: `SQABase` has no attribute `name`.
.filter(exp_sqa_class.name == experiment_name)
.one_or_none()
)
if gs_sqa is None:
raise ValueError(
f"Experiment {experiment_name} does not have a generation strategy "
"attached to it."
)
return decoder.generation_strategy_from_sqa(gs_sqa=gs_sqa)
def _load_generation_strategy_by_id(gs_id: int, decoder: Decoder) -> GenerationStrategy:
"""Finds a generation strategy stored by a given ID and restores it."""
with session_scope() as session:
gs_sqa = session.query(SQAGenerationStrategy).filter_by(id=gs_id).one_or_none()
if gs_sqa is None:
raise ValueError(f"Generation strategy with ID #{gs_id} not found.")
return decoder.generation_strategy_from_sqa(gs_sqa=gs_sqa)
def _load_experiment(experiment_name: str, decoder: Decoder) -> Experiment:
"""Load experiment by name, using given Decoder instance.
1) Get SQLAlchemy object from DB.
2) Convert to corresponding Ax object.
"""
# Convert SQA to user-facing class outside of session scope to avoid timeouts
return decoder.experiment_from_sqa(experiment_sqa=_get_experiment_sqa(
experiment_name=experiment_name))
class DBSettings(NamedTuple):
"""
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.
"""
creator: Optional[Callable] = None
decoder: Decoder = Decoder(config=SQAConfig())
encoder: Encoder = Encoder(config=SQAConfig())
url: Optional[str] = None
def load_generation_strategy_by_experiment_name(
experiment_name: str,
config: Optional[SQAConfig] = None) -> GenerationStrategy:
"""Finds a generation strategy attached to an experiment specified by a name
and restores it from its corresponding SQA object.
"""
config = config or SQAConfig()
decoder = Decoder(config=config)
return _load_generation_strategy_by_experiment_name(
experiment_name=experiment_name, decoder=decoder)
def save_or_update_trial(
experiment: Experiment, trial: BaseTrial, config: Optional[SQAConfig] = None
) -> None:
"""Add new trial to the experiment, or update if already exists
(using default SQAConfig)."""
config = config or SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
_save_or_update_trial(
experiment=experiment, trial=trial, encoder=encoder, decoder=decoder
)
def save_experiment(experiment: Experiment, config: Optional[SQAConfig] = None) -> None:
"""Save experiment (using default SQAConfig)."""
if not isinstance(experiment, Experiment):
raise ValueError("Can only save instances of Experiment")
if not experiment.has_name:
raise ValueError("Experiment name must be set prior to saving.")
config = config or SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
_save_experiment(experiment=experiment, encoder=encoder, decoder=decoder)
def load_generation_strategy_by_id(
gs_id: int,
config: Optional[SQAConfig] = None,
experiment: Optional[Experiment] = None,
reduced_state: bool = False,
) -> GenerationStrategy:
"""Finds a generation strategy stored by a given ID and restores it."""
config = config or SQAConfig()
decoder = Decoder(config=config)
return _load_generation_strategy_by_id(
gs_id=gs_id, decoder=decoder, experiment=experiment, reduced_state=reduced_state
)
def test_sqa_storage(self):
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
ax_client = AxClient(db_settings=db_settings)
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
for _ in range(5):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(
trial_index=trial_index, raw_data=branin(*parameters.values())
)
gs = ax_client.generation_strategy
ax_client = AxClient(db_settings=db_settings)
ax_client.load_experiment_from_database("test_experiment")
# Trial #4 was completed after the last time the generation strategy
# generated candidates, so pre-save generation strategy was not
# "aware" of completion of trial #4. Post-restoration generation
# strategy is aware of it, however, since it gets restored with most
# up-to-date experiment data. Do adding trial #4 to the seen completed
# trials of pre-storage GS to check their equality otherwise.
gs._seen_trial_indices_by_status[TrialStatus.COMPLETED].add(4)
self.assertEqual(gs, ax_client.generation_strategy)
with self.assertRaises(ValueError):
# Overwriting existing experiment.
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
with self.assertRaises(ValueError):
# Overwriting existing experiment with overwrite flag with present
# DB settings. This should fail as we no longer allow overwriting
# experiments stored in the DB.
ax_client.create_experiment(
name="test_experiment",
parameters=[{"name": "x", "type": "range", "bounds": [-5.0, 10.0]}],
overwrite_existing_experiment=True,
)
# Original experiment should still be in DB and not have been overwritten.
self.assertEqual(len(ax_client.experiment.trials), 5)
def _load_experiment(experiment_name: str, decoder: Decoder) -> Experiment:
"""Load experiment by name, using given Decoder instance.
1) Get SQLAlchemy object from DB.
2) Convert to corresponding Ax object.
"""
with session_scope() as session:
sqa_experiment = (
session.query(SQAExperiment).filter_by(name=experiment_name).one_or_none()
)
if sqa_experiment is None:
raise ValueError(f"Experiment `{experiment_name}` not found.")
return decoder.experiment_from_sqa(sqa_experiment)
def _load_generation_strategy_by_id(
gs_id: int,
decoder: Decoder,
experiment: Optional[Experiment] = None,
reduced_state: bool = False,
) -> GenerationStrategy:
"""Finds a generation strategy stored by a given ID and restores it."""
gs_sqa = _get_generation_strategy_sqa(
gs_id=gs_id, decoder=decoder, reduced_state=reduced_state
)
return decoder.generation_strategy_from_sqa(
gs_sqa=gs_sqa, experiment=experiment, reduced_state=reduced_state
)
def _load_generation_strategy_by_experiment_name(
experiment_name: str, decoder: Decoder) -> GenerationStrategy:
"""Load a generation strategy attached to an experiment specified by a name,
using given Decoder instance.
1) Get SQLAlchemy object from DB.
2) Convert to corresponding Ax object.
"""
experiment_sqa = _get_experiment_sqa(experiment_name=experiment_name)
if experiment_sqa.generation_strategy is None:
raise ValueError(
f"Experiment {experiment_name} does not have a generation strategy "
"attached to it.")
gs_sqa = experiment_sqa.generation_strategy
return decoder.generation_strategy_from_sqa(gs_sqa=gs_sqa)
def test_save_load_experiment(self):
exp = get_experiment()
init_test_engine_and_session_factory(force_init=True)
db_settings = DBSettings(
encoder=Encoder(config=SQAConfig()),
decoder=Decoder(config=SQAConfig()),
creator=None,
)
save_experiment(exp, db_settings)
load_experiment(exp.name, db_settings)
simple_experiment = get_simple_experiment()
save_experiment(simple_experiment, db_settings)
with self.assertRaisesRegex(ValueError, "Service API only"):
load_experiment(simple_experiment.name, db_settings)
def test_suppress_all_storage_errors(self, mock_save_exp, _):
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
BareBonesTestScheduler(
experiment=self.branin_experiment, # Has runner and metrics.
generation_strategy=self.two_sobol_steps_GS,
options=SchedulerOptions(
init_seconds_between_polls=0.1, # Short between polls so test is fast.
suppress_storage_errors_after_retries=True,
),
db_settings=db_settings,
)
self.assertEqual(mock_save_exp.call_count, 3)
def save_generation_strategy(generation_strategy: GenerationStrategy,
config: Optional[SQAConfig] = None) -> int:
"""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.
"""
# Start up SQA encoder.
config = config or SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
return _save_generation_strategy(generation_strategy=generation_strategy,
encoder=encoder,
decoder=decoder)
def test_sqa_storage(self):
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
ax_client = AxClient(db_settings=db_settings)
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
for _ in range(5):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(
trial_index=trial_index, raw_data=branin(*parameters.values())
)
gs = ax_client.generation_strategy
ax_client = AxClient(db_settings=db_settings)
ax_client.load_experiment_from_database("test_experiment")
self.assertEqual(gs, ax_client.generation_strategy)
with self.assertRaises(ValueError):
# Overwriting existing experiment.
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
with self.assertRaises(ValueError):
# Overwriting existing experiment with overwrite flag with present
# DB settings. This should fail as we no longer allow overwriting
# experiments stored in the DB.
ax_client.create_experiment(
name="test_experiment",
parameters=[{"name": "x", "type": "range", "bounds": [-5.0, 10.0]}],
overwrite_existing_experiment=True,
)
# Original experiment should still be in DB and not have been overwritten.
self.assertEqual(len(ax_client.experiment.trials), 5)
def save_or_update_trials(
experiment: Experiment,
trials: List[BaseTrial],
config: Optional[SQAConfig] = None,
batch_size: Optional[int] = None,
) -> None:
"""Add new trials to the experiment, or update if already exists
(using default SQAConfig)."""
config = config or SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
_save_or_update_trials(
experiment=experiment,
trials=trials,
encoder=encoder,
decoder=decoder,
batch_size=batch_size,
)
def update_generation_strategy(
generation_strategy: GenerationStrategy,
generator_runs: List[GeneratorRun],
config: Optional[SQAConfig] = None,
batch_size: Optional[int] = None,
) -> None:
"""Update generation strategy's current step and attach generator runs
(using default SQAConfig)."""
config = config or SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
_update_generation_strategy(
generation_strategy=generation_strategy,
generator_runs=generator_runs,
encoder=encoder,
decoder=decoder,
batch_size=batch_size,
)
def test_sqa_storage(self):
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
ax_client = AxClient(db_settings=db_settings)
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
for _ in range(5):
parameters, trial_index = ax_client.get_next_trial()
ax_client.complete_trial(
trial_index=trial_index, raw_data=branin(*parameters.values())
)
gs = ax_client.generation_strategy
ax_client = AxClient(db_settings=db_settings)
ax_client.load_experiment_from_database("test_experiment")
self.assertEqual(gs, ax_client.generation_strategy)
with self.assertRaises(ValueError):
# Overwriting existing experiment.
ax_client.create_experiment(
name="test_experiment",
parameters=[
{"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
# Overwriting existing experiment with overwrite flag.
ax_client.create_experiment(
name="test_experiment",
parameters=[{"name": "x1", "type": "range", "bounds": [-5.0, 10.0]}],
overwrite_existing_experiment=True,
)
# There should be no trials, as we just put in a fresh experiment.
self.assertEqual(len(ax_client.experiment.trials), 0)
def load_experiment(
experiment_name: str,
config: Optional[SQAConfig] = None,
reduced_state: bool = False,
) -> Experiment:
"""Load experiment by name.
Args:
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).
"""
config = config or SQAConfig()
decoder = Decoder(config=config)
return _load_experiment(
experiment_name=experiment_name, decoder=decoder, reduced_state=reduced_state
)
def _load_experiment(
experiment_name: str,
decoder: Decoder,
reduced_state: bool = False,
load_trials_in_batches_of_size: Optional[int] = None,
) -> Experiment:
"""Load experiment by name, using given Decoder instance.
1) Get SQLAlchemy object from DB.
2) Convert to corresponding Ax object.
"""
# pyre-ignore Incompatible variable type [9]: exp_sqa_class is declared to have type
# `Type[SQAExperiment]` but is used as type `Type[ax.storage.sqa_store.db.SQABase]`
exp_sqa_class: Type[SQAExperiment] = decoder.config.class_to_sqa_class[
Experiment]
# pyre-ignore Incompatible variable type [9]: trial_sqa_class is decl. to have type
# `Type[SQATrial]` but is used as type `Type[ax.storage.sqa_store.db.SQABase]`
trial_sqa_class: Type[SQATrial] = decoder.config.class_to_sqa_class[Trial]
if reduced_state:
_get_experiment_sqa_func = _get_experiment_sqa_reduced_state
else:
imm_OC_and_SS = _get_experiment_immutable_opt_config_and_search_space(
experiment_name=experiment_name, exp_sqa_class=exp_sqa_class)
_get_experiment_sqa_func = (
_get_experiment_sqa_immutable_opt_config_and_search_space
if imm_OC_and_SS else _get_experiment_sqa)
experiment_sqa = _get_experiment_sqa_func(
experiment_name=experiment_name,
exp_sqa_class=exp_sqa_class,
trial_sqa_class=trial_sqa_class,
load_trials_in_batches_of_size=load_trials_in_batches_of_size,
)
return decoder.experiment_from_sqa(
experiment_sqa=experiment_sqa,
reduced_state=reduced_state,
)
def _load_generation_strategy_by_experiment_name(
experiment_name: str, decoder: Decoder
) -> GenerationStrategy:
"""Load a generation strategy attached to an experiment specified by a name,
using given Decoder instance.
1) Get SQLAlchemy object from DB.
2) Convert to corresponding Ax object.
"""
with session_scope() as session:
gs_sqa = (
session.query(SQAGenerationStrategy)
.join(SQAExperiment.generation_strategy)
.filter(SQAExperiment.name == experiment_name)
.one_or_none()
)
if gs_sqa is None:
raise ValueError(
f"Experiment {experiment_name} does not have a generation strategy "
"attached to it."
)
return decoder.generation_strategy_from_sqa(gs_sqa=gs_sqa)
def test_sqa_storage(self):
init_test_engine_and_session_factory(force_init=True)
config = SQAConfig()
encoder = Encoder(config=config)
decoder = Decoder(config=config)
db_settings = DBSettings(encoder=encoder, decoder=decoder)
ax = AxClient(db_settings=db_settings)
ax.create_experiment(
name="test_experiment",
parameters=[
{"name": "x1", "type": "range", "bounds": [-5.0, 10.0]},
{"name": "x2", "type": "range", "bounds": [0.0, 15.0]},
],
minimize=True,
)
for _ in range(5):
parameters, trial_index = ax.get_next_trial()
ax.complete_trial(
trial_index=trial_index, raw_data=branin(*parameters.values())
)
gs = ax.generation_strategy
ax = AxClient(db_settings=db_settings)
ax.load_experiment_from_database("test_experiment")
self.assertEqual(gs, ax.generation_strategy)