From c5e63ce018fa85fbae667481541ea4b441ea8d1f Mon Sep 17 00:00:00 2001 From: Cesar Cardoso Date: Wed, 17 Apr 2024 14:29:13 -0700 Subject: [PATCH] Support fixed features in Service API (#2372) Summary: Add the possibility of specifying some `FixedFeatures` as `fixed_features` in `AxClient.get_next_trial` and `AxClient.get_next_trials` which is currently only possible with the developer API. Reviewed By: saitcakmak Differential Revision: D56068035 --- ax/service/ax_client.py | 49 +++++++++++++++++++----------- ax/service/tests/test_ax_client.py | 20 +++++++++--- 2 files changed, 47 insertions(+), 22 deletions(-) diff --git a/ax/service/ax_client.py b/ax/service/ax_client.py index 7d1a9524a34..23666cde9be 100644 --- a/ax/service/ax_client.py +++ b/ax/service/ax_client.py @@ -495,7 +495,10 @@ def set_search_space( wrap_error_message_in=CHOLESKY_ERROR_ANNOTATION, ) def get_next_trial( - self, ttl_seconds: Optional[int] = None, force: bool = False + self, + ttl_seconds: Optional[int] = None, + force: bool = False, + fixed_features: Optional[FixedFeatures] = None, ) -> Tuple[TParameterization, int]: """ Generate trial with the next set of parameters to try in the iteration process. @@ -508,6 +511,9 @@ def get_next_trial( failed properly. force: If set to True, this function will bypass the global stopping strategy's decision and generate a new trial anyway. + fixed_features: A FixedFeatures object containing any + features that should be fixed at specified values during + generation. Returns: Tuple of trial parameterization, trial index @@ -530,7 +536,10 @@ def get_next_trial( try: trial = self.experiment.new_trial( - generator_run=self._gen_new_generator_run(), ttl_seconds=ttl_seconds + generator_run=self._gen_new_generator_run( + fixed_features=fixed_features + ), + ttl_seconds=ttl_seconds, ) except MaxParallelismReachedException as e: if self._early_stopping_strategy is not None: @@ -580,7 +589,10 @@ def get_current_trial_generation_limit(self) -> Tuple[int, bool]: return self.generation_strategy.current_generator_run_limit() def get_next_trials( - self, max_trials: int, ttl_seconds: Optional[int] = None + self, + max_trials: int, + ttl_seconds: Optional[int] = None, + fixed_features: Optional[FixedFeatures] = None, ) -> Tuple[Dict[int, TParameterization], bool]: """Generate as many trials as currently possible. @@ -597,6 +609,9 @@ def get_next_trials( ttl_seconds: If specified, will consider the trial failed after this many seconds. Used to detect dead trials that were not marked failed properly. + fixed_features: A FixedFeatures object containing any + features that should be fixed at specified values during + generation. Returns: two-item tuple of: - mapping from trial indices to parameterizations in those trials, @@ -616,7 +631,9 @@ def get_next_trials( trials_dict = {} for _ in range(max_trials): try: - params, trial_index = self.get_next_trial(ttl_seconds=ttl_seconds) + params, trial_index = self.get_next_trial( + ttl_seconds=ttl_seconds, fixed_features=fixed_features + ) trials_dict[trial_index] = params except OptimizationComplete as err: logger.info( @@ -1744,20 +1761,16 @@ def _save_generation_strategy_to_db_if_possible( suppress_all_errors=suppress_all_errors, ) - def _get_last_completed_trial_index(self) -> int: - # infer last completed trial as the trial_index to use - # TODO: use Experiment.completed_trials once D46484953 lands. - completed_indices = [ - t.index for t in self.experiment.trials_by_status[TrialStatus.COMPLETED] - ] - completed_indices.append(0) # handle case of no completed trials - return max(completed_indices) - - def _gen_new_generator_run(self, n: int = 1) -> GeneratorRun: + def _gen_new_generator_run( + self, n: int = 1, fixed_features: Optional[FixedFeatures] = None + ) -> GeneratorRun: """Generate new generator run for this experiment. Args: n: Number of arms to generate. + fixed_features: A FixedFeatures object containing any + features that should be fixed at specified values during + generation. """ # If random seed is not set for this optimization, context manager does # nothing; otherwise, it sets the random seed for torch, but only for the @@ -1766,10 +1779,12 @@ def _gen_new_generator_run(self, n: int = 1) -> GeneratorRun: # serious negative impact on the performance of the models that employ # stochasticity. - fixed_feats = InstantiationBase.make_fixed_observation_features( - fixed_features=FixedFeatures( - parameters={}, trial_index=self._get_last_completed_trial_index() + fixed_feats = ( + InstantiationBase.make_fixed_observation_features( + fixed_features=fixed_features ) + if fixed_features + else None ) with manual_seed(seed=self._random_seed): return not_none(self.generation_strategy).gen( diff --git a/ax/service/tests/test_ax_client.py b/ax/service/tests/test_ax_client.py index d232efafed8..549d511f8dc 100644 --- a/ax/service/tests/test_ax_client.py +++ b/ax/service/tests/test_ax_client.py @@ -65,6 +65,7 @@ observed_pareto, predicted_pareto, ) +from ax.service.utils.instantiation import FixedFeatures from ax.storage.sqa_store.db import init_test_engine_and_session_factory from ax.storage.sqa_store.decoder import Decoder from ax.storage.sqa_store.encoder import Encoder @@ -2847,11 +2848,20 @@ def test_gen_fixed_features(self) -> None: with mock.patch.object( GenerationStrategy, "gen", wraps=ax_client.generation_strategy.gen ) as mock_gen: - params, idx = ax_client.get_next_trial() - call_kwargs = mock_gen.call_args_list[0][1] - ff = call_kwargs["fixed_features"] - self.assertEqual(ff.parameters, {}) - self.assertEqual(ff.trial_index, 0) + with self.subTest("fixed_features is None"): + params, idx = ax_client.get_next_trial() + call_kwargs = mock_gen.call_args_list[0][1] + ff = call_kwargs["fixed_features"] + self.assertIsNone(ff) + with self.subTest("fixed_features is set"): + fixed_features = FixedFeatures( + parameters={"x": 0.0, "y": 5.0}, trial_index=0 + ) + params, idx = ax_client.get_next_trial(fixed_features=fixed_features) + call_kwargs = mock_gen.call_args_list[1][1] + ff = call_kwargs["fixed_features"] + self.assertEqual(ff.parameters, fixed_features.parameters) + self.assertEqual(ff.trial_index, 0) def test_get_optimization_trace_discard_infeasible_trials(self) -> None: ax_client = AxClient()