Revert D42270109: Multisect successfully blamed D42270109 for test or build failures

Summary:
This diff is reverting D42270109
D42270109: [Ax] Update the model in ModelSpec.fit if all non-data args are the same by saitcakmak has been identified to be causing the following test or build failures:

Tests affected:
- [pts:pts_test - pts.tests.test_client.PTSClientTest: test_plot_slice](https://www.internalfb.com/intern/test/562950033661552/)

Here's the Multisect link:
https://www.internalfb.com/multisect/1699980
Here are the tasks that are relevant to this breakage:

We're generating a revert to back out the changes in this diff, please note the backout may land if someone accepts it.

Reviewed By: Balandat

Differential Revision: D44089870

fbshipit-source-id: 594ac92b106ef46c878d843eda4f8f815d2132b0

Author: generatedunixname89002005232357

ax/modelbridge/base.py

@@ -94,7 +94,6 @@ def __init__(
         optimization_config: Optional[OptimizationConfig] = None,
         fit_out_of_design: bool = False,
         fit_abandoned: bool = False,
-        fit_on_init: bool = True,
     ) -> None:
         """
         Applies transforms and fits model.
@@ -123,19 +122,12 @@ def __init__(
             fit_abandoned: Whether data for abandoned arms or trials should be
                 included in model training data. If ``False``, only
                 non-abandoned points are returned.
-            fit_on_init: Whether to fit the model on initialization. This can
-                be used to skip model fitting when a fitted model is not needed.
-                To fit the model afterwards, use `_process_and_transform_data`
-                to get the transformed inputs and call `_fit_if_implemented` with
-                the transformed inputs.
         """
-        t_fit_start = time.monotonic()
+        t_fit_start = time.time()
         transforms = transforms or []
         # pyre-ignore: Cast is a Tranform
         transforms: List[Type[Transform]] = [Cast] + transforms
 
-        self.fit_time: float = 0.0
-        self.fit_time_since_gen: float = 0.0
         self._metric_names: Set[str] = set()
         self._training_data: List[Observation] = []
         self._optimization_config: Optional[OptimizationConfig] = optimization_config
@@ -146,89 +138,56 @@ def __init__(
         self._model_key: Optional[str] = None
         self._model_kwargs: Optional[Dict[str, Any]] = None
         self._bridge_kwargs: Optional[Dict[str, Any]] = None
-        self._model_space: SearchSpace = search_space.clone()
+
+        # pyre-fixme[4]: Attribute must be annotated.
+        self._model_space = search_space.clone()
         self._raw_transforms = transforms
         self._transform_configs: Optional[Dict[str, TConfig]] = transform_configs
         self._fit_out_of_design = fit_out_of_design
         self._fit_abandoned = fit_abandoned
-        self._experiment_has_immutable_search_space_and_opt_config: bool = (
-            experiment is not None and experiment.immutable_search_space_and_opt_config
-        )
+        imm = experiment and experiment.immutable_search_space_and_opt_config
+        # pyre-fixme[4]: Attribute must be annotated.
+        self._experiment_has_immutable_search_space_and_opt_config = imm
         if experiment is not None:
             if self._optimization_config is None:
                 self._optimization_config = experiment.optimization_config
             self._arms_by_signature = experiment.arms_by_signature
 
-        # Convert Data to Observations, transform observations & search space.
-        observations, search_space = self._process_and_transform_data(
-            experiment=experiment, data=data
-        )
+        # Convert Data to Observations
+        observations = self._prepare_observations(experiment=experiment, data=data)
 
-        # Set model status quo.
+        observations_raw = self._set_training_data(
+            observations=observations, search_space=search_space
+        )
+        # Set model status quo
         # NOTE: training data must be set before setting the status quo.
         self._set_status_quo(
             experiment=experiment,
             status_quo_name=status_quo_name,
             status_quo_features=status_quo_features,
         )
+        observations, search_space = self._transform_data(
+            observations=observations_raw,
+            search_space=search_space,
+            transforms=transforms,
+            transform_configs=transform_configs,
+        )
 
-        # Save model, apply terminal transform, and fit.
+        # Save model, apply terminal transform, and fit
         self.model = model
-        if fit_on_init:
-            self._fit_if_implemented(
-                search_space=search_space,
-                observations=observations,
-                time_so_far=time.monotonic() - t_fit_start,
-            )
-
-    def _fit_if_implemented(
-        self,
-        search_space: SearchSpace,
-        observations: List[Observation],
-        time_so_far: float,
-    ) -> None:
-        r"""Fits the model if `_fit` is implemented and stores fit time.
-
-        Args:
-            search_space: A transformed search space for fitting the model.
-            observations: The observations to fit the model with. These should
-                also be transformed.
-            time_so_far: Time spent in initializing the model up to
-                `_fit_if_implemented` call.
-        """
         try:
-            t_fit_start = time.monotonic()
             self._fit(
-                model=self.model,
+                model=model,
                 search_space=search_space,
                 observations=observations,
             )
-            self.fit_time += time.monotonic() - t_fit_start + time_so_far
-            self.fit_time_since_gen += self.fit_time
+            # pyre-fixme[4]: Attribute must be annotated.
+            self.fit_time = time.time() - t_fit_start
+            # pyre-fixme[4]: Attribute must be annotated.
+            self.fit_time_since_gen = float(self.fit_time)
         except NotImplementedError:
-            pass
-
-    def _process_and_transform_data(
-        self,
-        experiment: Optional[Experiment] = None,
-        data: Optional[Data] = None,
-    ) -> Tuple[List[Observation], SearchSpace]:
-        r"""Processes the data into observations and returns transformed
-        observations and the search space. This packages the following methods:
-        * self._prepare_observations
-        * self._set_training_data
-        * self._transform_data
-        """
-        observations = self._prepare_observations(experiment=experiment, data=data)
-        observations_raw = self._set_training_data(
-            observations=observations, search_space=self._model_space
-        )
-        return self._transform_data(
-            observations=observations_raw,
-            search_space=self._model_space,
-            transforms=self._raw_transforms,
-            transform_configs=self._transform_configs,
-        )
+            self.fit_time = 0.0
+            self.fit_time_since_gen = 0.0
 
     def _prepare_observations(
         self, experiment: Optional[Experiment], data: Optional[Data]
@@ -607,7 +566,7 @@ def update(self, new_data: Data, experiment: Experiment) -> None:
                 `update`.
             experiment: Experiment, in which this data was obtained.
         """
-        t_update_start = time.monotonic()
+        t_update_start = time.time()
         observations = self._prepare_observations(experiment=experiment, data=new_data)
         obs_raw = self._extend_training_data(observations=observations)
         observations, search_space = self._transform_data(
@@ -620,8 +579,8 @@ def update(self, new_data: Data, experiment: Experiment) -> None:
             search_space=search_space,
             observations=observations,
         )
-        self.fit_time += time.monotonic() - t_update_start
-        self.fit_time_since_gen += time.monotonic() - t_update_start
+        self.fit_time += time.time() - t_update_start
+        self.fit_time_since_gen += time.time() - t_update_start
 
     def _update(
         self,

ax/modelbridge/map_torch.py

@@ -63,7 +63,6 @@ def __init__(
         status_quo_features: Optional[ObservationFeatures] = None,
         optimization_config: Optional[OptimizationConfig] = None,
         fit_out_of_design: bool = False,
-        fit_on_init: bool = True,
         default_model_gen_options: Optional[TConfig] = None,
         map_data_limit_rows_per_metric: Optional[int] = None,
         map_data_limit_rows_per_group: Optional[int] = None,
@@ -94,11 +93,6 @@ def __init__(
                 the model.
             fit_out_of_design: If specified, all training data is returned.
                 Otherwise, only in design points are returned.
-            fit_on_init: Whether to fit the model on initialization. This can
-                be used to skip model fitting when a fitted model is not needed.
-                To fit the model afterwards, use `_process_and_transform_data`
-                to get the transformed inputs and call `_fit_if_implemented` with
-                the transformed inputs.
             default_model_gen_options: Options passed down to `model.gen(...)`.
             map_data_limit_rows_per_metric: Subsample the map data so that the
                 total number of rows per metric is limited by this value.
@@ -129,7 +123,6 @@ def __init__(
             status_quo_features=status_quo_features,
             optimization_config=optimization_config,
             fit_out_of_design=fit_out_of_design,
-            fit_on_init=fit_on_init,
             default_model_gen_options=default_model_gen_options,
         )
 

ax/modelbridge/model_spec.py

@@ -68,9 +68,6 @@ class ModelSpec(Base):
     # stored cross validation diagnostics set in cross validate
     _diagnostics: Optional[CVDiagnostics] = None
 
-    # Stored to check if the model can be safely updated in fit.
-    _last_fit_arg_ids: Optional[Dict[str, int]] = None
-
     def __post_init__(self) -> None:
         self.model_kwargs = self.model_kwargs or {}
         self.model_gen_kwargs = self.model_gen_kwargs or {}
@@ -128,28 +125,11 @@ def fit(
         # adding contents of `model_kwargs` passed to this method, to
         # `self.model_kwargs`.
         combined_model_kwargs = {**(self.model_kwargs or {}), **model_kwargs}
-        if self._fitted_model is not None and self._safe_to_update(
-            experiment=experiment, combined_model_kwargs=combined_model_kwargs
-        ):
-            # Update the data on the modelbridge and call `_fit`.
-            # This will skip model fitting if the data has not changed.
-            observations, search_space = self.fitted_model._process_and_transform_data(
-                experiment=experiment, data=data
-            )
-            self.fitted_model._fit_if_implemented(
-                search_space=search_space, observations=observations, time_so_far=0.0
-            )
-
-        else:
-            # Fit from scratch.
-            self._fitted_model = self.model_enum(
-                experiment=experiment,
-                data=data,
-                **combined_model_kwargs,
-            )
-            self._last_fit_arg_ids = self._get_fit_arg_ids(
-                experiment=experiment, combined_model_kwargs=combined_model_kwargs
-            )
+        self._fitted_model = self.model_enum(
+            experiment=experiment,
+            data=data,
+            **combined_model_kwargs,
+        )
 
     def cross_validate(
         self,
@@ -235,32 +215,6 @@ def copy(self) -> ModelSpec:
             model_cv_kwargs=deepcopy(self.model_cv_kwargs),
         )
 
-    def _safe_to_update(
-        self,
-        experiment: Experiment,
-        combined_model_kwargs: Dict[str, Any],
-    ) -> bool:
-        """Checks if the object id of any of the non-data fit arguments has changed.
-
-        This is a cheap way of checking that we're attempting to re-fit the same
-        model for the same experiment, which is a very reasonable expectation
-        since this all happens on the same `ModelSpec` instance.
-        """
-        return self._last_fit_arg_ids == self._get_fit_arg_ids(
-            experiment=experiment, combined_model_kwargs=combined_model_kwargs
-        )
-
-    def _get_fit_arg_ids(
-        self,
-        experiment: Experiment,
-        combined_model_kwargs: Dict[str, Any],
-    ) -> Dict[str, int]:
-        """Construct a dictionary mapping arg name to object id."""
-        return {
-            "experiment": id(experiment),
-            **{k: id(v) for k, v in combined_model_kwargs.items()},
-        }
-
     def _assert_fitted(self) -> None:
         """Helper that verifies a model was fitted, raising an error if not"""
         if self._fitted_model is None:

ax/modelbridge/tests/test_base_modelbridge.py

@@ -258,15 +258,6 @@ def test_ModelBridge(
             modelbridge.transform_observation_features([get_observation2().features])
         mock_tr.assert_called_with(modelbridge, [get_observation2trans().features])
 
-        # Test that fit is not called when fit_on_init = False.
-        mock_fit.reset_mock()
-        modelbridge = ModelBridge(
-            search_space=ss,
-            model=Model(),
-            fit_on_init=False,
-        )
-        self.assertEqual(mock_fit.call_count, 0)
-
     @mock.patch(
         "ax.modelbridge.base.observations_from_data",
         autospec=True,

ax/modelbridge/tests/test_generation_strategy.py

@@ -59,9 +59,7 @@ def setUp(self) -> None:
             f"{TorchModelBridge.__module__}.TorchModelBridge", spec=True
         )
         self.mock_torch_model_bridge = self.torch_model_bridge_patcher.start()
-        mock_mb = self.mock_torch_model_bridge.return_value
-        mock_mb.gen.return_value = self.gr
-        mock_mb._process_and_transform_data.return_value = (None, None)
+        self.mock_torch_model_bridge.return_value.gen.return_value = self.gr
 
         # Mock out slow TS.
         self.discrete_model_bridge_patcher = patch(
@@ -304,7 +302,6 @@ def test_sobol_GPEI_strategy(self) -> None:
                         "transforms": Cont_X_trans,
                         "fit_out_of_design": False,
                         "fit_abandoned": False,
-                        "fit_on_init": True,
                     },
                 )
                 ms = g._model_state_after_gen

ax/modelbridge/tests/test_model_spec.py

@@ -5,14 +5,12 @@
 # LICENSE file in the root directory of this source tree.
 
 import warnings
-from unittest import mock
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import Mock, patch
 
 from ax.core.observation import ObservationFeatures
 from ax.exceptions.core import UserInputError
 from ax.modelbridge.factory import get_sobol
 from ax.modelbridge.model_spec import FactoryFunctionModelSpec, ModelSpec
-from ax.modelbridge.modelbridge_utils import extract_search_space_digest
 from ax.modelbridge.registry import Models
 from ax.utils.common.testutils import TestCase
 from ax.utils.testing.core_stubs import get_branin_experiment
@@ -41,49 +39,23 @@ def test_construct(self) -> None:
         with self.assertRaises(NotImplementedError):
             ms.update(experiment=self.experiment, new_data=self.data)
 
-    @fast_botorch_optimize
-    # We can use `extract_search_space_digest` as a surrogate for executing
-    # the full TorchModelBridge._fit.
-    @mock.patch(
-        "ax.modelbridge.torch.extract_search_space_digest",
-        wraps=extract_search_space_digest,
-    )
-    def test_fit(self, wrapped_extract_ssd: Mock) -> None:
-        ms = ModelSpec(model_enum=Models.GPEI)
-        # This should fit the model as usual.
-        ms.fit(experiment=self.experiment, data=self.data)
-        wrapped_extract_ssd.assert_called_once()
-        self.assertIsNotNone(ms._last_fit_arg_ids)
-        self.assertEqual(ms._last_fit_arg_ids["experiment"], id(self.experiment))
-        # This should skip the model fit.
-        with mock.patch("ax.modelbridge.torch.logger") as mock_logger:
-            ms.fit(experiment=self.experiment, data=self.data)
-        mock_logger.info.assert_called_with(
-            "The observations are identical to the last set of observations "
-            "used to fit the model. Skipping model fitting."
-        )
-        wrapped_extract_ssd.assert_called_once()
-
     def test_model_key(self) -> None:
         ms = ModelSpec(model_enum=Models.GPEI)
         self.assertEqual(ms.model_key, "GPEI")
 
     @patch(f"{ModelSpec.__module__}.compute_diagnostics")
     @patch(f"{ModelSpec.__module__}.cross_validate", return_value=["fake-cv-result"])
-    def test_cross_validate_with_GP_model(
-        self, mock_cv: Mock, mock_diagnostics: Mock
-    ) -> None:
+    # pyre-fixme[3]: Return type must be annotated.
+    def test_cross_validate_with_GP_model(self, mock_cv: Mock, mock_diagnostics: Mock):
         mock_enum = Mock()
-        fake_mb = MagicMock()
-        fake_mb._process_and_transform_data = MagicMock(return_value=(None, None))
-        mock_enum.return_value = fake_mb
+        mock_enum.return_value = "fake-modelbridge"
         ms = ModelSpec(model_enum=mock_enum, model_cv_kwargs={"test_key": "test-value"})
         ms.fit(
             experiment=self.experiment,
             data=self.experiment.trials[0].fetch_data(),
         )
         cv_results, cv_diagnostics = ms.cross_validate()
-        mock_cv.assert_called_with(model=fake_mb, test_key="test-value")
+        mock_cv.assert_called_with(model="fake-modelbridge", test_key="test-value")
         mock_diagnostics.assert_called_with(["fake-cv-result"])
 
         self.assertIsNotNone(cv_results)
@@ -112,14 +84,15 @@ def test_cross_validate_with_GP_model(
 
             self.assertIsNotNone(cv_results)
             self.assertIsNotNone(cv_diagnostics)
-            mock_cv.assert_called_with(model=fake_mb, test_key="test-value")
+            mock_cv.assert_called_with(model="fake-modelbridge", test_key="test-value")
             mock_diagnostics.assert_called_with(["fake-cv-result"])
 
     @patch(f"{ModelSpec.__module__}.compute_diagnostics")
     @patch(f"{ModelSpec.__module__}.cross_validate", side_effect=NotImplementedError)
+    # pyre-fixme[3]: Return type must be annotated.
     def test_cross_validate_with_non_GP_model(
         self, mock_cv: Mock, mock_diagnostics: Mock
-    ) -> None:
+    ):
         mock_enum = Mock()
         mock_enum.return_value = "fake-modelbridge"
         ms = ModelSpec(model_enum=mock_enum, model_cv_kwargs={"test_key": "test-value"})