Replace isinstance(MultiTypeExperiment) checks with feature checks (facebook#5001)

Summary:
Pull Request resolved: facebook#5001

Move `filter_trials_by_type` and `get_trial_indices_for_statuses` from
`multi_type_experiment.py` to `experiment.py`, with backward-compatible
re-exports from the old module.

Replace all `isinstance(experiment, MultiTypeExperiment)` checks with
`experiment.default_trial_type is not None` feature checks in:
- `ax/orchestration/orchestrator.py`
- `ax/fb/realtime/utils.py`
- `ax/fb/axolotl/utils_chronos.py`

Replace `assert_is_instance(experiment, MultiTypeExperiment).method()`
calls with direct calls on base `Experiment` (which now has
`runner_for_trial_type` and `metrics_for_trial_type`).

This is Phase 6 of the MultiTypeExperiment pull-up plan.

Differential Revision: D94999020

Author: mpolson64

ax/core/experiment.py

@@ -2739,3 +2739,47 @@ def add_arm_and_prevent_naming_collision(
                 stacklevel=2,
             )
             new_trial.add_arm(none_throws(old_trial.arm).clone(clear_name=True))
+
+
+def filter_trials_by_type(
+    trials: Sequence[BaseTrial], trial_type: str | None
+) -> list[BaseTrial]:
+    """Filter trials by trial type if provided.
+
+    This filters trials by trial type if the experiment has multiple
+    trial types.
+
+    Args:
+        trials: Trials to filter.
+        trial_type: The trial type to filter by. If None, all trials are returned.
+
+    Returns:
+        Filtered trials.
+    """
+    if trial_type is not None:
+        return [t for t in trials if t.trial_type == trial_type]
+    return list(trials)
+
+
+def get_trial_indices_for_statuses(
+    experiment: Experiment, statuses: set[TrialStatus], trial_type: str | None = None
+) -> set[int]:
+    """Get trial indices for a set of statuses.
+
+    Args:
+        experiment: The experiment to get trial indices from.
+        statuses: Set of statuses to get trial indices for.
+        trial_type: If provided, only return indices for trials of this type.
+
+    Returns:
+        Set of trial indices for the given statuses.
+    """
+    return {
+        i
+        for i, t in experiment.trials.items()
+        if (t.status in statuses)
+        and (
+            (trial_type is None)
+            or ((trial_type is not None) and (t.trial_type == trial_type))
+        )
+    }

ax/core/multi_type_experiment.py

@@ -6,12 +6,14 @@
 
 # pyre-strict
 
-from collections.abc import Sequence
 from typing import Any, Self
 
 from ax.core.arm import Arm
-from ax.core.base_trial import BaseTrial, TrialStatus
-from ax.core.experiment import Experiment
+from ax.core.experiment import (
+    Experiment,
+    filter_trials_by_type,
+    get_trial_indices_for_statuses,
+)
 from ax.core.metric import Metric
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.runner import Runner
@@ -171,42 +173,9 @@ def remove_metric(self, metric_name: str) -> Self:
         return self
 
 
-def filter_trials_by_type(
-    trials: Sequence[BaseTrial], trial_type: str | None
-) -> list[BaseTrial]:
-    """Filter trials by trial type if provided.
-
-    This filters trials by trial type if the experiment is a
-    MultiTypeExperiment.
-
-    Args:
-        trials: Trials to filter.
-
-    Returns:
-        Filtered trials.
-    """
-    if trial_type is not None:
-        return [t for t in trials if t.trial_type == trial_type]
-    return list(trials)
-
-
-def get_trial_indices_for_statuses(
-    experiment: Experiment, statuses: set[TrialStatus], trial_type: str | None = None
-) -> set[int]:
-    """Get trial indices for a set of statuses.
-
-    Args:
-        statuses: Set of statuses to get trial indices for.
-
-    Returns:
-        Set of trial indices for the given statuses.
-    """
-    return {
-        i
-        for i, t in experiment.trials.items()
-        if (t.status in statuses)
-        and (
-            (trial_type is None)
-            or ((trial_type is not None) and (t.trial_type == trial_type))
-        )
-    }
+# Re-exported from ax.core.experiment for backward compatibility.
+__all__ = [
+    "MultiTypeExperiment",
+    "filter_trials_by_type",
+    "get_trial_indices_for_statuses",
+]

ax/orchestration/orchestrator.py

@@ -21,14 +21,13 @@
 from ax.adapter.adapter_utils import get_fixed_features_from_experiment
 from ax.adapter.base import Adapter
 from ax.core.base_trial import BaseTrial
-from ax.core.experiment import Experiment
-from ax.core.generator_run import GeneratorRun
-from ax.core.metric import Metric, MetricFetchE, MetricFetchResult
-from ax.core.multi_type_experiment import (
+from ax.core.experiment import (
+    Experiment,
     filter_trials_by_type,
     get_trial_indices_for_statuses,
-    MultiTypeExperiment,
 )
+from ax.core.generator_run import GeneratorRun
+from ax.core.metric import Metric, MetricFetchE, MetricFetchResult
 from ax.core.runner import Runner
 from ax.core.trial import Trial
 from ax.core.trial_status import TrialStatus
@@ -58,7 +57,7 @@
     set_ax_logger_levels,
 )
 from ax.utils.common.timeutils import current_timestamp_in_millis
-from pyre_extensions import assert_is_instance, none_throws
+from pyre_extensions import none_throws
 
 
 NOT_IMPLEMENTED_IN_BASE_CLASS_MSG = """ \
@@ -367,21 +366,21 @@ def options(self, options: OrchestratorOptions) -> None:
     def trial_type(self) -> str | None:
         """Trial type for the experiment this Orchestrator is running.
 
-        This returns None if the experiment is not a MultitypeExperiment
+        This returns None if the experiment does not have multiple trial types.
 
         Returns:
             Trial type for the experiment this Orchestrator is running if the
-            experiment is a MultiTypeExperiment and None otherwise.
+            experiment has multiple trial types and None otherwise.
         """
-        if isinstance(self.experiment, MultiTypeExperiment):
+        if self.experiment.default_trial_type is not None:
             return self.options.mt_experiment_trial_type
         return None
 
     @property
     def running_trials(self) -> list[BaseTrial]:
         """Currently running trials.
 
-        Note: if the experiment is a MultiTypeExperiment, then this will
+        Note: if the experiment has multiple trial types, then this will
         only fetch trials of type `orchestrator.trial_type`.
 
 
@@ -397,7 +396,7 @@ def running_trials(self) -> list[BaseTrial]:
     def trials(self) -> list[BaseTrial]:
         """All trials.
 
-        Note: if the experiment is a MultiTypeExperiment, then this will
+        Note: if the experiment has multiple trial types, then this will
         only fetch trials of type `orchestrator.trial_type`.
 
         Returns:
@@ -424,7 +423,7 @@ def running_trial_indices(self) -> set[int]:
     def failed_abandoned_trial_indices(self) -> set[int]:
         """Failed or abandoned trials.
 
-        Note: if the experiment is a MultiTypeExperiment, then this will
+        Note: if the experiment has multiple trial types, then this will
         only fetch trials of type `orchestrator.trial_type`.
 
         Returns:
@@ -441,7 +440,7 @@ def pending_trials(self) -> list[BaseTrial]:
         """Running or staged trials on the experiment this Orchestrator is
         running.
 
-        Note: if the experiment is a MultiTypeExperiment, then this will
+        Note: if the experiment has multiple trial types, then this will
         only fetch trials of type `orchestrator.trial_type`.
 
         Returns:
@@ -457,7 +456,7 @@ def pending_trials(self) -> list[BaseTrial]:
     def candidate_trials(self) -> list[BaseTrial]:
         """Candidate trials on the experiment this Orchestrator is running.
 
-        Note: if the experiment is a MultiTypeExperiment, then this will
+        Note: if the experiment has multiple trial types, then this will
         only fetch trials of type `orchestrator.trial_type`.
 
         Returns:
@@ -472,7 +471,7 @@ def candidate_trials(self) -> list[BaseTrial]:
     def trials_expecting_data(self) -> list[BaseTrial]:
         """Trials expecting data.
 
-        Note: if the experiment is a MultiTypeExperiment, then this will
+        Note: if the experiment has multiple trial types, then this will
         only fetch trials of type `orchestrator.trial_type`.
         """
         trials = []
@@ -488,9 +487,9 @@ def runner(self) -> Runner:
         instance.
         """
         if self.trial_type is not None:
-            runner = assert_is_instance(
-                self.experiment, MultiTypeExperiment
-            ).runner_for_trial_type(trial_type=none_throws(self.trial_type))
+            runner = self.experiment.runner_for_trial_type(
+                trial_type=none_throws(self.trial_type)
+            )
         else:
             runner = self.experiment.runner
         if runner is None:
@@ -1626,10 +1625,11 @@ def _validate_options(self, options: OrchestratorOptions) -> None:
                     "will be unable to fetch intermediate results with which to "
                     "evaluate early stopping criteria."
                 )
-        if isinstance(self.experiment, MultiTypeExperiment):
+        if self.experiment.default_trial_type is not None:
             if options.mt_experiment_trial_type is None:
                 raise UserInputError(
-                    "Must specify `mt_experiment_trial_type` for MultiTypeExperiment."
+                    "Must specify `mt_experiment_trial_type` for experiments "
+                    "with multiple trial types."
                 )
             if not self.experiment.supports_trial_type(
                 options.mt_experiment_trial_type
@@ -1640,8 +1640,8 @@ def _validate_options(self, options: OrchestratorOptions) -> None:
                 )
         elif options.mt_experiment_trial_type is not None:
             raise UserInputError(
-                "`mt_experiment_trial_type` must be None unless the experiment is a "
-                "MultiTypeExperiment."
+                "`mt_experiment_trial_type` must be None unless the experiment "
+                "has multiple trial types."
             )
 
     def _get_max_pending_trials(self) -> int:
@@ -2040,9 +2040,9 @@ def _fetch_and_process_trials_data_results(
         try:
             kwargs = deepcopy(self.options.fetch_kwargs)
             if self.trial_type is not None:
-                metrics = assert_is_instance(
-                    self.experiment, MultiTypeExperiment
-                ).metrics_for_trial_type(trial_type=none_throws(self.trial_type))
+                metrics = self.experiment.metrics_for_trial_type(
+                    trial_type=none_throws(self.trial_type)
+                )
                 kwargs["metrics"] = metrics
             results = self.experiment.fetch_trials_data_results(
                 trial_indices=trial_indices,

ax/orchestration/tests/test_orchestrator.py

@@ -2796,8 +2796,8 @@ def test_validate_options_not_none_mt_trial_type(
         # compatible with the type of experiment (single or multi-type)
         if msg is None:
             msg = (
-                "`mt_experiment_trial_type` must be None unless the experiment is a "
-                "MultiTypeExperiment."
+                "`mt_experiment_trial_type` must be None unless the experiment "
+                "has multiple trial types."
             )
         options = OrchestratorOptions(
             init_seconds_between_polls=0,  # No wait bw polls so test is fast.
@@ -3080,7 +3080,8 @@ def test_validate_options_not_none_mt_trial_type(
         # test if a MultiTypeExperiment with `mt_experiment_trial_type=None`
         self.orchestrator_options_kwargs["mt_experiment_trial_type"] = None
         super().test_validate_options_not_none_mt_trial_type(
-            msg="Must specify `mt_experiment_trial_type` for MultiTypeExperiment."
+            msg="Must specify `mt_experiment_trial_type` for experiments "
+            "with multiple trial types."
         )
 
     def test_run_n_trials_single_step_existing_experiment(