From c3cdd3b86612b3f49361ed919d6a5bcb307478aa Mon Sep 17 00:00:00 2001 From: Miles Olson Date: Mon, 9 Dec 2024 08:02:45 -0800 Subject: [PATCH] Implement get_best_arm, get_pareto_frontier (#3142) Summary: Implement "best point" functionality in Client. A few notes: * Renamed get_best_trial --> get_best_arm: this is more accurate to what it is actually doing (calculating the best in sample point) and renaming will allow us to continue using this method when we create BatchClient * If GS is not on a predictive step we return {} for the prediction term in both get_best_arm and get_pareto_frontier as well as log an error Differential Revision: D66702545 --- ax/preview/api/client.py | 162 ++++++++++++++++++++++++++-- ax/preview/api/tests/test_client.py | 153 ++++++++++++++++++++++++++ 2 files changed, 304 insertions(+), 11 deletions(-) diff --git a/ax/preview/api/client.py b/ax/preview/api/client.py index d8cb8edc28f..f3b71dd2f16 100644 --- a/ax/preview/api/client.py +++ b/ax/preview/api/client.py @@ -45,6 +45,7 @@ optimization_config_from_string, ) from ax.service.scheduler import Scheduler, SchedulerOptions +from ax.service.utils.best_point_mixin import BestPointMixin from ax.utils.common.logger import get_logger from ax.utils.common.random import with_rng_seed from pyre_extensions import assert_is_instance, none_throws @@ -667,29 +668,168 @@ def compute_analyses( return cards - def get_best_trial( + def get_best_arm( self, use_model_predictions: bool = True - ) -> tuple[int, TParameterization, TOutcome]: + ) -> tuple[str, TParameterization, TOutcome]: """ - Calculates the best in-sample trial. + Identifies the best parameterization tried in the experiment so far, the best + in-sample arm. + + If `use_model_predictions` is True, first attempts to do so with the model used + in optimization and its corresponding predictions if available. If + `use_model_predictions` is False or attempts to use the model fails, falls back + to the best raw objective based on the data fetched from the experiment. + + Arms which violate outcome constraints are not eligible to be the best arm. Returns: - - The index of the best trial - - The parameters of the best trial - - The metric values associated withthe best trial + - The name of the best arm + - The parameters of the best arm + - The metric values for the best arm. Uses model prediction if + use_model_predictions=True, otherwise returns observed data. """ - ... + + if len(self._none_throws_experiment().trials) < 1: + raise UnsupportedError( + "No trials have been run yet. Please run at least one trial before " + "calling get_best_arm." + ) + + # Note: Using BestPointMixin directly instead of inheriting to avoid exposing + # unwanted public methods + trial_index, parameters, _ = none_throws( + BestPointMixin._get_best_trial( + experiment=self._none_throws_experiment(), + # Requiring true GenerationStrategy here, ideally we will loosen this + # in the future + generation_strategy=assert_is_instance( + self._generation_strategy_or_choose(), GenerationStrategy + ), + use_model_predictions=use_model_predictions, + ) + ) + + arm = none_throws( + assert_is_instance( + self._none_throws_experiment().trials[trial_index], Trial + ).arm + ) + + if use_model_predictions: + try: + # pyre-fixme[6]: Core Ax allows users to specify TParameterization + # values as None but we do not allow this in the API. + prediction = self.predict(points=[parameters])[0] + except UnsupportedError: + logger.error( + "Model predictions are not available, returning empty prediction" + ) + + prediction = {} + else: + data_dict = ( + self._none_throws_experiment() + .lookup_data(trial_indices=[trial_index]) + .df.to_dict() + ) + + prediction = { + data_dict["metric_name"][i]: (data_dict["mean"][i], data_dict["sem"][i]) + for i in range(len(data_dict["metric_name"])) + } + + # pyre-fixme[7]: Core Ax allows users to specify TParameterization values as + # None but we do not allow this in the API. + return arm.name, parameters, prediction def get_pareto_frontier( self, use_model_predictions: bool = True - ) -> dict[int, tuple[TParameterization, TOutcome]]: + ) -> dict[str, tuple[TParameterization, TOutcome]]: """ Calculates the in-sample Pareto frontier. Returns: - A mapping of trial index to its parameterization and metric values. + A mapping of the arm name to the parameterization and predicted or observed + outcome. """ - ... + + if len(self._none_throws_experiment().trials) < 1: + raise UnsupportedError( + "No trials have been run yet. Please run at least one trial before " + "calling get_pareto_frontier." + ) + + frontier = BestPointMixin._get_pareto_optimal_parameters( + experiment=self._none_throws_experiment(), + # Requiring true GenerationStrategy here, ideally we will loosen this + # in the future + generation_strategy=assert_is_instance( + self._generation_strategy_or_choose(), GenerationStrategy + ), + use_model_predictions=use_model_predictions, + ) + + frontier_list = [*frontier.items()] + + arm_names = [ + none_throws( + assert_is_instance( + self._none_throws_experiment().trials[trial_index], Trial + ).arm + ).name + for trial_index, _ in frontier_list + ] + + if use_model_predictions: + try: + predictions = self.predict( + # pyre-fixme[6]: Core Ax allows users to specify TParameterization + # values as None but we do not allow this in the API. + points=[value[0] for _key, value in frontier_list] + ) + except UnsupportedError: + logger.error( + "Model predictions are not available, returning empty prediction" + ) + + predictions: list[TOutcome] = [{} for _ in frontier] + else: + predictions = [] + for trial_index in frontier.keys(): + data_dict = ( + self._none_throws_experiment() + .lookup_data(trial_indices=[trial_index]) + .df.to_dict() + ) + + predictions.append( + { + data_dict["metric_name"][i]: ( + data_dict["mean"][i], + data_dict["sem"][i], + ) + for i in range(len(data_dict["metric_name"])) + } + ) + + try: + predictions = self.predict( + # pyre-fixme[6]: Core Ax allows users to specify TParameterization + # values as None but we do not allow this in the API. + points=[value[0] for _key, value in frontier_list] + ) + except UnsupportedError: + logger.error( + "Model predictions are not available, returning empty prediction" + ) + predictions: list[TOutcome] = [{} for _ in frontier] + + # pyre-fixme[7]: Core Ax allows users to specify TParameterization + # values as None but we do not allow this in the API. + return { + arm_names[i]: (frontier_list[i][1][0], predictions[i]) + for i in range(len(frontier_list)) + } def predict( self, @@ -723,7 +863,7 @@ def predict( for parameters in points ] ) - except (UserInputError, AssertionError) as e: + except (NotImplementedError, AssertionError) as e: raise UnsupportedError( "Predicting with the GenerationStrategy's modelbridge failed. This " "could be because the current GenerationNode is not predictive -- try " diff --git a/ax/preview/api/tests/test_client.py b/ax/preview/api/tests/test_client.py index ca0ddaf8cd1..6df01eb8446 100644 --- a/ax/preview/api/tests/test_client.py +++ b/ax/preview/api/tests/test_client.py @@ -911,6 +911,159 @@ def test_compute_analyses(self) -> None: self.assertEqual(len(cards), 1) self.assertEqual(cards[0].name, "ParallelCoordinatesPlot") + def test_get_best_arm(self) -> None: + client = Client() + + client.configure_experiment( + experiment_config=ExperimentConfig( + parameters=[ + RangeParameterConfig( + name="x1", parameter_type=ParameterType.FLOAT, bounds=(-1, 1) + ), + ], + name="foo", + ) + ) + client.configure_optimization(objective="foo") + # Set num_initialization_trials=3 so we can reach a predictive GenerationNode + # quickly + client.configure_generation_strategy( + generation_strategy_config=GenerationStrategyConfig( + num_initialization_trials=3 + ) + ) + + with self.assertRaisesRegex(UnsupportedError, "No trials have been run yet"): + client.get_best_arm() + + for _ in range(3): + for index, parameters in client.get_next_trials(maximum_trials=1).items(): + client.complete_trial( + trial_index=index, + raw_data={ + "foo": assert_is_instance(parameters["x1"], float) ** 2, + }, + ) + + name, parameters, prediction = client.get_best_arm() + self.assertIn( + name, + [ + none_throws(assert_is_instance(trial, Trial).arm).name + for trial in client._none_throws_experiment().trials.values() + ], + ) + self.assertTrue( + client._none_throws_experiment().search_space.check_membership( + parameterization=parameters # pyre-ignore[6] + ) + ) + self.assertEqual(prediction, {}) # No prediction since we are still in Sobol + + # Run a non-Sobol trial + for index, parameters in client.get_next_trials(maximum_trials=1).items(): + client.complete_trial( + trial_index=index, + raw_data={ + "foo": assert_is_instance(parameters["x1"], float) ** 2, + }, + ) + name, parameters, prediction = client.get_best_arm() + self.assertIn( + name, + [ + none_throws(assert_is_instance(trial, Trial).arm).name + for trial in client._none_throws_experiment().trials.values() + ], + ) + self.assertTrue( + client._none_throws_experiment().search_space.check_membership( + parameterization=parameters # pyre-ignore[6] + ) + ) + self.assertEqual({*prediction.keys()}, {"foo"}) + + def test_get_pareto_frontier(self) -> None: + client = Client() + + client.configure_experiment( + experiment_config=ExperimentConfig( + parameters=[ + RangeParameterConfig( + name="x1", parameter_type=ParameterType.FLOAT, bounds=(-1, 1) + ), + ], + name="foo", + ) + ) + client.configure_optimization(objective="foo, bar") + # Set num_initialization_trials=3 so we can reach a predictive GenerationNode + # quickly + client.configure_generation_strategy( + generation_strategy_config=GenerationStrategyConfig( + num_initialization_trials=3 + ) + ) + + with self.assertRaisesRegex(UnsupportedError, "No trials have been run yet"): + client.get_pareto_frontier() + + for _ in range(3): + for index, parameters in client.get_next_trials(maximum_trials=1).items(): + client.complete_trial( + trial_index=index, + raw_data={ + "foo": assert_is_instance(parameters["x1"], float) ** 2, + "bar": 0.0, + }, + ) + + frontier = client.get_pareto_frontier(False) + for name, point in frontier.items(): + parameters, prediction = point + + self.assertIn( + name, + [ + none_throws(assert_is_instance(trial, Trial).arm).name + for trial in client._none_throws_experiment().trials.values() + ], + ) + self.assertTrue( + client._none_throws_experiment().search_space.check_membership( + parameterization=parameters # pyre-ignore[6] + ) + ) + self.assertEqual( + prediction, {} + ) # No prediction since we are still in Sobol + + # Run a non-Sobol trial + for index, parameters in client.get_next_trials(maximum_trials=1).items(): + client.complete_trial( + trial_index=index, + raw_data={ + "foo": assert_is_instance(parameters["x1"], float) ** 2, + "bar": 0.0, + }, + ) + frontier = client.get_pareto_frontier() + for name, point in frontier.items(): + parameters, prediction = point + self.assertIn( + name, + [ + none_throws(assert_is_instance(trial, Trial).arm).name + for trial in client._none_throws_experiment().trials.values() + ], + ) + self.assertTrue( + client._none_throws_experiment().search_space.check_membership( + parameterization=parameters # pyre-ignore[6] + ) + ) + self.assertEqual({*prediction.keys()}, {"foo", "bar"}) + @mock_botorch_optimize def test_predict(self) -> None: client = Client()