def test_load_pyfunc_loads_torch_model_using_pickle_module_specified_at_save_time( module_scoped_subclassed_model, model_path): custom_pickle_module = pickle mlflow.pytorch.save_model( path=model_path, pytorch_model=module_scoped_subclassed_model, pickle_module=custom_pickle_module, ) import_module_fn = importlib.import_module imported_modules = [] def track_module_imports(module_name): imported_modules.append(module_name) return import_module_fn(module_name) with mock.patch("importlib.import_module") as import_mock, mock.patch( "torch.load") as torch_load_mock: import_mock.side_effect = track_module_imports pyfunc.load_model(model_path) torch_load_mock.assert_called_with(mock.ANY, pickle_module=custom_pickle_module) assert custom_pickle_module.__name__ in imported_modules
def test_load_model_loads_torch_model_using_pickle_module_specified_at_save_time( module_scoped_subclassed_model, ): custom_pickle_module = pickle artifact_path = "pytorch_model" with mlflow.start_run(): mlflow.pytorch.log_model( artifact_path=artifact_path, pytorch_model=module_scoped_subclassed_model, pickle_module=custom_pickle_module, ) model_uri = "runs:/{run_id}/{artifact_path}".format( run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path) import_module_fn = importlib.import_module imported_modules = [] def track_module_imports(module_name): imported_modules.append(module_name) return import_module_fn(module_name) with mock.patch("importlib.import_module") as import_mock, mock.patch( "torch.load") as torch_load_mock: import_mock.side_effect = track_module_imports pyfunc.load_model(model_uri=model_uri) torch_load_mock.assert_called_with(mock.ANY, pickle_module=custom_pickle_module) assert custom_pickle_module.__name__ in imported_modules
def test_iris_data_model_can_be_loaded_and_evaluated_as_pyfunc(saved_tf_iris_model, model_path): mlflow.tensorflow.save_model( tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=model_path, ) pyfunc_wrapper = pyfunc.load_model(model_path) # can call predict with a df results_df = pyfunc_wrapper.predict(saved_tf_iris_model.inference_df) assert isinstance(results_df, pandas.DataFrame) pandas.testing.assert_frame_equal( results_df, saved_tf_iris_model.expected_results_df, check_less_precise=1 ) # can also call predict with a dict inp_data = {} for col_name in list(saved_tf_iris_model.inference_df): inp_data[col_name] = saved_tf_iris_model.inference_df[col_name].values results = pyfunc_wrapper.predict(inp_data) assert isinstance(results, dict) pandas.testing.assert_frame_equal( pandas.DataFrame(data=results), saved_tf_iris_model.expected_results_df, check_less_precise=1, ) # can not call predict with a list inp_list = [] for df_col_name in list(saved_tf_iris_model.inference_df): inp_list.append(saved_tf_iris_model.inference_df[df_col_name].values) with pytest.raises(TypeError): results = pyfunc_wrapper.predict(inp_list)
def serve(model_path: str, unix_socket_path: str, shutdown_if_no_requests_after_sec: int = 60): """ Serve mlflow pyfunc model via a unix file socket. Expects pickled model input, responds with either pickled output or pickled exception. :param model_path: Local model path (server model path). :param unix_socket_path: Unix file socket path. :param shutdown_if_no_requests_after_sec: :return: """ model_path = model_path or os.environ.get( scoring_server._SERVER_MODEL_PATH) unix_socket_path = unix_socket_path or os.environ.get( UNIX_SOCKET_PATH) or tempfile.NamedTemporaryFile().name model = load_model(model_path) # type: Model print(f'Serving model: {model_path}\n' f'at: {unix_socket_path}\n' f'Process id: {os.getpid()}') try: os.unlink(unix_socket_path) except OSError: if os.path.exists(unix_socket_path): raise server = Server( unix_socket_path, PredictRequestHandler, shutdown_if_no_requests_timeout=shutdown_if_no_requests_after_sec) server.model = model server.model_path = model_path server.serve_forever()
def test_model_save_load(fastai_model, model_path): model = fastai_model.model mlflow.fastai.save_model(fastai_learner=model, path=model_path) reloaded_model = mlflow.fastai.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) # Verify reloaded model computes same predictions as original model test_data = TabularList.from_df(fastai_model.inference_dataframe) model.data.add_test(test_data) reloaded_model.data.add_test(test_data) real_preds, real_target = map(lambda output: output.numpy(), model.get_preds(DatasetType.Test)) reloaded_preds, reloaded_target = map( lambda output: output.numpy(), reloaded_model.get_preds(DatasetType.Test)) np.testing.assert_array_almost_equal(real_preds, reloaded_preds) np.testing.assert_array_almost_equal(real_target, reloaded_target) model_wrapper = mlflow.fastai._FastaiModelWrapper(model) reloaded_model_wrapper = mlflow.fastai._FastaiModelWrapper(reloaded_model) model_result = model_wrapper.predict(fastai_model.inference_dataframe) reloaded_result = reloaded_model_wrapper.predict( fastai_model.inference_dataframe) pyfunc_result = reloaded_pyfunc.predict(fastai_model.inference_dataframe) compare_wrapper_results(model_result, reloaded_result) compare_wrapper_results(reloaded_result, pyfunc_result)
def test_load_pyfunc_succeeds_for_older_models_with_pyfunc_data_field( sklearn_knn_model, model_path ): """ This test verifies that scikit-learn models saved in older versions of MLflow are loaded successfully by ``mlflow.pyfunc.load_model``. These older models specify a pyfunc ``data`` field referring directly to a serialized scikit-learn model file. In contrast, newer models omit the ``data`` field. """ mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE, ) model_conf_path = os.path.join(model_path, "MLmodel") model_conf = Model.load(model_conf_path) pyfunc_conf = model_conf.flavors.get(pyfunc.FLAVOR_NAME) sklearn_conf = model_conf.flavors.get(mlflow.sklearn.FLAVOR_NAME) assert sklearn_conf is not None assert pyfunc_conf is not None pyfunc_conf[pyfunc.DATA] = sklearn_conf["pickled_model"] reloaded_knn_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_equal( sklearn_knn_model.model.predict(sklearn_knn_model.inference_data), reloaded_knn_pyfunc.predict(sklearn_knn_model.inference_data), )
def test_iris_data_model_can_be_loaded_and_evaluated_as_pyfunc( saved_tf_iris_model, model_path): mlflow.tensorflow.save_model( tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=model_path, ) pyfunc_wrapper = pyfunc.load_model(model_path) # can call predict with a df results_df = pyfunc_wrapper.predict(saved_tf_iris_model.inference_df) assert isinstance(results_df, pd.DataFrame) for key in results_df.keys(): assert np.array_equal(results_df[key], saved_tf_iris_model.raw_df[key]) # can also call predict with a dict inp_dict = {} for df_col_name in list(saved_tf_iris_model.inference_df): inp_dict[df_col_name] = saved_tf_iris_model.inference_df[ df_col_name].values results = pyfunc_wrapper.predict(inp_dict) assert isinstance(results, dict) for key in results.keys(): assert np.array_equal(results[key], saved_tf_iris_model.raw_df[key].tolist()) # can not call predict with a list inp_list = [] for df_col_name in list(saved_tf_iris_model.inference_df): inp_list.append(saved_tf_iris_model.inference_df[df_col_name].values) with pytest.raises(TypeError): results = pyfunc_wrapper.predict(inp_list)
def test_load_pyfunc_succeeds_when_data_is_model_file_instead_of_directory( module_scoped_subclassed_model, model_path, data): """ This test verifies that PyTorch models saved in older versions of MLflow are loaded successfully by ``mlflow.pytorch.load_model``. The ``data`` path associated with these older models is serialized PyTorch model file, as opposed to the current format: a directory containing a serialized model file and pickle module information. """ mlflow.pytorch.save_model(path=model_path, pytorch_model=module_scoped_subclassed_model) model_conf_path = os.path.join(model_path, "MLmodel") model_conf = Model.load(model_conf_path) pyfunc_conf = model_conf.flavors.get(pyfunc.FLAVOR_NAME) assert pyfunc_conf is not None model_data_path = os.path.join(model_path, pyfunc_conf[pyfunc.DATA]) assert os.path.exists(model_data_path) assert mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME in os.listdir( model_data_path) pyfunc_conf[pyfunc.DATA] = os.path.join( model_data_path, mlflow.pytorch._SERIALIZED_TORCH_MODEL_FILE_NAME) model_conf.save(model_conf_path) loaded_pyfunc = pyfunc.load_model(model_path) np.testing.assert_array_almost_equal( loaded_pyfunc.predict(data[0]), pd.DataFrame(_predict(model=module_scoped_subclassed_model, data=data)), decimal=4, )
def test_model_save_load_built_in_high_level_api( pd_model_built_in_high_level_api, model_path): model = pd_model_built_in_high_level_api.model test_dataset = pd_model_built_in_high_level_api.inference_dataframe mlflow.paddle.save_model(pd_model=model, path=model_path) reloaded_pd_model = mlflow.paddle.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) low_level_test_dataset = [x[0] for x in test_dataset] np.testing.assert_array_almost_equal( np.array(model.predict(test_dataset)).squeeze(), np.array(reloaded_pyfunc.predict( np.array(low_level_test_dataset))).squeeze(), decimal=5, ) np.testing.assert_array_almost_equal( np.array(reloaded_pd_model( np.array(low_level_test_dataset))).squeeze(), np.array(reloaded_pyfunc.predict( np.array(low_level_test_dataset))).squeeze(), decimal=5, )
def test_model_save_load(fastai_model, model_path): model = fastai_model.model mlflow.fastai.save_model(fastai_learner=model, path=model_path) reloaded_model = mlflow.fastai.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) # Verify reloaded model computes same predictions as original model dl_model = model.dls.test_dl(fastai_model.inference_dataframe) dl_reloaded_model = reloaded_model.dls.test_dl( fastai_model.inference_dataframe) real_preds, _ = map( lambda output: output.numpy() if output is not None else output, model.get_preds(dl=dl_model), ) reloaded_preds, _ = map( lambda output: output.numpy() if output is not None else output, reloaded_model.get_preds(dl=dl_reloaded_model), ) np.testing.assert_array_almost_equal(real_preds, reloaded_preds) model_wrapper = mlflow.fastai._FastaiModelWrapper(model) reloaded_model_wrapper = mlflow.fastai._FastaiModelWrapper(reloaded_model) model_result = model_wrapper.predict(fastai_model.inference_dataframe) reloaded_result = reloaded_model_wrapper.predict( fastai_model.inference_dataframe) pyfunc_result = reloaded_pyfunc.predict(fastai_model.inference_dataframe) np.testing.assert_array_almost_equal(model_result, reloaded_result) np.testing.assert_array_almost_equal(reloaded_result, pyfunc_result)
def _model(self) -> Any: """ The return object has to follow the API of mlflow.pyfunc.PythonModel. """ from mlflow import pyfunc return pyfunc.load_model(model_uri=self._model_uri)
def load_mlflow_model(data): with tempfile.TemporaryDirectory() as tmp: data = io.BytesIO(data) with zipfile.ZipFile(data) as artifact: artifact.extractall(path=tmp) model = load_model(tmp) return model
def main(argv): with mlflow.start_run(): args = parser.parse_args(argv[1:]) # Builds, trains and evaluates a tf.estimator. Then, exports it for inference, # logs the exported model with MLflow, and loads the fitted model back as a PyFunc. (x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data() # There are 13 features we are using for inference. feat_cols = [ tf.feature_column.numeric_column(key="features", shape=(x_train.shape[1], )) ] feat_spec = { "features": tf.placeholder("float", name="features", shape=[None, x_train.shape[1]]) } hidden_units = [50, 20] steps = args.steps regressor = tf.estimator.DNNRegressor(hidden_units=hidden_units, feature_columns=feat_cols) train_input_fn = tf.estimator.inputs.numpy_input_fn( {"features": x_train}, y_train, num_epochs=None, shuffle=True) regressor.train(train_input_fn, steps=steps) test_input_fn = tf.estimator.inputs.numpy_input_fn( {"features": x_test}, y_test, num_epochs=None, shuffle=True) # Compute mean squared error mse = regressor.evaluate(test_input_fn, steps=steps) # Building a receiver function for exporting receiver_fn = tf.estimator.export.build_raw_serving_input_receiver_fn( feat_spec) temp = tempfile.mkdtemp() try: # The model is automatically logged when export_saved_model() is called. saved_estimator_path = regressor.export_savedmodel( temp, receiver_fn).decode("utf-8") # Since the model was automatically logged as an artifact (more specifically # a MLflow Model), we don't need to use saved_estimator_path to load back the model. # MLflow takes care of it! pyfunc_model = pyfunc.load_model(mlflow.get_artifact_uri('model')) df = pd.DataFrame(data=x_test, columns=["features"] * x_train.shape[1]) # Checking the PyFunc's predictions are the same as the original model's predictions. predict_df = pyfunc_model.predict(df) predict_df['original_labels'] = y_test print(predict_df) finally: shutil.rmtree(temp)
def test_model_retrain_built_in_high_level_api( pd_model_built_in_high_level_api, model_path, model_retrain_path, get_dataset_built_in_high_level_api, ): model = pd_model_built_in_high_level_api.model mlflow.paddle.save_model(pd_model=model, path=model_path, training=True) training_dataset, test_dataset = get_dataset_built_in_high_level_api model_retrain = paddle.Model(UCIHousing()) model_retrain = mlflow.paddle.load_model(model_uri=model_path, model=model_retrain) optim = paddle.optimizer.Adam(learning_rate=0.015, parameters=model.parameters()) model_retrain.prepare(optim, paddle.nn.MSELoss()) model_retrain.fit(training_dataset, epochs=6, batch_size=8, verbose=1) mlflow.paddle.save_model(pd_model=model_retrain, path=model_retrain_path, training=False) with pytest.raises(TypeError, match="This model can't be loaded"): mlflow.paddle.load_model(model_uri=model_retrain_path, model=model_retrain) error_model = 0 error_model_type = type(error_model) with pytest.raises( TypeError, match="Invalid object type `{}` for `model`, must be `paddle.Model`" .format(error_model_type), ): mlflow.paddle.load_model(model_uri=model_retrain_path, model=error_model) reloaded_pd_model = mlflow.paddle.load_model(model_uri=model_retrain_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_retrain_path) low_level_test_dataset = [x[0] for x in test_dataset] np.testing.assert_array_almost_equal( np.array(model_retrain.predict(test_dataset)).squeeze(), np.array(reloaded_pyfunc.predict( np.array(low_level_test_dataset))).squeeze(), decimal=5, ) np.testing.assert_array_almost_equal( np.array(reloaded_pd_model( np.array(low_level_test_dataset))).squeeze(), np.array(reloaded_pyfunc.predict( np.array(low_level_test_dataset))).squeeze(), decimal=5, )
def test_iris_data_model_can_be_loaded_and_evaluated_as_pyfunc(saved_tf_iris_model, model_path): mlflow.tensorflow.save_model(tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=model_path) pyfunc_wrapper = pyfunc.load_model(model_path) results_df = pyfunc_wrapper.predict(saved_tf_iris_model.inference_df) for key in results_df.keys(): assert(np.array_equal(results_df[key], saved_tf_iris_model.raw_df[key]))
def test_iris_data_model_can_be_loaded_and_evaluated_as_pyfunc(saved_tf_iris_model, model_path): mlflow.tensorflow.save_model(tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=model_path) pyfunc_wrapper = pyfunc.load_model(model_path) results_df = pyfunc_wrapper.predict(saved_tf_iris_model.inference_df) pandas.testing.assert_frame_equal( results_df, saved_tf_iris_model.expected_results_df, check_less_precise=1)
def test_transformer_model_export(spark_model_transformer, model_path, spark_custom_env): sparkm.save_model(spark_model_transformer.model, path=model_path, conda_env=spark_custom_env) # score and compare the reloaded sparkml model reloaded_model = sparkm.load_model(model_uri=model_path) preds_df = reloaded_model.transform(spark_model_transformer.spark_df) preds = [x.features for x in preds_df.select("features").collect()] assert spark_model_transformer.predictions == preds # 2. score and compare reloaded pyfunc m = pyfunc.load_model(model_path) preds2 = m.predict(spark_model_transformer.spark_df.toPandas()) assert spark_model_transformer.predictions == preds2
def __init__(self): #next(os.walk('mlruns\\0'))[1][0] #self.pyfunc_model = pyfunc.load_pyfunc("mlruns\\0\\"+next(os.walk('mlruns\\0'))[1][0]+"\\artifacts\\model") artifact_dir = "" print(os.listdir(".")) for (root, dirs, files) in os.walk('mlruns/0', topdown=True): artifact_dir = dirs[0] break self.pyfunc_model = pyfunc.load_model( "s3://tiger-mlflow/1/afa14f38393e4ca5a1a600dbfb43ae73/artifacts/model/" )
def test_iris_data_model_can_be_loaded_and_evaluated_as_pyfunc( saved_tf_iris_model, model_path): mlflow.tensorflow.save_model( tf_saved_model_dir=saved_tf_iris_model.path, tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags, tf_signature_def_key=saved_tf_iris_model.signature_def_key, path=model_path) pyfunc_wrapper = pyfunc.load_model(model_path) results_df = pyfunc_wrapper.predict(saved_tf_iris_model.inference_df) assert results_df.equals(saved_tf_iris_model.expected_results_df)
def test_model_pyfunc_save_load(prophet_model, model_path): model = prophet_model.model mlflow.prophet.save_model(pr_model=model, path=model_path) loaded_pyfunc = pyfunc.load_model(model_uri=model_path) horizon_df = future_horizon_df(model, FORECAST_HORIZON) np.testing.assert_array_equal( generate_forecast(model, FORECAST_HORIZON), loaded_pyfunc.predict(horizon_df)[TARGET_FIELD_NAME], )
def __init__(self): print("Create controller with configuration file !!!!!") f = open("conf.txt", "r") conf = f.read() conf = conf.split(',\n') self.name = conf[0] self.version = conf[1] self.source = conf[2] f.close() print(self.name, ' ', self.version, ' ', self.source) self.pyfunc_model = pyfunc.load_model('tensor')
def test_categorical_model_can_be_loaded_and_evaluated_as_pyfunc( saved_tf_categorical_model, model_path): mlflow.tensorflow.save_model(tf_saved_model_dir=saved_tf_categorical_model.path, tf_meta_graph_tags=saved_tf_categorical_model.meta_graph_tags, tf_signature_def_key=saved_tf_categorical_model.signature_def_key, path=model_path) pyfunc_wrapper = pyfunc.load_model(model_path) results_df = pyfunc_wrapper.predict(saved_tf_categorical_model.inference_df) # Precision is less accurate for the categorical model when we load back the saved model. pandas.testing.assert_frame_equal( results_df, saved_tf_categorical_model.expected_results_df, check_less_precise=3)
def _test_model_save_load(statsmodels_model, model_path, *predict_args): mlflow.statsmodels.save_model(statsmodels_model=statsmodels_model.model, path=model_path) reloaded_model = mlflow.statsmodels.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) if hasattr(statsmodels_model.model, "predict"): np.testing.assert_array_almost_equal( statsmodels_model.model.predict(*predict_args), reloaded_model.predict(*predict_args), ) np.testing.assert_array_almost_equal( reloaded_model.predict(*predict_args), reloaded_pyfunc.predict(statsmodels_model.inference_dataframe), )
def main(run_id, model_run_id): mlflow.set_experiment('restaurant-reviews-allergy') model_path = f'runs:/{model_run_id}/model' sentiment_model = load_model(model_path) chunks = download_chunked_data(run_id, 'chunks', format='pkl') scored_chunks = [sentiment_model.predict(chunk) for chunk in chunks] df = pd.concat(scored_chunks) mlflow.log_param('run_id', run_id) mlflow.log_param('model_run_id', model_run_id) mlflow.set_tag('step', 'append_sentiment_labels') logger = MlflowArtifactLogger() logger.add_artifact(df, 'data_with_sentiment.pkl') logger.log_artifacts('')
def test_sklearn_model_save_load(xgb_sklearn_model, model_path): model = xgb_sklearn_model.model mlflow.xgboost.save_model(xgb_model=model, path=model_path) reloaded_model = mlflow.xgboost.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_almost_equal( model.predict(xgb_sklearn_model.inference_dataframe), reloaded_model.predict(xgb_sklearn_model.inference_dataframe), ) np.testing.assert_array_almost_equal( reloaded_model.predict(xgb_sklearn_model.inference_dataframe), reloaded_pyfunc.predict(xgb_sklearn_model.inference_dataframe), )
def test_model_save_load(cb_model, model_path): model, inference_dataframe = cb_model mlflow.catboost.save_model(cb_model=model, path=model_path) loaded_model = mlflow.catboost.load_model(model_uri=model_path) np.testing.assert_array_almost_equal( model.predict(inference_dataframe), loaded_model.predict(inference_dataframe), ) loaded_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_almost_equal( loaded_model.predict(inference_dataframe), loaded_pyfunc.predict(inference_dataframe), )
def load_model(path): """ Loads the model object in a specific path (pyfunc) Parameters ---------- path : str path where the model object will be loaded. Returns ------- None """ model = pyfunc.load_model(path) return model
def test_model_export(spark_model_iris, model_path, spark_custom_env): sparkm.save_model(spark_model_iris.model, path=model_path, conda_env=spark_custom_env) # 1. score and compare reloaded sparkml model reloaded_model = sparkm.load_model(model_uri=model_path) preds_df = reloaded_model.transform(spark_model_iris.spark_df) preds1 = [x.prediction for x in preds_df.select("prediction").collect()] assert spark_model_iris.predictions == preds1 m = pyfunc.load_model(model_path) # 2. score and compare reloaded pyfunc preds2 = m.predict(spark_model_iris.pandas_df) assert spark_model_iris.predictions == preds2 # 3. score and compare reloaded pyfunc Spark udf preds3 = score_model_as_udf(model_uri=model_path, pandas_df=spark_model_iris.pandas_df) assert spark_model_iris.predictions == preds3 assert os.path.exists(sparkm.DFS_TMP)
def test_model_save_load(lgb_model, model_path): model = lgb_model.model mlflow.lightgbm.save_model(lgb_model=model, path=model_path) reloaded_model = mlflow.lightgbm.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_almost_equal( model.predict(lgb_model.inference_dataframe), reloaded_model.predict(lgb_model.inference_dataframe), ) np.testing.assert_array_almost_equal( reloaded_model.predict(lgb_model.inference_dataframe), reloaded_pyfunc.predict(lgb_model.inference_dataframe), )
def test_model_save_load(sklearn_knn_model, model_path): knn_model = sklearn_knn_model.model mlflow.sklearn.save_model(sk_model=knn_model, path=model_path) reloaded_knn_model = mlflow.sklearn.load_model(model_uri=model_path) reloaded_knn_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_equal( knn_model.predict(sklearn_knn_model.inference_data), reloaded_knn_model.predict(sklearn_knn_model.inference_data), ) np.testing.assert_array_equal( reloaded_knn_model.predict(sklearn_knn_model.inference_data), reloaded_knn_pyfunc.predict(sklearn_knn_model.inference_data), )