import inspect import json import logging import os from pathlib import Path from typing import Any, NamedTuple from unittest import mock import numpy as np import pandas as pd import pyspark import pytest import yaml from packaging.version import Version from pyspark.ml.classification import LogisticRegression from pyspark.ml.feature import VectorAssembler from pyspark.ml.pipeline import Pipeline from sklearn import datasets import mlflow import mlflow.pyfunc.scoring_server as pyfunc_scoring_server import mlflow.tracking import mlflow.utils.file_utils from mlflow import pyfunc from mlflow.entities.model_registry import ModelVersion from mlflow.environment_variables import MLFLOW_DFS_TMP from mlflow.exceptions import MlflowException from mlflow.models import Model, ModelSignature from mlflow.models.utils import _read_example from mlflow.spark import _add_code_from_conf_to_system_path from mlflow.store.artifact.databricks_models_artifact_repo import DatabricksModelsArtifactRepository from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository from mlflow.store.artifact.unity_catalog_models_artifact_repo import ( UnityCatalogModelsArtifactRepository, ) from mlflow.tracking.artifact_utils import _download_artifact_from_uri from mlflow.types import DataType from mlflow.types.schema import ColSpec, Schema from mlflow.utils.environment import _get_pip_deps, _mlflow_conda_env from mlflow.utils.file_utils import TempDir from mlflow.utils.model_utils import _get_flavor_configuration from tests.helper_functions import ( _assert_pip_requirements, _compare_conda_env_requirements, _compare_logged_code_paths, _mlflow_major_version_string, assert_register_model_called_with_local_model_path, score_model_in_sagemaker_docker_container, ) from tests.pyfunc.test_spark import get_spark_session, score_model_as_udf from tests.store.artifact.constants import MODELS_ARTIFACT_REPOSITORY _logger = logging.getLogger(__name__) PYSPARK_VERSION = Version(pyspark.__version__) @pytest.fixture def spark_custom_env(tmp_path): conda_env = os.path.join(tmp_path, "conda_env.yml") additional_pip_deps = ["/opt/mlflow", f"pyspark=={PYSPARK_VERSION}", "pytest"] if PYSPARK_VERSION < Version("3.4"): additional_pip_deps.extend([ # Versions of PySpark < 3.4 are incompatible with pandas >= 2 "pandas<2", # pandas<2.0 is incompatible with numpy>=2.0 "numpy<2.0", ]) _mlflow_conda_env(conda_env, additional_pip_deps=additional_pip_deps) return conda_env class SparkModelWithData(NamedTuple): model: Any spark_df: Any pandas_df: Any predictions: Any def _get_spark_session_with_retry(max_tries=3): conf = pyspark.SparkConf() for attempt in range(max_tries): try: return get_spark_session(conf) except Exception as e: if attempt >= max_tries - 1: raise _logger.exception( f"Attempt {attempt} to create a SparkSession failed ({e!r}), retrying..." ) # Specify `autouse=True` to ensure that a context is created # before any tests are executed. This ensures that the Hadoop filesystem # does not create its own SparkContext. @pytest.fixture(scope="module") def spark(): if Version(pyspark.__version__) < Version("3.1"): # A workaround for this issue: # https://stackoverflow.com/questions/62109276/errorjava-lang-unsupportedoperationexception-for-pyspark-pandas-udf-documenta spark_home = ( os.environ.get("SPARK_HOME") if "SPARK_HOME" in os.environ else os.path.dirname(pyspark.__file__) ) conf_dir = os.path.join(spark_home, "conf") os.makedirs(conf_dir, exist_ok=True) with open(os.path.join(conf_dir, "spark-defaults.conf"), "w") as f: conf = """ spark.driver.extraJavaOptions="-Dio.netty.tryReflectionSetAccessible=true" spark.executor.extraJavaOptions="-Dio.netty.tryReflectionSetAccessible=true" """ f.write(conf) with _get_spark_session_with_retry() as spark: yield spark def iris_pandas_df(): iris = datasets.load_iris() X = iris.data y = iris.target feature_names = ["0", "1", "2", "3"] df = pd.DataFrame(X, columns=feature_names) # to make spark_udf work df["label"] = pd.Series(y) return df @pytest.fixture(scope="module") def iris_df(spark): pdf = iris_pandas_df() feature_names = list(pdf.drop("label", axis=1).columns) iris_spark_df = spark.createDataFrame(pdf) return feature_names, pdf, iris_spark_df @pytest.fixture(scope="module") def iris_signature(): return ModelSignature( inputs=Schema([ ColSpec(name="0", type=DataType.double), ColSpec(name="1", type=DataType.double), ColSpec(name="2", type=DataType.double), ColSpec(name="3", type=DataType.double), ]), outputs=Schema([ColSpec(type=DataType.double)]), ) @pytest.fixture(scope="module") def spark_model_iris(iris_df): feature_names, iris_pandas_df, iris_spark_df = iris_df assembler = VectorAssembler(inputCols=feature_names, outputCol="features") lr = LogisticRegression(maxIter=50, regParam=0.1, elasticNetParam=0.8) pipeline = Pipeline(stages=[assembler, lr]) # Fit the model model = pipeline.fit(iris_spark_df) preds_df = model.transform(iris_spark_df) preds = [x.prediction for x in preds_df.select("prediction").collect()] return SparkModelWithData( model=model, spark_df=iris_spark_df, pandas_df=iris_pandas_df, predictions=preds ) @pytest.fixture(scope="module") def spark_model_transformer(iris_df): feature_names, iris_pandas_df, iris_spark_df = iris_df assembler = VectorAssembler(inputCols=feature_names, outputCol="features") # Fit the model preds_df = assembler.transform(iris_spark_df) preds = [x.features for x in preds_df.select("features").collect()] return SparkModelWithData( model=assembler, spark_df=iris_spark_df, pandas_df=iris_pandas_df, predictions=preds ) @pytest.fixture(scope="module") def spark_model_estimator(iris_df): feature_names, iris_pandas_df, iris_spark_df = iris_df assembler = VectorAssembler(inputCols=feature_names, outputCol="features") features_df = assembler.transform(iris_spark_df) lr = LogisticRegression(maxIter=50, regParam=0.1, elasticNetParam=0.8) # Fit the model model = lr.fit(features_df) preds_df = model.transform(features_df) preds = [x.prediction for x in preds_df.select("prediction").collect()] return SparkModelWithData( model=model, spark_df=features_df, pandas_df=iris_pandas_df, predictions=preds ) @pytest.fixture def model_path(tmp_path): return os.path.join(tmp_path, "model") @pytest.mark.usefixtures("spark") def test_hadoop_filesystem(tmp_path): # copy local dir to and back from HadoopFS and make sure the results match from mlflow.spark import _HadoopFileSystem as FS test_dir_0 = os.path.join(tmp_path, "expected") test_file_0 = os.path.join(test_dir_0, "root", "file_0") test_dir_1 = os.path.join(test_dir_0, "root", "subdir") test_file_1 = os.path.join(test_dir_1, "file_1") os.makedirs(os.path.dirname(test_file_0)) with open(test_file_0, "w") as f: f.write("test0") os.makedirs(os.path.dirname(test_file_1)) with open(test_file_1, "w") as f: f.write("test1") remote = "/tmp/mlflow/test0" # File should not be copied in this case assert os.path.abspath(test_dir_0) == FS.maybe_copy_from_local_file(test_dir_0, remote) FS.copy_from_local_file(test_dir_0, remote, remove_src=False) local = os.path.join(tmp_path, "actual") FS.copy_to_local_file(remote, local, remove_src=True) assert sorted(os.listdir(os.path.join(local, "root"))) == sorted([ "subdir", "file_0", ".file_0.crc", ]) assert sorted(os.listdir(os.path.join(local, "root", "subdir"))) == sorted([ "file_1", ".file_1.crc", ]) # compare the files with open(os.path.join(test_dir_0, "root", "file_0")) as expected_f: with open(os.path.join(local, "root", "file_0")) as actual_f: assert expected_f.read() == actual_f.read() with open(os.path.join(test_dir_0, "root", "subdir", "file_1")) as expected_f: with open(os.path.join(local, "root", "subdir", "file_1")) as actual_f: assert expected_f.read() == actual_f.read() # make sure we cleanup assert not os.path.exists(FS._remote_path(remote).toString()) # skip file: prefix FS.copy_from_local_file(test_dir_0, remote, remove_src=False) assert os.path.exists(FS._remote_path(remote).toString()) # skip file: prefix FS.delete(remote) assert not os.path.exists(FS._remote_path(remote).toString()) # skip file: prefix def test_model_export(spark_model_iris, model_path, spark_custom_env): mlflow.spark.save_model(spark_model_iris.model, path=model_path, conda_env=spark_custom_env) # 1. score and compare reloaded sparkml model reloaded_model = mlflow.spark.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(MLFLOW_DFS_TMP.get()) def test_model_export_with_signature_and_examples(spark_model_iris, iris_signature): features_df = spark_model_iris.pandas_df.drop("label", axis=1) example_ = features_df.head(3) for signature in (None, iris_signature): for example in (None, example_): with TempDir() as tmp: path = tmp.path("model") mlflow.spark.save_model( spark_model_iris.model, path=path, signature=signature, input_example=example ) mlflow_model = Model.load(path) if example is None and signature is None: assert mlflow_model.signature is None else: assert mlflow_model.signature == iris_signature if example is None: assert mlflow_model.saved_input_example_info is None else: assert all((_read_example(mlflow_model, path) == example).all()) def test_model_export_raise_when_example_is_spark_dataframe(spark, spark_model_iris, model_path): features_df = spark_model_iris.pandas_df.drop("label", axis=1) example = spark.createDataFrame(features_df.head(3)) with pytest.raises(MlflowException, match="Examples can not be provided as Spark Dataframe."): mlflow.spark.save_model(spark_model_iris.model, path=model_path, input_example=example) def test_log_model_with_signature_and_examples(spark_model_iris, iris_signature): features_df = spark_model_iris.pandas_df.drop("label", axis=1) example_ = features_df.head(3) artifact_path = "model" for signature in (None, iris_signature): for example in (None, example_): with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, signature=signature, input_example=example, ) mlflow_model = Model.load(model_info.model_uri) if example is None and signature is None: assert mlflow_model.signature is None else: assert mlflow_model.signature == iris_signature if example is None: assert mlflow_model.saved_input_example_info is None else: assert all((_read_example(mlflow_model, model_info.model_uri) == example).all()) def test_estimator_model_export(spark_model_estimator, model_path, spark_custom_env): mlflow.spark.save_model( spark_model_estimator.model, path=model_path, conda_env=spark_custom_env ) # score and compare the reloaded sparkml model reloaded_model = mlflow.spark.load_model(model_uri=model_path) preds_df = reloaded_model.transform(spark_model_estimator.spark_df) preds = [x.prediction for x in preds_df.select("prediction").collect()] assert spark_model_estimator.predictions == preds # 2. score and compare reloaded pyfunc m = pyfunc.load_model(model_path) preds2 = m.predict(spark_model_estimator.spark_df.toPandas()) assert spark_model_estimator.predictions == preds2 def test_transformer_model_export(spark_model_transformer, model_path, spark_custom_env): mlflow.spark.save_model( spark_model_transformer.model, path=model_path, conda_env=spark_custom_env ) # score and compare the reloaded sparkml model reloaded_model = mlflow.spark.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 @pytest.mark.skipif( PYSPARK_VERSION.is_devrelease, reason="this test does not support PySpark dev version." ) def test_model_deployment(spark_model_iris, model_path, spark_custom_env, monkeypatch): mlflow.spark.save_model( spark_model_iris.model, path=model_path, conda_env=spark_custom_env, ) scoring_response = score_model_in_sagemaker_docker_container( model_uri=model_path, data=spark_model_iris.pandas_df, content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON, flavor=mlflow.pyfunc.FLAVOR_NAME, ) from mlflow.deployments import PredictionsResponse np.testing.assert_array_almost_equal( spark_model_iris.predictions, PredictionsResponse.from_json(scoring_response.content).get_predictions( predictions_format="ndarray" ), decimal=4, ) @pytest.mark.skipif( "dev" in pyspark.__version__, reason="The dev version of pyspark built from the source doesn't exist on PyPI or Anaconda", ) def test_sagemaker_docker_model_scoring_with_default_conda_env(spark_model_iris, model_path): mlflow.spark.save_model( spark_model_iris.model, path=model_path, extra_pip_requirements=["/opt/mlflow"] ) scoring_response = score_model_in_sagemaker_docker_container( model_uri=model_path, data=spark_model_iris.pandas_df, content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON, flavor=mlflow.pyfunc.FLAVOR_NAME, ) deployed_model_preds = np.array(json.loads(scoring_response.content)["predictions"]) np.testing.assert_array_almost_equal( deployed_model_preds, spark_model_iris.predictions, decimal=4 ) @pytest.mark.parametrize("should_start_run", [False, True]) @pytest.mark.parametrize("use_dfs_tmpdir", [False, True]) def test_sparkml_model_log( tmp_path, tmp_sqlite_uri, spark_model_iris, should_start_run, use_dfs_tmpdir ): old_tracking_uri = mlflow.get_tracking_uri() dfs_tmpdir = None if use_dfs_tmpdir else tmp_path.joinpath("test") try: mlflow.set_tracking_uri(tmp_sqlite_uri) if should_start_run: mlflow.start_run() artifact_path = "model" model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, dfs_tmpdir=dfs_tmpdir, ) reloaded_model = mlflow.spark.load_model( model_uri=model_info.model_uri, dfs_tmpdir=dfs_tmpdir ) preds_df = reloaded_model.transform(spark_model_iris.spark_df) preds = [x.prediction for x in preds_df.select("prediction").collect()] assert spark_model_iris.predictions == preds finally: mlflow.end_run() mlflow.set_tracking_uri(old_tracking_uri) @pytest.mark.parametrize( ("registry_uri", "artifact_repo_class"), [ ("databricks-uc", UnityCatalogModelsArtifactRepository), ("databricks", DatabricksModelsArtifactRepository), ], ) def test_load_spark_model_from_models_uri( tmp_path, spark_model_estimator, registry_uri, artifact_repo_class ): model_dir = str(tmp_path.joinpath("spark_model")) model_name = "mycatalog.myschema.mymodel" fake_model_version = ModelVersion(name=model_name, version=str(3), creation_timestamp=0) with ( mock.patch(f"{MODELS_ARTIFACT_REPOSITORY}.get_underlying_uri") as mock_get_underlying_uri, mock.patch.object( artifact_repo_class, "download_artifacts", return_value=model_dir ) as mock_download_artifacts, mock.patch("mlflow.get_registry_uri", return_value=registry_uri), mock.patch.object( mlflow.tracking._model_registry.client.ModelRegistryClient, "get_model_version_by_alias", return_value=fake_model_version, ) as get_model_version_by_alias_mock, ): mlflow.spark.save_model( path=model_dir, spark_model=spark_model_estimator.model, ) mock_get_underlying_uri.return_value = "nonexistentscheme://fakeuri" mlflow.spark.load_model(f"models:/{model_name}/1") # Assert that we downloaded both the MLmodel file and the whole model itself using # the models:/ URI kwargs = ( {"lineage_header_info": None} if artifact_repo_class is UnityCatalogModelsArtifactRepository else {} ) mock_download_artifacts.assert_called_once_with("", None, **kwargs) mock_download_artifacts.reset_mock() mlflow.spark.load_model(f"models:/{model_name}@Champion") mock_download_artifacts.assert_called_once_with("", None, **kwargs) assert get_model_version_by_alias_mock.called_with(model_name, "Champion") @pytest.mark.parametrize("should_start_run", [False, True]) @pytest.mark.parametrize("use_dfs_tmpdir", [False, True]) def test_sparkml_estimator_model_log( tmp_path, tmp_sqlite_uri, spark_model_estimator, should_start_run, use_dfs_tmpdir ): old_tracking_uri = mlflow.get_tracking_uri() dfs_tmpdir = None if use_dfs_tmpdir else tmp_path.joinpath("test") try: mlflow.set_tracking_uri(tmp_sqlite_uri) if should_start_run: mlflow.start_run() artifact_path = "model" model_info = mlflow.spark.log_model( spark_model_estimator.model, artifact_path=artifact_path, dfs_tmpdir=dfs_tmpdir, ) reloaded_model = mlflow.spark.load_model( model_uri=model_info.model_uri, dfs_tmpdir=dfs_tmpdir ) preds_df = reloaded_model.transform(spark_model_estimator.spark_df) preds = [x.prediction for x in preds_df.select("prediction").collect()] assert spark_model_estimator.predictions == preds finally: mlflow.end_run() mlflow.set_tracking_uri(old_tracking_uri) def test_log_model_calls_register_model(tmp_path, spark_model_iris): artifact_path = "model" dfs_tmp_dir = tmp_path.joinpath("test") register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model") with mlflow.start_run(), register_model_patch: model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, dfs_tmpdir=dfs_tmp_dir, registered_model_name="AdsModel1", ) assert_register_model_called_with_local_model_path( register_model_mock=mlflow.tracking._model_registry.fluent._register_model, model_uri=model_info.model_uri, registered_model_name="AdsModel1", ) def test_log_model_no_registered_model_name(tmp_path, spark_model_iris): artifact_path = "model" dfs_tmp_dir = os.path.join(tmp_path, "test") register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model") with mlflow.start_run(), register_model_patch: mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, dfs_tmpdir=dfs_tmp_dir, ) mlflow.tracking._model_registry.fluent._register_model.assert_not_called() def test_log_model_skips_maybe_save_for_acled_artifact_uri(tmp_path): """_maybe_save_model should not be called for Databricks ACL-protected artifact URIs (dbfs:/databricks/mlflow-tracking/...) since Spark cannot write to them directly. Calling it wastes ~6s per model on a guaranteed Py4JError before falling back. """ acled_uri = "dbfs:/databricks/mlflow-tracking/abc123/run456/artifacts" class FakePipelineModel: def __init__(self, stages=None): pass mock_model = FakePipelineModel() with ( mock.patch("mlflow.spark._validate_model"), mock.patch("mlflow.spark._is_spark_connect_model", return_value=False), mock.patch("mlflow.spark._maybe_save_model") as mock_maybe_save, mock.patch("mlflow.get_artifact_uri", return_value=acled_uri), mock.patch("mlflow.spark._should_use_mlflowdbfs", return_value=False), mock.patch("mlflow.models.Model._log_v2") as mock_log_v2, mock.patch("pyspark.ml.PipelineModel", FakePipelineModel), mlflow.start_run(), ): mlflow.spark.log_model( mock_model, artifact_path="model", dfs_tmpdir=str(tmp_path), ) mock_maybe_save.assert_not_called() mock_log_v2.assert_called_once() def test_sparkml_model_load_from_remote_uri_succeeds(spark_model_iris, model_path, mock_s3_bucket): mlflow.spark.save_model(spark_model=spark_model_iris.model, path=model_path) artifact_root = f"s3://{mock_s3_bucket}" artifact_path = "model" artifact_repo = S3ArtifactRepository(artifact_root) artifact_repo.log_artifacts(model_path, artifact_path=artifact_path) model_uri = artifact_root + "/" + artifact_path reloaded_model = mlflow.spark.load_model(model_uri=model_uri) preds_df = reloaded_model.transform(spark_model_iris.spark_df) preds = [x.prediction for x in preds_df.select("prediction").collect()] assert spark_model_iris.predictions == preds def test_sparkml_model_save_persists_specified_conda_env_in_mlflow_model_directory( spark_model_iris, model_path, spark_custom_env ): mlflow.spark.save_model( spark_model=spark_model_iris.model, path=model_path, conda_env=spark_custom_env ) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != spark_custom_env with open(spark_custom_env) as f: spark_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == spark_custom_env_parsed def test_sparkml_model_save_persists_requirements_in_mlflow_model_directory( spark_model_iris, model_path, spark_custom_env ): mlflow.spark.save_model( spark_model=spark_model_iris.model, path=model_path, conda_env=spark_custom_env ) saved_pip_req_path = os.path.join(model_path, "requirements.txt") _compare_conda_env_requirements(spark_custom_env, saved_pip_req_path) def test_log_model_with_pip_requirements(spark_model_iris, tmp_path): expected_mlflow_version = _mlflow_major_version_string() # Path to a requirements file req_file = tmp_path.joinpath("requirements.txt") req_file.write_text("a") with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path="model", pip_requirements=str(req_file) ) _assert_pip_requirements(model_info.model_uri, [expected_mlflow_version, "a"], strict=True) # List of requirements with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path="model", pip_requirements=[f"-r {req_file}", "b"] ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, "a", "b"], strict=True ) # Constraints file with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path="model", pip_requirements=[f"-c {req_file}", "b"] ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, "b", "-c constraints.txt"], ["a"], strict=True, ) def test_log_model_with_extra_pip_requirements(spark_model_iris, tmp_path): expected_mlflow_version = _mlflow_major_version_string() default_reqs = mlflow.spark.get_default_pip_requirements() # Path to a requirements file req_file = tmp_path.joinpath("requirements.txt") req_file.write_text("a") with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path="model", extra_pip_requirements=str(req_file) ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, *default_reqs, "a"] ) # List of requirements with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path="model", extra_pip_requirements=[f"-r {req_file}", "b"], ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, *default_reqs, "a", "b"] ) # Constraints file with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path="model", extra_pip_requirements=[f"-c {req_file}", "b"], ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, *default_reqs, "b", "-c constraints.txt"], ["a"], ) def test_sparkml_model_save_accepts_conda_env_as_dict(spark_model_iris, model_path): conda_env = dict(mlflow.spark.get_default_conda_env()) conda_env["dependencies"].append("pytest") mlflow.spark.save_model( spark_model=spark_model_iris.model, path=model_path, conda_env=conda_env ) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == conda_env def test_sparkml_model_log_persists_specified_conda_env_in_mlflow_model_directory( spark_model_iris, model_path, spark_custom_env ): artifact_path = "model" with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, conda_env=spark_custom_env, ) model_path = _download_artifact_from_uri(artifact_uri=model_info.model_uri) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != spark_custom_env with open(spark_custom_env) as f: spark_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == spark_custom_env_parsed def test_sparkml_model_log_persists_requirements_in_mlflow_model_directory( spark_model_iris, model_path, spark_custom_env ): artifact_path = "model" with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, conda_env=spark_custom_env, ) model_path = _download_artifact_from_uri(artifact_uri=model_info.model_uri) saved_pip_req_path = os.path.join(model_path, "requirements.txt") _compare_conda_env_requirements(spark_custom_env, saved_pip_req_path) def test_sparkml_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies( spark_model_iris, model_path ): mlflow.spark.save_model(spark_model=spark_model_iris.model, path=model_path) _assert_pip_requirements(model_path, mlflow.spark.get_default_pip_requirements()) def test_sparkml_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies( spark_model_iris, ): artifact_path = "model" with mlflow.start_run(): model_info = mlflow.spark.log_model(spark_model_iris.model, artifact_path=artifact_path) _assert_pip_requirements(model_info.model_uri, mlflow.spark.get_default_pip_requirements()) def test_pyspark_version_is_logged_without_dev_suffix(spark_model_iris): expected_mlflow_version = _mlflow_major_version_string() unsuffixed_version = "2.4.0" for dev_suffix in [".dev0", ".dev", ".dev1", "dev.a", ".devb"]: with mock.patch("importlib_metadata.version", return_value=unsuffixed_version + dev_suffix): with mlflow.start_run(): model_info = mlflow.spark.log_model(spark_model_iris.model, artifact_path="model") _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, f"pyspark=={unsuffixed_version}"] ) for unaffected_version in ["2.0", "2.3.4", "2"]: with mock.patch("importlib_metadata.version", return_value=unaffected_version): pip_deps = _get_pip_deps(mlflow.spark.get_default_conda_env()) assert any(x == f"pyspark=={unaffected_version}" for x in pip_deps) def test_model_is_recorded_when_using_direct_save(spark_model_iris): # Patch `is_local_uri` to enforce direct model serialization to DFS with mock.patch("mlflow.spark.is_local_uri", return_value=False): with mlflow.start_run(): mlflow.spark.log_model(spark_model_iris.model, artifact_path="model") current_tags = mlflow.get_run(mlflow.active_run().info.run_id).data.tags assert mlflow.utils.mlflow_tags.MLFLOW_LOGGED_MODELS in current_tags @pytest.mark.parametrize( ( "artifact_uri", "db_runtime_version", "mlflowdbfs_disabled", "mlflowdbfs_available", "dbutils_available", "expected_uri", "expect_log_v2", ), [ ( "dbfs:/databricks/mlflow-tracking/a/b", "12.0", "", True, True, "mlflowdbfs:///artifacts?run_id={}&path=/model/sparkml", False, ), ( "dbfs:/databricks/mlflow-tracking/a/b", "12.0", "false", True, True, "mlflowdbfs:///artifacts?run_id={}&path=/model/sparkml", False, ), # ACL-protected paths where mlflowdbfs is unavailable/disabled always route through # Model._log_v2 because _maybe_save_model is skipped via is_databricks_acled_artifacts_uri. # In real Databricks, _maybe_save_model always fails with Py4JError for these paths anyway. ( "dbfs:/databricks/mlflow-tracking/a/b", "12.0", "false", True, False, None, True, ), ( "dbfs:/databricks/mlflow-tracking/a/b", "12.0", "", False, True, None, True, ), ( "dbfs:/databricks/mlflow-tracking/a/b", "", "", True, True, None, True, ), ( "dbfs:/databricks/mlflow-tracking/a/b", "12.0", "true", True, True, None, True, ), ("dbfs:/root/a/b", "12.0", "", True, True, "dbfs:/root/a/b/model/sparkml", False), ("s3://mybucket/a/b", "12.0", "", True, True, "s3://mybucket/a/b/model/sparkml", False), ], ) def test_model_logged_via_mlflowdbfs_when_appropriate( monkeypatch, spark_model_iris, artifact_uri, db_runtime_version, mlflowdbfs_disabled, mlflowdbfs_available, dbutils_available, expected_uri, expect_log_v2, ): def mock_spark_session_load(path): raise Exception("MlflowDbfsClient operation failed!") mock_spark_session = mock.Mock() mock_read_spark_session = mock.Mock() mock_read_spark_session.load = mock_spark_session_load from mlflow.utils.databricks_utils import _get_dbutils as og_getdbutils def mock_get_dbutils(): # _get_dbutils is called during run creation and model logging; to avoid breaking run # creation, we only mock the output if _get_dbutils is called during spark model logging caller_fn_name = inspect.stack()[1].function if caller_fn_name == "_should_use_mlflowdbfs": if dbutils_available: return mock.Mock() else: raise Exception("dbutils not available") else: return og_getdbutils() with ( mock.patch( "mlflow.utils._spark_utils._get_active_spark_session", return_value=mock_spark_session ), mock.patch("mlflow.get_artifact_uri", return_value=artifact_uri), mock.patch( "mlflow.spark._HadoopFileSystem.is_filesystem_available", return_value=mlflowdbfs_available, ), mock.patch("mlflow.utils.databricks_utils.MlflowCredentialContext", autospec=True), mock.patch("mlflow.utils.databricks_utils._get_dbutils", mock_get_dbutils), mock.patch.object(spark_model_iris.model, "save") as mock_save, mock.patch("mlflow.models.infer_pip_requirements", return_value=[]) as mock_infer, mock.patch("mlflow.models.Model._log_v2") as mock_log_v2, ): with mlflow.start_run(): if db_runtime_version: monkeypatch.setenv("DATABRICKS_RUNTIME_VERSION", db_runtime_version) monkeypatch.setenv("DISABLE_MLFLOWDBFS", mlflowdbfs_disabled) mlflow.spark.log_model(spark_model_iris.model, artifact_path="model") if expect_log_v2: # ACL-protected paths where mlflowdbfs is unavailable skip _maybe_save_model # entirely and fall through to Model._log_v2. In production, _maybe_save_model # always raises Py4JError for these paths, so skipping it is correct. mock_log_v2.assert_called_once() mock_save.assert_not_called() else: mock_save.assert_called_once_with( expected_uri.format(mlflow.active_run().info.run_id) ) if expected_uri.startswith("mlflowdbfs"): # If mlflowdbfs is used, infer_pip_requirements should load the model from the # remote model path instead of a local tmp path. assert ( mock_infer.call_args[0][0] == "dbfs:/databricks/mlflow-tracking/a/b/model/sparkml" ) @pytest.mark.parametrize("dummy_read_shows_mlflowdbfs_available", [True, False]) def test_model_logging_uses_mlflowdbfs_if_appropriate_when_hdfs_check_fails( monkeypatch, spark_model_iris, dummy_read_shows_mlflowdbfs_available ): def mock_spark_session_load(path): if dummy_read_shows_mlflowdbfs_available: raise Exception("MlflowdbfsClient operation failed!") else: raise Exception("mlflowdbfs filesystem not found") mock_read_spark_session = mock.Mock() mock_read_spark_session.load = mock_spark_session_load mock_spark_session = mock.Mock() mock_spark_session.read = mock_read_spark_session from mlflow.utils.databricks_utils import _get_dbutils as og_getdbutils def mock_get_dbutils(): # _get_dbutils is called during run creation and model logging; to avoid breaking run # creation, we only mock the output if _get_dbutils is called during spark model logging caller_fn_name = inspect.stack()[1].function if caller_fn_name == "_should_use_mlflowdbfs": return mock.Mock() else: return og_getdbutils() with ( mock.patch( "mlflow.utils._spark_utils._get_active_spark_session", return_value=mock_spark_session, ), mock.patch( "mlflow.get_artifact_uri", return_value="dbfs:/databricks/mlflow-tracking/a/b", ), mock.patch( "mlflow.spark._HadoopFileSystem.is_filesystem_available", side_effect=Exception("MlflowDbfsClient operation failed!"), ), mock.patch("mlflow.utils.databricks_utils.MlflowCredentialContext", autospec=True), mock.patch( "mlflow.utils.databricks_utils._get_dbutils", mock_get_dbutils, ), mock.patch.object(spark_model_iris.model, "save") as mock_save, mock.patch("mlflow.models.Model._log_v2") as mock_log_v2, ): with mlflow.start_run(): monkeypatch.setenv("DATABRICKS_RUNTIME_VERSION", "12.0") mlflow.spark.log_model(spark_model_iris.model, artifact_path="model") run_id = mlflow.active_run().info.run_id if dummy_read_shows_mlflowdbfs_available: mock_save.assert_called_once_with( f"mlflowdbfs:///artifacts?run_id={run_id}&path=/model/sparkml" ) else: # mlflowdbfs unavailable + ACL-protected path: _maybe_save_model is skipped, # Model._log_v2 is called directly. In production, _maybe_save_model always # raises Py4JError for these ACL-protected paths, so skipping it is correct. mock_log_v2.assert_called_once() mock_save.assert_not_called() def test_log_model_with_code_paths(spark_model_iris): artifact_path = "model" with ( mlflow.start_run(), mock.patch( "mlflow.spark._add_code_from_conf_to_system_path", wraps=_add_code_from_conf_to_system_path, ) as add_mock, ): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, code_paths=[__file__] ) _compare_logged_code_paths(__file__, model_info.model_uri, mlflow.spark.FLAVOR_NAME) mlflow.spark.load_model(model_info.model_uri) add_mock.assert_called() def test_virtualenv_subfield_points_to_correct_path(spark_model_iris, model_path): mlflow.spark.save_model(spark_model_iris.model, path=model_path) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) python_env_path = Path(model_path, pyfunc_conf[pyfunc.ENV]["virtualenv"]) assert python_env_path.exists() assert python_env_path.is_file() def test_model_save_load_with_metadata(spark_model_iris, model_path): mlflow.spark.save_model( spark_model_iris.model, path=model_path, metadata={"metadata_key": "metadata_value"} ) reloaded_model = mlflow.pyfunc.load_model(model_uri=model_path) assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value" def test_model_log_with_metadata(spark_model_iris): with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path="model", metadata={"metadata_key": "metadata_value"}, ) reloaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri) assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value" _df_input_example = iris_pandas_df().drop("label", axis=1).iloc[[0]] @pytest.mark.parametrize( "input_example", # array and dict input examples are not supported any more as they # won't be converted to pandas dataframe when saving example [_df_input_example], ) def test_model_log_with_signature_inference(spark_model_iris, input_example): artifact_path = "model" with mlflow.start_run(): model_info = mlflow.spark.log_model( spark_model_iris.model, artifact_path=artifact_path, input_example=input_example ) mlflow_model = Model.load(model_info.model_uri) input_columns = mlflow_model.signature.inputs.inputs assert all(col.type == DataType.double for col in input_columns) column_names = [col.name for col in input_columns] if isinstance(input_example, list): assert column_names == [0, 1, 2, 3] else: assert column_names == ["0", "1", "2", "3"] assert mlflow_model.signature.outputs == Schema([ColSpec(type=DataType.double)]) def test_log_model_with_vector_input_type_signature(spark, spark_model_estimator): from pyspark.ml.functions import vector_to_array from mlflow.types.schema import SparkMLVector model = spark_model_estimator.model with mlflow.start_run(): model_info = mlflow.spark.log_model( model, artifact_path="model", signature=ModelSignature( inputs=Schema([ ColSpec(name="features", type=SparkMLVector()), ]), outputs=Schema([ColSpec(type=DataType.double)]), ), ) model_uri = model_info.model_uri model_meta = Model.load(model_uri) input_type = model_meta.signature.inputs.input_dict()["features"].type assert isinstance(input_type, SparkMLVector) pyfunc_model = pyfunc.load_model(model_uri) infer_data = spark_model_estimator.spark_df.withColumn( "features", vector_to_array("features") ).toPandas() preds = pyfunc_model.predict(infer_data) assert spark_model_estimator.predictions == preds