Only consider completed trials for best point (#1819)

Summary: Pull Request resolved: #1819

Reviewed By: bernardbeckerman, dme65

Differential Revision: D48985230

fbshipit-source-id: 0d6121fa3bfdec2238f3698d501747b7219cfec7

ax/service/tests/test_best_point_utils.py

@@ -125,7 +125,7 @@ def test_best_raw_objective_point(self) -> None:
         self.assertEqual(get_best_parameters(exp, Models), None)
         exp.new_trial(
             generator_run=GeneratorRun(arms=[Arm(parameters={"x1": 5.0, "x2": 5.0})])
-        ).run()
+        ).run().complete()
         exp.fetch_data()
         # pyre-fixme[16]: Optional type has no attribute `clone`.
         opt_conf = exp.optimization_config.clone()
@@ -205,7 +205,7 @@ def test_best_raw_objective_point_scalarized(self) -> None:
         self.assertEqual(get_best_parameters(exp, Models), None)
         exp.new_trial(
             generator_run=GeneratorRun(arms=[Arm(parameters={"x1": 5.0, "x2": 5.0})])
-        ).run()
+        ).run().complete()
         exp.fetch_data()
         self.assertEqual(get_best_raw_objective_point(exp)[0], {"x1": 5.0, "x2": 5.0})
 
@@ -223,7 +223,7 @@ def test_best_raw_objective_point_scalarized_multi(self) -> None:
         self.assertEqual(get_best_parameters(exp, Models), None)
         exp.new_trial(
             generator_run=GeneratorRun(arms=[Arm(parameters={"x1": 5.0, "x2": 5.0})])
-        ).run()
+        ).run().complete()
         exp.fetch_data()
         self.assertEqual(get_best_raw_objective_point(exp)[0], {"x1": 5.0, "x2": 5.0})
 

ax/service/utils/best_point.py

@@ -100,8 +100,15 @@ def get_best_raw_objective_point_with_trial_index(
             logger.warning(
                 "No status quo provided; relative constraints will be ignored."
             )
+
+    # Only COMPLETED trials should be considered when identifying the best point
+    completed_indices = {
+        t.index for t in experiment.trials_by_status[TrialStatus.COMPLETED]
+    }
+    completed_df = dat.df[dat.df["trial_index"].isin(completed_indices)]
+
     feasible_df = _filter_feasible_rows(
-        df=dat.df,
+        df=completed_df,
         optimization_config=optimization_config,
     )
     objective = optimization_config.objective