Move add_trial_type, update_runner, supports_trial_type to base Experiment (#5003)

Summary:
Pull Request resolved: #5003

Phase 3 of moving MultiTypeExperiment features into base Experiment.

Moves `add_trial_type` and `update_runner` from MultiTypeExperiment to the base Experiment class, making them available to all experiments.

Updates `supports_trial_type` to unify the logic: for multi-type experiments (where `default_trial_type` is set), only trial types registered in `_trial_type_to_runner` are supported. For single-type experiments, `None` is supported along with `SHORT_RUN` and `LONG_RUN` for backward compatibility with generation strategies that use those trial types.

Removes the corresponding overrides from MultiTypeExperiment — all three methods are now inherited from the base class.

Reviewed By: saitcakmak

Differential Revision: D94988577

Author: mpolson64

ax/core/experiment.py

@@ -2033,6 +2033,34 @@ def default_trials(self) -> set[int]:
             if trial.trial_type == self.default_trial_type
         }
 
+    def add_trial_type(self, trial_type: str, runner: Runner | None = None) -> Self:
+        """Add a new trial type to be supported by this experiment.
+
+        Args:
+            trial_type: The new trial type to be added.
+            runner: The default runner for trials of this type.
+        """
+        if self.supports_trial_type(trial_type):
+            raise ValueError(f"Experiment already contains trial_type `{trial_type}`")
+
+        if runner is not None:
+            self._trial_type_to_runner[trial_type] = runner
+
+        return self
+
+    def update_runner(self, trial_type: str, runner: Runner) -> Self:
+        """Update the default runner for an existing trial type.
+
+        Args:
+            trial_type: The trial type whose runner should be updated.
+            runner: The new runner for trials of this type.
+        """
+        if not self.supports_trial_type(trial_type):
+            raise ValueError(f"Experiment does not contain trial_type `{trial_type}`")
+        self._trial_type_to_runner[trial_type] = runner
+        self._runner = runner
+        return self
+
     def runner_for_trial_type(self, trial_type: str | None) -> Runner | None:
         """The default runner to use for a given trial type.
 
@@ -2047,14 +2075,20 @@ def runner_for_trial_type(self, trial_type: str | None) -> Runner | None:
     def supports_trial_type(self, trial_type: str | None) -> bool:
         """Whether this experiment allows trials of the given type.
 
-        The base experiment class only supports None. For experiments
-        with multiple trial types, use the MultiTypeExperiment class.
+        For experiments with a ``default_trial_type`` (multi-type experiments),
+        only trial types registered in ``_trial_type_to_runner`` are supported.
+        For single-type experiments, ``None`` is always supported, along with
+        ``SHORT_RUN`` and ``LONG_RUN`` for backward compatibility with
+        generation strategies that use those trial types.
         """
+        if self._default_trial_type is not None:
+            return trial_type in self._trial_type_to_runner
         return (
             trial_type is None
             or trial_type == Keys.SHORT_RUN
             or trial_type == Keys.LONG_RUN
             or trial_type == Keys.LILO_LABELING
+            or trial_type in self._trial_type_to_runner
         )
 
     def attach_trial(

ax/core/multi_type_experiment.py

@@ -104,19 +104,6 @@ def __init__(
                     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.
-
-        Args:
-            trial_type: The new trial_type to be added.
-            runner: The default runner for trials of this type.
-        """
-        if self.supports_trial_type(trial_type):
-            raise ValueError(f"Experiment already contains trial_type `{trial_type}`")
-
-        self._trial_type_to_runner[trial_type] = runner
-        return self
-
     # pyre does not support inferring the type of property setter decorators
     # or the `.fset` attribute on properties.
     # pyre-fixme[56]: Pyre was not able to infer the type of the decorator.
@@ -130,20 +117,6 @@ def optimization_config(self, optimization_config: OptimizationConfig) -> None:
                 self.default_trial_type
             )
 
-    def update_runner(self, trial_type: str, runner: Runner) -> Self:
-        """Update the default runner for an existing trial_type.
-
-        Args:
-            trial_type: The new trial_type to be added.
-            runner: The new runner for trials of this type.
-        """
-        if not self.supports_trial_type(trial_type):
-            raise ValueError(f"Experiment does not contain trial_type `{trial_type}`")
-
-        self._trial_type_to_runner[trial_type] = runner
-        self._runner = runner
-        return self
-
     def add_tracking_metric(
         self,
         metric: Metric,
@@ -232,13 +205,6 @@ def _fetch_trial_data(
         # Invoke parent's fetch method using only metrics for this trial_type
         return super()._fetch_trial_data(trial.index, metrics=metrics, **kwargs)
 
-    def supports_trial_type(self, trial_type: str | None) -> bool:
-        """Whether this experiment allows trials of the given type.
-
-        Only trial types defined in the trial_type_to_runner are allowed.
-        """
-        return trial_type in self._trial_type_to_runner.keys()
-
 
 def filter_trials_by_type(
     trials: Sequence[BaseTrial], trial_type: str | None