def test_use_repl_context_if_available(tmpdir):
# Simulate a case where `dbruntime.databricks_repl_context.get_context` is unavailable.
with pytest.raises(ModuleNotFoundError,
match="No module named 'dbruntime'"):
from dbruntime.databricks_repl_context import get_context # pylint: disable=unused-import
command_context_mock = mock.MagicMock()
command_context_mock.jobId().get.return_value = "job_id"
with mock.patch(
"mlflow.utils.databricks_utils._get_command_context",
return_value=command_context_mock) as mock_get_command_context:
assert databricks_utils.get_job_id() == "job_id"
mock_get_command_context.assert_called_once()
# Create a fake databricks_repl_context module
tmpdir.mkdir("dbruntime").join("databricks_repl_context.py").write("""
def get_context():
pass
""")
sys.path.append(tmpdir.strpath)
# Simulate a case where the REPL context object is not initialized.
with mock.patch(
"dbruntime.databricks_repl_context.get_context",
return_value=None,
) as mock_get_context, mock.patch(
"mlflow.utils.databricks_utils._get_command_context",
return_value=command_context_mock) as mock_get_command_context:
assert databricks_utils.get_job_id() == "job_id"
mock_get_command_context.assert_called_once()
with mock.patch(
"dbruntime.databricks_repl_context.get_context",
return_value=mock.MagicMock(jobId="job_id"),
) as mock_get_context, mock.patch(
"mlflow.utils.databricks_utils._get_dbutils") as mock_dbutils:
assert databricks_utils.get_job_id() == "job_id"
mock_get_context.assert_called_once()
mock_dbutils.assert_not_called()
with mock.patch(
"dbruntime.databricks_repl_context.get_context",
return_value=mock.MagicMock(notebookId="notebook_id"),
) as mock_get_context, mock.patch(
"mlflow.utils.databricks_utils._get_property_from_spark_context"
) as mock_spark_context:
assert databricks_utils.get_notebook_id() == "notebook_id"
mock_get_context.assert_called_once()
mock_spark_context.assert_not_called()
with mock.patch(
"dbruntime.databricks_repl_context.get_context",
return_value=mock.MagicMock(isInCluster=True),
) as mock_get_context, mock.patch(
"mlflow.utils._spark_utils._get_active_spark_session"
) as mock_spark_session:
assert databricks_utils.is_in_cluster()
mock_get_context.assert_called_once()
mock_spark_session.assert_not_called()
def _load_model(model_uri, dfs_tmpdir_base=None):
from pyspark.ml.pipeline import PipelineModel
if dfs_tmpdir_base is None:
dfs_tmpdir_base = DFS_TMP
dfs_tmpdir = _tmp_path(dfs_tmpdir_base)
if databricks_utils.is_in_cluster() and databricks_utils.is_dbfs_fuse_available():
return _load_model_databricks(model_uri, dfs_tmpdir)
model_uri = _HadoopFileSystem.maybe_copy_from_uri(model_uri, dfs_tmpdir)
return PipelineModel.load(model_uri)
def request_headers(self):
request_headers = {}
if databricks_utils.is_in_databricks_notebook():
request_headers["notebook_id"] = databricks_utils.get_notebook_id()
if databricks_utils.is_in_databricks_job():
request_headers["job_id"] = databricks_utils.get_job_id()
request_headers["job_run_id"] = databricks_utils.get_job_run_id()
request_headers["job_type"] = databricks_utils.get_job_type()
if databricks_utils.is_in_cluster():
request_headers["cluster_id"] = databricks_utils.get_cluster_id()
return request_headers
def save_model(
spark_model,
path,
mlflow_model=None,
conda_env=None,
dfs_tmpdir=None,
sample_input=None,
signature: ModelSignature = None,
input_example: ModelInputExample = None,
pip_requirements=None,
extra_pip_requirements=None,
):
"""
Save a Spark MLlib Model to a local path.
By default, this function saves models using the Spark MLlib persistence mechanism.
Additionally, if a sample input is specified using the ``sample_input`` parameter, the model
is also serialized in MLeap format and the MLeap flavor is added.
:param spark_model: Spark model to be saved - MLflow can only save descendants of
pyspark.ml.Model which implement MLReadable and MLWritable.
:param path: Local path where the model is to be saved.
:param mlflow_model: MLflow model config this flavor is being added to.
:param conda_env: Either a dictionary representation of a Conda environment or the path to a
Conda environment yaml file. If provided, this decsribes the environment
this model should be run in. At minimum, it should specify the dependencies
contained in :func:`get_default_conda_env()`. If `None`, the default
:func:`get_default_conda_env()` environment is added to the model.
The following is an *example* dictionary representation of a Conda
environment::
{
'name': 'mlflow-env',
'channels': ['defaults'],
'dependencies': [
'python=3.7.0',
'pyspark=2.3.0'
]
}
:param 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``.
:param sample_input: A sample input that is used to add the MLeap flavor to the model.
This must be a PySpark DataFrame that the model can evaluate. If
``sample_input`` is ``None``, the MLeap flavor is not added.
:param signature: :py:class:`ModelSignature <mlflow.models.ModelSignature>`
describes model input and output :py:class:`Schema <mlflow.types.Schema>`.
The model signature can be :py:func:`inferred <mlflow.models.infer_signature>`
from datasets with valid model input (e.g. the training dataset with target
column omitted) and valid model output (e.g. model predictions generated on
the training dataset), for example:
.. code-block:: python
from mlflow.models.signature import infer_signature
train = df.drop_column("target_label")
predictions = ... # compute model predictions
signature = infer_signature(train, predictions)
:param input_example: Input example provides one or several instances of valid
model input. The example can be used as a hint of what data to feed the
model. The given example will be converted to a Pandas DataFrame and then
serialized to json using the Pandas split-oriented format. Bytes are
base64-encoded.
:param pip_requirements: {{ pip_requirements }}
:param extra_pip_requirements: {{ extra_pip_requirements }}
.. code-block:: python
:caption: Example
from mlflow import spark
from pyspark.ml.pipeline.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
if not isinstance(spark_model, PipelineModel):
spark_model = PipelineModel([spark_model])
if mlflow_model is None:
mlflow_model = Model()
# 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 = DFS_TMP
tmp_path = _tmp_path(dfs_tmpdir)
spark_model.save(tmp_path)
sparkml_data_path = os.path.abspath(
os.path.join(path, _SPARK_MODEL_PATH_SUB))
# 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,
sample_input=sample_input,
conda_env=conda_env,
signature=signature,
input_example=input_example,
pip_requirements=pip_requirements,
extra_pip_requirements=extra_pip_requirements,
)
def in_context(self):
return (databricks_utils.is_in_cluster()
or databricks_utils.is_in_databricks_notebook()
or databricks_utils.is_in_databricks_job())
def in_context(self):
return databricks_utils.is_in_cluster()
