mlflow--mlflow
wehub-resource-sync
107 行
3.6 KiB
Python
107 行
3.6 KiB
Python
import logging
|
|
import threading
|
|
import time
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
import mlflow.spark.autologging as autologging_module
|
|
from mlflow.spark.autologging import _stop_listen_for_spark_activity
|
|
|
|
|
|
@pytest.fixture
|
|
def autolog_caplog(caplog):
|
|
# "mlflow" sets propagate=False, so its records never reach caplog's root handler.
|
|
# Enable propagation so caplog captures them. Don't attach caplog.handler directly:
|
|
# pytest >= 9.1 attaches it to non-propagating loggers, which double-counts records.
|
|
mlflow_logger = logging.getLogger("mlflow")
|
|
original_propagate = mlflow_logger.propagate
|
|
mlflow_logger.propagate = True
|
|
try:
|
|
with caplog.at_level(logging.WARNING, logger="mlflow.spark.autologging"):
|
|
yield caplog
|
|
finally:
|
|
mlflow_logger.propagate = original_propagate
|
|
|
|
|
|
def _make_spark_context(shutdown_side_effect=None):
|
|
gw = MagicMock()
|
|
if shutdown_side_effect is not None:
|
|
gw.shutdown_callback_server.side_effect = shutdown_side_effect
|
|
sc = MagicMock()
|
|
sc._gateway = gw
|
|
return sc
|
|
|
|
|
|
def test_fast_shutdown_returns_without_warning(autolog_caplog):
|
|
sc = _make_spark_context()
|
|
_stop_listen_for_spark_activity(sc)
|
|
sc._gateway.shutdown_callback_server.assert_called_once()
|
|
assert not any("did not complete" in r.message for r in autolog_caplog.records)
|
|
|
|
|
|
def test_hanging_shutdown_times_out_and_logs_warning(autolog_caplog):
|
|
hang = threading.Event()
|
|
|
|
def _hang():
|
|
hang.wait()
|
|
|
|
sc = _make_spark_context(shutdown_side_effect=_hang)
|
|
try:
|
|
with patch.object(autologging_module, "_CALLBACK_SERVER_SHUTDOWN_TIMEOUT_SECONDS", 0.2):
|
|
_stop_listen_for_spark_activity(sc)
|
|
warnings = [r for r in autolog_caplog.records if "did not complete" in r.message]
|
|
assert len(warnings) == 1
|
|
assert "CLOSE_WAIT" in warnings[0].message
|
|
finally:
|
|
hang.set()
|
|
|
|
|
|
def test_hanging_shutdown_does_not_block_caller():
|
|
hang = threading.Event()
|
|
|
|
def _hang():
|
|
hang.wait()
|
|
|
|
sc = _make_spark_context(shutdown_side_effect=_hang)
|
|
try:
|
|
with patch.object(autologging_module, "_CALLBACK_SERVER_SHUTDOWN_TIMEOUT_SECONDS", 0.2):
|
|
start = time.monotonic()
|
|
_stop_listen_for_spark_activity(sc)
|
|
elapsed = time.monotonic() - start
|
|
assert elapsed < 2.0
|
|
finally:
|
|
hang.set()
|
|
|
|
|
|
def test_exception_during_shutdown_logs_error_not_timeout(autolog_caplog):
|
|
sc = _make_spark_context(shutdown_side_effect=Exception("connection reset"))
|
|
with patch.object(autologging_module, "_CALLBACK_SERVER_SHUTDOWN_TIMEOUT_SECONDS", 0.2):
|
|
_stop_listen_for_spark_activity(sc)
|
|
error_warnings = [r for r in autolog_caplog.records if "Failed to shut down" in r.message]
|
|
assert len(error_warnings) == 1
|
|
assert "connection reset" in error_warnings[0].message
|
|
assert not any("did not complete" in r.message for r in autolog_caplog.records)
|
|
|
|
|
|
def test_shutdown_thread_is_daemon():
|
|
created_threads = []
|
|
original_init = threading.Thread.__init__
|
|
|
|
def capturing_init(self, *args, **kwargs):
|
|
original_init(self, *args, **kwargs)
|
|
created_threads.append(self)
|
|
|
|
hang = threading.Event()
|
|
sc = _make_spark_context(shutdown_side_effect=lambda: hang.wait())
|
|
try:
|
|
with (
|
|
patch.object(threading.Thread, "__init__", capturing_init),
|
|
patch.object(autologging_module, "_CALLBACK_SERVER_SHUTDOWN_TIMEOUT_SECONDS", 0.1),
|
|
):
|
|
_stop_listen_for_spark_activity(sc)
|
|
assert created_threads
|
|
assert all(t.daemon for t in created_threads)
|
|
finally:
|
|
hang.set()
|