Enable auto dependency inference in spark flavor

[harupy]

Signed-off-by: harupy hkawamura0130@gmail.com

What changes are proposed in this pull request?

Enable auto dependency inference in spark flavor.

How is this patch tested?

Existing tests

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: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• 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

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

[WeichenXu123]

if the case not in databricks runtime ?

[harupy]

@WeichenXu123

If not, a new spark session is created in the following code:

mlflow/mlflow/spark.py

Lines 687 to 700 in c4b8e84

	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 = (
	pyspark.sql.SparkSession.builder.config("spark.python.worker.reuse", True)
	.config("spark.databricks.io.cache.enabled", False)
	# In Spark 3.1 and above, we need to set this conf explicitly to enable creating
	# a SparkSession on the workers
	.config("spark.executor.allowSparkContext", "true")
	.master("local[1]")
	.getOrCreate()
	)

[WeichenXu123]

Shall we add try/catch here and add fallback handling ?
e.g. the case User upgrade/downgrade the builtin mlflow of databrick runtime and found the dbruntime.spark_connection.initialize_spark_connection API does not exist.

[dbczumar]

Agree with @WeichenXu123 , I think a try/catch is a good idea here.

[harupy]

This approach breaks when initialize_spark_connection is renamed or moved to a different module.

[harupy]

Should we make the spark initialization process modifiable via monkey-patching or an environment variable?

[WeichenXu123]

Add a MLR test to prevent it break.

[dbczumar]

@harupy LGTM once https://github.com/mlflow/mlflow/pull/4759/files#r699873541 is addressed. Great work!

[harupy]

@WeichenXu123 @dbczumar Added try-catch. We could optimistically run the subsequent process after the failure, but mlflow.spark.load_pyfunc will also propbably fail If initialize_spark_connection fails.

[dbczumar]

@harupy Makes sense that it will likely fail. Can we raise a clearer error message for debugging purposes?

[harupy]

Suggested change

	try:
	# pylint: disable=import-error
	from dbruntime.spark_connection import initialize_spark_connection
	initialize_spark_connection()
	except Exception:
	pass
	try:
	# pylint: disable=import-error
	from dbruntime.spark_connection import initialize_spark_connection
	initialize_spark_connection()
	except Exception as e:
	raise Exception("Attempted to initialize a spark session to load the model but failed") from e

[dbczumar]

Sounds good to me!