Make get_best_raw_objective_point_with_trial_index return a TModelPredictArm (#3753)

Summary:
Pull Request resolved: #3753

In every usage of `get_best_raw_objective_point_with_trial_index`, its third return value is either ignored or changed into a TModelPredictArm.

This PR
* makes it return a TModelPredictArm, which makes some downstream functions more unitary.
* Fixes some type annotations that were too broad (the TModelPredictArm won't be None)
* Updates docstrings

Reviewed By: Balandat

Differential Revision: D74401853

fbshipit-source-id: bd184c2e6f0c5e673d08d6fb5278cf2029063cfc

Author: esantorella

ax/service/managed_loop.py

@@ -260,17 +260,10 @@ def get_best_point(self) -> tuple[TParameterization, TModelPredictArm | None]:
             return parameterizations, predictions
 
         # Could not find through model, default to using raw objective.
-        _, parameterization, values = get_best_raw_objective_point_with_trial_index(
-            experiment=self.experiment
-        )
-        # For values, grab just the means to conform to TModelPredictArm format.
-        return (
-            parameterization,
-            (
-                {k: v[0] for k, v in values.items()},  # v[0] is mean
-                {k: {k: v[1] * v[1]} for k, v in values.items()},  # v[1] is sem
-            ),
+        _, parameterization, predict_arm = (
+            get_best_raw_objective_point_with_trial_index(experiment=self.experiment)
         )
+        return parameterization, predict_arm
 
     def get_current_model(self) -> Adapter | None:
         """Obtain the most recently used model in optimization."""

ax/service/utils/best_point.py

@@ -76,13 +76,17 @@ def get_best_raw_objective_point_with_trial_index(
     experiment: Experiment,
     optimization_config: OptimizationConfig | None = None,
     trial_indices: Iterable[int] | None = None,
-) -> tuple[int, TParameterization, dict[str, tuple[float, float]]]:
+) -> tuple[int, TParameterization, TModelPredictArm]:
     """Given an experiment, identifies the arm that had the best raw objective,
     based on the data fetched from the experiment.
 
-    Note: This function will error with an invalid configuration. If you would
+    TModelPredictArm is of the form:
+        ({metric_name: mean}, {metric_name_1: {metric_name_2: cov_1_2}})
+
+    Note: This function will error with invalid inputs. If you would
     prefer for error logs rather than exceptions, use
-    `get_best_by_raw_objective_with_trial_index`.
+    `get_best_by_raw_objective_with_trial_index`, which returns None if
+    inputs are invalid.
 
     Args:
         experiment: Experiment, on which to identify best raw objective arm.
@@ -92,8 +96,7 @@ def get_best_raw_objective_point_with_trial_index(
             retrieve data from all available trials.
 
     Returns:
-        Tuple of parameterization and a mapping from metric name to a tuple of
-            the corresponding objective mean and SEM.
+        Tuple of trial index, parameterization, and model predictions for it.
     """
     optimization_config = optimization_config or experiment.optimization_config
     if optimization_config is None:
@@ -171,8 +174,9 @@ def get_best_raw_objective_point_with_trial_index(
         row["metric_name"]: (row["mean"], row["sem"])
         for _, row in objective_rows.iterrows()
     }
+    predict_arm = _raw_values_to_model_predict_arm(values=vals)
 
-    return best_trial_index, none_throws(best_arm).parameters, vals
+    return best_trial_index, none_throws(best_arm).parameters, predict_arm
 
 
 def _extract_best_arm_from_gr(
@@ -320,7 +324,7 @@ def get_best_by_raw_objective_with_trial_index(
     experiment: Experiment,
     optimization_config: OptimizationConfig | None = None,
     trial_indices: Iterable[int] | None = None,
-) -> tuple[int, TParameterization, TModelPredictArm | None] | None:
+) -> tuple[int, TParameterization, TModelPredictArm] | None:
     """Given an experiment, identifies the arm that had the best raw objective,
     based on the data fetched from the experiment.
 
@@ -341,11 +345,7 @@ def get_best_by_raw_objective_with_trial_index(
         Tuple of trial index, parameterization, and model predictions for it.
     """
     try:
-        (
-            trial_index,
-            parameterization,
-            values,
-        ) = get_best_raw_objective_point_with_trial_index(
+        result = get_best_raw_objective_point_with_trial_index(
             experiment=experiment,
             optimization_config=optimization_config,
             trial_indices=trial_indices,
@@ -356,11 +356,7 @@ def get_best_by_raw_objective_with_trial_index(
             f"'{err}'. Returning None."
         )
         return None
-    return (
-        trial_index,
-        parameterization,
-        _raw_values_to_model_predict_arm(values),
-    )
+    return result
 
 
 def get_pareto_optimal_parameters(