[ML-12132] Fix persistence of Spark models on passthrough-enabled environments

[smurching]

What changes are proposed in this pull request?

Fixes mlflow.spark.log_model, mlflow.spark.load_model APIs on passthrough-enabled environments against ACL'd artifact locations.

Previously, we had the following behavior:

• mlflow.spark.save_model: saves spark model to temporary directory on Spark's default distributed FS (e.g. DBFS on Databricks), then copies the persisted spark model directory to the desired local path using HDFS APIs called via Py4j. The latter step fails in passthrough environments as HDFS APIs aren't allowlisted
• mlflow.spark.log_model: attempts to write directly to the destination directory, which fails for ACL'd artifact locations. Then, falls back to calling save_model & log_artifacts, which fails as described above
• mlflow.spark.load_model: attempts to load model directly from the artifact location via Spark. If this fails, downloads model files locally, uploads them to Spark's distributed FS via HDFS APIs (fails due to lack of allowlisting), and then loads Spark model from the distributed FS

This PR makes the following changes when running on Databricks, for both passthrough-enabled & non-passthrough envs:

• mlflow.spark.save_model: copy model files locally using DBFS FUSE instead of denylisted HDFS APIs. This also fixes log_model
• mlflow.spark.load_model: upload locally-downloaded model files to Spark's distributed FS via DBFS FUSE instead of denylisted HDFS APIs.

How is this patch tested?

Manually

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• 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/projects: MLproject format, project running backends
• area/scoring: Local serving, model deployment tools, spark UDFs
• area/server-infra: MLflow server, JavaScript dev server
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, JavaScript, plotting
• 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

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[edgan8]

Why do we need the abspath check?

[smurching]

Good Q, it's because persisting a spark model to a path like "/my-directory" actually saves to "dbfs:/my-directory".

I'll add this explanation/clarifying comment to the code, lmk if it makes sense:

    # If we're running on a Databricks cluster, we can detect whether the tempdir we're saving our
    # Spark model to is on DBFS, by checking if it either starts with dbfs:/
    # e.g. "dbfs:/my-directory" or is an absolute path like "/my-directory" (in which case Spark
    # saves to the default distributed FS => DBFS). If we're using a DBFS
    # tempdir, copy to local FS via the FUSE mount. Otherwise, use HDFS APIs to copy the model to
    # the local FS.

[edgan8]

Can you add a docstring to this method so that it's more obvious that this comment just applies to the immediate if statement below and not to the whole method?

[edgan8]

if dbfs_uri is a posixpath abspath there's no way it can be a dbfs:/ uri right?

[edgan8]

When does the file get written to tmp_path fuse so we can move it? Are we assuming that the FUSE mount has been synced with HDFS? What if the sync hasn't finished yet?

[smurching]

In this save_model method, the goal is to save the SparkML model to a local filesystem path

The flow here is:

• Save SparkML model to a temporary dbfs location (dbfs:/tmp/my-model)
• Move that persisted Spark model to the local filesystem using FUSE and a file move (move /dbs/tmp/my-model to the /my/desired/local/path)

We leverage FUSE to copy from DBFS -> local as there's no allowlisted JVM-side, HDFS API for copying from dbfs:/... to the local filesystem

[smurching]

The FUSE access has the same consistency properties as standard HDFS-style access to the same DBFS path. That is, writing to dbfs:/tmp/my-model and reading /dbfs/tmp/my-model is equivalent from a blob-storage-eventual-consistency perspective to writing to dbfs:/tmp/my-model and subsequently reading dbfs:/tmp/my-model

[edgan8]

Can we factor out this copytree functionality into a more generic helper function?

[edgan8]

This is no longer specific to fuse, it's just "absolute_dir_path" right?

[smurching]

👍 yep that's right! I also added a unit test for this method, although it's an internal one used only in the mlflow.spark module, because it's somewhat complex

[smurching]

Note: should be a posixpath join since HDFS URIs are posixpath style (e.g. we don't want to include backward slashes \ here on platforms like Windows)

[edgan8]

Nice refactoring! Can you update the PR description to clarify that this is changing dbfs access from HDFS to FUSE for all clusters (not just the passthrough ones that motivated this)?