Improve ergonomics around metric_name --> signature conversion in Adapter and Transform (#5045)

Summary:

D93520819 removed Metric instances from the OptimizationConfig which made it more verbose to get any individual Metric's signature, which we use for bookkeeping throughout the adapter stack.

Previously, functions like extract_objective_weights, extract_outcome_constraints, and extract_objective_thresholds accepted an Experiment object solely to call experiment.get_metric(name).signature. This diff replaces that experiment parameter with a lightweight dict[str, str] (metric_name_to_signature) across the adapter and transform layers.

Ultimately this reduces coupling to Experiment in the modeling layer, making these functions easier to test and reuse in contexts where a full Experiment object isn't available.

Reviewed By: Balandat

Differential Revision: D96855090

Author: mpolson64

ax/adapter/adapter_utils.py

@@ -9,10 +9,10 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Callable, Iterable, Mapping, MutableMapping, Sequence
+from collections.abc import Iterable, Mapping, MutableMapping, Sequence
 from copy import deepcopy
 from logging import Logger
-from typing import Any, cast, SupportsFloat, TYPE_CHECKING
+from typing import Any, Callable, cast, SupportsFloat, TYPE_CHECKING
 
 import numpy as np
 import numpy.typing as npt
@@ -206,7 +206,7 @@ def extract_objective_thresholds(
     objective_thresholds: TRefPoint,
     objective: Objective,
     outcomes: list[str],
-    experiment: Experiment,
+    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray | None:
     """Extracts objective thresholds' values, in the order of `outcomes`.
 
@@ -222,7 +222,7 @@ def extract_objective_thresholds(
         objective_thresholds: Objective thresholds to extract values from.
         objective: The corresponding Objective, for validation purposes.
         outcomes: n-length list of names of metrics.
-        experiment: The experiment, used to map metric names to signatures.
+        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns:
         (n,) array of thresholds
@@ -232,7 +232,7 @@ def extract_objective_thresholds(
 
     objective_threshold_dict = {}
     for ot in objective_thresholds:
-        ot_signature = experiment.get_metric(ot.metric_names[0]).signature
+        ot_signature = metric_name_to_signature[ot.metric_names[0]]
         if ot.relative:
             raise ValueError(
                 f"Objective {ot_signature} has a relative threshold that "
@@ -242,7 +242,7 @@ def extract_objective_thresholds(
 
     # Check that all thresholds correspond to a metric.
     obj_metric_signatures = [
-        experiment.get_metric(name).signature for name in objective.metric_names
+        metric_name_to_signature[name] for name in objective.metric_names
     ]
     if set(objective_threshold_dict.keys()).difference(set(obj_metric_signatures)):
         raise ValueError(
@@ -259,7 +259,9 @@ def extract_objective_thresholds(
 
 
 def extract_objective_weights(
-    objective: Objective, outcomes: list[str], experiment: Experiment
+    objective: Objective,
+    outcomes: list[str],
+    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray:
     """Extract a weights for objectives.
 
@@ -277,7 +279,7 @@ def extract_objective_weights(
     Args:
         objective: Objective to extract weights from.
         outcomes: n-length list of metric signatures.
-        experiment: The experiment, used to map metric names to signatures.
+        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns:
         n-length array of weights.
@@ -287,13 +289,15 @@ def extract_objective_weights(
     # metric_weights returns sign-encoded (name, weight) tuples for all
     # objective types (single, scalarized, multi).
     for obj_metric_name, obj_weight in objective.metric_weights:
-        sig = experiment.get_metric(obj_metric_name).signature
+        sig = metric_name_to_signature[obj_metric_name]
         objective_weights[outcomes.index(sig)] = obj_weight
     return objective_weights
 
 
 def extract_objective_weight_matrix(
-    objective: Objective, outcomes: list[str], experiment: Experiment
+    objective: Objective,
+    outcomes: list[str],
+    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray:
     """Extract a 2D weight matrix for objectives.
 
@@ -308,6 +312,7 @@ def extract_objective_weight_matrix(
     Args:
         objective: Objective to extract weights from.
         outcomes: n-length list of signatures of metrics.
+        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns:
         ``(n_objectives, n)`` array of weights.
@@ -319,19 +324,23 @@ def extract_objective_weight_matrix(
                 extract_objective_weights(
                     objective=Objective(expression=f"{weight} * {name}"),
                     outcomes=outcomes,
-                    experiment=experiment,
+                    metric_name_to_signature=metric_name_to_signature,
                 )
             )
         return np.stack(rows, axis=0)
     else:
         # Single row – covers Objective and ScalarizedObjective
-        return extract_objective_weights(objective, outcomes, experiment).reshape(1, -1)
+        return extract_objective_weights(
+            objective=objective,
+            outcomes=outcomes,
+            metric_name_to_signature=metric_name_to_signature,
+        ).reshape(1, -1)
 
 
 def extract_outcome_constraints(
     outcome_constraints: list[OutcomeConstraint],
     outcomes: list[str],
-    experiment: Experiment,
+    metric_name_to_signature: Mapping[str, str],
 ) -> TBounds:
     if len(outcome_constraints) == 0:
         return None
@@ -342,10 +351,10 @@ def extract_outcome_constraints(
         s = 1 if c.op == ComparisonOp.LEQ else -1
         if isinstance(c, ScalarizedOutcomeConstraint):
             for c_metric_name, c_weight in c.metric_weights:
-                j = outcomes.index(experiment.get_metric(c_metric_name).signature)
+                j = outcomes.index(metric_name_to_signature[c_metric_name])
                 A[i, j] = s * c_weight
         else:
-            j = outcomes.index(experiment.get_metric(c.metric_names[0]).signature)
+            j = outcomes.index(metric_name_to_signature[c.metric_names[0]])
             A[i, j] = s
         b[i, 0] = s * c.bound
     return (A, b)
@@ -658,18 +667,18 @@ def get_pareto_frontier_and_configs(
     objective_weights = extract_objective_weight_matrix(
         objective=optimization_config.objective,
         outcomes=adapter.outcomes,
-        experiment=adapter._experiment,
+        metric_name_to_signature=adapter.metric_name_to_signature,
     )
     outcome_constraints = extract_outcome_constraints(
         outcome_constraints=optimization_config.outcome_constraints,
         outcomes=adapter.outcomes,
-        experiment=adapter._experiment,
+        metric_name_to_signature=adapter.metric_name_to_signature,
     )
     obj_t = extract_objective_thresholds(
         objective_thresholds=optimization_config.objective_thresholds,
         objective=optimization_config.objective,
         outcomes=adapter.outcomes,
-        experiment=adapter._experiment,
+        metric_name_to_signature=adapter.metric_name_to_signature,
     )
     if obj_t is not None:
         obj_t = array_to_tensor(obj_t)
@@ -1128,7 +1137,7 @@ def observation_features_to_array(
 def feasible_hypervolume(
     optimization_config: MultiObjectiveOptimizationConfig,
     values: dict[str, npt.NDArray],
-    experiment: Experiment,
+    metric_name_to_signature: Mapping[str, str],
 ) -> npt.NDArray:
     """Compute the feasible hypervolume each iteration.
 
@@ -1137,26 +1146,26 @@ def feasible_hypervolume(
         values: Dictionary from metric name to array of value at each
             iteration (each array is `n`-dim). If optimization config contains
             outcome constraints, values for them must be present in `values`.
-        experiment: The experiment, used to map metric names to signatures.
+        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns: Array of feasible hypervolumes.
     """
     # Get objective at each iteration
     obj_threshold_dict = {
-        experiment.get_metric(ot.metric_names[0]).signature: ot.bound
+        metric_name_to_signature[ot.metric_names[0]]: ot.bound
         for ot in optimization_config.objective_thresholds
     }
     obj_metric_names = optimization_config.objective.metric_names
-    obj_metrics = [experiment.get_metric(name) for name in obj_metric_names]
-    f_vals = np.hstack([values[m.signature].reshape(-1, 1) for m in obj_metrics])
-    obj_thresholds = np.array([obj_threshold_dict[m.signature] for m in obj_metrics])
+    obj_metric_sigs = [metric_name_to_signature[name] for name in obj_metric_names]
+    f_vals = np.hstack([values[sig].reshape(-1, 1) for sig in obj_metric_sigs])
+    obj_thresholds = np.array([obj_threshold_dict[sig] for sig in obj_metric_sigs])
     # Set infeasible points to be the objective threshold
     for oc in optimization_config.outcome_constraints:
         if oc.relative:
             raise ValueError(
                 "Benchmark aggregation does not support relative constraints"
             )
-        oc_sig = experiment.get_metric(oc.metric_names[0]).signature
+        oc_sig = metric_name_to_signature[oc.metric_names[0]]
         g = values[oc_sig]
         feas = g <= oc.bound if oc.op == ComparisonOp.LEQ else g >= oc.bound
         f_vals[~feas] = obj_thresholds

ax/adapter/base.py

@@ -190,6 +190,9 @@ def __init__(
         )
         self._experiment_properties: dict[str, Any] = experiment._properties
         self._experiment: Experiment = experiment
+        self._metric_name_to_signature: dict[str, str] = {
+            name: metric.signature for name, metric in self._experiment.metrics.items()
+        }
 
         if self._optimization_config is None:
             self._optimization_config = experiment.optimization_config
@@ -525,6 +528,11 @@ def metric_signatures(self) -> set[str]:
         """Metric signatures present in training data."""
         return self._metric_signatures
 
+    @property
+    def metric_name_to_signature(self) -> dict[str, str]:
+        """Mapping from metric names to their signatures."""
+        return self._metric_name_to_signature
+
     @property
     def model_space(self) -> SearchSpace:
         """SearchSpace used to fit model."""

ax/adapter/cross_validation.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from collections.abc import Callable, Iterable
+from collections.abc import Callable, Iterable, Mapping
 from logging import Logger
 from typing import cast, NamedTuple
 from warnings import warn
@@ -21,7 +21,6 @@
 from ax.adapter.data_utils import ExperimentData
 from ax.adapter.observation_utils import unwrap_observation_data
 from ax.adapter.torch import TorchAdapter
-from ax.core.experiment import Experiment
 from ax.core.observation import Observation, ObservationData, ObservationFeatures
 from ax.core.optimization_config import OptimizationConfig
 from ax.exceptions.core import UnsupportedError
@@ -574,7 +573,7 @@ def assess_model_fit(
 def has_good_opt_config_model_fit(
     optimization_config: OptimizationConfig,
     assess_model_fit_result: AssessModelFitResult,
-    experiment: Experiment,
+    metric_name_to_signature: Mapping[str, str],
 ) -> bool:
     """Assess model fit for given diagnostics results across the optimization
     config metrics
@@ -587,7 +586,7 @@ def has_good_opt_config_model_fit(
     Args:
         optimization_config: Objective/Outcome constraint metrics to assess
         assess_model_fit_result: Output of assess_model_fit
-        experiment: The experiment, used to map metric names to signatures.
+        metric_name_to_signature: Mapping from metric names to signatures.
 
     Returns:
         Two dictionaries, one for good metrics, one for bad metrics, each
@@ -597,7 +596,7 @@ def has_good_opt_config_model_fit(
     # Bad fit criteria: Any objective metrics are poorly fit
     # TODO[]: Incl. outcome constraints in assessment
     has_good_opt_config_fit = all(
-        experiment.get_metric(name).signature
+        metric_name_to_signature[name]
         in assess_model_fit_result.good_fit_metrics_to_fisher_score
         for name in optimization_config.objective.metric_names
     )

ax/adapter/discrete.py

@@ -164,17 +164,19 @@ def _gen(
             outcome_constraints = None
         else:
             validate_transformed_optimization_config(
-                optimization_config, self.outcomes, experiment=self._experiment
+                optimization_config,
+                self.outcomes,
+                metric_name_to_signature=self.metric_name_to_signature,
             )
             objective_weights = extract_objective_weights(
                 objective=optimization_config.objective,
                 outcomes=self.outcomes,
-                experiment=self._experiment,
+                metric_name_to_signature=self.metric_name_to_signature,
             )
             outcome_constraints = extract_outcome_constraints(
                 outcome_constraints=optimization_config.outcome_constraints,
                 outcomes=self.outcomes,
-                experiment=self._experiment,
+                metric_name_to_signature=self.metric_name_to_signature,
             )
 
         # Get fixed features

ax/adapter/tests/test_adapter_utils.py

@@ -43,7 +43,6 @@
 from ax.utils.common.hash_utils import compute_lilo_input_hash
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import (
-    get_branin_experiment,
     get_experiment_with_observations,
     get_hierarchical_search_space,
     get_search_space_for_range_values,
@@ -79,11 +78,8 @@ def test_feasible_hypervolume(self) -> None:
                 ),
             ],
         )
-        experiment = Experiment(
-            search_space=SearchSpace(parameters=[]),
-            optimization_config=optimization_config,
-            tracking_metrics=[mc],
-        )
+        # For plain Metric objects, signature == name.
+        metric_name_to_signature = {m.name: m.name for m in [ma, mb, mc]}
         feas_hv = feasible_hypervolume(
             optimization_config,
             values={
@@ -112,7 +108,7 @@ def test_feasible_hypervolume(self) -> None:
                     ]
                 ),
             },
-            experiment=experiment,
+            metric_name_to_signature=metric_name_to_signature,
         )
         self.assertEqual(list(feas_hv), [0.0, 0.0, 1.0, 1.0])
 
@@ -544,24 +540,34 @@ def test_can_map_to_binary(self) -> None:
     def test_extract_objective_weight_matrix(self) -> None:
         m1, m2, m3 = Metric(name="m1"), Metric(name="m2"), Metric(name="m3")
         outcomes = ["m1", "m2", "m3"]
-        experiment = get_branin_experiment()
-        experiment.add_metric(m1)
-        experiment.add_metric(m2)
-        experiment.add_metric(m3)
+        # For plain Metric objects, signature == name.
+        metric_name_to_signature = {name: name for name in outcomes}
 
         # Single Objective: one row, nonzero only in matching column.
         obj = Objective(metric=m1, minimize=False)
-        result = extract_objective_weight_matrix(obj, outcomes, experiment)
+        result = extract_objective_weight_matrix(
+            objective=obj,
+            outcomes=outcomes,
+            metric_name_to_signature=metric_name_to_signature,
+        )
         np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0]])
 
         # Minimization flips the sign.
         obj_min = Objective(metric=m2, minimize=True)
-        result = extract_objective_weight_matrix(obj_min, outcomes, experiment)
+        result = extract_objective_weight_matrix(
+            objective=obj_min,
+            outcomes=outcomes,
+            metric_name_to_signature=metric_name_to_signature,
+        )
         np.testing.assert_array_equal(result, [[0.0, -1.0, 0.0]])
 
         # ScalarizedObjective: single row with multiple nonzero entries.
         scal = ScalarizedObjective(metrics=[m1, m3], weights=[0.3, 0.7], minimize=False)
-        result = extract_objective_weight_matrix(scal, outcomes, experiment)
+        result = extract_objective_weight_matrix(
+            objective=scal,
+            outcomes=outcomes,
+            metric_name_to_signature=metric_name_to_signature,
+        )
         np.testing.assert_array_almost_equal(result, [[0.3, 0.0, 0.7]])
 
         # MultiObjective: one row per sub-objective.
@@ -571,7 +577,11 @@ def test_extract_objective_weight_matrix(self) -> None:
                 Objective(metric=m3, minimize=True),
             ]
         )
-        result = extract_objective_weight_matrix(multi, outcomes, experiment)
+        result = extract_objective_weight_matrix(
+            objective=multi,
+            outcomes=outcomes,
+            metric_name_to_signature=metric_name_to_signature,
+        )
         np.testing.assert_array_equal(result, [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0]])
 
     def test_get_fresh_pairwise_trial_indices(self) -> None: