def testExperimentUpdateTrial(self): save_experiment(self.experiment) trial = self.experiment.trials[0] trial.mark_staged() update_trial(experiment=self.experiment, trial=trial) loaded_experiment = load_experiment(self.experiment.name) self.assertEqual(trial, loaded_experiment.trials[0]) trial._run_metadata = {"foo": "bar"} update_trial(experiment=self.experiment, trial=trial) loaded_experiment = load_experiment(self.experiment.name) self.assertEqual(trial, loaded_experiment.trials[0]) self.experiment.attach_data(get_data(trial_index=trial.index)) update_trial(experiment=self.experiment, trial=trial) loaded_experiment = load_experiment(self.experiment.name) self.assertEqual(self.experiment, loaded_experiment) trial = self.experiment.new_batch_trial( generator_run=get_generator_run()) save_new_trial(experiment=self.experiment, trial=trial) self.experiment.attach_data(get_data(trial_index=trial.index)) update_trial(experiment=self.experiment, trial=trial) loaded_experiment = load_experiment(self.experiment.name) self.assertEqual(self.experiment, loaded_experiment)
def test_factorial_thompson_strategy(self, mock_update, mock_gen, mock_discrete): exp = get_branin_experiment() factorial_thompson_generation_strategy = GenerationStrategy(steps=[ GenerationStep(model=Models.FACTORIAL, num_arms=1), GenerationStep(model=Models.THOMPSON, num_arms=-1), ]) self.assertEqual(factorial_thompson_generation_strategy.name, "Factorial+Thompson") self.assertEqual( factorial_thompson_generation_strategy.generator_changes, [1]) for i in range(2): data = get_data() if i > 0 else None factorial_thompson_generation_strategy.gen(experiment=exp, new_data=data) exp.new_batch_trial().add_arm(Arm(parameters={"x1": i, "x2": i})) if i < 1: mock_discrete.assert_called() args, kwargs = mock_discrete.call_args self.assertIsInstance(kwargs.get("model"), FullFactorialGenerator) exp.new_batch_trial() else: mock_discrete.assert_called() args, kwargs = mock_discrete.call_args self.assertIsInstance( kwargs.get("model"), (ThompsonSampler, EmpiricalBayesThompsonSampler), )
def test_hierarchical_search_space(self): experiment = get_hierarchical_search_space_experiment() self.assertTrue(experiment.search_space.is_hierarchical) self.sobol_GS.gen(experiment=experiment) for _ in range(10): # During each iteration, check that all transformed observation features # contain all parameters of the flat search space. with patch.object(RandomModelBridge, "_fit") as mock_model_fit, patch.object( RandomModelBridge, "gen"): self.sobol_GS.gen(experiment=experiment) mock_model_fit.assert_called_once() obs_feats = mock_model_fit.call_args[1].get( "observation_features") all_parameter_names = ( experiment.search_space._all_parameter_names.copy()) # One of the parameter names is modified by transforms (because it's # one-hot encoded). all_parameter_names.remove("model") all_parameter_names.add("model_OH_PARAM_") for obsf in obs_feats: for p_name in all_parameter_names: self.assertIn(p_name, obsf.parameters) trial = (experiment.new_trial(generator_run=self.sobol_GS.gen( experiment=experiment)).mark_running( no_runner_required=True).mark_completed()) experiment.attach_data( get_data( metric_name="m1", trial_index=trial.index, num_non_sq_arms=1, include_sq=False, )) experiment.attach_data( get_data( metric_name="m2", trial_index=trial.index, num_non_sq_arms=1, include_sq=False, ))
def testEmptyMetrics(self): empty_experiment = Experiment(name="test_experiment", search_space=get_search_space()) self.assertEqual(empty_experiment.num_trials, 0) with self.assertRaises(ValueError): empty_experiment.fetch_data() batch = empty_experiment.new_batch_trial() self.assertEqual(empty_experiment.num_trials, 1) with self.assertRaises(ValueError): batch.fetch_data() empty_experiment.add_tracking_metric(Metric(name="some_metric")) empty_experiment.attach_data(get_data()) self.assertFalse(empty_experiment.fetch_data().df.empty)
def testEmptyMetrics(self): empty_experiment = Experiment(name="test_experiment", search_space=get_search_space()) self.assertEqual(empty_experiment.num_trials, 0) with self.assertRaises(ValueError): empty_experiment.fetch_data() batch = empty_experiment.new_batch_trial() batch.mark_running(no_runner_required=True) self.assertEqual(empty_experiment.num_trials, 1) with self.assertRaises(ValueError): batch.fetch_data() empty_experiment.add_tracking_metric(Metric(name="ax_test_metric")) self.assertTrue(empty_experiment.fetch_data().df.empty) empty_experiment.attach_data(get_data()) batch.mark_completed() self.assertFalse(empty_experiment.fetch_data().df.empty)
def get_modelbridge(mock_gen_arms, mock_observations_from_data, status_quo_name: Optional[str] = None) -> ModelBridge: exp = get_experiment() modelbridge = ModelBridge( search_space=get_search_space(), model=FullFactorialGenerator(), experiment=exp, data=get_data(), status_quo_name=status_quo_name, ) modelbridge._predict = mock.MagicMock( "ax.modelbridge.base.ModelBridge._predict", autospec=True, return_value=[get_observation().data], ) return modelbridge
def test_factorial_thompson_strategy(self): exp = get_branin_experiment() factorial_thompson_generation_strategy = GenerationStrategy(steps=[ GenerationStep(model=Models.FACTORIAL, num_arms=1), GenerationStep(model=Models.THOMPSON, num_arms=-1), ]) self.assertEqual(factorial_thompson_generation_strategy.name, "Factorial+Thompson") self.assertEqual( factorial_thompson_generation_strategy.model_transitions, [1]) mock_model_bridge = self.mock_discrete_model_bridge.return_value # Initial factorial batch. exp.new_batch_trial( factorial_thompson_generation_strategy.gen(experiment=exp)) args, kwargs = mock_model_bridge._set_kwargs_to_save.call_args self.assertEqual(kwargs.get("model_key"), "Factorial") # Subsequent Thompson sampling batch. exp.new_batch_trial( factorial_thompson_generation_strategy.gen(experiment=exp, data=get_data())) args, kwargs = mock_model_bridge._set_kwargs_to_save.call_args self.assertEqual(kwargs.get("model_key"), "Thompson")
class TestGenerationStrategy(TestCase): def setUp(self): self.gr = GeneratorRun(arms=[Arm(parameters={"x1": 1, "x2": 2})]) # Mock out slow GPEI. self.torch_model_bridge_patcher = patch( f"{TorchModelBridge.__module__}.TorchModelBridge", spec=True ) self.mock_torch_model_bridge = self.torch_model_bridge_patcher.start() self.mock_torch_model_bridge.return_value.gen.return_value = self.gr # Mock out slow TS. self.discrete_model_bridge_patcher = patch( f"{DiscreteModelBridge.__module__}.DiscreteModelBridge", spec=True ) self.mock_discrete_model_bridge = self.discrete_model_bridge_patcher.start() self.mock_discrete_model_bridge.return_value.gen.return_value = self.gr # Mock in `Models` registry self.registry_setup_dict_patcher = patch.dict( f"{Models.__module__}.MODEL_KEY_TO_MODEL_SETUP", { "Factorial": MODEL_KEY_TO_MODEL_SETUP["Factorial"]._replace( bridge_class=self.mock_discrete_model_bridge ), "Thompson": MODEL_KEY_TO_MODEL_SETUP["Thompson"]._replace( bridge_class=self.mock_discrete_model_bridge ), "GPEI": MODEL_KEY_TO_MODEL_SETUP["GPEI"]._replace( bridge_class=self.mock_torch_model_bridge ), }, ) self.mock_in_registry = self.registry_setup_dict_patcher.start() def tearDown(self): self.torch_model_bridge_patcher.stop() self.discrete_model_bridge_patcher.stop() self.registry_setup_dict_patcher.stop() def test_validation(self): # num_trials can be positive or -1. with self.assertRaises(ValueError): GenerationStrategy( steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-10), ] ) # only last num_trials can be -1. with self.assertRaises(ValueError): GenerationStrategy( steps=[ GenerationStep(model=Models.SOBOL, num_trials=-1), GenerationStep(model=Models.GPEI, num_trials=10), ] ) exp = Experiment( name="test", search_space=SearchSpace(parameters=[get_choice_parameter()]) ) factorial_thompson_generation_strategy = GenerationStrategy( steps=[ GenerationStep(model=Models.FACTORIAL, num_trials=1), GenerationStep(model=Models.THOMPSON, num_trials=2), ] ) self.assertTrue(factorial_thompson_generation_strategy._uses_registered_models) self.assertFalse( factorial_thompson_generation_strategy.uses_non_registered_models ) with self.assertRaises(ValueError): factorial_thompson_generation_strategy.gen(exp) self.assertEqual(GenerationStep(model=sum, num_trials=1).model_name, "sum") def test_custom_callables_for_models(self): exp = get_branin_experiment() sobol_factory_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=get_sobol, num_trials=-1)] ) self.assertFalse(sobol_factory_generation_strategy._uses_registered_models) self.assertTrue(sobol_factory_generation_strategy.uses_non_registered_models) gr = sobol_factory_generation_strategy.gen(experiment=exp, n=1) self.assertEqual(len(gr.arms), 1) def test_string_representation(self): gs1 = GenerationStrategy( steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-1), ] ) self.assertEqual( str(gs1), ( "GenerationStrategy(name='Sobol+GPEI', steps=[Sobol for 5 trials," " GPEI for subsequent trials])" ), ) gs2 = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=-1)] ) self.assertEqual( str(gs2), "GenerationStrategy(name='Sobol', steps=[Sobol for all trials])" ) def test_equality(self): gs1 = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-1), ], ) gs2 = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-1), ], ) self.assertEqual(gs1, gs2) # Clone_reset() doesn't clone exactly, so they won't be equal. gs3 = gs1.clone_reset() self.assertEqual(gs1, gs3) def test_restore_from_generator_run(self): gs = GenerationStrategy( steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-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 test_min_observed(self): # We should fail to transition the next model if there is not # enough data observed. exp = get_branin_experiment(get_branin_experiment()) gs = GenerationStrategy( steps=[ GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5), GenerationStep(model=Models.GPEI, num_trials=1), ] ) self.assertFalse(gs.uses_non_registered_models) for _ in range(5): exp.new_trial(gs.gen(exp)) with self.assertRaises(DataRequiredError): gs.gen(exp) def test_do_not_enforce_min_observations(self): # We should be able to move on to the next model if there is not # enough data observed if `enforce_num_trials` setting is False, in which # case the previous model should be used until there is enough data. exp = get_branin_experiment() gs = GenerationStrategy( steps=[ GenerationStep( model=Models.SOBOL, num_trials=1, min_trials_observed=5, enforce_num_trials=False, ), GenerationStep(model=Models.GPEI, num_trials=1), ] ) for _ in range(2): gs.gen(exp) # Make sure Sobol is used to generate the 6th point. self.assertIsInstance(gs._model, RandomModelBridge) def test_sobol_GPEI_strategy(self): exp = get_branin_experiment() sobol_GPEI = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=2), ], ) self.assertEqual(sobol_GPEI.name, "Sobol+GPEI") self.assertEqual(sobol_GPEI.model_transitions, [5]) # exp.new_trial(generator_run=sobol_GPEI.gen(exp)).run() for i in range(7): g = sobol_GPEI.gen(exp) exp.new_trial(generator_run=g).run() if i > 4: self.mock_torch_model_bridge.assert_called() else: self.assertEqual(g._model_key, "Sobol") self.assertEqual( g._model_kwargs, { "seed": None, "deduplicate": False, "init_position": i, "scramble": True, "generated_points": None, }, ) self.assertEqual( g._bridge_kwargs, { "optimization_config": None, "status_quo_features": None, "status_quo_name": None, "transform_configs": None, "transforms": Cont_X_trans, "fit_out_of_design": False, }, ) self.assertEqual(g._model_state_after_gen, {"init_position": i + 1}) # Check completeness error message when GS should be done. with self.assertRaises(GenerationStrategyCompleted): g = sobol_GPEI.gen(exp) def test_sobol_GPEI_strategy_keep_generating(self): exp = get_branin_experiment() sobol_GPEI_generation_strategy = GenerationStrategy( steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-1), ] ) self.assertEqual(sobol_GPEI_generation_strategy.name, "Sobol+GPEI") self.assertEqual(sobol_GPEI_generation_strategy.model_transitions, [5]) exp.new_trial(generator_run=sobol_GPEI_generation_strategy.gen(exp)).run() for i in range(1, 15): g = sobol_GPEI_generation_strategy.gen(exp) exp.new_trial(generator_run=g).run() if i > 4: self.assertIsInstance( sobol_GPEI_generation_strategy.model, TorchModelBridge ) @patch(f"{Experiment.__module__}.Experiment.fetch_data", return_value=get_data()) def test_factorial_thompson_strategy(self, _): exp = get_branin_experiment() factorial_thompson_generation_strategy = GenerationStrategy( steps=[ GenerationStep(model=Models.FACTORIAL, num_trials=1), GenerationStep(model=Models.THOMPSON, num_trials=-1), ] ) self.assertEqual( factorial_thompson_generation_strategy.name, "Factorial+Thompson" ) self.assertEqual(factorial_thompson_generation_strategy.model_transitions, [1]) mock_model_bridge = self.mock_discrete_model_bridge.return_value # Initial factorial batch. exp.new_batch_trial(factorial_thompson_generation_strategy.gen(experiment=exp)) args, kwargs = mock_model_bridge._set_kwargs_to_save.call_args self.assertEqual(kwargs.get("model_key"), "Factorial") # Subsequent Thompson sampling batch. exp.new_batch_trial(factorial_thompson_generation_strategy.gen(experiment=exp)) args, kwargs = mock_model_bridge._set_kwargs_to_save.call_args self.assertEqual(kwargs.get("model_key"), "Thompson") def test_clone_reset(self): ftgs = GenerationStrategy( steps=[ GenerationStep(model=Models.FACTORIAL, num_trials=1), GenerationStep(model=Models.THOMPSON, num_trials=2), ] ) ftgs._curr = ftgs._steps[1] self.assertEqual(ftgs._curr.index, 1) self.assertEqual(ftgs.clone_reset()._curr.index, 0) def test_kwargs_passed(self): gs = GenerationStrategy( steps=[ GenerationStep( model=Models.SOBOL, num_trials=1, model_kwargs={"scramble": False} ) ] ) exp = get_branin_experiment() gs.gen(exp) self.assertFalse(gs._model.model.scramble) def test_sobol_GPEI_strategy_batches(self): mock_GPEI_gen = self.mock_torch_model_bridge.return_value.gen mock_GPEI_gen.return_value = GeneratorRun( arms=[ Arm(parameters={"x1": 1, "x2": 2}), Arm(parameters={"x1": 3, "x2": 4}), ] ) exp = get_branin_experiment() sobol_GPEI_generation_strategy = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep(model=Models.SOBOL, num_trials=1), GenerationStep(model=Models.GPEI, num_trials=6), ], ) self.assertEqual(sobol_GPEI_generation_strategy.name, "Sobol+GPEI") self.assertEqual(sobol_GPEI_generation_strategy.model_transitions, [1]) gr = sobol_GPEI_generation_strategy.gen(exp, n=2) exp.new_batch_trial(generator_run=gr).run() for i in range(1, 8): if i == 7: # Check completeness error message. with self.assertRaises(GenerationStrategyCompleted): g = sobol_GPEI_generation_strategy.gen(exp, n=2) else: g = sobol_GPEI_generation_strategy.gen(exp, n=2) exp.new_batch_trial(generator_run=g).run() self.assertIsInstance(sobol_GPEI_generation_strategy.model, TorchModelBridge) def test_with_factory_function(self): """Checks that generation strategy works with custom factory functions. No information about the model should be saved on generator run.""" def get_sobol(search_space: SearchSpace) -> RandomModelBridge: return RandomModelBridge( search_space=search_space, model=SobolGenerator(), transforms=Cont_X_trans, ) exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=get_sobol, num_trials=5)] ) g = sobol_generation_strategy.gen(exp) self.assertIsInstance(sobol_generation_strategy.model, RandomModelBridge) self.assertIsNone(g._model_key) self.assertIsNone(g._model_kwargs) self.assertIsNone(g._bridge_kwargs) def test_store_experiment(self): exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=5)] ) self.assertIsNone(sobol_generation_strategy._experiment) sobol_generation_strategy.gen(exp) self.assertIsNotNone(sobol_generation_strategy._experiment) def test_trials_as_df(self): exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=5)] ) # No trials yet, so the DF will be None. self.assertIsNone(sobol_generation_strategy.trials_as_df) # Now the trial should appear in the DF. trial = exp.new_trial(sobol_generation_strategy.gen(experiment=exp)) self.assertFalse(sobol_generation_strategy.trials_as_df.empty) self.assertEqual( sobol_generation_strategy.trials_as_df.head()["Trial Status"][0], "CANDIDATE", ) # Changes in trial status should be reflected in the DF. trial._status = TrialStatus.RUNNING self.assertEqual( sobol_generation_strategy.trials_as_df.head()["Trial Status"][0], "RUNNING" ) def test_max_parallelism_reached(self): exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=5, max_parallelism=1)] ) exp.new_trial( generator_run=sobol_generation_strategy.gen(experiment=exp) ).mark_running(no_runner_required=True) with self.assertRaises(MaxParallelismReachedException): sobol_generation_strategy.gen(experiment=exp) @patch(f"{RandomModelBridge.__module__}.RandomModelBridge.update") @patch(f"{Experiment.__module__}.Experiment.fetch_trials_data") def test_use_update(self, mock_fetch_trials_data, mock_update): exp = get_branin_experiment() sobol_gs_with_update = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=-1, use_update=True)] ) # Try without passing data (generation strategy fetches data from 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): trial = exp.new_trial( generator_run=sobol_gs_with_update.gen(experiment=exp) ) self.assertEqual( mock_fetch_trials_data.call_args[1].get("trial_indices"), {i} ) trial._status = TrialStatus.COMPLETED # Try with passing data. sobol_gs_with_update.gen( experiment=exp, data=get_branin_data(trial_indices=range(4)) ) # 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} )
class TestGenerationStrategy(TestCase): def setUp(self): self.gr = GeneratorRun(arms=[Arm(parameters={"x1": 1, "x2": 2})]) # Mock out slow GPEI. self.torch_model_bridge_patcher = patch( f"{TorchModelBridge.__module__}.TorchModelBridge", spec=True) self.mock_torch_model_bridge = self.torch_model_bridge_patcher.start() self.mock_torch_model_bridge.return_value.gen.return_value = self.gr # Mock out slow TS. self.discrete_model_bridge_patcher = patch( f"{DiscreteModelBridge.__module__}.DiscreteModelBridge", spec=True) self.mock_discrete_model_bridge = self.discrete_model_bridge_patcher.start( ) self.mock_discrete_model_bridge.return_value.gen.return_value = self.gr # Mock in `Models` registry self.registry_setup_dict_patcher = patch.dict( f"{Models.__module__}.MODEL_KEY_TO_MODEL_SETUP", { "Factorial": MODEL_KEY_TO_MODEL_SETUP["Factorial"]._replace( bridge_class=self.mock_discrete_model_bridge), "Thompson": MODEL_KEY_TO_MODEL_SETUP["Thompson"]._replace( bridge_class=self.mock_discrete_model_bridge), "GPEI": MODEL_KEY_TO_MODEL_SETUP["GPEI"]._replace( bridge_class=self.mock_torch_model_bridge), }, ) self.mock_in_registry = self.registry_setup_dict_patcher.start() # model bridges are mocked, which makes kwargs' validation difficult, # so for now we will skip it in the generation strategy tests. # NOTE: Starting with Python3.8 this is not a problem as `autospec=True` # ensures that the mocks have correct signatures, but in earlier # versions kwarg validation on mocks does not really work. self.step_model_kwargs = {"silently_filter_kwargs": True} self.hss_experiment = get_hierarchical_search_space_experiment() self.sobol_GPEI_GS = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep( model=Models.SOBOL, num_trials=5, model_kwargs=self.step_model_kwargs, ), GenerationStep(model=Models.GPEI, num_trials=2, model_kwargs=self.step_model_kwargs), ], ) self.sobol_GS = GenerationStrategy(steps=[ GenerationStep( Models.SOBOL, num_trials=-1, should_deduplicate=True, ) ]) def tearDown(self): self.torch_model_bridge_patcher.stop() self.discrete_model_bridge_patcher.stop() self.registry_setup_dict_patcher.stop() def test_name(self): self.sobol_GS.name = "SomeGSName" self.assertEqual(self.sobol_GS.name, "SomeGSName") def test_validation(self): # num_trials can be positive or -1. with self.assertRaises(UserInputError): GenerationStrategy(steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-10), ]) # only last num_trials can be -1. with self.assertRaises(UserInputError): GenerationStrategy(steps=[ GenerationStep(model=Models.SOBOL, num_trials=-1), GenerationStep(model=Models.GPEI, num_trials=10), ]) exp = Experiment( name="test", search_space=SearchSpace(parameters=[get_choice_parameter()])) factorial_thompson_generation_strategy = GenerationStrategy(steps=[ GenerationStep(model=Models.FACTORIAL, num_trials=1), GenerationStep(model=Models.THOMPSON, num_trials=2), ]) self.assertTrue( factorial_thompson_generation_strategy._uses_registered_models) self.assertFalse( factorial_thompson_generation_strategy.uses_non_registered_models) with self.assertRaises(ValueError): factorial_thompson_generation_strategy.gen(exp) self.assertEqual( GenerationStep(model=sum, num_trials=1).model_name, "sum") with self.assertRaisesRegex(UserInputError, "Maximum parallelism should be"): GenerationStrategy(steps=[ GenerationStep( model=Models.SOBOL, num_trials=5, max_parallelism=-1), GenerationStep(model=Models.GPEI, num_trials=-1), ]) def test_custom_callables_for_models(self): exp = get_branin_experiment() sobol_factory_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=get_sobol, num_trials=-1)]) self.assertFalse( sobol_factory_generation_strategy._uses_registered_models) self.assertTrue( sobol_factory_generation_strategy.uses_non_registered_models) gr = sobol_factory_generation_strategy.gen(experiment=exp, n=1) self.assertEqual(len(gr.arms), 1) def test_string_representation(self): gs1 = GenerationStrategy(steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-1), ]) self.assertEqual( str(gs1), ("GenerationStrategy(name='Sobol+GPEI', steps=[Sobol for 5 trials," " GPEI for subsequent trials])"), ) gs2 = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=-1)]) self.assertEqual( str(gs2), "GenerationStrategy(name='Sobol', steps=[Sobol for all trials])") def test_equality(self): gs1 = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-1), ], ) gs2 = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep(model=Models.SOBOL, num_trials=5), GenerationStep(model=Models.GPEI, num_trials=-1), ], ) self.assertEqual(gs1, gs2) # Clone_reset() doesn't clone exactly, so they won't be equal. gs3 = gs1.clone_reset() self.assertEqual(gs1, gs3) 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 test_min_observed(self): # We should fail to transition the next model if there is not # enough data observed. exp = get_branin_experiment(get_branin_experiment()) gs = GenerationStrategy(steps=[ GenerationStep( model=Models.SOBOL, num_trials=5, min_trials_observed=5), GenerationStep(model=Models.GPEI, num_trials=1), ]) self.assertFalse(gs.uses_non_registered_models) for _ in range(5): exp.new_trial(gs.gen(exp)) with self.assertRaises(DataRequiredError): gs.gen(exp) def test_do_not_enforce_min_observations(self): # We should be able to move on to the next model if there is not # enough data observed if `enforce_num_trials` setting is False, in which # case the previous model should be used until there is enough data. exp = get_branin_experiment() gs = GenerationStrategy(steps=[ GenerationStep( model=Models.SOBOL, num_trials=1, min_trials_observed=5, enforce_num_trials=False, ), GenerationStep(model=Models.GPEI, num_trials=1), ]) for _ in range(2): gs.gen(exp) # Make sure Sobol is used to generate the 6th point. self.assertIsInstance(gs._model, RandomModelBridge) def test_sobol_GPEI_strategy(self): exp = get_branin_experiment() self.assertEqual(self.sobol_GPEI_GS.name, "Sobol+GPEI") self.assertEqual(self.sobol_GPEI_GS.model_transitions, [5]) for i in range(7): g = self.sobol_GPEI_GS.gen(exp) exp.new_trial(generator_run=g).run() self.assertEqual(len(self.sobol_GPEI_GS._generator_runs), i + 1) if i > 4: self.mock_torch_model_bridge.assert_called() else: self.assertEqual(g._model_key, "Sobol") self.assertEqual( g._model_kwargs, { "seed": None, "deduplicate": False, "init_position": i, "scramble": True, "generated_points": None, "fallback_to_sample_polytope": False, }, ) self.assertEqual( g._bridge_kwargs, { "optimization_config": None, "status_quo_features": None, "status_quo_name": None, "transform_configs": None, "transforms": Cont_X_trans, "fit_out_of_design": False, "fit_abandoned": False, }, ) self.assertEqual(g._model_state_after_gen, {"init_position": i + 1}) # Check completeness error message when GS should be done. with self.assertRaises(GenerationStrategyCompleted): g = self.sobol_GPEI_GS.gen(exp) def test_sobol_GPEI_strategy_keep_generating(self): exp = get_branin_experiment() sobol_GPEI_generation_strategy = GenerationStrategy(steps=[ GenerationStep( model=Models.SOBOL, num_trials=5, model_kwargs=self.step_model_kwargs, ), GenerationStep( model=Models.GPEI, num_trials=-1, model_kwargs=self.step_model_kwargs, ), ]) self.assertEqual(sobol_GPEI_generation_strategy.name, "Sobol+GPEI") self.assertEqual(sobol_GPEI_generation_strategy.model_transitions, [5]) exp.new_trial( generator_run=sobol_GPEI_generation_strategy.gen(exp)).run() for i in range(1, 15): g = sobol_GPEI_generation_strategy.gen(exp) exp.new_trial(generator_run=g).run() if i > 4: self.assertIsInstance(sobol_GPEI_generation_strategy.model, TorchModelBridge) def test_sobol_strategy(self): exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy(steps=[ GenerationStep( model=Models.SOBOL, num_trials=5, max_parallelism=10, use_update=False, enforce_num_trials=False, ) ]) for i in range(1, 6): sobol_generation_strategy.gen(exp, n=1) self.assertEqual(len(sobol_generation_strategy._generator_runs), i) @patch(f"{Experiment.__module__}.Experiment.fetch_data", return_value=get_data()) def test_factorial_thompson_strategy(self, _): exp = get_branin_experiment() factorial_thompson_generation_strategy = GenerationStrategy(steps=[ GenerationStep( model=Models.FACTORIAL, num_trials=1, model_kwargs=self.step_model_kwargs, ), GenerationStep( model=Models.THOMPSON, num_trials=-1, model_kwargs=self.step_model_kwargs, ), ]) self.assertEqual(factorial_thompson_generation_strategy.name, "Factorial+Thompson") self.assertEqual( factorial_thompson_generation_strategy.model_transitions, [1]) mock_model_bridge = self.mock_discrete_model_bridge.return_value # Initial factorial batch. exp.new_batch_trial( factorial_thompson_generation_strategy.gen(experiment=exp)) args, kwargs = mock_model_bridge._set_kwargs_to_save.call_args self.assertEqual(kwargs.get("model_key"), "Factorial") # Subsequent Thompson sampling batch. exp.new_batch_trial( factorial_thompson_generation_strategy.gen(experiment=exp)) args, kwargs = mock_model_bridge._set_kwargs_to_save.call_args self.assertEqual(kwargs.get("model_key"), "Thompson") def test_clone_reset(self): ftgs = GenerationStrategy(steps=[ GenerationStep(model=Models.FACTORIAL, num_trials=1), GenerationStep(model=Models.THOMPSON, num_trials=2), ]) ftgs._curr = ftgs._steps[1] self.assertEqual(ftgs._curr.index, 1) self.assertEqual(ftgs.clone_reset()._curr.index, 0) def test_kwargs_passed(self): gs = GenerationStrategy(steps=[ GenerationStep(model=Models.SOBOL, num_trials=1, model_kwargs={"scramble": False}) ]) exp = get_branin_experiment() gs.gen(exp) self.assertFalse(gs._model.model.scramble) def test_sobol_GPEI_strategy_batches(self): mock_GPEI_gen = self.mock_torch_model_bridge.return_value.gen mock_GPEI_gen.return_value = GeneratorRun(arms=[ Arm(parameters={ "x1": 1, "x2": 2 }), Arm(parameters={ "x1": 3, "x2": 4 }), ]) exp = get_branin_experiment() sobol_GPEI_generation_strategy = GenerationStrategy( name="Sobol+GPEI", steps=[ GenerationStep( model=Models.SOBOL, num_trials=1, model_kwargs=self.step_model_kwargs, ), GenerationStep(model=Models.GPEI, num_trials=6, model_kwargs=self.step_model_kwargs), ], ) self.assertEqual(sobol_GPEI_generation_strategy.name, "Sobol+GPEI") self.assertEqual(sobol_GPEI_generation_strategy.model_transitions, [1]) gr = sobol_GPEI_generation_strategy.gen(exp, n=2) exp.new_batch_trial(generator_run=gr).run() for i in range(1, 8): if i == 7: # Check completeness error message. with self.assertRaises(GenerationStrategyCompleted): g = sobol_GPEI_generation_strategy.gen(exp, n=2) else: g = sobol_GPEI_generation_strategy.gen(exp, n=2) exp.new_batch_trial(generator_run=g).run() self.assertIsInstance(sobol_GPEI_generation_strategy.model, TorchModelBridge) def test_with_factory_function(self): """Checks that generation strategy works with custom factory functions. No information about the model should be saved on generator run.""" def get_sobol(search_space: SearchSpace) -> RandomModelBridge: return RandomModelBridge( search_space=search_space, model=SobolGenerator(), transforms=Cont_X_trans, ) exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=get_sobol, num_trials=5)]) g = sobol_generation_strategy.gen(exp) self.assertIsInstance(sobol_generation_strategy.model, RandomModelBridge) self.assertIsNone(g._model_key) self.assertIsNone(g._model_kwargs) self.assertIsNone(g._bridge_kwargs) def test_store_experiment(self): exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=5)]) self.assertIsNone(sobol_generation_strategy._experiment) sobol_generation_strategy.gen(exp) self.assertIsNotNone(sobol_generation_strategy._experiment) def test_trials_as_df(self): exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy( steps=[GenerationStep(model=Models.SOBOL, num_trials=5)]) # No trials yet, so the DF will be None. self.assertIsNone(sobol_generation_strategy.trials_as_df) # Now the trial should appear in the DF. trial = exp.new_trial(sobol_generation_strategy.gen(experiment=exp)) self.assertFalse(sobol_generation_strategy.trials_as_df.empty) self.assertEqual( sobol_generation_strategy.trials_as_df.head()["Trial Status"][0], "CANDIDATE", ) # Changes in trial status should be reflected in the DF. trial._status = TrialStatus.RUNNING self.assertEqual( sobol_generation_strategy.trials_as_df.head()["Trial Status"][0], "RUNNING") def test_max_parallelism_reached(self): exp = get_branin_experiment() sobol_generation_strategy = GenerationStrategy(steps=[ GenerationStep(model=Models.SOBOL, num_trials=5, max_parallelism=1) ]) exp.new_trial(generator_run=sobol_generation_strategy.gen( experiment=exp)).mark_running(no_runner_required=True) with self.assertRaises(MaxParallelismReachedException): sobol_generation_strategy.gen(experiment=exp) @patch(f"{RandomModelBridge.__module__}.RandomModelBridge.update") @patch(f"{Experiment.__module__}.Experiment.lookup_data") 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_deduplication(self): tiny_parameters = [ FixedParameter( name="x1", parameter_type=ParameterType.FLOAT, value=1.0, ), ChoiceParameter( name="x2", parameter_type=ParameterType.FLOAT, values=[float(x) for x in range(2)], ), ] tiny_search_space = SearchSpace( parameters=cast(List[Parameter], tiny_parameters)) exp = get_branin_experiment(search_space=tiny_search_space) sobol = GenerationStrategy( name="Sobol", steps=[ GenerationStep( model=Models.SOBOL, num_trials=-1, model_kwargs=self.step_model_kwargs, should_deduplicate=True, ), ], ) for _ in range(2): g = sobol.gen(exp) exp.new_trial(generator_run=g).run() self.assertEqual(len(exp.arms_by_signature), 2) with self.assertRaisesRegex(GenerationStrategyRepeatedPoints, "exceeded `MAX_GEN_DRAWS`"): g = sobol.gen(exp) 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=(NUM_ROUNDS - 1) * SECOND_STEP_PARALLELISM, 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) # Optimization should now be complete. ( num_trials_to_gen, opt_complete, ) = sobol_gs_with_parallelism_limits.current_generator_run_limit() self.assertTrue(opt_complete) self.assertEqual(num_trials_to_gen, 0) # 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_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_hierarchical_search_space(self): experiment = get_hierarchical_search_space_experiment() self.assertTrue(experiment.search_space.is_hierarchical) self.sobol_GS.gen(experiment=experiment) for _ in range(10): # During each iteration, check that all transformed observation features # contain all parameters of the flat search space. with patch.object(RandomModelBridge, "_fit") as mock_model_fit, patch.object( RandomModelBridge, "gen"): self.sobol_GS.gen(experiment=experiment) mock_model_fit.assert_called_once() obs_feats = mock_model_fit.call_args[1].get( "observation_features") all_parameter_names = ( experiment.search_space._all_parameter_names.copy()) # One of the parameter names is modified by transforms (because it's # one-hot encoded). all_parameter_names.remove("model") all_parameter_names.add("model_OH_PARAM_") for obsf in obs_feats: for p_name in all_parameter_names: self.assertIn(p_name, obsf.parameters) trial = (experiment.new_trial(generator_run=self.sobol_GS.gen( experiment=experiment)).mark_running( no_runner_required=True).mark_completed()) experiment.attach_data( get_data( metric_name="m1", trial_index=trial.index, num_non_sq_arms=1, include_sq=False, )) experiment.attach_data( get_data( metric_name="m2", trial_index=trial.index, num_non_sq_arms=1, include_sq=False, )) # ------------- Testing helpers (put tests above this line) ------------- def _run_GS_for_N_rounds(self, gs: GenerationStrategy, exp: Experiment, num_rounds: int) -> List[int]: could_gen = [] for _ in range(num_rounds): ( num_trials_to_gen, opt_complete, ) = gs.current_generator_run_limit() self.assertFalse(opt_complete) could_gen.append(num_trials_to_gen) trials = [] for _ in range(num_trials_to_gen): gr = gs.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() return could_gen