Address loading in pipeline that uses unloaded custom objective class code

evalml/pipelines/__init__.py

@@ -11,4 +11,4 @@
     list_model_types,
     load_pipeline,
     save_pipeline
-)
+)

evalml/pipelines/utils.py

@@ -1,4 +1,4 @@
-import pickle
+import cloudpickle
 
 from .classification import (
     LogisticRegressionPipeline,
@@ -85,7 +85,7 @@ def save_pipeline(pipeline, file_path):
         None
     """
     with open(file_path, 'wb') as f:
-        pickle.dump(pipeline, f)
+        cloudpickle.dump(pipeline, f)
 
 
 def load_pipeline(file_path):
@@ -98,4 +98,4 @@ def load_pipeline(file_path):
         Pipeline obj
     """
     with open(file_path, 'rb') as f:
-        return pickle.load(f)
+        return cloudpickle.load(f)

evalml/tests/pipeline_tests/test_pipelines.py

@@ -1,11 +1,8 @@
-import errno
 import os
-import shutil
 
 import pandas as pd
 import pytest
 
-import evalml.tests as tests
 from evalml.model_types import ModelTypes
 from evalml.objectives import FraudCost, Precision
 from evalml.pipelines import LogisticRegressionPipeline
@@ -17,8 +14,6 @@
 )
 from evalml.problem_types import ProblemTypes
 
-CACHE = os.path.join(os.path.dirname(tests.__file__), '.cache')
-
 
 def test_list_model_types():
     assert set(list_model_types(ProblemTypes.BINARY)) == set([ModelTypes.RANDOM_FOREST, ModelTypes.XGBOOST, ModelTypes.LINEAR_MODEL])
@@ -33,21 +28,9 @@ def test_get_pipelines():
         get_pipelines(problem_type=ProblemTypes.REGRESSION, model_types=["random_forest", "xgboost"])
 
 
-@pytest.fixture
-def path_management():
-    path = CACHE
-    try:
-        os.makedirs(path)
-    except OSError as e:
-        if e.errno != errno.EEXIST:  # EEXIST corresponds to FileExistsError
-            raise e
-    yield path
-    shutil.rmtree(path)
-
-
-def test_serialization(X_y, trained_model, path_management):
+def test_serialization(X_y, tmpdir):
     X, y = X_y
-    path = os.path.join(path_management, 'pipe.pkl')
+    path = os.path.join(tmpdir, 'pipe.pkl')
     objective = Precision()
 
     pipeline = LogisticRegressionPipeline(objective=objective, penalty='l2', C=1.0, impute_strategy='mean', number_features=len(X[0]))
@@ -56,6 +39,26 @@ def test_serialization(X_y, trained_model, path_management):
     assert pipeline.score(X, y) == load_pipeline(path).score(X, y)
 
 
+@pytest.fixture
+def pickled_pipeline_path(X_y, tmpdir):
+    X, y = X_y
+    path = os.path.join(tmpdir, 'pickled_pipe.pkl')
+    MockPrecision = type('MockPrecision', (Precision,), {})
+    objective = MockPrecision()
+    pipeline = LogisticRegressionPipeline(objective=objective, penalty='l2', C=1.0, impute_strategy='mean', number_features=len(X[0]))
+    pipeline.fit(X, y)
+    save_pipeline(pipeline, path)
+    return path
+
+
+def test_load_pickled_pipeline_with_custom_objective(X_y, pickled_pipeline_path):
+    X, y = X_y
+    objective = Precision()
+    pipeline = LogisticRegressionPipeline(objective=objective, penalty='l2', C=1.0, impute_strategy='mean', number_features=len(X[0]))
+    pipeline.fit(X, y)
+    assert load_pipeline(pickled_pipeline_path).score(X, y) == pipeline.score(X, y)
+
+
 def test_reproducibility(X_y):
     X, y = X_y
     X = pd.DataFrame(X)

evalml/utils/__init__.py

@@ -1,3 +1,3 @@
 # flake8:noqa
 from .logging_utils import Logger
-from .convert_time import convert_to_seconds
+from .convert_time import convert_to_seconds