def test_restore_from_generator_run(self):
gs = GenerationStrategy(
steps=[GenerationStep(model=Models.SOBOL, num_trials=5)])
# No generator runs on GS, so can't restore from one.
with self.assertRaises(ValueError):
gs._restore_model_from_generator_run()
exp = get_branin_experiment(with_batch=True)
gs.gen(experiment=exp)
model = gs.model
# Create a copy of the generation strategy and check that when
# we restore from last generator run, the model will be set
# correctly and that `_seen_trial_indices_by_status` is filled.
new_gs = GenerationStrategy(
steps=[GenerationStep(model=Models.SOBOL, num_trials=5)])
new_gs._experiment = exp
new_gs._generator_runs = gs._generator_runs
self.assertIsNone(new_gs._seen_trial_indices_by_status)
new_gs._restore_model_from_generator_run()
self.assertEqual(gs._seen_trial_indices_by_status,
exp.trial_indices_by_status)
# Model should be reset, but it should be the same model with same data.
self.assertIsNot(model, new_gs.model)
self.assertEqual(model.__class__,
new_gs.model.__class__) # Model bridge.
self.assertEqual(model.model.__class__,
new_gs.model.model.__class__) # Model.
self.assertEqual(model._training_data, new_gs.model._training_data)
def generation_strategy_from_json(
generation_strategy_json: Dict[str, Any]) -> GenerationStrategy:
"""Load generation strategy from JSON."""
steps = object_from_json(generation_strategy_json.pop("steps"))
gs = GenerationStrategy(steps=steps,
name=generation_strategy_json.pop("name"))
gs._db_id = object_from_json(generation_strategy_json.pop("db_id"))
gs._experiment = object_from_json(
generation_strategy_json.pop("experiment"))
gs._curr = gs._steps[generation_strategy_json.pop("curr_index")]
gs._generator_runs = object_from_json(
generation_strategy_json.pop("generator_runs"))
if generation_strategy_json.pop(
"had_initialized_model"): # pragma: no cover
# If model in the current step was not directly from the `Models` enum,
# pass its type to `restore_model_from_generator_run`, which will then
# attempt to use this type to recreate the model.
if type(gs._curr.model) != Models:
models_enum = type(gs._curr.model)
assert issubclass(models_enum, Models)
# pyre-ignore[6]: `models_enum` typing hackiness
gs._restore_model_from_generator_run(models_enum=models_enum)
return gs
gs._restore_model_from_generator_run()
return gs
def generation_strategy_from_sqa(
self, gs_sqa: SQAGenerationStrategy) -> GenerationStrategy:
"""Convert SQALchemy generation strategy to Ax `GenerationStrategy`."""
steps = object_from_json(gs_sqa.steps)
gs = GenerationStrategy(name=gs_sqa.name, steps=steps)
gs._curr = gs._steps[gs_sqa.curr_index]
gs._generator_runs = [
self.generator_run_from_sqa(gr) for gr in gs_sqa.generator_runs
]
if len(gs._generator_runs) > 0:
# Generation strategy had an initialized model.
# pyre-ignore[16]: SQAGenerationStrategy does not have `experiment` attr.
gs._experiment = self.experiment_from_sqa(gs_sqa.experiment)
# If model in the current step was not directly from the `Models` enum,
# pass its type to `restore_model_from_generator_run`, which will then
# attempt to use this type to recreate the model.
if type(gs._curr.model) != Models:
models_enum = type(gs._curr.model)
assert issubclass(models_enum, Models)
# pyre-ignore[6]: `models_enum` typing hackiness
gs._restore_model_from_generator_run(models_enum=models_enum)
else:
gs._restore_model_from_generator_run()
gs._db_id = gs_sqa.id
return gs
def generation_strategy_from_sqa(
self,
gs_sqa: SQAGenerationStrategy,
experiment: Optional[Experiment] = None,
reduced_state: bool = False,
) -> GenerationStrategy:
"""Convert SQALchemy generation strategy to Ax `GenerationStrategy`."""
steps = object_from_json(
gs_sqa.steps,
decoder_registry=self.config.json_decoder_registry,
class_decoder_registry=self.config.json_class_decoder_registry,
)
gs = GenerationStrategy(name=gs_sqa.name, steps=steps)
gs._curr = gs._steps[gs_sqa.curr_index]
immutable_ss_and_oc = (experiment.immutable_search_space_and_opt_config
if experiment is not None else False)
if reduced_state and gs_sqa.generator_runs:
# Only fully load the last of the generator runs, load the rest with
# reduced state.
gs._generator_runs = [
self.generator_run_from_sqa(
generator_run_sqa=gr,
reduced_state=True,
immutable_search_space_and_opt_config=immutable_ss_and_oc,
) for gr in gs_sqa.generator_runs[:-1]
]
gs._generator_runs.append(
self.generator_run_from_sqa(
generator_run_sqa=gs_sqa.generator_runs[-1],
reduced_state=False,
immutable_search_space_and_opt_config=immutable_ss_and_oc,
))
else:
gs._generator_runs = [
self.generator_run_from_sqa(
generator_run_sqa=gr,
reduced_state=False,
immutable_search_space_and_opt_config=immutable_ss_and_oc,
) for gr in gs_sqa.generator_runs
]
if len(gs._generator_runs) > 0:
# Generation strategy had an initialized model.
if experiment is None:
raise SQADecodeError(
"Cannot decode a generation strategy with a non-zero number of "
"generator runs without an experiment.")
gs._experiment = experiment
# If model in the current step was not directly from the `Models` enum,
# pass its type to `restore_model_from_generator_run`, which will then
# attempt to use this type to recreate the model.
if type(gs._curr.model) != Models:
models_enum = type(gs._curr.model)
assert issubclass(models_enum, ModelRegistryBase)
# pyre-ignore[6]: `models_enum` typing hackiness
gs._restore_model_from_generator_run(models_enum=models_enum)
else:
gs._restore_model_from_generator_run()
gs.db_id = gs_sqa.id
return gs
def test_restore_from_generator_run(self):
gs = GenerationStrategy(steps=[
GenerationStep(model=Models.SOBOL, num_arms=5),
GenerationStep(model=Models.GPEI, num_arms=-1),
])
with self.assertRaises(ValueError):
gs._restore_model_from_generator_run()
gs.gen(experiment=get_branin_experiment())
model = gs.model
gs._restore_model_from_generator_run()
# Model should be reset.
self.assertIsNot(model, gs.model)
def generation_strategy_from_sqa(
self, gs_sqa: SQAGenerationStrategy) -> GenerationStrategy:
"""Convert SQALchemy generation strategy to Ax `GenerationStrategy`."""
steps = object_from_json(gs_sqa.steps)
gs = GenerationStrategy(name=gs_sqa.name, steps=steps)
gs._curr = gs._steps[gs_sqa.curr_index]
gs._generator_runs = [
self.generator_run_from_sqa(gr) for gr in gs_sqa.generator_runs
]
if len(gs._generator_runs) > 0:
# Generation strategy had an initialized model.
# pyre-ignore[16]: SQAGenerationStrategy does not have `experiment` attr.
gs._experiment = self.experiment_from_sqa(gs_sqa.experiment)
gs._restore_model_from_generator_run()
gs._db_id = gs_sqa.id
return gs
def generation_strategy_from_json(
generation_strategy_json: Dict[str, Any]
) -> GenerationStrategy:
"""Load generation strategy from JSON."""
steps = object_from_json(generation_strategy_json.pop("steps"))
gs = GenerationStrategy(steps=steps, name=generation_strategy_json.pop("name"))
gs._experiment = object_from_json(generation_strategy_json.pop("experiment"))
gs._generated = generation_strategy_json.pop("generated")
gs._observed = generation_strategy_json.pop("observed")
gs._data = object_from_json(generation_strategy_json.pop("data"))
gs._curr = gs._steps[generation_strategy_json.pop("current_step_index")]
gs._generator_runs = object_from_json(
generation_strategy_json.pop("generator_runs")
)
if generation_strategy_json.pop("had_initialized_model"): # pragma: no cover
gs._restore_model_from_generator_run()
return gs
def generation_strategy_from_json(
experiment: Experiment, generation_strategy_json: Dict[str, Any]
) -> GenerationStrategy:
"""Load generation strategy from JSON."""
gs = GenerationStrategy(
steps=object_from_json(generation_strategy_json.pop("steps")),
name=generation_strategy_json.pop("name"),
)
gs._generated = generation_strategy_json.pop("generated")
gs._observed = generation_strategy_json.pop("observed")
gs._data = object_from_json(generation_strategy_json.pop("data"))
gs._curr = object_from_json(generation_strategy_json.pop("curr"))
gs._last_generator_run = object_from_json(
generation_strategy_json.pop("last_generator_run")
)
if generation_strategy_json.pop("had_initialized_model"): # pragma: no cover
gs._restore_model_from_generator_run(experiment=experiment)
return gs
