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