Add ability to use built-in pickle for saving AutoMLSearch

docs/source/release_notes.rst

@@ -17,6 +17,7 @@ Release Notes
         * Added support for showing a Individual Conditional Expectations plot when graphing Partial Dependence :pr:`2386`
         * Exposed ``thread_count`` for Catboost estimators as ``n_jobs`` parameter :pr:`2410`
         * Updated Objectives API to allow for sample weighting :pr:`2433`
+        * Added ability to use built-in pickle for saving AutoMLSearch :pr:`2463`
     * Fixes
         * Deleted unreachable line from ``IterativeAlgorithm`` :pr:`2464`
     * Changes

evalml/automl/automl_search.py

@@ -1,4 +1,5 @@
 import copy
+import pickle
 import sys
 import time
 import traceback
@@ -1307,31 +1308,50 @@ def best_pipeline(self):
 
         return self._best_pipeline
 
-    def save(self, file_path, pickle_protocol=cloudpickle.DEFAULT_PROTOCOL):
+    def save(
+        self,
+        file_path,
+        pickle_type="cloudpickle",
+        pickle_protocol=cloudpickle.DEFAULT_PROTOCOL,
+    ):
         """Saves AutoML object at file path
 
         Arguments:
             file_path (str): location to save file
+            pickle_type {"pickle", "cloudpickle"}: the pickling library to use.
             pickle_protocol (int): the pickle data stream format.
 
         Returns:
             None
         """
+        if pickle_type == "cloudpickle":
+            pkl_lib = cloudpickle
+        elif pickle_type == "pickle":
+            pkl_lib = pickle
+        else:
+            raise ValueError(
+                f"`pickle_type` must be either 'pickle' or 'cloudpickle'. Received {pickle_type}"
+            )
+
         with open(file_path, "wb") as f:
-            cloudpickle.dump(self, f, protocol=pickle_protocol)
+            pkl_lib.dump(self, f, protocol=pickle_protocol)
 
     @staticmethod
-    def load(file_path):
+    def load(
+        file_path,
+        pickle_type="cloudpickle",
+    ):
         """Loads AutoML object at file path
 
         Arguments:
             file_path (str): location to find file to load
+            pickle_type {"pickle", "cloudpickle"}: the pickling library to use. Currently not used since the standard pickle library can handle cloudpickles.
 
         Returns:
             AutoSearchBase object
         """
         with open(file_path, "rb") as f:
-            return cloudpickle.load(f)
+            return pickle.load(f)
 
     def train_pipelines(self, pipelines):
         """Train a list of pipelines on the training data.

evalml/automl/pipeline_search_plots.py

@@ -32,6 +32,7 @@ def __init__(self, results, objective):
         self.best_score_by_iter_fig = self._go.FigureWidget(data, layout)
         self.best_score_by_iter_fig.update_layout(showlegend=False)
         self.update(results, objective)
+        self._go = None
 
     def update(self, results, objective):
         if len(results["search_order"]) > 0 and len(results["pipeline_results"]) > 0:

evalml/tests/automl_tests/test_automl.py

@@ -650,7 +650,8 @@ def test_automl_allowed_component_graphs_algorithm(
         assert actual.parameters == expected.parameters
 
 
-def test_automl_serialization(X_y_binary, tmpdir):
+@pytest.mark.parametrize("pickle_type", ["cloudpickle", "pickle", "invalid"])
+def test_automl_serialization(pickle_type, X_y_binary, tmpdir):
     X, y = X_y_binary
     path = os.path.join(str(tmpdir), "automl.pkl")
     num_max_iterations = 5
@@ -663,35 +664,50 @@ def test_automl_serialization(X_y_binary, tmpdir):
         n_jobs=1,
     )
     automl.search()
-    automl.save(path)
-    loaded_automl = automl.load(path)
 
-    for i in range(num_max_iterations):
-        assert (
-            automl.get_pipeline(i).__class__ == loaded_automl.get_pipeline(i).__class__
-        )
-        assert (
-            automl.get_pipeline(i).parameters
-            == loaded_automl.get_pipeline(i).parameters
+    if automl.search_iteration_plot:
+        # Testing pickling of SearchIterationPlot object
+        automl.search_iteration_plot = automl.plot.search_iteration_plot(
+            interactive_plot=True
         )
 
-        for id_, pipeline_results in automl.results["pipeline_results"].items():
-            loaded_ = loaded_automl.results["pipeline_results"][id_]
-            for name in pipeline_results:
-                # Use np to check percent_better_than_baseline because of (possible) nans
-                if name == "percent_better_than_baseline_all_objectives":
-                    for objective_name, value in pipeline_results[name].items():
+    if pickle_type == "invalid":
+        with pytest.raises(
+            ValueError,
+            match="`pickle_type` must be either 'pickle' or 'cloudpickle'. Received invalid",
+        ):
+            automl.save(path, pickle_type=pickle_type)
+    else:
+        automl.save(path, pickle_type=pickle_type)
+        loaded_automl = automl.load(path)
+
+        for i in range(num_max_iterations):
+            assert (
+                automl.get_pipeline(i).__class__
+                == loaded_automl.get_pipeline(i).__class__
+            )
+            assert (
+                automl.get_pipeline(i).parameters
+                == loaded_automl.get_pipeline(i).parameters
+            )
+
+            for id_, pipeline_results in automl.results["pipeline_results"].items():
+                loaded_ = loaded_automl.results["pipeline_results"][id_]
+                for name in pipeline_results:
+                    # Use np to check percent_better_than_baseline because of (possible) nans
+                    if name == "percent_better_than_baseline_all_objectives":
+                        for objective_name, value in pipeline_results[name].items():
+                            np.testing.assert_almost_equal(
+                                value, loaded_[name][objective_name]
+                            )
+                    elif name == "percent_better_than_baseline":
                         np.testing.assert_almost_equal(
-                            value, loaded_[name][objective_name]
+                            pipeline_results[name], loaded_[name]
                         )
-                elif name == "percent_better_than_baseline":
-                    np.testing.assert_almost_equal(
-                        pipeline_results[name], loaded_[name]
-                    )
-                else:
-                    assert pipeline_results[name] == loaded_[name]
+                    else:
+                        assert pipeline_results[name] == loaded_[name]
 
-    pd.testing.assert_frame_equal(automl.rankings, loaded_automl.rankings)
+        pd.testing.assert_frame_equal(automl.rankings, loaded_automl.rankings)
 
 
 @patch("cloudpickle.dump")