wehub-resource-sync
53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from pyspark.sql import SparkSession
|
|
from sklearn.datasets import load_iris
|
|
from sklearn.linear_model import LogisticRegression
|
|
|
|
import mlflow
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def spark():
|
|
spark = SparkSession.builder.remote("local[2]").getOrCreate()
|
|
yield spark
|
|
spark.stop()
|
|
|
|
|
|
def test_spark_udf_spark_connect(spark):
|
|
X, y = load_iris(return_X_y=True)
|
|
model = LogisticRegression().fit(X, y)
|
|
with mlflow.start_run():
|
|
info = mlflow.sklearn.log_model(model, name="model")
|
|
sdf = spark.createDataFrame(pd.DataFrame(X, columns=list("abcd")))
|
|
udf = mlflow.pyfunc.spark_udf(spark, info.model_uri, env_manager="local")
|
|
result = sdf.select(udf(*sdf.columns).alias("preds")).toPandas()
|
|
np.testing.assert_almost_equal(result["preds"].to_numpy(), model.predict(X))
|
|
|
|
|
|
@pytest.mark.parametrize("env_manager", ["conda", "virtualenv"])
|
|
def test_spark_udf_spark_connect_unsupported_env_manager(spark, tmp_path, env_manager):
|
|
with pytest.raises(
|
|
mlflow.MlflowException,
|
|
match=f"Environment manager {env_manager!r} is not supported",
|
|
):
|
|
mlflow.pyfunc.spark_udf(spark, str(tmp_path), env_manager=env_manager)
|
|
|
|
|
|
def test_spark_udf_spark_connect_with_model_logging(spark, db_uri):
|
|
X, y = load_iris(return_X_y=True, as_frame=True)
|
|
model = LogisticRegression().fit(X, y)
|
|
|
|
mlflow.set_tracking_uri(db_uri)
|
|
mlflow.set_experiment("test")
|
|
with mlflow.start_run():
|
|
signature = mlflow.models.infer_signature(X, y)
|
|
model_info = mlflow.sklearn.log_model(model, name="model", signature=signature)
|
|
|
|
udf = mlflow.pyfunc.spark_udf(spark, model_info.model_uri, env_manager="local")
|
|
X_test = X.head(5)
|
|
sdf = spark.createDataFrame(X_test)
|
|
preds = sdf.select(udf(*X_test.columns).alias("preds")).toPandas()["preds"]
|
|
np.testing.assert_array_almost_equal(preds, model.predict(X_test))
|