Inference trace and Best Point Recommendation (BPR) bugfix (facebook#4128)

Summary:

This diff addresses two issues in the computation of inference trace:
1. The generation strategy is copied inside run_optimization_with_orchestrator --> we retrieve the traces on an unused generation strategy --> get_best_point defaults to the best raw observation on ALL obserations

2. Relevant data not filtered in the fallback option for get_best_parameters_from_model_predictions_with_trial_index

Both of these individually lead to the inference trace being incorrect - the first to the best raw value of ALL trials, the second to the best predicted across ALL trials.

Changes:
- (re?)-Moved copying of generation strategy
- Added argument use_model_only_if_good to force model-based BPR even if model fit is bad
- Clearer sequencing in get_best_parameters_from_model_predictions_with_trial_index

Differential Revision: D80019803

Author: Carl Hvarfner

ax/benchmark/benchmark.py

@@ -342,6 +342,9 @@ def get_best_parameters(
         experiment=experiment,
         generation_strategy=generation_strategy,
         trial_indices=trial_indices,
+        # disables the model quality check which determines whether to evaluate
+        # inference trace or raw observations when retrieving the best point
+        use_model_only_if_good=False,
     )
     if result is None:
         # This can happen if no points are predicted to satisfy all outcome
@@ -507,7 +510,7 @@ def run_optimization_with_orchestrator(
 
     orchestrator = Orchestrator(
         experiment=experiment,
-        generation_strategy=method.generation_strategy.clone_reset(),
+        generation_strategy=method.generation_strategy,
         options=orchestrator_options,
     )
 
@@ -562,6 +565,8 @@ def benchmark_replication(
     Return:
         ``BenchmarkResult`` object.
     """
+    # Reset the generation strategy to ensure that it is in an unused state.
+    method.generation_strategy = method.generation_strategy.clone_reset()
     experiment = run_optimization_with_orchestrator(
         problem=problem,
         method=method,

ax/benchmark/tests/test_benchmark.py

@@ -490,6 +490,7 @@ def test_run_optimization_with_orchestrator(self) -> None:
                 none_throws(runner.simulated_backend_runner).simulator._verbose_logging
             )
 
+        method.generation_strategy = method.generation_strategy.clone_reset()
         with self.subTest("Logs not produced by default"), self.assertNoLogs(
             level=logging.INFO, logger=logger
         ), self.assertNoLogs(logger=logger):
@@ -617,9 +618,9 @@ def test_early_stopping(self) -> None:
             self.assertEqual(max_run, {0: 4, 1: 2, 2: 2, 3: 2})
 
     def test_replication_variable_runtime(self) -> None:
-        method = get_async_benchmark_method(max_pending_trials=1)
         for map_data in [False, True]:
             with self.subTest(map_data=map_data):
+                method = get_async_benchmark_method(max_pending_trials=1)
                 problem = get_async_benchmark_problem(
                     map_data=map_data,
                     step_runtime_fn=lambda params: params["x0"] + 1,
@@ -1195,6 +1196,7 @@ def test_get_opt_trace_by_cumulative_epochs(self) -> None:
             ):
                 get_opt_trace_by_steps(experiment=experiment)
 
+        method.generation_strategy = method.generation_strategy.clone_reset()
         with self.subTest("Constrained"):
             problem = get_benchmark_problem("constrained_gramacy_observed_noise")
             experiment = self.run_optimization_with_orchestrator(

ax/service/tests/test_best_point_utils.py

@@ -10,7 +10,7 @@
 from itertools import product
 from typing import Any
 from unittest import mock
-from unittest.mock import patch, PropertyMock
+from unittest.mock import Mock, patch, PropertyMock
 
 import numpy as np
 
@@ -919,6 +919,209 @@ def test_get_best_point_with_model_prediction(
         self.assertEqual(best_params, params)
         self.assertEqual(predictions, ({"y": mock.ANY}, {"y": {"y": mock.ANY}}))
 
+    @mock_botorch_optimize
+    def test_get_best_parameters_from_model_predictions_with_trial_index(
+        self,
+    ) -> None:
+        # Setup experiment
+        exp = get_branin_experiment()
+        gs = choose_generation_strategy_legacy(
+            search_space=exp.search_space,
+            num_initialization_trials=3,
+            suggested_model_override=Generators.BOTORCH_MODULAR,
+        )
+
+        # Add some trials with data
+        for _ in range(4):
+            generator_run = gs.gen_single_trial(experiment=exp, n=1)
+            trial = exp.new_trial(generator_run=generator_run)
+            trial.run().mark_completed()
+            exp.attach_data(exp.fetch_data())
+
+        # Test 1: No adapter (None) - should fall back to generator run
+        with self.subTest("No adapter - fallback to generator run"):
+            result = get_best_parameters_from_model_predictions_with_trial_index(
+                experiment=exp, adapter=None
+            )
+            self.assertIsNotNone(result)
+            trial_index, params, _ = none_throws(result)
+            self.assertIsInstance(trial_index, int)
+            self.assertIsInstance(params, dict)
+
+        # Test 2: Non-TorchAdapter - should fall back to generator run
+        # Then, the recommendation should be in-sample
+        with self.subTest("Non-TorchAdapter - fallback to generator run"):
+            non_torch_adapter = Mock()  # Not a TorchAdapter
+            result = none_throws(
+                get_best_parameters_from_model_predictions_with_trial_index(
+                    experiment=exp, adapter=non_torch_adapter
+                )
+            )
+            arm_params = result[1]
+            self.assertTrue(
+                arm_params in [v.parameters for v in exp.arms_by_name.values()]
+            )
+
+        # Test 3: TorchAdapter with use_model_only_if_good=False ->
+        # skip model fit check
+        with self.subTest("TorchAdapter with use_model_only_if_good=False"):
+            with patch.object(
+                TorchAdapter,
+                "model_best_point",
+                return_value=(
+                    exp.trials[0].arms[0],
+                    ({"branin": 1.0}, {"branin": {"branin": 0.1}}),
+                ),
+            ) as mock_model_best_point, patch(
+                "ax.service.utils.best_point.cross_validate"
+            ) as mock_cv:
+                result = get_best_parameters_from_model_predictions_with_trial_index(
+                    experiment=exp, adapter=gs.adapter, use_model_only_if_good=False
+                )
+
+                # Should not call cross_validate when consider_model_fit=False
+                mock_cv.assert_not_called()
+                mock_model_best_point.assert_called_once()
+                self.assertIsNotNone(result)
+
+        # Test 4: TorchAdapter with good model fit - should use adapter
+        with self.subTest("TorchAdapter with good model fit"):
+            with patch.object(
+                TorchAdapter,
+                "model_best_point",
+                return_value=(
+                    exp.trials[0].arms[0],
+                    ({"branin": 1.0}, {"branin": {"branin": 0.1}}),
+                ),
+            ) as mock_model_best_point, patch(
+                "ax.service.utils.best_point.assess_model_fit",
+                return_value=AssessModelFitResult(
+                    good_fit_metrics_to_fisher_score={"branin": 1.0},
+                    bad_fit_metrics_to_fisher_score={},
+                ),
+            ), self.assertLogs(logger=best_point_logger, level="INFO") as lg:
+                result = get_best_parameters_from_model_predictions_with_trial_index(
+                    experiment=exp, adapter=gs.adapter, use_model_only_if_good=True
+                )
+
+                mock_model_best_point.assert_called_once()
+
+                # Should log that model fit is acceptable
+                self.assertTrue(
+                    any("Model fit is acceptable" in log for log in lg.output),
+                    msg=lg.output,
+                )
+
+                self.assertIsNotNone(result)
+
+        # Test 5: TorchAdapter with bad model fit - should fall back to raw data
+        with self.subTest("TorchAdapter with bad model fit"):
+            with patch.object(
+                TorchAdapter,
+                "model_best_point",
+                return_value=(
+                    exp.trials[0].arms[0],
+                    ({"branin": 1.0}, {"branin": {"branin": 0.1}}),
+                ),
+            ) as mock_model_best_point, patch(
+                "ax.service.utils.best_point.assess_model_fit",
+                return_value=AssessModelFitResult(
+                    good_fit_metrics_to_fisher_score={},
+                    bad_fit_metrics_to_fisher_score={"branin": 0.1},
+                ),
+            ), patch(
+                "ax.service.utils.best_point.get_best_by_raw_objective_with_trial_index",
+                return_value=(0, {"x1": 1.0, "x2": 2.0}, ({"branin": 5.0}, {})),
+            ) as mock_raw_best, self.assertLogs(
+                logger=best_point_logger, level="WARN"
+            ) as lg:
+                result = get_best_parameters_from_model_predictions_with_trial_index(
+                    experiment=exp, adapter=gs.adapter, use_model_only_if_good=True
+                )
+
+                # Should not call model_best_point when model fit is bad
+                mock_model_best_point.assert_not_called()
+                # Should call raw objective fallbacak
+                mock_raw_best.assert_called_once()
+
+                # Should log warning about poor model fit
+                self.assertTrue(
+                    any("Model fit is poor" in log for log in lg.output), msg=lg.output
+                )
+
+                self.assertIsNotNone(result)
+
+        # Test 6: TorchAdapter with model_best_point returning None -> fall back to GR
+        with self.subTest("TorchAdapter with model_best_point returning None"):
+            with patch.object(
+                TorchAdapter, "model_best_point", return_value=None
+            ) as mock_model_best_point, patch(
+                "ax.service.utils.best_point.assess_model_fit",
+                return_value=AssessModelFitResult(
+                    good_fit_metrics_to_fisher_score={"branin": 1.0},
+                    bad_fit_metrics_to_fisher_score={},
+                ),
+            ):
+                result = get_best_parameters_from_model_predictions_with_trial_index(
+                    experiment=exp, adapter=gs.adapter, use_model_only_if_good=True
+                )
+
+                mock_model_best_point.assert_called_once()
+                # Should still return a result from generator run fallback
+                self.assertIsNotNone(result)
+
+        # Test 7: No generator run available - should return None
+        with self.subTest("No generator run available"):
+            # Create experiment with no generator runs
+            empty_exp = get_branin_experiment()
+            empty_exp.new_trial().run().mark_completed()
+
+            result = get_best_parameters_from_model_predictions_with_trial_index(
+                experiment=empty_exp, adapter=None
+            )
+            self.assertIsNone(result)
+
+        # Test 10: Trial indices subset - should work with subset of data
+        with self.subTest("Trial indices subset"):
+            trial_indices = [0, 1]  # Only use first two trials
+
+            with patch.object(
+                TorchAdapter,
+                "model_best_point",
+                return_value=(
+                    exp.trials[0].arms[0],
+                    ({"branin": 1.0}, {"branin": {"branin": 0.1}}),
+                ),
+            ):
+                result = get_best_parameters_from_model_predictions_with_trial_index(
+                    experiment=exp,
+                    adapter=gs.adapter,
+                    trial_indices=trial_indices,
+                    use_model_only_if_good=False,
+                )
+
+                self.assertIsNotNone(result)
+
+        # Test 11: Noisy data with poor model fit - should log additional warning
+        with self.subTest("Noisy data with poor model fit"):
+            with patch(
+                "ax.service.utils.best_point.assess_model_fit",
+                return_value=AssessModelFitResult(
+                    good_fit_metrics_to_fisher_score={},
+                    bad_fit_metrics_to_fisher_score={"branin": 0.1},
+                ),
+            ), patch(
+                "ax.service.utils.best_point._is_all_noiseless",
+                return_value=False,  # Simulate noisy data
+            ), patch(
+                "ax.service.utils.best_point"
+                ".get_best_by_raw_objective_with_trial_index",
+                return_value=(0, {"x1": 1.0, "x2": 2.0}, ({"branin": 5.0}, {})),
+            ):
+                result = get_best_parameters_from_model_predictions_with_trial_index(
+                    experiment=exp, adapter=gs.adapter, use_model_only_if_good=True
+                )
+
 
 def _repeat_elements(list_to_replicate: list[Any], n_repeats: int) -> pd.Series:
     return pd.Series([item for item in list_to_replicate for _ in range(n_repeats)])