Remove unused refit_on_update kwarg

Summary:
This has no usage since the `update` method has been deprecated. Removing it to clean up the API.

Updated the decoders to make sure we can load previous GS.

Reviewed By: Balandat

Differential Revision: D56796465

fbshipit-source-id: e759b382a50f8644f76e1ea44f671a672646d977

Author: saitcakmak

ax/modelbridge/factory.py

@@ -245,7 +245,6 @@ def get_botorch(
     acqf_constructor: TAcqfConstructor = get_qLogNEI,
     acqf_optimizer: TOptimizer = scipy_optimizer,  # pyre-ignore[9]
     refit_on_cv: bool = False,
-    refit_on_update: bool = True,
     optimization_config: Optional[OptimizationConfig] = None,
 ) -> TorchModelBridge:
     """Instantiates a BotorchModel."""
@@ -266,7 +265,6 @@ def get_botorch(
             acqf_constructor=acqf_constructor,
             acqf_optimizer=acqf_optimizer,
             refit_on_cv=refit_on_cv,
-            refit_on_update=refit_on_update,
             optimization_config=optimization_config,
         ),
     )

ax/modelbridge/tests/test_registry.py

@@ -153,7 +153,6 @@ def test_enum_sobol_GPEI(self) -> None:
                     "value": f"{botorch_defaults}.recommend_best_observed_point",
                 },
                 "refit_on_cv": False,
-                "refit_on_update": True,
                 "warm_start_refitting": True,
                 "use_input_warping": False,
                 "use_loocv_pseudo_likelihood": False,

ax/models/tests/test_alebo.py

@@ -286,7 +286,6 @@ def test_ALEBO(self) -> None:
         self.assertTrue(torch.equal(B, m.B))
         self.assertEqual(m.laplace_nsamp, 5)
         self.assertEqual(m.fit_restarts, 1)
-        self.assertEqual(m.refit_on_update, True)
         self.assertEqual(m.refit_on_cv, False)
         self.assertEqual(m.warm_start_refitting, False)
 

ax/models/torch/alebo.py

@@ -792,7 +792,7 @@ def alebo_acqf_optimizer(
                     if base_X_pending is not None
                     else candidates
                 )
-        logger.info(f"Generated sequential candidate {i+1} of {n}")
+        logger.info(f"Generated sequential candidate {i + 1} of {n}")
     if acq_has_X_pend:
         acq_function.set_X_pending(base_X_pending)
     return candidates, torch.stack(acq_value_list)
@@ -828,7 +828,6 @@ def __init__(
         self.laplace_nsamp = laplace_nsamp
         self.fit_restarts = fit_restarts
         super().__init__(
-            refit_on_update=True,  # Important to not get stuck in local opt.
             refit_on_cv=False,
             warm_start_refitting=False,
             acqf_constructor=ei_or_nei,  # pyre-ignore

ax/models/torch/botorch.py

@@ -125,8 +125,6 @@ class BotorchModel(TorchModel):
             signature as described below.
         refit_on_cv: If True, refit the model for each fold when performing
             cross-validation.
-        refit_on_update: If True, refit the model after updating the training
-            data using the `update` method.
         warm_start_refitting: If True, start model refitting from previous
             model parameters in order to speed up the fitting process.
         prior: An optional dictionary that contains the specification of GP model prior.
@@ -251,7 +249,6 @@ def __init__(
         acqf_optimizer: TOptimizer = scipy_optimizer,
         best_point_recommender: TBestPointRecommender = recommend_best_observed_point,
         refit_on_cv: bool = False,
-        refit_on_update: bool = True,
         warm_start_refitting: bool = True,
         use_input_warping: bool = False,
         use_loocv_pseudo_likelihood: bool = False,
@@ -276,7 +273,6 @@ def __init__(
         # pyre-fixme[4]: Attribute must be annotated.
         self._kwargs = kwargs
         self.refit_on_cv = refit_on_cv
-        self.refit_on_update = refit_on_update
         self.warm_start_refitting = warm_start_refitting
         self.use_input_warping = use_input_warping
         self.use_loocv_pseudo_likelihood = use_loocv_pseudo_likelihood

ax/models/torch/botorch_modular/model.py

@@ -140,7 +140,6 @@ class BoTorchModel(TorchModel, Base):
             based on the data provided.
         surrogate: In liu of SurrogateSpecs, an instance of `Surrogate` may be
             provided to be used as the sole Surrogate for all outcomes
-        refit_on_update: Unused.
         refit_on_cv: Whether to reoptimize model parameters during call to
             `BoTorchmodel.cross_validate`.
         warm_start_refit: Whether to load parameters from either the provided
@@ -169,7 +168,6 @@ def __init__(
         acquisition_options: Optional[Dict[str, Any]] = None,
         botorch_acqf_class: Optional[Type[AcquisitionFunction]] = None,
         # TODO: [T168715924] Revisit these "refit" arguments.
-        refit_on_update: bool = True,
         refit_on_cv: bool = False,
         warm_start_refit: bool = True,
     ) -> None:
@@ -216,7 +214,6 @@ def __init__(
         self.acquisition_options = acquisition_options or {}
         self._botorch_acqf_class = botorch_acqf_class
 
-        self.refit_on_update = refit_on_update
         self.refit_on_cv = refit_on_cv
         self.warm_start_refit = warm_start_refit
 

ax/models/torch/botorch_moo.py

@@ -206,7 +206,6 @@ def __init__(
         best_point_recommender: TBestPointRecommender = recommend_best_observed_point,
         frontier_evaluator: TFrontierEvaluator = pareto_frontier_evaluator,
         refit_on_cv: bool = False,
-        refit_on_update: bool = True,
         warm_start_refitting: bool = False,
         use_input_warping: bool = False,
         use_loocv_pseudo_likelihood: bool = False,
@@ -222,7 +221,6 @@ def __init__(
         # pyre-fixme[4]: Attribute must be annotated.
         self._kwargs = kwargs
         self.refit_on_cv = refit_on_cv
-        self.refit_on_update = refit_on_update
         self.warm_start_refitting = warm_start_refitting
         self.use_input_warping = use_input_warping
         self.use_loocv_pseudo_likelihood = use_loocv_pseudo_likelihood

ax/models/torch/fully_bayesian.py

@@ -523,7 +523,6 @@ def __init__(
         acqf_optimizer: TOptimizer = scipy_optimizer,
         best_point_recommender: TBestPointRecommender = recommend_best_observed_point,
         refit_on_cv: bool = False,
-        refit_on_update: bool = True,
         warm_start_refitting: bool = True,
         use_input_warping: bool = False,
         # use_saas is deprecated. TODO: remove
@@ -553,8 +552,6 @@ def __init__(
                 signature as described below.
             refit_on_cv: If True, refit the model for each fold when performing
                 cross-validation.
-            refit_on_update: If True, refit the model after updating the training
-                data using the `update` method.
             warm_start_refitting: If True, start model refitting from previous
                 model parameters in order to speed up the fitting process.
             use_input_warping: A boolean indicating whether to use input warping
@@ -581,7 +578,6 @@ def __init__(
             acqf_optimizer=acqf_optimizer,
             best_point_recommender=best_point_recommender,
             refit_on_cv=refit_on_cv,
-            refit_on_update=refit_on_update,
             warm_start_refitting=warm_start_refitting,
             use_input_warping=use_input_warping,
             num_samples=num_samples,
@@ -619,7 +615,6 @@ def __init__(
         best_point_recommender: TBestPointRecommender = recommend_best_observed_point,
         frontier_evaluator: TFrontierEvaluator = pareto_frontier_evaluator,
         refit_on_cv: bool = False,
-        refit_on_update: bool = True,
         warm_start_refitting: bool = False,
         use_input_warping: bool = False,
         num_samples: int = 256,
@@ -646,7 +641,6 @@ def __init__(
             best_point_recommender=best_point_recommender,
             frontier_evaluator=frontier_evaluator,
             refit_on_cv=refit_on_cv,
-            refit_on_update=refit_on_update,
             warm_start_refitting=warm_start_refitting,
             use_input_warping=use_input_warping,
             num_samples=num_samples,

ax/models/torch/tests/test_model.py

@@ -188,7 +188,6 @@ def test_init(self) -> None:
         self.assertEqual(model.botorch_acqf_class, qExpectedImprovement)
 
         # Check defaults for refitting settings.
-        self.assertTrue(model.refit_on_update)
         self.assertFalse(model.refit_on_cv)
         self.assertTrue(model.warm_start_refit)
 
@@ -197,11 +196,9 @@ def test_init(self) -> None:
             surrogate=self.surrogate,
             acquisition_class=self.acquisition_class,
             acquisition_options=self.acquisition_options,
-            refit_on_update=False,
             refit_on_cv=True,
             warm_start_refit=False,
         )
-        self.assertFalse(mdl2.refit_on_update)
         self.assertTrue(mdl2.refit_on_cv)
         self.assertFalse(mdl2.warm_start_refit)
 

ax/storage/json_store/decoder.py

@@ -683,6 +683,7 @@ def generation_step_from_json(
         generation_step_json
     )
     kwargs = generation_step_json.pop("model_kwargs", None)
+    kwargs.pop("fit_on_update", None)  # Remove deprecated fit_on_update.
     gen_kwargs = generation_step_json.pop("model_gen_kwargs", None)
     completion_criteria = (
         object_from_json(
@@ -741,6 +742,7 @@ def model_spec_from_json(
 ) -> ModelSpec:
     """Load ModelSpec from JSON."""
     kwargs = model_spec_json.pop("model_kwargs", None)
+    kwargs.pop("fit_on_update", None)  # Remove deprecated fit_on_update.
     gen_kwargs = model_spec_json.pop("model_gen_kwargs", None)
     return ModelSpec(
         model_enum=object_from_json(

ax/storage/json_store/encoders.py

@@ -579,7 +579,6 @@ def botorch_model_to_dict(model: BoTorchModel) -> Dict[str, Any]:
             model.surrogate_specs if len(model.surrogate_specs) > 0 else None
         ),
         "botorch_acqf_class": model._botorch_acqf_class,
-        "refit_on_update": model.refit_on_update,
         "refit_on_cv": model.refit_on_cv,
         "warm_start_refit": model.warm_start_refit,
     }

ax/storage/json_store/tests/test_json_store.py

@@ -18,6 +18,7 @@
 from ax.core.runner import Runner
 from ax.exceptions.core import AxStorageWarning
 from ax.exceptions.storage import JSONDecodeError, JSONEncodeError
+from ax.modelbridge.generation_node import GenerationStep
 from ax.modelbridge.generation_strategy import GenerationStrategy
 from ax.modelbridge.registry import Models
 from ax.storage.json_store.decoder import (
@@ -668,3 +669,23 @@ def test_objective_backwards_compatibility(self) -> None:
         self.assertNotEqual(objective, objective_loaded)
         self.assertTrue(objective_loaded.minimize)
         self.assertTrue(objective_loaded.metric.lower_is_better)
+
+    def test_generation_step_backwards_compatibility(self) -> None:
+        # Test that we can load a generation step with fit_on_update.
+        json = {
+            "__type": "GenerationStep",
+            "model": {"__type": "Models", "name": "BOTORCH_MODULAR"},
+            "num_trials": 5,
+            "min_trials_observed": 0,
+            "completion_criteria": [],
+            "max_parallelism": None,
+            "use_update": False,
+            "enforce_num_trials": True,
+            "model_kwargs": {"fit_on_update": False, "other_kwarg": 5},
+            "model_gen_kwargs": {},
+            "index": -1,
+            "should_deduplicate": False,
+        }
+        generation_step = object_from_json(json)
+        self.assertIsInstance(generation_step, GenerationStep)
+        self.assertEqual(generation_step.model_kwargs, {"other_kwarg": 5})

ax/utils/common/constants.py

@@ -73,7 +73,6 @@ class Keys(str, Enum):
     QMC = "qmc"
     RAW_INNER_SAMPLES = "raw_inner_samples"
     RAW_SAMPLES = "raw_samples"
-    REFIT_ON_UPDATE = "refit_on_update"
     SAMPLER = "sampler"
     SEED_INNER = "seed_inner"
     SEQUENTIAL = "sequential"

tutorials/modular_botax.ipynb

@@ -236,7 +236,6 @@
     "    acquisition_class=None,\n",
     "    # Less common model settings shown with default values, refer\n",
     "    # to `BoTorchModel` documentation for detail\n",
-    "    refit_on_update=True,\n",
     "    refit_on_cv=False,\n",
     "    warm_start_refit=True,\n",
     ")"