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:
- Moved copying of generation strategy to the level`benchmark_replication`, since results need to be computed on the used `generation_strategy` and not an empty copy. This means that `run_optimization_with_orchestrator` no longer `clone_and_reset`'s the GS.
- Clearer sequencing in get_best_parameters_from_model_predictions_with_trial_index
- Removed model fit quality check as part of BPR

Previous, redacted changes:
- Added argument use_model_only_if_good to force model-based BPR even if model fit is bad

Differential Revision: D80019803

Author: Carl Hvarfner

ax/benchmark/benchmark.py

@@ -51,8 +51,10 @@
 from ax.core.utils import get_model_times
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 from ax.service.orchestrator import Orchestrator
-from ax.service.utils.best_point import get_trace
-from ax.service.utils.best_point_mixin import BestPointMixin
+from ax.service.utils.best_point import (
+    get_best_parameters_from_model_predictions_with_trial_index,
+    get_trace,
+)
 from ax.service.utils.orchestrator_options import OrchestratorOptions, TrialType
 from ax.utils.common.logger import DEFAULT_LOG_LEVEL, get_logger
 from ax.utils.common.random import with_rng_seed
@@ -338,9 +340,9 @@ def get_best_parameters(
             best point.
         trial_indices: Use data from only these trials. If None, use all data.
     """
-    result = BestPointMixin._get_best_trial(
+    result = get_best_parameters_from_model_predictions_with_trial_index(
         experiment=experiment,
-        generation_strategy=generation_strategy,
+        adapter=generation_strategy.adapter,
         trial_indices=trial_indices,
     )
     if result is None:
@@ -507,7 +509,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 +564,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

@@ -491,6 +491,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):
@@ -618,9 +619,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,
@@ -1196,6 +1197,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(