wehub-resource-sync
1291 lines
54 KiB
Python
1291 lines
54 KiB
Python
"""
|
|
The ``mlflow.spark`` module provides an API for logging and loading Spark MLlib models. This module
|
|
exports Spark MLlib models with the following flavors:
|
|
|
|
Spark MLlib (native) format
|
|
Allows models to be loaded as Spark Transformers for scoring in a Spark session.
|
|
Models with this flavor can be loaded as PySpark PipelineModel objects in Python.
|
|
This is the main flavor and is always produced.
|
|
:py:mod:`mlflow.pyfunc`
|
|
Supports deployment outside of Spark by instantiating a SparkContext and reading
|
|
input data as a Spark DataFrame prior to scoring. Also supports deployment in Spark
|
|
as a Spark UDF. Models with this flavor can be loaded as Python functions
|
|
for performing inference. This flavor is always produced.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import posixpath
|
|
import re
|
|
import shutil
|
|
from typing import Any
|
|
from urllib.parse import urlparse
|
|
|
|
import yaml
|
|
from packaging.version import Version
|
|
|
|
import mlflow
|
|
from mlflow import environment_variables, pyfunc
|
|
from mlflow.environment_variables import MLFLOW_DFS_TMP
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models import Model, ModelInputExample, ModelSignature
|
|
from mlflow.models.model import MLMODEL_FILE_NAME
|
|
from mlflow.models.signature import _LOG_MODEL_INFER_SIGNATURE_WARNING_TEMPLATE
|
|
from mlflow.models.utils import _Example, _save_example
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.store.artifact.databricks_artifact_repo import DatabricksArtifactRepository
|
|
from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
|
|
from mlflow.tracking.artifact_utils import (
|
|
_download_artifact_from_uri,
|
|
_get_root_uri_and_artifact_path,
|
|
)
|
|
from mlflow.types.schema import SparkMLVector
|
|
from mlflow.utils import _get_fully_qualified_class_name, databricks_utils
|
|
from mlflow.utils.autologging_utils import autologging_integration, safe_patch
|
|
from mlflow.utils.class_utils import _get_class_from_string
|
|
from mlflow.utils.docstring_utils import LOG_MODEL_PARAM_DOCS, format_docstring
|
|
from mlflow.utils.environment import (
|
|
_CONDA_ENV_FILE_NAME,
|
|
_CONSTRAINTS_FILE_NAME,
|
|
_PYTHON_ENV_FILE_NAME,
|
|
_REQUIREMENTS_FILE_NAME,
|
|
_mlflow_conda_env,
|
|
_process_conda_env,
|
|
_process_pip_requirements,
|
|
_PythonEnv,
|
|
_validate_env_arguments,
|
|
)
|
|
from mlflow.utils.file_utils import (
|
|
TempDir,
|
|
get_total_file_size,
|
|
shutil_copytree_without_file_permissions,
|
|
write_to,
|
|
)
|
|
from mlflow.utils.model_utils import (
|
|
_add_code_from_conf_to_system_path,
|
|
_validate_and_copy_code_paths,
|
|
)
|
|
from mlflow.utils.requirements_utils import _get_pinned_requirement
|
|
from mlflow.utils.uri import (
|
|
append_to_uri_path,
|
|
dbfs_hdfs_uri_to_fuse_path,
|
|
generate_tmp_dfs_path,
|
|
get_databricks_profile_uri_from_artifact_uri,
|
|
is_databricks_acled_artifacts_uri,
|
|
is_local_uri,
|
|
is_valid_dbfs_uri,
|
|
)
|
|
|
|
FLAVOR_NAME = "spark"
|
|
|
|
_SPARK_MODEL_PATH_SUB = "sparkml"
|
|
_MLFLOWDBFS_SCHEME = "mlflowdbfs"
|
|
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_default_pip_requirements(is_spark_connect_model=False):
|
|
"""
|
|
Returns:
|
|
A list of default pip requirements for MLflow Models produced by this flavor.
|
|
Calls to :func:`save_model()` and :func:`log_model()` produce a pip environment
|
|
that, at minimum, contains these requirements.
|
|
"""
|
|
import pyspark
|
|
|
|
# Strip the suffix from `dev` versions of PySpark, which are not
|
|
# available for installation from Anaconda or PyPI
|
|
pyspark_req_str = "pyspark[connect]" if is_spark_connect_model else "pyspark"
|
|
pyspark_req = re.sub(r"(\.?)dev.*$", "", _get_pinned_requirement(pyspark_req_str))
|
|
reqs = [pyspark_req]
|
|
if Version(pyspark.__version__) < Version("3.4"):
|
|
# Versions of PySpark < 3.4 are incompatible with pandas >= 2
|
|
reqs.append("pandas<2")
|
|
|
|
if is_spark_connect_model:
|
|
reqs.extend([
|
|
# Spark connect ML Model uses spark torch distributor to train model
|
|
_get_pinned_requirement("torch"),
|
|
# Spark connect ML Model saves feature transformers as sklearn transformer format.
|
|
_get_pinned_requirement("scikit-learn", module="sklearn"),
|
|
# Spark connect ML evaluators depend on torcheval package.
|
|
_get_pinned_requirement("torcheval"),
|
|
])
|
|
return reqs
|
|
|
|
|
|
def get_default_conda_env(is_spark_connect_model=False):
|
|
"""
|
|
Returns:
|
|
The default Conda environment for MLflow Models produced by calls to
|
|
:func:`save_model()` and :func:`log_model()`. This Conda environment
|
|
contains the current version of PySpark that is installed on the caller's
|
|
system. ``dev`` versions of PySpark are replaced with stable versions in
|
|
the resulting Conda environment (e.g., if you are running PySpark version
|
|
``2.4.5.dev0``, invoking this method produces a Conda environment with a
|
|
dependency on PySpark version ``2.4.5``).
|
|
"""
|
|
return _mlflow_conda_env(
|
|
additional_pip_deps=get_default_pip_requirements(
|
|
is_spark_connect_model=is_spark_connect_model
|
|
)
|
|
)
|
|
|
|
|
|
@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name="pyspark"))
|
|
def log_model(
|
|
spark_model,
|
|
artifact_path,
|
|
conda_env=None,
|
|
code_paths=None,
|
|
dfs_tmpdir=None,
|
|
registered_model_name=None,
|
|
signature: ModelSignature = None,
|
|
input_example: ModelInputExample = None,
|
|
await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
|
|
pip_requirements=None,
|
|
extra_pip_requirements=None,
|
|
metadata=None,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Log a Spark MLlib model as an MLflow artifact for the current run. This uses the
|
|
MLlib persistence format and produces an MLflow Model with the Spark flavor.
|
|
|
|
Note: If no run is active, it will instantiate a run to obtain a run_id.
|
|
|
|
Args:
|
|
spark_model: Spark model to be saved - MLflow can only save descendants of
|
|
pyspark.ml.Model or pyspark.ml.Transformer which implement
|
|
MLReadable and MLWritable.
|
|
|
|
.. Note:: The provided Spark model's `transform` method must generate one column
|
|
named with "prediction", the column is used as MLflow pyfunc model output.
|
|
Most Spark models generate the output column with "prediction" name that
|
|
contains prediction labels by default.
|
|
To set probability column as the output column for probabilistic
|
|
classification models, you need to set "probabilityCol" param to "prediction"
|
|
and set "predictionCol" param to "".
|
|
(e.g. `model.setProbabilityCol("prediction").setPredictionCol("")`)
|
|
artifact_path: Run relative artifact path.
|
|
conda_env: {{ conda_env }}
|
|
code_paths: {{ code_paths }}
|
|
dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
|
|
filesystem if running in local mode. The model is written in this
|
|
destination and then copied into the model's artifact directory. This is
|
|
necessary as Spark ML models read from and write to DFS if running on a
|
|
cluster. If this operation completes successfully, all temporary files
|
|
created on the DFS are removed. Defaults to ``/tmp/mlflow``.
|
|
For models defined in `pyspark.ml.connect` module, this param is ignored.
|
|
registered_model_name: If given, create a model version under
|
|
``registered_model_name``, also creating a registered model if one
|
|
with the given name does not exist.
|
|
signature: A Model Signature object that describes the input and output Schema of the
|
|
model. The model signature can be inferred using `infer_signature` function
|
|
of `mlflow.models.signature`.
|
|
Note if your Spark model contains Spark ML vector type input or output column,
|
|
you should create ``SparkMLVector`` vector type for the column,
|
|
`infer_signature` function can also infer ``SparkMLVector`` vector type correctly
|
|
from Spark Dataframe input / output.
|
|
When loading a Spark ML model with ``SparkMLVector`` vector type input as MLflow
|
|
pyfunc model, it accepts ``Array[double]`` type input. MLflow internally converts
|
|
the array into Spark ML vector and then invoke Spark model for inference. Similarly,
|
|
if the model has vector type output, MLflow internally converts Spark ML vector
|
|
output data into ``Array[double]`` type inference result.
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.models import infer_signature
|
|
from pyspark.sql.functions import col
|
|
from pyspark.ml.classification import LogisticRegression
|
|
from pyspark.ml.functions import array_to_vector
|
|
import pandas as pd
|
|
import mlflow
|
|
|
|
train_df = spark.createDataFrame(
|
|
[([3.0, 4.0], 0), ([5.0, 6.0], 1)], schema="features array<double>, label long"
|
|
).select(array_to_vector("features").alias("features"), col("label"))
|
|
lor = LogisticRegression(maxIter=2)
|
|
lor.setPredictionCol("").setProbabilityCol("prediction")
|
|
lor_model = lor.fit(train_df)
|
|
|
|
test_df = train_df.select("features")
|
|
prediction_df = lor_model.transform(train_df).select("prediction")
|
|
|
|
signature = infer_signature(test_df, prediction_df)
|
|
|
|
with mlflow.start_run() as run:
|
|
model_info = mlflow.spark.log_model(
|
|
lor_model,
|
|
"model",
|
|
signature=signature,
|
|
)
|
|
|
|
# The following signature is outputted:
|
|
# inputs:
|
|
# ['features': SparkML vector (required)]
|
|
# outputs:
|
|
# ['prediction': SparkML vector (required)]
|
|
print(model_info.signature)
|
|
|
|
loaded = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
|
|
test_dataset = pd.DataFrame({"features": [[1.0, 2.0]]})
|
|
|
|
# `loaded.predict` accepts `Array[double]` type input column,
|
|
# and generates `Array[double]` type output column.
|
|
print(loaded.predict(test_dataset))
|
|
|
|
input_example: {{ input_example }}
|
|
await_registration_for: Number of seconds to wait for the model version to finish
|
|
being created and is in ``READY`` status. By default, the function
|
|
waits for five minutes. Specify 0 or None to skip waiting.
|
|
pip_requirements: {{ pip_requirements }}
|
|
extra_pip_requirements: {{ extra_pip_requirements }}
|
|
metadata: {{ metadata }}
|
|
kwargs: Extra arguments to pass to the Spark model ``save`` method.
|
|
|
|
Returns:
|
|
A :py:class:`ModelInfo <mlflow.models.model.ModelInfo>` instance that contains the
|
|
metadata of the logged model.
|
|
|
|
.. code-block:: python
|
|
:caption: Example
|
|
|
|
from pyspark.ml import Pipeline
|
|
from pyspark.ml.classification import LogisticRegression
|
|
from pyspark.ml.feature import HashingTF, Tokenizer
|
|
|
|
training = spark.createDataFrame(
|
|
[
|
|
(0, "a b c d e spark", 1.0),
|
|
(1, "b d", 0.0),
|
|
(2, "spark f g h", 1.0),
|
|
(3, "hadoop mapreduce", 0.0),
|
|
],
|
|
["id", "text", "label"],
|
|
)
|
|
tokenizer = Tokenizer(inputCol="text", outputCol="words")
|
|
hashingTF = HashingTF(inputCol=tokenizer.getOutputCol(), outputCol="features")
|
|
lr = LogisticRegression(maxIter=10, regParam=0.001)
|
|
pipeline = Pipeline(stages=[tokenizer, hashingTF, lr])
|
|
model = pipeline.fit(training)
|
|
mlflow.spark.log_model(model, "spark-model")
|
|
"""
|
|
_validate_model(spark_model)
|
|
from pyspark.ml import PipelineModel
|
|
|
|
if _is_spark_connect_model(spark_model):
|
|
# TODO: Use `Model.log` once `mlflowdbfs` supports logged model artifacts.
|
|
# `mlflowdbfs` doesn't support logged model artifacts yet, so we use `Model._log_v2`.
|
|
return Model._log_v2(
|
|
artifact_path=artifact_path,
|
|
flavor=mlflow.spark,
|
|
spark_model=spark_model,
|
|
conda_env=conda_env,
|
|
code_paths=code_paths,
|
|
registered_model_name=registered_model_name,
|
|
signature=signature,
|
|
input_example=input_example,
|
|
await_registration_for=await_registration_for,
|
|
pip_requirements=pip_requirements,
|
|
extra_pip_requirements=extra_pip_requirements,
|
|
metadata=metadata,
|
|
)
|
|
|
|
if not isinstance(spark_model, PipelineModel):
|
|
spark_model = PipelineModel([spark_model])
|
|
run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
|
|
run_root_artifact_uri = mlflow.get_artifact_uri()
|
|
remote_model_path = None
|
|
if _should_use_mlflowdbfs(run_root_artifact_uri):
|
|
remote_model_path = append_to_uri_path(
|
|
run_root_artifact_uri, artifact_path, _SPARK_MODEL_PATH_SUB
|
|
)
|
|
mlflowdbfs_path = _mlflowdbfs_path(run_id, artifact_path)
|
|
with databricks_utils.MlflowCredentialContext(
|
|
get_databricks_profile_uri_from_artifact_uri(run_root_artifact_uri)
|
|
):
|
|
try:
|
|
spark_model.save(mlflowdbfs_path)
|
|
except Exception as e:
|
|
raise MlflowException("failed to save spark model via mlflowdbfs") from e
|
|
|
|
# If the artifact URI is a local filesystem path, defer to Model.log() to persist the model,
|
|
# since Spark may not be able to write directly to the driver's filesystem. For example,
|
|
# writing to `file:/uri` will write to the local filesystem from each executor, which will
|
|
# be incorrect on multi-node clusters.
|
|
# If the artifact URI is not a local filesystem path we attempt to write directly to the
|
|
# artifact repo via Spark. If this fails, we defer to Model.log().
|
|
elif (
|
|
is_local_uri(run_root_artifact_uri)
|
|
or databricks_utils.is_in_databricks_serverless_runtime()
|
|
or databricks_utils.is_in_databricks_shared_cluster_runtime()
|
|
or is_databricks_acled_artifacts_uri(run_root_artifact_uri)
|
|
or not _maybe_save_model(
|
|
spark_model,
|
|
append_to_uri_path(run_root_artifact_uri, artifact_path),
|
|
)
|
|
):
|
|
dfs_tmpdir = dfs_tmpdir or MLFLOW_DFS_TMP.get()
|
|
_check_databricks_uc_volume_tmpdir_availability(dfs_tmpdir)
|
|
# TODO: Use `Model.log` once `mlflowdbfs` supports logged model artifacts.
|
|
# `mlflowdbfs` doesn't support logged model artifacts yet, so we use `Model._log_v2`.
|
|
return Model._log_v2(
|
|
artifact_path=artifact_path,
|
|
flavor=mlflow.spark,
|
|
spark_model=spark_model,
|
|
conda_env=conda_env,
|
|
code_paths=code_paths,
|
|
dfs_tmpdir=dfs_tmpdir,
|
|
registered_model_name=registered_model_name,
|
|
signature=signature,
|
|
input_example=input_example,
|
|
await_registration_for=await_registration_for,
|
|
pip_requirements=pip_requirements,
|
|
extra_pip_requirements=extra_pip_requirements,
|
|
metadata=metadata,
|
|
)
|
|
# Otherwise, override the default model log behavior and save model directly to artifact repo
|
|
mlflow_model = Model(artifact_path=artifact_path, run_id=run_id)
|
|
with TempDir() as tmp:
|
|
tmp_model_metadata_dir = tmp.path()
|
|
_save_model_metadata(
|
|
tmp_model_metadata_dir,
|
|
spark_model,
|
|
mlflow_model,
|
|
conda_env,
|
|
code_paths,
|
|
signature=signature,
|
|
input_example=input_example,
|
|
pip_requirements=pip_requirements,
|
|
extra_pip_requirements=extra_pip_requirements,
|
|
remote_model_path=remote_model_path,
|
|
)
|
|
mlflow.tracking.fluent.log_artifacts(tmp_model_metadata_dir, artifact_path)
|
|
mlflow.tracking.fluent._record_logged_model(mlflow_model)
|
|
if registered_model_name is not None:
|
|
mlflow.register_model(
|
|
f"runs:/{run_id}/{artifact_path}",
|
|
registered_model_name,
|
|
await_registration_for,
|
|
)
|
|
return mlflow_model.get_model_info()
|
|
|
|
|
|
def _mlflowdbfs_path(run_id, artifact_path):
|
|
if artifact_path.startswith("/"):
|
|
raise MlflowException(
|
|
f"artifact_path should be relative, found: {artifact_path}",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
return "{}:///artifacts?run_id={}&path=/{}".format(
|
|
_MLFLOWDBFS_SCHEME, run_id, posixpath.join(artifact_path, _SPARK_MODEL_PATH_SUB)
|
|
)
|
|
|
|
|
|
def _maybe_save_model(spark_model, model_dir):
|
|
from py4j.protocol import Py4JError
|
|
|
|
try:
|
|
spark_model.save(posixpath.join(model_dir, _SPARK_MODEL_PATH_SUB))
|
|
return True
|
|
except Py4JError:
|
|
return False
|
|
|
|
|
|
class _HadoopFileSystem:
|
|
"""
|
|
Interface to org.apache.hadoop.fs.FileSystem.
|
|
|
|
Spark ML models expect to read from and write to Hadoop FileSystem when running on a cluster.
|
|
Since MLflow works on local directories, we need this interface to copy the files between
|
|
the current DFS and local dir.
|
|
"""
|
|
|
|
def __init__(self):
|
|
raise Exception("This class should not be instantiated")
|
|
|
|
_filesystem = None
|
|
_conf = None
|
|
|
|
@classmethod
|
|
def _jvm(cls):
|
|
from pyspark import SparkContext
|
|
|
|
return SparkContext._gateway.jvm
|
|
|
|
@classmethod
|
|
def _fs(cls):
|
|
if not cls._filesystem:
|
|
cls._filesystem = cls._jvm().org.apache.hadoop.fs.FileSystem.get(cls._conf())
|
|
return cls._filesystem
|
|
|
|
@classmethod
|
|
def _conf(cls):
|
|
from pyspark import SparkContext
|
|
|
|
sc = SparkContext.getOrCreate()
|
|
return sc._jsc.hadoopConfiguration()
|
|
|
|
@classmethod
|
|
def _local_path(cls, path):
|
|
return cls._jvm().org.apache.hadoop.fs.Path(os.path.abspath(path))
|
|
|
|
@classmethod
|
|
def _remote_path(cls, path):
|
|
return cls._jvm().org.apache.hadoop.fs.Path(path)
|
|
|
|
@classmethod
|
|
def _stats(cls):
|
|
return cls._jvm().org.apache.hadoop.fs.FileSystem.getGlobalStorageStatistics()
|
|
|
|
@classmethod
|
|
def copy_to_local_file(cls, src, dst, remove_src):
|
|
cls._fs().copyToLocalFile(remove_src, cls._remote_path(src), cls._local_path(dst))
|
|
|
|
@classmethod
|
|
def copy_from_local_file(cls, src, dst, remove_src):
|
|
cls._fs().copyFromLocalFile(remove_src, cls._local_path(src), cls._remote_path(dst))
|
|
|
|
@classmethod
|
|
def qualified_local_path(cls, path):
|
|
return cls._fs().makeQualified(cls._local_path(path)).toString()
|
|
|
|
@classmethod
|
|
def maybe_copy_from_local_file(cls, src, dst):
|
|
"""
|
|
Conditionally copy the file to the Hadoop DFS.
|
|
The file is copied iff the configuration has distributed filesystem.
|
|
|
|
Returns:
|
|
If copied, return new target location, otherwise return (absolute) source path.
|
|
"""
|
|
local_path = cls._local_path(src)
|
|
qualified_local_path = cls._fs().makeQualified(local_path).toString()
|
|
if qualified_local_path == "file:" + local_path.toString():
|
|
return local_path.toString()
|
|
cls.copy_from_local_file(src, dst, remove_src=False)
|
|
_logger.info("Copied SparkML model to %s", dst)
|
|
return dst
|
|
|
|
@classmethod
|
|
def _try_file_exists(cls, dfs_path):
|
|
try:
|
|
return cls._fs().exists(dfs_path)
|
|
except Exception as ex:
|
|
# Log a debug-level message, since existence checks may raise exceptions
|
|
# in normal operating circumstances that do not warrant warnings
|
|
_logger.debug(
|
|
"Unexpected exception while checking if model uri is visible on DFS: %s", ex
|
|
)
|
|
return False
|
|
|
|
@classmethod
|
|
def maybe_copy_from_uri(cls, src_uri, dst_path, local_model_path=None):
|
|
"""
|
|
Conditionally copy the file to the Hadoop DFS from the source uri.
|
|
In case the file is already on the Hadoop DFS do nothing.
|
|
|
|
Returns:
|
|
If copied, return new target location, otherwise return source uri.
|
|
"""
|
|
try:
|
|
# makeQualified throws if wrong schema / uri
|
|
dfs_path = cls._fs().makeQualified(cls._remote_path(src_uri))
|
|
if cls._try_file_exists(dfs_path):
|
|
_logger.info("File '%s' is already on DFS, copy is not necessary.", src_uri)
|
|
return src_uri
|
|
except Exception:
|
|
_logger.info("URI '%s' does not point to the current DFS.", src_uri)
|
|
_logger.info("File '%s' not found on DFS. Will attempt to upload the file.", src_uri)
|
|
return cls.maybe_copy_from_local_file(
|
|
local_model_path or _download_artifact_from_uri(src_uri), dst_path
|
|
)
|
|
|
|
@classmethod
|
|
def delete(cls, path):
|
|
cls._fs().delete(cls._remote_path(path), True)
|
|
|
|
@classmethod
|
|
def is_filesystem_available(cls, scheme):
|
|
return scheme in [stats.getScheme() for stats in cls._stats().iterator()]
|
|
|
|
|
|
def _should_use_mlflowdbfs(root_uri):
|
|
# The `mlflowdbfs` scheme does not appear in the available schemes returned from
|
|
# the Hadoop FileSystem API until a read call has been issued.
|
|
from mlflow.utils._spark_utils import _get_active_spark_session
|
|
|
|
if (
|
|
databricks_utils.is_in_databricks_serverless_runtime()
|
|
or databricks_utils.is_in_databricks_shared_cluster_runtime()
|
|
or not is_valid_dbfs_uri(root_uri)
|
|
or not is_databricks_acled_artifacts_uri(root_uri)
|
|
or not databricks_utils.is_in_databricks_runtime()
|
|
or (environment_variables._DISABLE_MLFLOWDBFS.get() or "").lower() == "true"
|
|
):
|
|
return False
|
|
|
|
try:
|
|
databricks_utils._get_dbutils()
|
|
except Exception:
|
|
# If dbutils is unavailable, indicate that mlflowdbfs is unavailable
|
|
# because usage of mlflowdbfs depends on dbutils
|
|
return False
|
|
|
|
mlflowdbfs_read_exception_str = None
|
|
try:
|
|
_get_active_spark_session().read.load("mlflowdbfs:///artifact?run_id=foo&path=/bar")
|
|
except Exception as e:
|
|
# The load invocation is expected to throw an exception.
|
|
mlflowdbfs_read_exception_str = str(e)
|
|
|
|
try:
|
|
return _HadoopFileSystem.is_filesystem_available(_MLFLOWDBFS_SCHEME)
|
|
except Exception:
|
|
# The HDFS filesystem logic used to determine mlflowdbfs availability on Databricks
|
|
# clusters may not work on certain Databricks cluster types due to unavailability of
|
|
# the _HadoopFileSystem.is_filesystem_available() API. As a temporary workaround,
|
|
# we check the contents of the expected exception raised by a dummy mlflowdbfs
|
|
# read for evidence that mlflowdbfs is available. If "MlflowdbfsClient" is present
|
|
# in the exception contents, we can safely assume that mlflowdbfs is available because
|
|
# `MlflowdbfsClient` is exclusively used by mlflowdbfs for performing MLflow
|
|
# file storage operations
|
|
#
|
|
# TODO: Remove this logic once the _HadoopFileSystem.is_filesystem_available() check
|
|
# below is determined to work on all Databricks cluster types
|
|
return "MlflowdbfsClient" in (mlflowdbfs_read_exception_str or "")
|
|
|
|
|
|
def _save_model_metadata(
|
|
dst_dir,
|
|
spark_model,
|
|
mlflow_model,
|
|
conda_env,
|
|
code_paths,
|
|
signature=None,
|
|
input_example=None,
|
|
pip_requirements=None,
|
|
extra_pip_requirements=None,
|
|
remote_model_path=None,
|
|
):
|
|
"""
|
|
Saves model metadata into the passed-in directory.
|
|
If mlflowdbfs is not used, the persisted metadata assumes that a model can be
|
|
loaded from a relative path to the metadata file (currently hard-coded to "sparkml").
|
|
If mlflowdbfs is used, remote_model_path should be provided, and the model needs to
|
|
be loaded from the remote_model_path.
|
|
"""
|
|
import pyspark
|
|
|
|
is_spark_connect_model = _is_spark_connect_model(spark_model)
|
|
if signature is not None:
|
|
mlflow_model.signature = signature
|
|
if input_example is not None:
|
|
_save_example(mlflow_model, input_example, dst_dir)
|
|
|
|
code_dir_subpath = _validate_and_copy_code_paths(code_paths, dst_dir)
|
|
mlflow_model.add_flavor(
|
|
FLAVOR_NAME,
|
|
pyspark_version=pyspark.__version__,
|
|
model_data=_SPARK_MODEL_PATH_SUB,
|
|
code=code_dir_subpath,
|
|
model_class=_get_fully_qualified_class_name(spark_model),
|
|
)
|
|
pyfunc.add_to_model(
|
|
mlflow_model,
|
|
loader_module="mlflow.spark",
|
|
data=_SPARK_MODEL_PATH_SUB,
|
|
conda_env=_CONDA_ENV_FILE_NAME,
|
|
python_env=_PYTHON_ENV_FILE_NAME,
|
|
code=code_dir_subpath,
|
|
)
|
|
if size := get_total_file_size(dst_dir):
|
|
mlflow_model.model_size_bytes = size
|
|
mlflow_model.save(os.path.join(dst_dir, MLMODEL_FILE_NAME))
|
|
|
|
if conda_env is None:
|
|
if pip_requirements is None:
|
|
default_reqs = get_default_pip_requirements(is_spark_connect_model)
|
|
if remote_model_path:
|
|
_logger.info(
|
|
"Inferring pip requirements by reloading the logged model from the databricks "
|
|
"artifact repository, which can be time-consuming. To speed up, explicitly "
|
|
"specify the conda_env or pip_requirements when calling log_model()."
|
|
)
|
|
# To ensure `_load_pyfunc` can successfully load the model during the dependency
|
|
# inference, `mlflow_model.save` must be called beforehand to save an MLmodel file.
|
|
inferred_reqs = mlflow.models.infer_pip_requirements(
|
|
remote_model_path or dst_dir,
|
|
FLAVOR_NAME,
|
|
fallback=default_reqs,
|
|
)
|
|
default_reqs = sorted(set(inferred_reqs).union(default_reqs))
|
|
else:
|
|
default_reqs = None
|
|
conda_env, pip_requirements, pip_constraints = _process_pip_requirements(
|
|
default_reqs,
|
|
pip_requirements,
|
|
extra_pip_requirements,
|
|
)
|
|
else:
|
|
conda_env, pip_requirements, pip_constraints = _process_conda_env(conda_env)
|
|
|
|
with open(os.path.join(dst_dir, _CONDA_ENV_FILE_NAME), "w") as f:
|
|
yaml.safe_dump(conda_env, stream=f, default_flow_style=False)
|
|
|
|
# Save `constraints.txt` if necessary
|
|
if pip_constraints:
|
|
write_to(os.path.join(dst_dir, _CONSTRAINTS_FILE_NAME), "\n".join(pip_constraints))
|
|
|
|
# Save `requirements.txt`
|
|
write_to(os.path.join(dst_dir, _REQUIREMENTS_FILE_NAME), "\n".join(pip_requirements))
|
|
|
|
_PythonEnv.current().to_yaml(os.path.join(dst_dir, _PYTHON_ENV_FILE_NAME))
|
|
|
|
|
|
def _validate_model(spark_model):
|
|
from pyspark.ml import Model as PySparkModel
|
|
from pyspark.ml import Transformer as PySparkTransformer
|
|
from pyspark.ml.util import MLReadable, MLWritable
|
|
|
|
if _is_spark_connect_model(spark_model):
|
|
return
|
|
|
|
if (
|
|
(
|
|
not isinstance(spark_model, PySparkModel)
|
|
and not isinstance(spark_model, PySparkTransformer)
|
|
)
|
|
or not isinstance(spark_model, MLReadable)
|
|
or not isinstance(spark_model, MLWritable)
|
|
):
|
|
raise MlflowException(
|
|
"Cannot serialize this model. MLflow can only save descendants of pyspark.ml.Model "
|
|
"or pyspark.ml.Transformer that implement MLWritable and MLReadable.",
|
|
INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
|
|
def _is_spark_connect_model(spark_model):
|
|
"""
|
|
Return whether the spark model is spark connect ML model
|
|
"""
|
|
try:
|
|
from pyspark.ml.connect import Model as ConnectModel
|
|
|
|
return isinstance(spark_model, ConnectModel)
|
|
except ImportError:
|
|
# pyspark < 3.5 does not support Spark connect ML model
|
|
return False
|
|
|
|
|
|
def _is_uc_volume_uri(url):
|
|
parsed_url = urlparse(url)
|
|
return parsed_url.scheme in ["", "dbfs"] and parsed_url.path.startswith("/Volumes")
|
|
|
|
|
|
def _check_databricks_uc_volume_tmpdir_availability(dfs_tmpdir):
|
|
if (
|
|
databricks_utils.is_in_databricks_serverless_runtime()
|
|
or databricks_utils.is_in_databricks_shared_cluster_runtime()
|
|
):
|
|
if not dfs_tmpdir or not _is_uc_volume_uri(dfs_tmpdir):
|
|
raise MlflowException(
|
|
"UC volume path must be provided to save, log or load SparkML models "
|
|
"in Databricks shared or serverless clusters. "
|
|
"Specify environment variable 'MLFLOW_DFS_TMP' "
|
|
"or 'dfs_tmpdir' argument that uses a UC volume path starting with '/Volumes/...' "
|
|
"when saving, logging or loading a model."
|
|
)
|
|
|
|
|
|
@format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name="pyspark"))
|
|
def save_model(
|
|
spark_model,
|
|
path,
|
|
mlflow_model=None,
|
|
conda_env=None,
|
|
code_paths=None,
|
|
dfs_tmpdir=None,
|
|
signature: ModelSignature = None,
|
|
input_example: ModelInputExample = None,
|
|
pip_requirements=None,
|
|
extra_pip_requirements=None,
|
|
metadata=None,
|
|
**kwargs,
|
|
):
|
|
"""
|
|
Save a Spark MLlib Model to a local path.
|
|
|
|
By default, this function saves models using the Spark MLlib persistence mechanism.
|
|
|
|
Args:
|
|
spark_model: Spark model to be saved - MLflow can only save descendants of
|
|
pyspark.ml.Model or pyspark.ml.Transformer which implement
|
|
MLReadable and MLWritable.
|
|
path: Local path where the model is to be saved.
|
|
mlflow_model: MLflow model config this flavor is being added to.
|
|
conda_env: {{ conda_env }}
|
|
code_paths: {{ code_paths }}
|
|
dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
|
|
filesystem if running in local mode. The model is be written in this
|
|
destination and then copied to the requested local path. This is necessary
|
|
as Spark ML models read from and write to DFS if running on a cluster. All
|
|
temporary files created on the DFS are removed if this operation
|
|
completes successfully. Defaults to ``/tmp/mlflow``.
|
|
signature: See the document of argument ``signature`` in :py:func:`mlflow.spark.log_model`.
|
|
input_example: {{ input_example }}
|
|
pip_requirements: {{ pip_requirements }}
|
|
extra_pip_requirements: {{ extra_pip_requirements }}
|
|
metadata: {{ metadata }}
|
|
kwargs: {{ kwargs }}
|
|
|
|
.. code-block:: python
|
|
:caption: Example
|
|
|
|
from mlflow import spark
|
|
from pyspark.ml.pipeline import PipelineModel
|
|
|
|
# your pyspark.ml.pipeline.PipelineModel type
|
|
model = ...
|
|
mlflow.spark.save_model(model, "spark-model")
|
|
"""
|
|
_validate_model(spark_model)
|
|
_validate_env_arguments(conda_env, pip_requirements, extra_pip_requirements)
|
|
|
|
from pyspark.ml import PipelineModel
|
|
|
|
from mlflow.utils._spark_utils import _get_active_spark_session
|
|
|
|
is_spark_connect_model = _is_spark_connect_model(spark_model)
|
|
|
|
if not is_spark_connect_model and not isinstance(spark_model, PipelineModel):
|
|
spark_model = PipelineModel([spark_model])
|
|
if mlflow_model is None:
|
|
mlflow_model = Model()
|
|
if metadata is not None:
|
|
mlflow_model.metadata = metadata
|
|
|
|
# for automatic signature inference, we use an inline implementation rather than the
|
|
# `_infer_signature_from_input_example` API because we need to convert model predictions from a
|
|
# list into a Pandas series for signature inference.
|
|
if signature is None and input_example is not None:
|
|
input_ex = _Example(input_example).inference_data
|
|
try:
|
|
spark = _get_active_spark_session()
|
|
if spark is not None:
|
|
input_example_spark_df = spark.createDataFrame(input_ex)
|
|
# `_infer_spark_model_signature` mutates the model. Copy the model to preserve the
|
|
# original model.
|
|
try:
|
|
spark_model = spark_model.copy()
|
|
except Exception:
|
|
_logger.debug(
|
|
"Failed to copy the model, using the original model.", exc_info=True
|
|
)
|
|
signature = mlflow.pyspark.ml._infer_spark_model_signature(
|
|
spark_model, input_example_spark_df
|
|
)
|
|
except Exception as e:
|
|
if environment_variables._MLFLOW_TESTING.get():
|
|
raise
|
|
_logger.warning(_LOG_MODEL_INFER_SIGNATURE_WARNING_TEMPLATE, repr(e))
|
|
_logger.debug("", exc_info=True)
|
|
elif signature is False:
|
|
signature = None
|
|
|
|
sparkml_data_path = os.path.abspath(os.path.join(path, _SPARK_MODEL_PATH_SUB))
|
|
|
|
if is_spark_connect_model:
|
|
spark_model.saveToLocal(sparkml_data_path)
|
|
else:
|
|
# Spark ML stores the model on DFS if running on a cluster
|
|
# Save it to a DFS temp dir first and copy it to local path
|
|
if dfs_tmpdir is None:
|
|
dfs_tmpdir = MLFLOW_DFS_TMP.get()
|
|
|
|
_check_databricks_uc_volume_tmpdir_availability(dfs_tmpdir)
|
|
tmp_path = generate_tmp_dfs_path(dfs_tmpdir)
|
|
spark_model.save(tmp_path)
|
|
|
|
if databricks_utils.is_in_databricks_runtime() and _is_uc_volume_uri(tmp_path):
|
|
# The temp DFS path is a UC volume path.
|
|
# Use UC volume fuse mount to read data.
|
|
tmp_path_fuse = urlparse(tmp_path).path
|
|
shutil.move(src=tmp_path_fuse, dst=sparkml_data_path)
|
|
else:
|
|
# We're copying the Spark model from DBFS to the local filesystem if (a) the temporary
|
|
# DFS URI we saved the Spark model to is a DBFS URI ("dbfs:/my-directory"), or (b) if
|
|
# we're running on a Databricks cluster and the URI is schemeless (e.g. looks like a
|
|
# filesystem absolute path like "/my-directory")
|
|
copying_from_dbfs = is_valid_dbfs_uri(tmp_path) or (
|
|
databricks_utils.is_in_cluster() and posixpath.abspath(tmp_path) == tmp_path
|
|
)
|
|
if copying_from_dbfs and databricks_utils.is_dbfs_fuse_available():
|
|
tmp_path_fuse = dbfs_hdfs_uri_to_fuse_path(tmp_path)
|
|
shutil.move(src=tmp_path_fuse, dst=sparkml_data_path)
|
|
else:
|
|
_HadoopFileSystem.copy_to_local_file(tmp_path, sparkml_data_path, remove_src=True)
|
|
|
|
_save_model_metadata(
|
|
dst_dir=path,
|
|
spark_model=spark_model,
|
|
mlflow_model=mlflow_model,
|
|
conda_env=conda_env,
|
|
code_paths=code_paths,
|
|
signature=signature,
|
|
input_example=input_example,
|
|
pip_requirements=pip_requirements,
|
|
extra_pip_requirements=extra_pip_requirements,
|
|
)
|
|
|
|
|
|
def _load_model_databricks_dbfs(dfs_tmpdir, local_model_path):
|
|
from pyspark.ml.pipeline import PipelineModel
|
|
|
|
# Spark ML expects the model to be stored on DFS
|
|
# Copy the model to a temp DFS location first. We cannot delete this file, as
|
|
# Spark may read from it at any point.
|
|
fuse_dfs_tmpdir = dbfs_hdfs_uri_to_fuse_path(dfs_tmpdir)
|
|
os.makedirs(fuse_dfs_tmpdir)
|
|
# Workaround for inability to use shutil.copytree with DBFS FUSE due to permission-denied
|
|
# errors on passthrough-enabled clusters when attempting to copy permission bits for directories
|
|
shutil_copytree_without_file_permissions(src_dir=local_model_path, dst_dir=fuse_dfs_tmpdir)
|
|
return PipelineModel.load(dfs_tmpdir)
|
|
|
|
|
|
def _load_model_databricks_uc_volume(dfs_tmpdir, local_model_path):
|
|
from pyspark.ml.pipeline import PipelineModel
|
|
|
|
# Copy the model to a temp DFS location first. We cannot delete this file, as
|
|
# Spark may read from it at any point.
|
|
fuse_dfs_tmpdir = urlparse(dfs_tmpdir).path
|
|
shutil.copytree(src=local_model_path, dst=fuse_dfs_tmpdir)
|
|
return PipelineModel.load(dfs_tmpdir)
|
|
|
|
|
|
def _load_model(model_uri, dfs_tmpdir_base=None, local_model_path=None):
|
|
from pyspark.ml.pipeline import PipelineModel
|
|
|
|
dfs_tmpdir = generate_tmp_dfs_path(dfs_tmpdir_base or MLFLOW_DFS_TMP.get())
|
|
|
|
_check_databricks_uc_volume_tmpdir_availability(dfs_tmpdir)
|
|
if (
|
|
databricks_utils.is_in_databricks_serverless_runtime()
|
|
or databricks_utils.is_in_databricks_shared_cluster_runtime()
|
|
):
|
|
return _load_model_databricks_uc_volume(
|
|
dfs_tmpdir, local_model_path or _download_artifact_from_uri(model_uri)
|
|
)
|
|
if databricks_utils.is_in_cluster() and databricks_utils.is_dbfs_fuse_available():
|
|
return _load_model_databricks_dbfs(
|
|
dfs_tmpdir, local_model_path or _download_artifact_from_uri(model_uri)
|
|
)
|
|
model_uri = _HadoopFileSystem.maybe_copy_from_uri(model_uri, dfs_tmpdir, local_model_path)
|
|
return PipelineModel.load(model_uri)
|
|
|
|
|
|
def _load_spark_connect_model(model_class, local_path):
|
|
return _get_class_from_string(model_class).loadFromLocal(local_path)
|
|
|
|
|
|
def load_model(model_uri, dfs_tmpdir=None, dst_path=None):
|
|
"""
|
|
Load the Spark MLlib model from the path.
|
|
|
|
Args:
|
|
model_uri: The location, in URI format, of the MLflow model, for example:
|
|
|
|
- ``/Users/me/path/to/local/model``
|
|
- ``relative/path/to/local/model``
|
|
- ``s3://my_bucket/path/to/model``
|
|
- ``runs:/<mlflow_run_id>/run-relative/path/to/model``
|
|
- ``models:/<model_name>/<model_version>``
|
|
- ``models:/<model_name>/<stage>``
|
|
|
|
For more information about supported URI schemes, see
|
|
`Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
|
|
artifact-locations>`_.
|
|
dfs_tmpdir: Temporary directory path on Distributed (Hadoop) File System (DFS) or local
|
|
filesystem if running in local mode. The model is loaded from this
|
|
destination. Defaults to ``/tmp/mlflow``.
|
|
dst_path: The local filesystem path to which to download the model artifact.
|
|
This directory must already exist. If unspecified, a local output
|
|
path will be created.
|
|
|
|
Returns:
|
|
pyspark.ml.pipeline.PipelineModel
|
|
|
|
.. code-block:: python
|
|
:caption: Example
|
|
|
|
import mlflow
|
|
|
|
model = mlflow.spark.load_model("spark-model")
|
|
# Prepare test documents, which are unlabeled (id, text) tuples.
|
|
test = spark.createDataFrame(
|
|
[(4, "spark i j k"), (5, "l m n"), (6, "spark hadoop spark"), (7, "apache hadoop")],
|
|
["id", "text"],
|
|
)
|
|
# Make predictions on test documents
|
|
prediction = model.transform(test)
|
|
"""
|
|
# This MUST be called prior to appending the model flavor to `model_uri` in order
|
|
# for `artifact_path` to take on the correct value for model loading via mlflowdbfs.
|
|
root_uri, artifact_path = _get_root_uri_and_artifact_path(model_uri)
|
|
|
|
local_mlflow_model_path = _download_artifact_from_uri(
|
|
artifact_uri=model_uri, output_path=dst_path
|
|
)
|
|
flavor_conf = Model.load(local_mlflow_model_path).flavors[FLAVOR_NAME]
|
|
_add_code_from_conf_to_system_path(local_mlflow_model_path, flavor_conf)
|
|
|
|
model_class = flavor_conf.get("model_class")
|
|
if model_class is not None and model_class.startswith("pyspark.ml.connect."):
|
|
spark_model_local_path = os.path.join(local_mlflow_model_path, flavor_conf["model_data"])
|
|
return _load_spark_connect_model(model_class, spark_model_local_path)
|
|
|
|
if _should_use_mlflowdbfs(model_uri) and (
|
|
run_id := DatabricksArtifactRepository._extract_run_id(model_uri)
|
|
):
|
|
from pyspark.ml.pipeline import PipelineModel
|
|
|
|
mlflowdbfs_path = _mlflowdbfs_path(run_id, artifact_path)
|
|
with databricks_utils.MlflowCredentialContext(
|
|
get_databricks_profile_uri_from_artifact_uri(root_uri)
|
|
):
|
|
return PipelineModel.load(mlflowdbfs_path)
|
|
|
|
sparkml_model_uri = append_to_uri_path(model_uri, flavor_conf["model_data"])
|
|
local_sparkml_model_path = os.path.join(local_mlflow_model_path, flavor_conf["model_data"])
|
|
return _load_model(
|
|
model_uri=sparkml_model_uri,
|
|
dfs_tmpdir_base=dfs_tmpdir,
|
|
local_model_path=local_sparkml_model_path,
|
|
)
|
|
|
|
|
|
def _load_pyfunc(path):
|
|
"""
|
|
Load PyFunc implementation. Called by ``pyfunc.load_model``.
|
|
|
|
Args:
|
|
path: Local filesystem path to the MLflow Model with the ``spark`` flavor.
|
|
"""
|
|
from mlflow.utils._spark_utils import (
|
|
_create_local_spark_session_for_loading_spark_model,
|
|
_get_active_spark_session,
|
|
)
|
|
|
|
model_meta_path = os.path.join(os.path.dirname(path), MLMODEL_FILE_NAME)
|
|
model_meta = Model.load(model_meta_path)
|
|
|
|
model_class = model_meta.flavors[FLAVOR_NAME].get("model_class")
|
|
if model_class is not None and model_class.startswith("pyspark.ml.connect."):
|
|
# Note:
|
|
# Spark connect ML models don't require a spark session for running inference.
|
|
spark = None
|
|
spark_model = _load_spark_connect_model(model_class, path)
|
|
|
|
else:
|
|
# NOTE: The `_create_local_spark_session_for_loading_spark_model()` call below may change
|
|
# settings of the active session which we do not intend to do here.
|
|
# In particular, setting master to local[1] can break distributed clusters.
|
|
# To avoid this problem, we explicitly check for an active session. This is not ideal but
|
|
# there is no good workaround at the moment.
|
|
spark = _get_active_spark_session()
|
|
if spark is None:
|
|
# NB: If there is no existing Spark context, create a new local one.
|
|
# NB: We're disabling caching on the new context since we do not need it and we want to
|
|
# avoid overwriting cache of underlying Spark cluster when executed on a Spark Worker
|
|
# (e.g. as part of spark_udf).
|
|
spark = _create_local_spark_session_for_loading_spark_model()
|
|
|
|
spark_model = _load_model(model_uri=path)
|
|
|
|
return _PyFuncModelWrapper(spark, spark_model, signature=model_meta.signature)
|
|
|
|
|
|
def _find_and_set_features_col_as_vector_if_needed(spark_df, spark_model):
|
|
"""
|
|
Finds the `featuresCol` column in spark_model and
|
|
then tries to cast that column to `vector` type.
|
|
This method is noop if the `featuresCol` is already of type `vector`
|
|
or if it can't be cast to `vector` type
|
|
Note:
|
|
If a spark ML pipeline contains a single Estimator stage, it requires
|
|
the input dataframe to contain features column of vector type.
|
|
But the autologging for pyspark ML casts vector column to array<double> type
|
|
for parity with the pd Dataframe. The following fix is required, which transforms
|
|
that features column back to vector type so that the pipeline stages can correctly work.
|
|
A valid scenario is if the auto-logged input example is directly used
|
|
for prediction, which would otherwise fail without this transformation.
|
|
|
|
Args:
|
|
spark_df: Input dataframe that contains `featuresCol`
|
|
spark_model: A pipeline model or a single transformer that contains `featuresCol` param
|
|
|
|
Returns:
|
|
A spark dataframe that contains features column of `vector` type.
|
|
"""
|
|
from pyspark.ml.linalg import Vectors, VectorUDT
|
|
from pyspark.sql import types as t
|
|
from pyspark.sql.functions import udf
|
|
|
|
def _find_stage_with_features_col(stage):
|
|
if stage.hasParam("featuresCol"):
|
|
|
|
def _array_to_vector(input_array):
|
|
return Vectors.dense(input_array)
|
|
|
|
array_to_vector_udf = udf(f=_array_to_vector, returnType=VectorUDT())
|
|
features_col_name = stage.extractParamMap().get(stage.featuresCol)
|
|
features_col_type = [
|
|
_field
|
|
for _field in spark_df.schema.fields
|
|
if _field.name == features_col_name
|
|
and _field.dataType
|
|
in [t.ArrayType(t.DoubleType(), True), t.ArrayType(t.DoubleType(), False)]
|
|
]
|
|
if len(features_col_type) == 1:
|
|
return spark_df.withColumn(
|
|
features_col_name, array_to_vector_udf(features_col_name)
|
|
)
|
|
return spark_df
|
|
|
|
if hasattr(spark_model, "stages"):
|
|
for stage in reversed(spark_model.stages):
|
|
return _find_stage_with_features_col(stage)
|
|
return _find_stage_with_features_col(spark_model)
|
|
|
|
|
|
class _PyFuncModelWrapper:
|
|
"""
|
|
Wrapper around Spark MLlib PipelineModel providing interface for scoring pandas DataFrame.
|
|
"""
|
|
|
|
def __init__(self, spark, spark_model, signature):
|
|
self.spark = spark
|
|
self.spark_model = spark_model
|
|
self.signature = signature
|
|
|
|
def get_raw_model(self):
|
|
"""
|
|
Returns the underlying model.
|
|
"""
|
|
return self.spark_model
|
|
|
|
def predict(
|
|
self,
|
|
pandas_df,
|
|
params: dict[str, Any] | None = None,
|
|
):
|
|
"""
|
|
Generate predictions given input data in a pandas DataFrame.
|
|
|
|
Args:
|
|
pandas_df: pandas DataFrame containing input data.
|
|
params: Additional parameters to pass to the model for inference.
|
|
|
|
Returns:
|
|
List with model predictions.
|
|
"""
|
|
if _is_spark_connect_model(self.spark_model):
|
|
# Spark connect ML model directly appends prediction result column to input pandas
|
|
# dataframe. To make input dataframe intact, make a copy first.
|
|
# TODO: apache/spark master has made a change to do shallow copy before
|
|
# calling `spark_model.transform`, so once spark 4.0 releases, we can
|
|
# remove this line.
|
|
pandas_df = pandas_df.copy(deep=False)
|
|
# Assuming the model output column name is "prediction".
|
|
# Spark model uses "prediction" as default model inference output column name.
|
|
return self.spark_model.transform(pandas_df)["prediction"]
|
|
|
|
# Convert List[np.float64] / np.array[np.float64] type to List[float] type,
|
|
# otherwise it will break `spark.createDataFrame` column type inferring.
|
|
if self.signature and self.signature.inputs:
|
|
for col_spec in self.signature.inputs.inputs:
|
|
if isinstance(col_spec.type, SparkMLVector):
|
|
col_name = col_spec.name or pandas_df.columns[0]
|
|
|
|
pandas_df[col_name] = pandas_df[col_name].map(
|
|
lambda array: [float(elem) for elem in array]
|
|
)
|
|
|
|
spark_df = self.spark.createDataFrame(pandas_df)
|
|
|
|
# Convert Array[Double] column to spark ML vector type according to signature
|
|
if self.signature and self.signature.inputs:
|
|
for col_spec in self.signature.inputs.inputs:
|
|
if isinstance(col_spec.type, SparkMLVector):
|
|
from pyspark.ml.functions import array_to_vector
|
|
|
|
col_name = col_spec.name or spark_df.columns[0]
|
|
spark_df = spark_df.withColumn(col_name, array_to_vector(col_name))
|
|
|
|
# For the case of no signature or signature logged by old version MLflow,
|
|
# the signature does not support spark ML vector type, in this case,
|
|
# automatically infer vector type input columns and do the conversion
|
|
# using `_find_and_set_features_col_as_vector_if_needed` utility function.
|
|
spark_df = _find_and_set_features_col_as_vector_if_needed(spark_df, self.spark_model)
|
|
|
|
prediction_column = mlflow.pyspark.ml._check_or_set_model_prediction_column(
|
|
self.spark_model, spark_df
|
|
)
|
|
prediction_df = self.spark_model.transform(spark_df).select(prediction_column)
|
|
|
|
# If signature output schema exists and it contains vector type columns,
|
|
# Convert spark ML vector type column to Array[Double] otherwise it will
|
|
# break enforce_schema checking
|
|
if self.signature and self.signature.outputs:
|
|
for col_spec in self.signature.outputs.inputs:
|
|
if isinstance(col_spec.type, SparkMLVector):
|
|
from pyspark.ml.functions import vector_to_array
|
|
|
|
col_name = col_spec.name or prediction_df.columns[0]
|
|
prediction_df = prediction_df.withColumn(col_name, vector_to_array(col_name))
|
|
return [x.prediction for x in prediction_df.collect()]
|
|
|
|
|
|
@autologging_integration(FLAVOR_NAME)
|
|
def autolog(disable=False, silent=False):
|
|
"""
|
|
Enables (or disables) and configures logging of Spark datasource paths, versions
|
|
(if applicable), and formats when they are read. This method is not threadsafe and assumes a
|
|
`SparkSession
|
|
<https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.SparkSession.html>`_
|
|
already exists with the
|
|
`mlflow-spark JAR
|
|
<https://www.mlflow.org/docs/latest/tracking.html#spark>`_
|
|
attached. It should be called on the Spark driver, not on the executors (i.e. do not call
|
|
this method within a function parallelized by Spark).
|
|
The mlflow-spark JAR used must match the Scala version of Spark. Please see the
|
|
`Maven Repository
|
|
<https://mvnrepository.com/artifact/org.mlflow/mlflow-spark>`_
|
|
for available versions. This API requires Spark 3.0 or above.
|
|
|
|
Datasource information is cached in memory and logged to all subsequent MLflow runs,
|
|
including the active MLflow run (if one exists when the data is read). Note that autologging of
|
|
Spark ML (MLlib) models is not currently supported via this API. Datasource autologging is
|
|
best-effort, meaning that if Spark is under heavy load or MLflow logging fails for any reason
|
|
(e.g., if the MLflow server is unavailable), logging may be dropped.
|
|
|
|
For any unexpected issues with autologging, check Spark driver and executor logs in addition
|
|
to stderr & stdout generated from your MLflow code - datasource information is pulled from
|
|
Spark, so logs relevant to debugging may show up amongst the Spark logs.
|
|
|
|
.. Note:: Spark datasource autologging only supports logging to MLflow runs in a single thread
|
|
|
|
.. code-block:: python
|
|
:caption: Example
|
|
|
|
import mlflow.spark
|
|
import os
|
|
import shutil
|
|
from pyspark.sql import SparkSession
|
|
|
|
# Create and persist some dummy data
|
|
# Note: the 2.12 in 'org.mlflow:mlflow-spark_2.12:2.16.2' below indicates the Scala
|
|
# version, please match this with that of Spark. The 2.16.2 indicates the mlflow version.
|
|
# Note: On environments like Databricks with pre-created SparkSessions,
|
|
# ensure the org.mlflow:mlflow-spark_2.12:2.16.2 is attached as a library to
|
|
# your cluster
|
|
spark = (
|
|
SparkSession.builder
|
|
.config(
|
|
"spark.jars.packages",
|
|
"org.mlflow:mlflow-spark_2.12:2.16.2",
|
|
)
|
|
.master("local[*]")
|
|
.getOrCreate()
|
|
)
|
|
df = spark.createDataFrame(
|
|
[(4, "spark i j k"), (5, "l m n"), (6, "spark hadoop spark"), (7, "apache hadoop")],
|
|
["id", "text"],
|
|
)
|
|
import tempfile
|
|
|
|
tempdir = tempfile.mkdtemp()
|
|
df.write.csv(os.path.join(tempdir, "my-data-path"), header=True)
|
|
# Enable Spark datasource autologging.
|
|
mlflow.spark.autolog()
|
|
loaded_df = spark.read.csv(
|
|
os.path.join(tempdir, "my-data-path"), header=True, inferSchema=True
|
|
)
|
|
# Call toPandas() to trigger a read of the Spark datasource. Datasource info
|
|
# (path and format) is logged to the current active run, or the
|
|
# next-created MLflow run if no run is currently active
|
|
with mlflow.start_run() as active_run:
|
|
pandas_df = loaded_df.toPandas()
|
|
|
|
Args:
|
|
disable: If ``True``, disables the Spark datasource autologging integration.
|
|
If ``False``, enables the Spark datasource autologging integration.
|
|
silent: If ``True``, suppress all event logs and warnings from MLflow during Spark
|
|
datasource autologging. If ``False``, show all events and warnings during Spark
|
|
datasource autologging.
|
|
"""
|
|
from pyspark import __version__ as pyspark_version
|
|
from pyspark.sql import SparkSession
|
|
|
|
from mlflow.spark.autologging import (
|
|
_listen_for_spark_activity,
|
|
_stop_listen_for_spark_activity,
|
|
)
|
|
from mlflow.utils import databricks_utils
|
|
from mlflow.utils._spark_utils import _get_active_spark_session
|
|
|
|
if (
|
|
databricks_utils.is_in_databricks_serverless_runtime()
|
|
or databricks_utils.is_in_databricks_shared_cluster_runtime()
|
|
):
|
|
if disable:
|
|
return
|
|
raise MlflowException(
|
|
"MLflow Spark dataset autologging is not supported on Databricks shared clusters "
|
|
"or Databricks serverless clusters."
|
|
)
|
|
|
|
# Check if environment variable PYSPARK_PIN_THREAD is set to false.
|
|
# The "Pin thread" concept was introduced since Pyspark 3.0.0 and set to default to true
|
|
# since Pyspark 3.2.0 (https://issues.apache.org/jira/browse/SPARK-35303). When pin thread
|
|
# is enabled, Pyspark manages Python and JVM threads in a 1:1, meaning that when one thread
|
|
# is terminated, the corresponding thread in the other side will be terminated as well.
|
|
# However, this causes an issue in spark autologging as our event listener thread may be
|
|
# terminated before receiving the datasource event.
|
|
# Hence, we have to disable it and decouple the thread management between Python and JVM.
|
|
if (
|
|
Version(pyspark_version) >= Version("3.2.0")
|
|
and os.environ.get("PYSPARK_PIN_THREAD", "").lower() != "false"
|
|
):
|
|
_logger.warning(
|
|
"With Pyspark >= 3.2, PYSPARK_PIN_THREAD environment variable must be set to false "
|
|
"for Spark datasource autologging to work."
|
|
)
|
|
|
|
def __init__(original, self, *args, **kwargs):
|
|
original(self, *args, **kwargs)
|
|
|
|
_listen_for_spark_activity(self._sc)
|
|
|
|
safe_patch(FLAVOR_NAME, SparkSession, "__init__", __init__, manage_run=False)
|
|
|
|
def patched_session_stop(original, self, *args, **kwargs):
|
|
_stop_listen_for_spark_activity(self.sparkContext)
|
|
original(self, *args, **kwargs)
|
|
|
|
safe_patch(FLAVOR_NAME, SparkSession, "stop", patched_session_stop, manage_run=False)
|
|
|
|
active_session = _get_active_spark_session()
|
|
if active_session is not None:
|
|
# We know SparkContext exists here already, so get it
|
|
sc = active_session.sparkContext
|
|
|
|
if disable:
|
|
_stop_listen_for_spark_activity(sc)
|
|
else:
|
|
_listen_for_spark_activity(sc)
|