Remove bifrucation around MultiTypeExperiment (#4874)

Summary:

With recent changes to experiment we no longer need this bifructation.

Next diff will remove places where we construct MultiTypeExperiment, and the one after will deprecate the class entirely

Differential Revision: D91920991

Author: mpolson64

ax/orchestration/orchestrator.py

@@ -27,7 +27,6 @@
 from ax.core.multi_type_experiment import (
     filter_trials_by_type,
     get_trial_indices_for_statuses,
-    MultiTypeExperiment,
 )
 from ax.core.runner import Runner
 from ax.core.trial import Trial
@@ -367,17 +366,11 @@ def options(self, options: OrchestratorOptions) -> None:
     def trial_type(self) -> str:
         """Trial type for the experiment this Orchestrator is running.
 
-        This returns None if the experiment is not a MultitypeExperiment
-
         Returns:
-            Trial type for the experiment this Orchestrator is running if the
-            experiment is a MultiTypeExperiment and None otherwise.
+            Trial type for the experiment this Orchestrator is running.
+            Defaults to Keys.DEFAULT_TRIAL_TYPE if not specified.
         """
-        if isinstance(self.experiment, MultiTypeExperiment):
-            return (
-                self.options.mt_experiment_trial_type or Keys.DEFAULT_TRIAL_TYPE.value
-            )
-        return Keys.DEFAULT_TRIAL_TYPE.value
+        return self.options.mt_experiment_trial_type or Keys.DEFAULT_TRIAL_TYPE.value
 
     @property
     def running_trials(self) -> list[BaseTrial]:
@@ -1619,23 +1612,14 @@ 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 options.mt_experiment_trial_type is None:
-                raise UserInputError(
-                    "Must specify `mt_experiment_trial_type` for MultiTypeExperiment."
-                )
+        if options.mt_experiment_trial_type is not None:
             if not self.experiment.supports_trial_type(
                 options.mt_experiment_trial_type
             ):
                 raise ValueError(
                     "Experiment does not support trial type "
                     f"{options.mt_experiment_trial_type}."
                 )
-        elif options.mt_experiment_trial_type is not None:
-            raise UserInputError(
-                "`mt_experiment_trial_type` must be None unless the experiment is a "
-                "MultiTypeExperiment."
-            )
 
     def _get_max_pending_trials(self) -> int:
         """Returns the maximum number of pending trials specified in the options, or

ax/orchestration/tests/test_orchestrator.py

@@ -2736,12 +2736,9 @@ def test_validate_options_not_none_mt_trial_type(
         self, msg: str | None = None
     ) -> None:
         # test that error is raised if `mt_experiment_trial_type` is not
-        # compatible with the type of experiment (single or multi-type)
+        # a supported trial type for this experiment
         if msg is None:
-            msg = (
-                "`mt_experiment_trial_type` must be None unless the experiment is a "
-                "MultiTypeExperiment."
-            )
+            msg = "Experiment does not support trial type type1."
         options = OrchestratorOptions(
             init_seconds_between_polls=0,  # No wait bw polls so test is fast.
             batch_size=10,
@@ -2752,7 +2749,7 @@ def test_validate_options_not_none_mt_trial_type(
             ),
         )
         gs = self.two_sobol_steps_GS
-        with self.assertRaisesRegex(UserInputError, msg):
+        with self.assertRaisesRegex(ValueError, msg):
             Orchestrator(
                 experiment=self.branin_experiment,
                 generation_strategy=gs,
@@ -3010,10 +3007,11 @@ def test_fetch_and_process_trials_data_results_failed_non_objective(
     def test_validate_options_not_none_mt_trial_type(
         self, msg: str | None = None
     ) -> None:
-        # test if a MultiTypeExperiment with `mt_experiment_trial_type=None`
-        self.orchestrator_options_kwargs["mt_experiment_trial_type"] = None
+        # test that error is raised if `mt_experiment_trial_type` is not
+        # a supported trial type for this experiment (using an invalid type)
+        self.orchestrator_options_kwargs["mt_experiment_trial_type"] = "invalid_type"
         super().test_validate_options_not_none_mt_trial_type(
-            msg="Must specify `mt_experiment_trial_type` for MultiTypeExperiment."
+            msg="Experiment does not support trial type invalid_type."
         )
 
     def test_run_n_trials_single_step_existing_experiment(

ax/service/ax_client.py

@@ -24,7 +24,6 @@
 from ax.core.evaluations_to_data import raw_evaluations_to_data
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
-from ax.core.multi_type_experiment import MultiTypeExperiment
 from ax.core.objective import MultiObjective, Objective
 from ax.core.observation import ObservationFeatures
 from ax.core.runner import Runner
@@ -458,15 +457,11 @@ def add_tracking_metrics(
             for metric_name in metric_names
         ]
 
-        if isinstance(self.experiment, MultiTypeExperiment):
-            experiment = assert_is_instance(self.experiment, MultiTypeExperiment)
-            experiment.add_tracking_metrics(
-                metrics=metric_objects,
-                metrics_to_trial_types=metrics_to_trial_types,
-                canonical_names=canonical_names,
-            )
-        else:
-            self.experiment.add_tracking_metrics(metrics=metric_objects)
+        self.experiment.add_tracking_metrics(
+            metrics=metric_objects,
+            metrics_to_trial_types=metrics_to_trial_types,
+            **({"canonical_names": canonical_names} if canonical_names else {}),
+        )
 
     @copy_doc(Experiment.remove_tracking_metric)
     def remove_tracking_metric(self, metric_name: str) -> None:

ax/service/tests/test_report_utils.py

@@ -199,9 +199,9 @@ def test_exp_to_df_with_failure(self) -> None:
         self.assertEqual(f"{fail_reason}...", df["reason"].iloc[0])
 
     def test_exp_to_df(self) -> None:
-        # MultiTypeExperiment should fail
+        # Experiments with multiple trial types should fail
         exp = get_multi_type_experiment()
-        with self.assertRaisesRegex(ValueError, "MultiTypeExperiment"):
+        with self.assertRaisesRegex(ValueError, "multiple trial types"):
             exp_to_df(exp=exp)
 
         # exp with no trials should return empty results

ax/service/utils/report_utils.py

@@ -32,7 +32,6 @@
 from ax.core.generator_run import GeneratorRunType
 from ax.core.map_metric import MapMetric
 from ax.core.metric import Metric
-from ax.core.multi_type_experiment import MultiTypeExperiment
 from ax.core.objective import MultiObjective, ScalarizedObjective
 from ax.core.optimization_config import (
     MultiObjectiveOptimizationConfig,
@@ -787,8 +786,11 @@ def exp_to_df(
         )
 
     # Accept Experiment and SimpleExperiment
-    if isinstance(exp, MultiTypeExperiment):
-        raise ValueError("Cannot transform MultiTypeExperiments to DataFrames.")
+    # Reject experiments with multiple trial types as they need special handling
+    if len(exp._trial_type_to_runner) > 1:
+        raise ValueError(
+            "Cannot transform experiments with multiple trial types to DataFrames."
+        )
 
     key_components = ["trial_index", "arm_name"]