Deprecate Trial.runner (#4460)

Summary:
Pull Request resolved: #4460

Upstream functionality from MultiTypeExperiment._runner_by_trial_type

Full details in: T222906773 and [doc](https://docs.google.com/document/d/1u0J0VA2VzMkJO4n-J-8ngR0PNdd4CgCmePp8HUFRuTA/edit?tab=t.z51rdls4788k), but a summary is:

 {F1982739301}

Reviewed By: mpolson64

Differential Revision: D83001393

fbshipit-source-id: e06b70c9fb085e4f2b9fd7666dc90f0017cd8c2b

Author: eonofrey

ax/core/base_trial.py

@@ -129,8 +129,6 @@ def __init__(
         self._run_metadata: dict[str, Any] = {}
         self._stop_metadata: dict[str, Any] = {}
 
-        self._runner: Runner | None = None
-
         # Counter to maintain how many arms have been named by this BatchTrial
         self._num_arms_created = 0
 
@@ -246,12 +244,30 @@ def did_not_complete(self) -> bool:
     @property
     def runner(self) -> Runner | None:
         """The runner object defining how to deploy the trial."""
-        return self._runner
+        return self.experiment.runner_for_trial_type(self.trial_type)
 
     @runner.setter
-    @immutable_once_run
     def runner(self, runner: Runner | None) -> None:
-        self._runner = runner
+        raise UnsupportedError(
+            "Setting runner on individual trials is no longer supported. "
+            "Use experiment-level runners instead."
+        )
+
+    @property
+    def _runner(self) -> Runner | None:
+        """Private runner access is not supported."""
+        raise UnsupportedError(
+            "Accessing _runner on individual trials is no longer supported. "
+            "Use trial.runner instead, which gets the runner from the experiment."
+        )
+
+    @_runner.setter
+    def _runner(self, runner: Runner | None) -> None:
+        """Private runner setting is not supported."""
+        raise UnsupportedError(
+            "Setting _runner on individual trials is no longer supported. "
+            "Use experiment-level runners instead."
+        )
 
     @property
     def deployed_name(self) -> str | None:
@@ -305,13 +321,6 @@ def _add_generator_run(self, generator_run: GeneratorRun) -> None:
         # 4. TODO: Capture which generator run the arms we are about to add this
         # this trial, came from.
 
-    def assign_runner(self) -> BaseTrial:
-        """Assigns default experiment runner if trial doesn't already have one."""
-        runner = self.experiment.runner_for_trial(self)
-        if runner is not None:
-            self._runner = runner.clone()
-        return self
-
     def update_run_metadata(self, metadata: dict[str, Any]) -> dict[str, Any]:
         """Updates the run metadata dict stored on this trial and returns the
         updated dict."""
@@ -338,15 +347,12 @@ def run(self) -> BaseTrial:
         if self.status != TrialStatus.CANDIDATE:
             raise ValueError("Can only run a candidate trial.")
 
-        # Default to experiment runner if trial doesn't have one
-        self.assign_runner()
-
-        if self._runner is None:
-            raise ValueError("No runner set on trial or experiment.")
+        if self.runner is None:
+            raise ValueError("No runner set on experiment.")
 
-        self.update_run_metadata(none_throws(self._runner).run(self))
+        self.update_run_metadata(none_throws(self.runner).run(self))
 
-        if none_throws(self._runner).staging_required:
+        if none_throws(self.runner).staging_required:
             self.mark_staged()
         else:
             self.mark_running()
@@ -379,11 +385,9 @@ def stop(self, new_status: TrialStatus, reason: str | None = None) -> BaseTrial:
                 "COMPLETED, ABANDONED or EARLY_STOPPED."
             )
 
-        # Default to experiment runner if trial doesn't have one
-        self.assign_runner()
-        if self._runner is None:
-            raise ValueError("No runner set on trial or experiment.")
-        runner = none_throws(self._runner)
+        if self.runner is None:
+            raise ValueError("No runner set on experiment.")
+        runner = none_throws(self.runner)
 
         self._stop_metadata = runner.stop(self, reason=reason)
         self.mark_as(new_status)
@@ -592,7 +596,7 @@ def mark_running(
 
         prev_step = (
             TrialStatus.STAGED
-            if self._runner is not None and self._runner.staging_required
+            if self.runner is not None and self.runner.staging_required
             else TrialStatus.CANDIDATE
         )
         prev_step_str = "staged" if prev_step == TrialStatus.STAGED else "candidate"
@@ -897,7 +901,6 @@ def _update_trial_attrs_on_clone(
         new_trial._run_metadata = deepcopy(self._run_metadata)
         new_trial._stop_metadata = deepcopy(self._stop_metadata)
         new_trial._num_arms_created = self._num_arms_created
-        new_trial.runner = self._runner.clone() if self._runner else None
 
         # Set status and reason accordingly.
         if self.status == TrialStatus.CANDIDATE:

ax/core/experiment.py

@@ -100,6 +100,7 @@ def __init__(
         default_data_type: DataType | None = None,
         auxiliary_experiments_by_purpose: None
         | (dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]]) = None,
+        default_trial_type: str | None = None,
     ) -> None:
         """Inits Experiment.
 
@@ -130,7 +131,7 @@ def __init__(
 
         self._name = name
         self.description = description
-        self.runner = runner
+        self._runner = runner
         self.is_test: bool = is_test
 
         self._data_by_trial: dict[int, OrderedDict[int, Data]] = {}
@@ -141,6 +142,12 @@ def __init__(
         self._trials: dict[int, BaseTrial] = {}
         self._properties: dict[str, Any] = properties or {}
         self._default_data_type: DataType = default_data_type or DataType.DATA
+
+        # Initialize trial type to runner mapping
+        self._default_trial_type = default_trial_type
+        self._trial_type_to_runner: dict[str | None, Runner | None] = {
+            default_trial_type: runner
+        }
         # Used to keep track of whether any trials on the experiment
         # specify a TTL. Since trials need to be checked for their TTL's
         # expiration often, having this attribute helps avoid unnecessary
@@ -391,6 +398,26 @@ def status_quo(self, status_quo: Arm | None) -> None:
 
         self._status_quo = status_quo
 
+    @property
+    def runner(self) -> Runner | None:
+        """Default runner used for trials on this experiment."""
+        return self._runner
+
+    @runner.setter
+    def runner(self, runner: Runner | None) -> None:
+        """Set the default runner and update trial type mapping."""
+        self._runner = runner
+        if runner is not None:
+            self._trial_type_to_runner[self._default_trial_type] = runner
+        else:
+            self._trial_type_to_runner = {None: None}
+
+    @runner.deleter
+    def runner(self) -> None:
+        """Delete the runner."""
+        self._runner = None
+        self._trial_type_to_runner = {None: None}
+
     @property
     def parameters(self) -> dict[str, Parameter]:
         """The parameters in the experiment's search space."""
@@ -1327,7 +1354,7 @@ def stop_trial_runs(
             reasons = [None] * len(trials)
 
         for trial, reason in zip(trials, reasons):
-            runner = self.runner_for_trial(trial=trial)
+            runner = self.runner_for_trial_type(trial_type=trial.trial_type)
             if runner is None:
                 raise RunnerNotFoundError(
                     "Unable to stop trial runs: Runner not configured "
@@ -1336,17 +1363,6 @@ def stop_trial_runs(
             runner.stop(trial=trial, reason=reason)
             trial.mark_early_stopped()
 
-    def reset_runners(self, runner: Runner) -> None:
-        """Replace all candidate trials runners.
-
-        Args:
-            runner: New runner to replace with.
-        """
-        for trial in self._trials.values():
-            if trial.status == TrialStatus.CANDIDATE:
-                trial.runner = runner
-        self.runner = runner
-
     def _attach_trial(self, trial: BaseTrial, index: int | None = None) -> int:
         """Attach a trial to this experiment.
 
@@ -1648,15 +1664,18 @@ def default_trial_type(self) -> str | None:
         In the base experiment class this is always None. For experiments
         with multiple trial types, use the MultiTypeExperiment class.
         """
-        return None
+        return self._default_trial_type
 
-    def runner_for_trial(self, trial: BaseTrial) -> Runner | None:
-        """The default runner to use for a given trial.
+    def runner_for_trial_type(self, trial_type: str | None) -> Runner | None:
+        """The default runner to use for a given trial type.
 
-        In the base experiment class, this is always the default experiment runner.
-        For experiments with multiple trial types, use the MultiTypeExperiment class.
+        Looks up the appropriate runner for this trial type in the trial_type_to_runner.
         """
-        return trial._runner if trial._runner else self.runner
+        if not self.supports_trial_type(trial_type):
+            raise ValueError(f"Trial type `{trial_type}` is not supported.")
+        if (runner := self._trial_type_to_runner.get(trial_type)) is None:
+            return self.runner  # return the default runner
+        return runner
 
     def supports_trial_type(self, trial_type: str | None) -> bool:
         """Whether this experiment allows trials of the given type.

ax/core/multi_type_experiment.py

@@ -71,13 +71,6 @@ def __init__(
             default_data_type: Enum representing the data type this experiment uses.
         """
 
-        self._default_trial_type = default_trial_type
-
-        # Map from trial type to default runner of that type
-        self._trial_type_to_runner: dict[str, Runner | None] = {
-            default_trial_type: default_runner
-        }
-
         # Specifies which trial type each metric belongs to
         self._metric_to_trial_type: dict[str, str] = {}
 
@@ -99,6 +92,8 @@ def __init__(
             properties=properties,
             default_data_type=default_data_type,
             tracking_metrics=tracking_metrics,
+            runner=default_runner,
+            default_trial_type=default_trial_type,
         )
 
     def add_trial_type(self, trial_type: str, runner: Runner) -> "MultiTypeExperiment":
@@ -138,6 +133,7 @@ def update_runner(self, trial_type: str, runner: Runner) -> "MultiTypeExperiment
             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(
@@ -159,7 +155,7 @@ def add_tracking_metric(
             raise ValueError(f"`{trial_type}` is not a supported trial type.")
 
         super().add_tracking_metric(metric)
-        self._metric_to_trial_type[metric.name] = 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
@@ -307,26 +303,6 @@ def default_trial_type(self) -> str | None:
         """Default trial type assigned to trials in this experiment."""
         return self._default_trial_type
 
-    def runner_for_trial(self, trial: BaseTrial) -> Runner | None:
-        """The default runner to use for a given trial.
-
-        Looks up the appropriate runner for this trial type in the trial_type_to_runner.
-        """
-        return (
-            trial._runner
-            if trial._runner
-            else self.runner_for_trial_type(trial_type=none_throws(trial.trial_type))
-        )
-
-    def runner_for_trial_type(self, trial_type: str) -> Runner | None:
-        """The default runner to use for a given trial type.
-
-        Looks up the appropriate runner for this trial type in the trial_type_to_runner.
-        """
-        if not self.supports_trial_type(trial_type):
-            raise ValueError(f"Trial type `{trial_type}` is not supported.")
-        return self._trial_type_to_runner[trial_type]
-
     def metrics_for_trial_type(self, trial_type: str) -> list[Metric]:
         """The default runner to use for a given trial type.
 
@@ -347,11 +323,6 @@ def supports_trial_type(self, trial_type: str | None) -> bool:
         """
         return trial_type in self._trial_type_to_runner.keys()
 
-    def reset_runners(self, runner: Runner) -> None:
-        raise NotImplementedError(
-            "MultiTypeExperiment does not support resetting all runners."
-        )
-
 
 def filter_trials_by_type(
     trials: Sequence[BaseTrial], trial_type: str | None

ax/core/tests/test_batch_trial.py

@@ -92,7 +92,7 @@ def test_UndefinedSetters(self) -> None:
             self.batch.status = TrialStatus.RUNNING
 
     def test_BasicSetter(self) -> None:
-        self.batch.runner = SyntheticRunner()
+        self.experiment.runner = SyntheticRunner()
         self.assertIsNotNone(self.batch.runner)
 
     def test_AddArm(self) -> None:
@@ -209,7 +209,7 @@ def test_BatchLifecycle(self) -> None:
         with patch.object(SyntheticRunner, "staging_required", staging_mock):
             mock_runner = SyntheticRunner()
             staging_mock.return_value = True
-            self.batch.runner = mock_runner
+            self.experiment.runner = mock_runner
             self.batch.run()
             self.assertEqual(self.batch.status, TrialStatus.STAGED)
             # Check that the trial statuses mapping on experiment has been updated.
@@ -231,7 +231,7 @@ def test_BatchLifecycle(self) -> None:
             with self.assertRaises(TrialMutationError):
                 self.batch.add_arms_and_weights(arms=self.arms, weights=self.weights)
 
-            with self.assertRaises(TrialMutationError):
+            with self.assertRaises(UnsupportedError):
                 self.batch.runner = None
 
             # Cannot run batch that was already run
@@ -317,7 +317,7 @@ def test_AbandonBatchTrial(self) -> None:
         self.assertEqual(self.batch.abandoned_reason, reason)
 
     def test_FailedBatchTrial(self) -> None:
-        self.batch.runner = SyntheticRunner()
+        self.experiment.runner = SyntheticRunner()
         self.batch.run()
         self.batch.mark_failed()
 
@@ -332,7 +332,7 @@ def test_StaleBatchTrial(self) -> None:
         self.assertIsNotNone(self.batch.time_completed)
 
     def test_EarlyStoppedBatchTrial(self) -> None:
-        self.batch.runner = SyntheticRunner()
+        self.experiment.runner = SyntheticRunner()
         self.batch.run()
         self.batch.attach_batch_trial_data(
             raw_data={
@@ -431,7 +431,7 @@ def test_Runner(self) -> None:
         with self.assertRaises(ValueError):
             self.batch.mark_running()
 
-        self.batch.runner = SyntheticRunner()
+        self.experiment.runner = SyntheticRunner()
         self.batch.run()
         self.assertEqual(self.batch.deployed_name, "test_0")
         self.assertNotEqual(len(self.batch.run_metadata.keys()), 0)

ax/core/tests/test_experiment.py

@@ -798,11 +798,10 @@ def test_ExperimentRunner(self) -> None:
         # pyre-fixme[6]: For 1st param expected `Optional[str]` but got `Dict[str,
         #  bool]`.
         new_runner = SyntheticRunner(dummy_metadata=identifier)
-
-        self.experiment.reset_runners(new_runner)
         # Don't update trials that have been run.
         self.assertEqual(batch.runner, original_runner)
         # Update default runner
+        self.experiment.runner = new_runner
         self.assertEqual(self.experiment.runner, new_runner)
         # Update candidate trial runners.
         self.assertEqual(self.experiment.trials[1].runner, new_runner)

ax/core/tests/test_multi_type_experiment.py

@@ -141,16 +141,12 @@ def test_BadBehavior(self) -> None:
         batch = self.experiment.new_batch_trial()
         batch._trial_type = "type3"  # Force override trial type
         with self.assertRaises(ValueError):
-            self.experiment.runner_for_trial(batch)
+            self.experiment.runner_for_trial_type(batch.trial_type)
 
         # Try making trial with unsupported trial type
         with self.assertRaises(ValueError):
             self.experiment.new_batch_trial(trial_type="type3")
 
-        # Try resetting runners.
-        with self.assertRaises(NotImplementedError):
-            self.experiment.reset_runners(SyntheticRunner())
-
     def test_setting_opt_config(self) -> None:
         self.assertDictEqual(
             self.experiment._metric_to_trial_type, {"m1": "type1", "m2": "type2"}