63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
import os
|
|
import shutil
|
|
import uuid
|
|
|
|
from mlflow.utils._spark_utils import _get_active_spark_session
|
|
from mlflow.utils.databricks_utils import (
|
|
get_databricks_nfs_temp_dir,
|
|
is_databricks_connect,
|
|
is_in_databricks_runtime,
|
|
is_in_databricks_serverless_runtime,
|
|
)
|
|
|
|
# Set spark config "spark.mlflow.nfs.rootDir" to specify a NFS (network file system) directory
|
|
# which is shared with all spark cluster nodes.
|
|
# This will help optimize routine of distributing spark driver files to remote workers.
|
|
# None represent no NFS directory available.
|
|
# Note:
|
|
# 1. If NFS directory set, you must ensure all spark cluster nodes using the same hardware and
|
|
# installed the same OS with the same environment configured, because mlflow uses NFS directory
|
|
# to distribute driver side virtual environment to remote workers if NFS available, heterogeneous
|
|
# cluster nodes might cause issues under the case.
|
|
# 2. The NFS directory must be mounted before importing mlflow.
|
|
# 3. For databricks users, don't set this config, databricks product sets up internal NFS service
|
|
# automatically.
|
|
_NFS_CACHE_ROOT_DIR = None
|
|
|
|
|
|
def get_nfs_cache_root_dir():
|
|
if is_in_databricks_runtime():
|
|
spark_sess = _get_active_spark_session()
|
|
if is_in_databricks_serverless_runtime():
|
|
# Databricks Serverless runtime VM can't access NFS.
|
|
nfs_enabled = False
|
|
else:
|
|
nfs_enabled = spark_sess and (
|
|
spark_sess.conf.get("spark.databricks.mlflow.nfs.enabled", "true").lower() == "true"
|
|
)
|
|
if nfs_enabled:
|
|
try:
|
|
return get_databricks_nfs_temp_dir()
|
|
except Exception:
|
|
nfs_root_dir = "/local_disk0/.ephemeral_nfs"
|
|
# Test whether the NFS directory is writable.
|
|
test_path = os.path.join(nfs_root_dir, uuid.uuid4().hex)
|
|
try:
|
|
os.makedirs(test_path)
|
|
return nfs_root_dir
|
|
except Exception:
|
|
# For databricks cluster enabled Table ACL, we have no permission to access NFS
|
|
# directory, in this case, return None, meaning NFS is not available.
|
|
return None
|
|
finally:
|
|
shutil.rmtree(test_path, ignore_errors=True)
|
|
else:
|
|
return None
|
|
else:
|
|
spark_session = _get_active_spark_session()
|
|
if is_databricks_connect(spark_session):
|
|
# Remote spark connect client can't access Databricks Serverless cluster NFS.
|
|
return None
|
|
if spark_session is not None:
|
|
return spark_session.conf.get("spark.mlflow.nfs.rootDir", None)
|