281 lines
10 KiB
Python
281 lines
10 KiB
Python
import concurrent.futures
|
|
import logging
|
|
import sys
|
|
import threading
|
|
import uuid
|
|
|
|
from py4j.java_gateway import CallbackServerParameters
|
|
|
|
from mlflow import MlflowClient
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.spark import FLAVOR_NAME
|
|
from mlflow.tracking.context.abstract_context import RunContextProvider
|
|
from mlflow.tracking.fluent import _get_latest_active_run
|
|
from mlflow.utils import _truncate_and_ellipsize
|
|
from mlflow.utils.autologging_utils import (
|
|
ExceptionSafeClass,
|
|
autologging_is_disabled,
|
|
)
|
|
from mlflow.utils.databricks_utils import get_repl_id as get_databricks_repl_id
|
|
from mlflow.utils.validation import MAX_TAG_VAL_LENGTH
|
|
|
|
_JAVA_PACKAGE = "org.mlflow.spark.autologging"
|
|
_SPARK_TABLE_INFO_TAG_NAME = "sparkDatasourceInfo"
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
_lock = threading.Lock()
|
|
_table_infos = []
|
|
_spark_table_info_listener = None
|
|
|
|
# Queue & singleton consumer thread for logging Spark datasource info asynchronously
|
|
_metric_queue = []
|
|
_thread_pool = concurrent.futures.ThreadPoolExecutor(
|
|
max_workers=1, thread_name_prefix="MlflowSparkAutologging"
|
|
)
|
|
|
|
|
|
# Exposed for testing
|
|
def _get_current_listener():
|
|
return _spark_table_info_listener
|
|
|
|
|
|
def _get_table_info_string(path, version, data_format):
|
|
if data_format == "delta":
|
|
return f"path={path},version={version},format={data_format}"
|
|
return f"path={path},format={data_format}"
|
|
|
|
|
|
def _merge_tag_lines(existing_tag, new_table_info):
|
|
if existing_tag is None:
|
|
return new_table_info
|
|
if new_table_info in existing_tag:
|
|
return existing_tag
|
|
return "\n".join([existing_tag, new_table_info])
|
|
|
|
|
|
def add_table_info_to_context_provider(path, version, data_format):
|
|
with _lock:
|
|
_table_infos.append((path, version, data_format))
|
|
|
|
|
|
def clear_table_infos():
|
|
"""Clear the table info accumulated SparkAutologgingContext.
|
|
|
|
This is currently only used in unit tests.
|
|
"""
|
|
with _lock:
|
|
global _table_infos
|
|
_table_infos = []
|
|
|
|
|
|
def _get_spark_major_version(sc):
|
|
spark_version_parts = sc.version.split(".")
|
|
spark_major_version = None
|
|
if len(spark_version_parts) > 0:
|
|
spark_major_version = int(spark_version_parts[0])
|
|
return spark_major_version
|
|
|
|
|
|
def _get_jvm_event_publisher(spark_context):
|
|
"""
|
|
Get JVM-side object implementing the following methods:
|
|
- init() for initializing JVM state needed for autologging (e.g. attaching a SparkListener
|
|
to watch for datasource reads)
|
|
- register(subscriber) for registering subscribers to receive datasource events
|
|
"""
|
|
jvm = spark_context._gateway.jvm
|
|
qualified_classname = "{}.{}".format(_JAVA_PACKAGE, "MlflowAutologEventPublisher")
|
|
return getattr(jvm, qualified_classname)
|
|
|
|
|
|
def _generate_datasource_tag_value(table_info_string):
|
|
return _truncate_and_ellipsize(table_info_string, MAX_TAG_VAL_LENGTH)
|
|
|
|
|
|
def _set_run_tag_async(run_id, path, version, data_format):
|
|
_thread_pool.submit(
|
|
_set_run_tag, run_id=run_id, path=path, version=version, data_format=data_format
|
|
)
|
|
|
|
|
|
def _set_run_tag(run_id, path, version, data_format):
|
|
client = MlflowClient()
|
|
table_info_string = _get_table_info_string(path, version, data_format)
|
|
existing_run = client.get_run(run_id)
|
|
existing_tag = existing_run.data.tags.get(_SPARK_TABLE_INFO_TAG_NAME)
|
|
new_table_info = _merge_tag_lines(existing_tag, table_info_string)
|
|
new_tag_value = _generate_datasource_tag_value(new_table_info)
|
|
client.set_tag(run_id, _SPARK_TABLE_INFO_TAG_NAME, new_tag_value)
|
|
|
|
|
|
def _stop_listen_for_spark_activity(spark_context):
|
|
gw = spark_context._gateway
|
|
try:
|
|
gw.shutdown_callback_server()
|
|
except Exception as e:
|
|
_logger.warning("Failed to shut down Spark callback server for autologging: %s", e)
|
|
|
|
|
|
def _listen_for_spark_activity(spark_context):
|
|
global _spark_table_info_listener
|
|
if _get_current_listener() is not None:
|
|
return
|
|
|
|
if _get_spark_major_version(spark_context) < 3:
|
|
raise MlflowException("Spark autologging unsupported for Spark versions < 3")
|
|
|
|
gw = spark_context._gateway
|
|
params = gw.callback_server_parameters
|
|
callback_server_params = CallbackServerParameters(
|
|
address=params.address,
|
|
port=params.port,
|
|
daemonize=True,
|
|
daemonize_connections=True,
|
|
eager_load=params.eager_load,
|
|
ssl_context=params.ssl_context,
|
|
accept_timeout=params.accept_timeout,
|
|
read_timeout=params.read_timeout,
|
|
auth_token=params.auth_token,
|
|
)
|
|
callback_server_started = gw.start_callback_server(callback_server_params)
|
|
|
|
try:
|
|
event_publisher = _get_jvm_event_publisher(spark_context)
|
|
event_publisher.init(1)
|
|
_spark_table_info_listener = PythonSubscriber()
|
|
event_publisher.register(_spark_table_info_listener)
|
|
except Exception as e:
|
|
if callback_server_started:
|
|
try:
|
|
gw.shutdown_callback_server()
|
|
except Exception as e:
|
|
_logger.warning(
|
|
"Failed to shut down Spark callback server for autologging: %s", str(e)
|
|
)
|
|
_spark_table_info_listener = None
|
|
raise MlflowException(
|
|
"Exception while attempting to initialize JVM-side state for Spark datasource "
|
|
"autologging. Note that Spark datasource autologging only works with Spark 3.0 "
|
|
"and above. Please create a new Spark session with required Spark version and "
|
|
"ensure you have the mlflow-spark JAR attached to your Spark session as described "
|
|
f"in https://mlflow.org/docs/latest/tracking/autolog.html#spark Exception:\n{e}"
|
|
)
|
|
|
|
# Register context provider for Spark autologging
|
|
from mlflow.tracking.context.registry import _run_context_provider_registry
|
|
|
|
_run_context_provider_registry.register(SparkAutologgingContext)
|
|
|
|
_logger.info("Autologging successfully enabled for spark.")
|
|
|
|
|
|
def _get_repl_id():
|
|
"""
|
|
Get a unique REPL ID for a PythonSubscriber instance. This is used to distinguish between
|
|
REPLs in multitenant, REPL-aware environments where multiple Python processes may share the
|
|
same Spark JVM (e.g. in Databricks). In such environments, we pull the REPL ID from Spark
|
|
local properties, and expect that the PythonSubscriber for the current Python process only
|
|
receives events for datasource reads triggered by the current process.
|
|
"""
|
|
repl_id = get_databricks_repl_id()
|
|
if repl_id:
|
|
return repl_id
|
|
main_file = sys.argv[0] if len(sys.argv) > 0 else "<console>"
|
|
return f"PythonSubscriber[{main_file}][{uuid.uuid4().hex}]"
|
|
|
|
|
|
class PythonSubscriber(metaclass=ExceptionSafeClass):
|
|
"""
|
|
Subscriber, intended to be instantiated once per Python process, that logs Spark table
|
|
information propagated from Java to the current MLflow run, starting a run if necessary.
|
|
class implements a Java interface (org.mlflow.spark.autologging.MlflowAutologEventSubscriber,
|
|
defined in the mlflow-spark package) that's called-into by autologging logic in the JVM in order
|
|
to propagate Spark datasource read events to Python.
|
|
|
|
This class leverages the Py4j callback mechanism to receive callbacks from the JVM, see
|
|
https://www.py4j.org/advanced_topics.html#implementing-java-interfaces-from-python-callback for
|
|
more information.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._repl_id = _get_repl_id()
|
|
|
|
def toString(self):
|
|
# For debugging
|
|
return f"PythonSubscriber<replId={self.replId()}>"
|
|
|
|
def ping(self):
|
|
return None
|
|
|
|
def notify(self, path, version, data_format):
|
|
try:
|
|
self._notify(path, version, data_format)
|
|
except Exception as e:
|
|
_logger.error(
|
|
"Unexpected exception %s while attempting to log Spark datasource "
|
|
"info. Exception:\n",
|
|
e,
|
|
)
|
|
|
|
def _notify(self, path, version, data_format):
|
|
"""
|
|
Method called by Scala SparkListener to propagate datasource read events to the current
|
|
Python process
|
|
"""
|
|
if autologging_is_disabled(FLAVOR_NAME):
|
|
return
|
|
# If there are active runs, simply set the tag on the latest active run
|
|
# Note that there's a TOCTOU race condition here - active_run() here can actually throw
|
|
# if the main thread happens to end the run & pop from the active run stack after we check
|
|
# the stack size but before we peek
|
|
|
|
# Note Spark datasource autologging is hard to support thread-local behavior,
|
|
# because the spark event listener callback (jvm side) does not have the python caller
|
|
# thread information, therefore the tag is set to the latest active run, ignoring threading
|
|
# information. This way, consistent behavior is kept with existing functionality for
|
|
# Spark in MLflow.
|
|
latest_active_run = _get_latest_active_run()
|
|
|
|
if latest_active_run:
|
|
_set_run_tag_async(latest_active_run.info.run_id, path, version, data_format)
|
|
else:
|
|
add_table_info_to_context_provider(path, version, data_format)
|
|
|
|
def replId(self):
|
|
return self._repl_id
|
|
|
|
class Java:
|
|
implements = [f"{_JAVA_PACKAGE}.MlflowAutologEventSubscriber"]
|
|
|
|
|
|
class SparkAutologgingContext(RunContextProvider):
|
|
"""
|
|
Context provider used when there's no active run. Accumulates datasource read information,
|
|
then logs that information to the next-created run. Note that this doesn't clear the accumulated
|
|
info when logging them to the next run, so it will be logged to any successive runs as well.
|
|
"""
|
|
|
|
def in_context(self):
|
|
return True
|
|
|
|
def tags(self):
|
|
# if autologging is disabled, then short circuit `tags()` and return empty dict.
|
|
if autologging_is_disabled(FLAVOR_NAME):
|
|
return {}
|
|
with _lock:
|
|
seen = set()
|
|
unique_infos = []
|
|
for info in _table_infos:
|
|
if info not in seen:
|
|
unique_infos.append(info)
|
|
seen.add(info)
|
|
if len(unique_infos) > 0:
|
|
tags = {
|
|
_SPARK_TABLE_INFO_TAG_NAME: _generate_datasource_tag_value(
|
|
"\n".join([_get_table_info_string(*info) for info in unique_infos])
|
|
)
|
|
}
|
|
else:
|
|
tags = {}
|
|
return tags
|