Remove isinstance(foo, MulitTypeExperiment) checks (#4740)

Summary:

As titled. Should be a no-op since these are still initialized as MultiTypeExperiment. Next diff will remove instantiations.

Differential Revision:
D90133237

Privacy Context Container: L1413903

Author: mpolson64

ax/core/experiment.py

@@ -1649,8 +1649,12 @@ def __repr__(self) -> str:
 
     @property
     def is_multi_type(self) -> bool:
-        """Whether this Experiment contains more than one trial type."""
-        return len(self._trial_type_to_runner) > 1
+        """
+        Whether this Experiment supports at least one non-None trial type.
+        """
+        supported_types = {*self._trial_type_to_runner.keys()}
+
+        return len(supported_types - {None}) > 0
 
     @property
     def default_trial_type(self) -> str | None:

ax/orchestration/orchestrator.py

@@ -362,7 +362,7 @@ def trial_type(self) -> str | None:
             Trial type for the experiment this Orchestrator is running if the
             experiment is a MultiTypeExperiment and None otherwise.
         """
-        if isinstance(self.experiment, MultiTypeExperiment):
+        if self.experiment.is_multi_type:
             return self.options.mt_experiment_trial_type
         return None
 
@@ -1576,7 +1576,7 @@ 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.is_multi_type:
             if options.mt_experiment_trial_type is None:
                 raise UserInputError(
                     "Must specify `mt_experiment_trial_type` for MultiTypeExperiment."

ax/orchestration/tests/test_orchestrator.py

@@ -1307,7 +1307,7 @@ def test_orchestrator_with_metric_with_new_data_after_completion(self) -> None:
         init_test_engine_and_session_factory(force_init=True)
         branin_gs = self.two_sobol_steps_GS
         # With runners & metrics, `Orchestrator.run_all_trials` should run.
-        if isinstance(self.branin_experiment, MultiTypeExperiment):
+        if self.branin_experiment.is_multi_type:
             self.branin_experiment.update_runner(
                 "type1", SyntheticRunnerWithPredictableStatusPolling()
             )
@@ -2209,7 +2209,7 @@ def test_generate_candidates_works_for_sobol(self) -> None:
         gs = get_online_sobol_mbm_generation_strategy()
 
         # this is a HITL experiment, so we don't want trials completing on their own.
-        if isinstance(self.branin_experiment, MultiTypeExperiment):
+        if self.branin_experiment.is_multi_type:
             self.branin_experiment.update_runner("type1", InfinitePollRunner())
         else:
             self.branin_experiment.runner = InfinitePollRunner()
@@ -2259,7 +2259,7 @@ def test_generate_candidates_can_remove_stale_candidates_with_ttl(
         gs = self.two_sobol_steps_GS
 
         # this is a HITL experiment, so we don't want trials completing on their own.
-        if isinstance(self.branin_experiment, MultiTypeExperiment):
+        if self.branin_experiment.is_multi_type:
             self.branin_experiment.update_runner("type1", InfinitePollRunner())
         else:
             self.branin_experiment.runner = InfinitePollRunner()
@@ -2312,7 +2312,7 @@ def test_generate_candidates_can_remove_stale_candidates(self) -> None:
         gs = self.two_sobol_steps_GS
 
         # this is a HITL experiment, so we don't want trials completing on their own.
-        if isinstance(self.branin_experiment, MultiTypeExperiment):
+        if self.branin_experiment.is_multi_type:
             self.branin_experiment.update_runner("type1", InfinitePollRunner())
         else:
             self.branin_experiment.runner = InfinitePollRunner()
@@ -2389,7 +2389,7 @@ def test_generate_candidates_does_not_fail_stale_candidates_if_fails_to_gen(
         gs = self.two_sobol_steps_GS
 
         # this is a HITL experiment, so we don't want trials completing on their own.
-        if isinstance(self.branin_experiment, MultiTypeExperiment):
+        if self.branin_experiment.is_multi_type:
             self.branin_experiment.update_runner("type1", InfinitePollRunner())
         else:
             self.branin_experiment.runner = InfinitePollRunner()

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
@@ -445,14 +444,10 @@ 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,
-            )
-        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,
+        )
 
     @copy_doc(Experiment.remove_tracking_metric)
     def remove_tracking_metric(self, metric_name: str) -> None:

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,7 +786,7 @@ def exp_to_df(
         )
 
     # Accept Experiment and SimpleExperiment
-    if isinstance(exp, MultiTypeExperiment):
+    if exp.is_multi_type:
         raise ValueError("Cannot transform MultiTypeExperiments to DataFrames.")
 
     key_components = ["trial_index", "arm_name"]

ax/storage/sqa_store/encoder.py

@@ -28,7 +28,6 @@
 from ax.core.experiment import Experiment
 from ax.core.generator_run import GeneratorRun
 from ax.core.metric import Metric
-from ax.core.multi_type_experiment import MultiTypeExperiment
 from ax.core.objective import MultiObjective, Objective, ScalarizedObjective
 from ax.core.optimization_config import (
     MultiObjectiveOptimizationConfig,
@@ -212,7 +211,7 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
             ]
             auxiliary_experiments_by_purpose[aux_exp_type] = aux_exp_jsons
         runners = []
-        if isinstance(experiment, MultiTypeExperiment):
+        if experiment.is_multi_type:
             experiment._properties[Keys.SUBCLASS] = "MultiTypeExperiment"
             for trial_type, runner in experiment._trial_type_to_runner.items():
                 runner_sqa = self.runner_to_sqa(none_throws(runner), trial_type)