Repository: mlflow/mlflow

Issue:
[FR] Support argument for `local_tmpdir` when loading spark model.
### Willingness to contribute

Yes. I would be willing to contribute this feature with guidance from the MLflow community.

### Proposal Summary

In the current implementation of `load_model` in `mlflow/spark.py`, `_download_artifact_from_uri` will be called to download the artifacts to a local path. The thing is that `load_model` doesn't support passing a value to the argument `output_path` in `_download_artifact_from_uri`. 

This `output_path` supports specifying a path for storing artifacts in local system, it will create a local path if the argument is unspecified. Sometimes, the auto created local path may not work well. In that case, a user specified path is needed.

### Motivation

> #### What is the use case for this feature?
Sometimes, the temp path created by `_download_artifact_from_uri` doesn't work well (e.g., not enough space, not proper auth, not proper path for customized file system...). As a result, `mlflow.spark.load_model` doesn't work well. We just need to specify a usable path to solve the issue.

For log_model, we can specify `dfs_tmpdir` in `mlflow.spark.log_model` to successfully log a spark model (default value is `/tmp/mlflow`, which doesn't work in our system). However, since we can't specifiy `local_tmpdir` or `output_path` in `mlflow.spark.load_model`, we can't load a spark model back.

> #### Why is this use case valuable to support for MLflow users in general?
It helps mlflow to support more file systems. 

With a specified `local_tmpdir`, users can avoid downloading the same artifcats in different programs. This will save both time and space, especially for large models.

> #### Why is this use case valuable to support for your project(s) or organization?
With this we can use mlflow to log and load model in our system.

> #### Why is it currently difficult to achieve this use case?
We can patch mlflow to achieve this use case, but it's not the optimized solution.

### Details

### Current implementation
```python
def load_model(model_uri, dfs_tmpdir=None):
    """
    Load the Spark MLlib model from the path.
    :param 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>`_.
    :param 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``.
    :return: pyspark.ml.pipeline.PipelineModel
    .. code-block:: python
        :caption: Example
        from mlflow import spark
        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)
    """
    if RunsArtifactRepository.is_runs_uri(model_uri):
        runs_uri = model_uri
        model_uri = RunsArtifactRepository.get_underlying_uri(model_uri)
        _logger.info("'%s' resolved as '%s'", runs_uri, model_uri)
    elif ModelsArtifactRepository.is_models_uri(model_uri):
        runs_uri = model_uri
        model_uri = ModelsArtifactRepository.get_underlying_uri(model_uri)
        _logger.info("'%s' resolved as '%s'", runs_uri, model_uri)
    # 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)

    flavor_conf = _get_flavor_configuration_from_uri(model_uri, FLAVOR_NAME)
    model_uri = append_to_uri_path(model_uri, flavor_conf["model_data"])
    local_model_path = _download_artifact_from_uri(model_uri)
    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)

    if _should_use_mlflowdbfs(model_uri):
        from pyspark.ml.pipeline import PipelineModel

        mlflowdbfs_path = _mlflowdbfs_path(
            DatabricksArtifactRepository._extract_run_id(model_uri), artifact_path
        )
        with databricks_utils.MlflowCredentialContext(
            get_databricks_profile_uri_from_artifact_uri(root_uri)
        ):
            return PipelineModel.load(mlflowdbfs_path)

    return _load_model(
        model_uri=model_uri, dfs_tmpdir_base=dfs_tmpdir, local_model_path=local_model_path
    )
```

### Proposal
```python
def load_model(model_uri, dfs_tmpdir=None, local_tmpdir=None):
    """
    Load the Spark MLlib model from the path.

...

    :param local_tmpdir: Temporary directory path on local filesystem.The model artifacts are 
                         download to this path. If unspecified, a local path will be created.

...

    local_model_path = _download_artifact_from_uri(model_uri, local_tmpdir)
    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)

...
```

### What component(s) does this bug affect?

- [X] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
@thinkall Your proposal sounds great! Feel free to file a pull request and let me know if you have any questions! :)

Resolving patch:
diff --git a/mlflow/spark.py b/mlflow/spark.py
--- a/mlflow/spark.py
+++ b/mlflow/spark.py
@@ -729,7 +729,7 @@ def _load_model(model_uri, dfs_tmpdir_base=None, local_model_path=None):
     return PipelineModel.load(model_uri)
 
 
-def load_model(model_uri, dfs_tmpdir=None):
+def load_model(model_uri, dfs_tmpdir=None, dst_path=None):
     """
     Load the Spark MLlib model from the path.
 
@@ -748,6 +748,9 @@ def load_model(model_uri, dfs_tmpdir=None):
     :param 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``.
+    :param 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.
     :return: pyspark.ml.pipeline.PipelineModel
 
     .. code-block:: python
@@ -778,7 +781,7 @@ def load_model(model_uri, dfs_tmpdir=None):
 
     flavor_conf = _get_flavor_configuration_from_uri(model_uri, FLAVOR_NAME)
     model_uri = append_to_uri_path(model_uri, flavor_conf["model_data"])
-    local_model_path = _download_artifact_from_uri(model_uri)
+    local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
     _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
 
     if _should_use_mlflowdbfs(model_uri):


Test patch: