Add trial_type support to base Experiment metric methods

Summary:
Phase 2 of moving MultiTypeExperiment features into base Experiment.

Updates the base Experiment metric management methods (`add_metric`, `update_metric`, `remove_metric`) to accept an optional `trial_type` parameter. When provided, metrics are associated with the specified trial type in `_trial_type_to_metric_names`. The `__init__` and `optimization_config` setter also now register metrics when `default_trial_type` is set.

The deprecated wrappers (`add_tracking_metric`, `add_tracking_metrics`, `update_tracking_metric`) now accept and pass through `trial_type` and `canonical_name` parameters.

On MultiTypeExperiment, overrides are simplified to delegate to the base class methods:
- `add_tracking_metric` delegates to `self.add_metric()`
- `add_tracking_metrics` override removed (inherited from base)
- `update_tracking_metric` delegates to `self.update_metric()`
- `remove_tracking_metric` replaced with `remove_metric` override

Differential Revision: D94986440

Author: mpolson64

ax/core/experiment.py

@@ -201,6 +201,10 @@ def __init__(
         # a naming collision ucurs.
         for m in [*(tracking_metrics or []), *(metrics or [])]:
             self._metrics[m.name] = m
+            if self._default_trial_type is not None:
+                self._trial_type_to_metric_names.setdefault(
+                    self._default_trial_type, set()
+                ).add(m.name)
 
         # call setters defined below
         self.status_quo = status_quo
@@ -574,6 +578,11 @@ def optimization_config(self, optimization_config: OptimizationConfig) -> None:
                     "but not found on experiment. Add it first with add_metric()."
                 )
         self._optimization_config = optimization_config
+        if self._default_trial_type is not None:
+            for metric_name in optimization_config.metric_names:
+                self._trial_type_to_metric_names.setdefault(
+                    self._default_trial_type, set()
+                ).add(metric_name)
 
     @property
     def is_moo_problem(self) -> bool:
@@ -826,59 +835,115 @@ def get_metric(self, name: str) -> Metric:
             )
         return self._metrics[name]
 
-    def add_metric(self, metric: Metric) -> Self:
+    def add_metric(self, metric: Metric, trial_type: str | None = None) -> Self:
         """Add a new metric to the experiment.
 
         Args:
             metric: Metric to be added.
+            trial_type: If provided, associates the metric with this trial type.
+                When ``None`` and a ``default_trial_type`` is set, defaults to
+                the default trial type.
         """
         if metric.name in self._metrics:
             raise ValueError(
                 f"Metric `{metric.name}` already defined on experiment. "
                 "Use `update_metric` to update an existing metric definition."
             )
+        if trial_type is None and self._default_trial_type is not None:
+            trial_type = self._default_trial_type
+        if trial_type is not None:
+            if not self.supports_trial_type(trial_type):
+                raise ValueError(f"`{trial_type}` is not a supported trial type.")
+            self._trial_type_to_metric_names.setdefault(trial_type, set()).add(
+                metric.name
+            )
         self._metrics[metric.name] = metric
         return self
 
-    def add_tracking_metric(self, metric: Metric) -> Self:
+    def add_tracking_metric(
+        self,
+        metric: Metric,
+        trial_type: str | None = None,
+        canonical_name: str | None = None,
+    ) -> Self:
         """*Deprecated.* Use ``add_metric`` instead."""
         warnings.warn(
             "add_tracking_metric is deprecated. Use add_metric instead.",
             DeprecationWarning,
             stacklevel=2,
         )
-        return self.add_metric(metric)
+        return self.add_metric(metric, trial_type=trial_type)
 
-    def add_tracking_metrics(self, metrics: list[Metric]) -> Experiment:
+    def add_tracking_metrics(
+        self,
+        metrics: list[Metric],
+        metrics_to_trial_types: dict[str, str] | None = None,
+        canonical_names: dict[str, str] | None = None,
+    ) -> Experiment:
         """*Deprecated.* Use ``add_metric`` instead."""
         warnings.warn(
             "add_tracking_metrics is deprecated. Use add_metric instead.",
             DeprecationWarning,
             stacklevel=2,
         )
+        metrics_to_trial_types = metrics_to_trial_types or {}
         for metric in metrics:
-            self.add_metric(metric)
+            canonical_name = (canonical_names or {}).get(metric.name)
+            self.add_tracking_metric(
+                metric=metric,
+                trial_type=metrics_to_trial_types.get(metric.name),
+                canonical_name=canonical_name,
+            )
         return self
 
-    def update_metric(self, metric: Metric) -> Self:
+    def update_metric(self, metric: Metric, trial_type: str | None = None) -> Self:
         """Redefine a metric that already exists on the experiment.
 
         Args:
             metric: New metric definition.
+            trial_type: If provided, reassociates the metric with this trial
+                type. When ``None``, keeps the metric's existing trial type.
         """
         if metric.name not in self._metrics:
             raise ValueError(f"Metric `{metric.name}` doesn't exist on experiment.")
+        if trial_type is not None:
+            if not self.supports_trial_type(trial_type):
+                raise ValueError(f"`{trial_type}` is not a supported trial type.")
+            oc = self._optimization_config
+            if (
+                oc is not None
+                and metric.name in oc.metric_names
+                and self._default_trial_type is not None
+                and trial_type != self._default_trial_type
+            ):
+                raise ValueError(
+                    f"Metric `{metric.name}` must remain a "
+                    f"`{self._default_trial_type}` metric because it is part of "
+                    "the optimization_config."
+                )
+            # Remove from any existing trial type set
+            for names in self._trial_type_to_metric_names.values():
+                names.discard(metric.name)
+            # Add to new trial type set
+            self._trial_type_to_metric_names.setdefault(trial_type, set()).add(
+                metric.name
+            )
         self._metrics[metric.name] = metric
         return self
 
-    def update_tracking_metric(self, metric: Metric) -> Experiment:
+    def update_tracking_metric(
+        self,
+        metric: Metric,
+        trial_type: str | None = None,
+        canonical_name: str | None = None,
+    ) -> Experiment:
         """*Deprecated.* Use ``update_metric`` instead."""
         warnings.warn(
             "update_tracking_metric is deprecated. Use update_metric instead.",
             DeprecationWarning,
             stacklevel=2,
         )
-        return self.update_metric(metric)
+        return self.update_metric(metric, trial_type=trial_type)
 
     def remove_metric(self, metric_name: str) -> Self:
         """Remove a metric from the experiment.
@@ -899,6 +964,9 @@ def remove_metric(self, metric_name: str) -> Self:
                 f"Metric `{metric_name}` is referenced by the optimization config "
                 "and cannot be removed. Update the optimization config first."
             )
+        # Clean up _trial_type_to_metric_names
+        for names in self._trial_type_to_metric_names.values():
+            names.discard(metric_name)
         del self._metrics[metric_name]
         return self
 

ax/core/multi_type_experiment.py

@@ -96,16 +96,13 @@ def __init__(
             default_data_type=default_data_type,
         )
 
-        # Ensure tracking metrics are registered in _metric_to_trial_type
-        # and _trial_type_to_metric_names.
-        # super().__init__ sets self._metrics directly, bypassing
-        # add_tracking_metric, so tracking metrics won't be in
-        # _metric_to_trial_type yet.
+        # Ensure tracking metrics are registered in _metric_to_trial_type.
+        # The base __init__ handles _trial_type_to_metric_names.
         for m in tracking_metrics or []:
             if m.name not in self._metric_to_trial_type:
-                tt = none_throws(self._default_trial_type)
-                self._metric_to_trial_type[m.name] = tt
-                self._trial_type_to_metric_names.setdefault(tt, set()).add(m.name)
+                self._metric_to_trial_type[m.name] = none_throws(
+                    self._default_trial_type
+                )
 
     def add_trial_type(self, trial_type: str, runner: Runner) -> Self:
         """Add a new trial_type to be supported by this experiment.
@@ -127,12 +124,11 @@ def add_trial_type(self, trial_type: str, runner: Runner) -> Self:
     def optimization_config(self, optimization_config: OptimizationConfig) -> None:
         # pyre-fixme[16]: `Optional` has no attribute `fset`.
         Experiment.optimization_config.fset(self, optimization_config)
+        # Base setter handles _trial_type_to_metric_names; update legacy dict.
         for metric_name in optimization_config.metric_names:
-            # Optimization config metrics are required to be the default trial type
-            # currently. TODO: remove that restriction (T202797235)
-            tt = none_throws(self.default_trial_type)
-            self._metric_to_trial_type[metric_name] = tt
-            self._trial_type_to_metric_names.setdefault(tt, set()).add(metric_name)
+            self._metric_to_trial_type[metric_name] = none_throws(
+                self.default_trial_type
+            )
 
     def update_runner(self, trial_type: str, runner: Runner) -> Self:
         """Update the default runner for an existing trial_type.
@@ -163,56 +159,12 @@ def add_tracking_metric(
         """
         if trial_type is None:
             trial_type = self._default_trial_type
-        if not self.supports_trial_type(trial_type):
-            raise ValueError(f"`{trial_type}` is not a supported trial type.")
-
-        super().add_tracking_metric(metric)
-        tt = none_throws(trial_type)
-        self._metric_to_trial_type[metric.name] = tt
-        self._trial_type_to_metric_names.setdefault(tt, set()).add(metric.name)
+        self.add_metric(metric, trial_type=trial_type)
+        self._metric_to_trial_type[metric.name] = none_throws(trial_type)
         if canonical_name is not None:
             self._metric_to_canonical_name[metric.name] = canonical_name
         return self
 
-    def add_tracking_metrics(
-        self,
-        metrics: list[Metric],
-        metrics_to_trial_types: dict[str, str] | None = None,
-        canonical_names: dict[str, str] | None = None,
-    ) -> Experiment:
-        """Add a list of new metrics to the experiment.
-
-        If any of the metrics are already defined on the experiment,
-        we raise an error and don't add any of them to the experiment
-
-        Args:
-            metrics: Metrics to be added.
-            metrics_to_trial_types: The mapping from metric names to corresponding
-                trial types for each metric. If provided, the metrics will be
-                added to their trial types. If not provided, then the default
-                trial type will be used.
-            canonical_names: A mapping of metric names to their
-                canonical names(The default metrics for which the metrics are
-                proxies.)
-
-        Returns:
-            The experiment with the added metrics.
-        """
-        metrics_to_trial_types = metrics_to_trial_types or {}
-        canonical_name = None
-        for metric in metrics:
-            if canonical_names is not None:
-                canonical_name = none_throws(canonical_names).get(metric.name, None)
-
-            self.add_tracking_metric(
-                metric=metric,
-                trial_type=metrics_to_trial_types.get(
-                    metric.name, self._default_trial_type
-                ),
-                canonical_name=canonical_name,
-            )
-        return self
-
     def update_tracking_metric(
         self,
         metric: Metric,
@@ -233,47 +185,17 @@ def update_tracking_metric(
             trial_type = self._metric_to_trial_type.get(
                 metric.name, self._default_trial_type
             )
-        oc = self.optimization_config
-        oc_metric_names = oc.metric_names if oc else set()
-        if metric.name in oc_metric_names and trial_type != self._default_trial_type:
-            raise ValueError(
-                f"Metric `{metric.name}` must remain a "
-                f"`{self._default_trial_type}` metric because it is part of the "
-                "optimization_config."
-            )
-        elif not self.supports_trial_type(trial_type):
-            raise ValueError(f"`{trial_type}` is not a supported trial type.")
-
-        super().update_tracking_metric(metric)
-        # Remove from old trial type set
-        old_tt = self._metric_to_trial_type.get(metric.name)
-        if old_tt is not None and old_tt in self._trial_type_to_metric_names:
-            self._trial_type_to_metric_names[old_tt].discard(metric.name)
-        # Add to new trial type set
-        tt = none_throws(trial_type)
-        self._metric_to_trial_type[metric.name] = tt
-        self._trial_type_to_metric_names.setdefault(tt, set()).add(metric.name)
+        self.update_metric(metric, trial_type=trial_type)
+        self._metric_to_trial_type[metric.name] = none_throws(trial_type)
         if canonical_name is not None:
             self._metric_to_canonical_name[metric.name] = canonical_name
         return self
 
-    @copy_doc(Experiment.remove_tracking_metric)
-    def remove_tracking_metric(self, metric_name: str) -> Self:
-        if metric_name not in self._metrics:
-            raise ValueError(f"Metric `{metric_name}` doesn't exist on experiment.")
-
-        # Clean up _trial_type_to_metric_names
-        old_tt = self._metric_to_trial_type.get(metric_name)
-        if old_tt is not None and old_tt in self._trial_type_to_metric_names:
-            self._trial_type_to_metric_names[old_tt].discard(metric_name)
-
-        # Required fields
-        del self._metrics[metric_name]
-        del self._metric_to_trial_type[metric_name]
-
-        # Optional
-        if metric_name in self._metric_to_canonical_name:
-            del self._metric_to_canonical_name[metric_name]
+    @copy_doc(Experiment.remove_metric)
+    def remove_metric(self, metric_name: str) -> Self:
+        super().remove_metric(metric_name)
+        self._metric_to_trial_type.pop(metric_name, None)
+        self._metric_to_canonical_name.pop(metric_name, None)
         return self
 
     @copy_doc(Experiment.fetch_data)