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_json(
generation_strategy_json: Dict[str, Any],
experiment: Optional[Experiment] = None) -> 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 = experiment or 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, ModelRegistryBase)
# 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 test_current_generator_run_limit_unlimited_second_step(self):
NUM_INIT_TRIALS = 5
SECOND_STEP_PARALLELISM = 3
NUM_ROUNDS = 4
exp = get_branin_experiment()
sobol_gs_with_parallelism_limits = GenerationStrategy(steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=NUM_INIT_TRIALS,
min_trials_observed=3,
),
GenerationStep(
model=Models.SOBOL,
num_trials=-1,
max_parallelism=SECOND_STEP_PARALLELISM,
),
])
sobol_gs_with_parallelism_limits._experiment = exp
could_gen = self._run_GS_for_N_rounds(
gs=sobol_gs_with_parallelism_limits,
exp=exp,
num_rounds=NUM_ROUNDS)
# We expect trials from first generation step + trials from remaining rounds in
# batches limited by parallelism setting in the second step.
self.assertEqual(
len(exp.trials),
NUM_INIT_TRIALS + (NUM_ROUNDS - 1) * SECOND_STEP_PARALLELISM,
)
self.assertTrue(all(t.status.is_completed
for t in exp.trials.values()))
self.assertEqual(could_gen, [NUM_INIT_TRIALS] +
[SECOND_STEP_PARALLELISM] * (NUM_ROUNDS - 1))
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_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_use_update(self, mock_lookup_data, mock_update):
exp = get_branin_experiment()
sobol_gs_with_update = GenerationStrategy(steps=[
GenerationStep(model=Models.SOBOL, num_trials=-1, use_update=True)
])
sobol_gs_with_update._experiment = exp
self.assertEqual(
sobol_gs_with_update._find_trials_completed_since_last_gen(),
set(),
)
with self.assertRaises(NotImplementedError):
# `BraninMetric` is available while running by default, which should
# raise an error when use with `use_update=True` on a generation step, as we
# have not yet properly addressed that edge case (for lack of use case).
sobol_gs_with_update.gen(experiment=exp)
core_stubs_module = get_branin_experiment.__module__
with patch(
f"{core_stubs_module}.BraninMetric.is_available_while_running",
return_value=False,
):
# Try without passing data (GS looks up data on experiment).
trial = exp.new_trial(generator_run=sobol_gs_with_update.gen(
experiment=exp))
mock_update.assert_not_called()
trial._status = TrialStatus.COMPLETED
for i in range(3):
gr = sobol_gs_with_update.gen(experiment=exp)
self.assertEqual(
mock_lookup_data.call_args[1].get("trial_indices"), {i})
trial = exp.new_trial(generator_run=gr)
trial._status = TrialStatus.COMPLETED
# `_seen_trial_indices_by_status` is set during `gen`, to the experiment's
# `trial_indices_by_Status` at the time of candidate generation.
self.assertNotEqual(
sobol_gs_with_update._seen_trial_indices_by_status,
exp.trial_indices_by_status,
)
# Try with passing data.
sobol_gs_with_update.gen(
experiment=exp, data=get_branin_data(trial_indices=range(4)))
# Now `_seen_trial_indices_by_status` should be set to experiment's,
self.assertEqual(
sobol_gs_with_update._seen_trial_indices_by_status,
exp.trial_indices_by_status,
)
# Only the data for the last completed trial should be considered new and passed
# to `update`.
self.assertEqual(
set(mock_update.call_args[1].get(
"new_data").df["trial_index"].values), {3})
# Try with passing same data as before; no update should be performed.
with patch.object(sobol_gs_with_update,
"_update_current_model") as mock_update:
sobol_gs_with_update.gen(
experiment=exp, data=get_branin_data(trial_indices=range(4)))
mock_update.assert_not_called()
def test_current_generator_run_limit(self):
NUM_INIT_TRIALS = 5
SECOND_STEP_PARALLELISM = 3
NUM_ROUNDS = 4
exp = get_branin_experiment()
sobol_gs_with_parallelism_limits = GenerationStrategy(steps=[
GenerationStep(
model=Models.SOBOL,
num_trials=NUM_INIT_TRIALS,
min_trials_observed=3,
),
GenerationStep(
model=Models.SOBOL,
num_trials=-1,
max_parallelism=SECOND_STEP_PARALLELISM,
),
])
sobol_gs_with_parallelism_limits._experiment = exp
could_gen = []
for _ in range(NUM_ROUNDS):
(
num_trials_to_gen,
opt_complete,
) = sobol_gs_with_parallelism_limits.current_generator_run_limit()
self.assertFalse(opt_complete)
could_gen.append(num_trials_to_gen)
trials = []
for _ in range(num_trials_to_gen):
gr = sobol_gs_with_parallelism_limits.gen(
experiment=exp,
pending_observations=get_pending(experiment=exp),
)
trials.append(
exp.new_trial(gr).mark_running(no_runner_required=True))
for trial in trials:
exp.attach_data(get_branin_data(trial_indices=[trial.index]))
trial.mark_completed()
# We expect trials from first generation step + trials from remaining rounds in
# batches limited by parallelism setting in the second step.
self.assertEqual(
len(exp.trials),
NUM_INIT_TRIALS + (NUM_ROUNDS - 1) * SECOND_STEP_PARALLELISM,
)
self.assertTrue(all(t.status.is_completed
for t in exp.trials.values()))
self.assertEqual(could_gen, [NUM_INIT_TRIALS] +
[SECOND_STEP_PARALLELISM] * (NUM_ROUNDS - 1))
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
