wehub-resource-sync
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
import warnings
|
|
from threading import Thread
|
|
|
|
from mlflow import MlflowClient
|
|
from mlflow.entities import Metric
|
|
from mlflow.utils.autologging_utils.logging_and_warnings import (
|
|
_WarningsController,
|
|
)
|
|
from mlflow.utils.autologging_utils.metrics_queue import (
|
|
_metrics_queue,
|
|
_metrics_queue_lock,
|
|
flush_metrics_queue,
|
|
)
|
|
|
|
|
|
def test_flush_metrics_queue_is_thread_safe():
|
|
"""
|
|
Autologging augments TensorBoard event logging hooks with MLflow `log_metric` API
|
|
calls. To prevent these API calls from blocking TensorBoard event logs, `log_metric`
|
|
API calls are scheduled via `_flush_queue` on a background thread. Accordingly, this test
|
|
verifies that `_flush_queue` is thread safe.
|
|
"""
|
|
|
|
client = MlflowClient()
|
|
run = client.create_run(experiment_id="0")
|
|
metric_queue_item = (run.info.run_id, Metric("foo", 0.1, 100, 1))
|
|
_metrics_queue.append(metric_queue_item)
|
|
|
|
# Verify that, if another thread holds a lock on the metric queue leveraged by
|
|
# _flush_queue, _flush_queue terminates and does not modify the queue
|
|
_metrics_queue_lock.acquire()
|
|
flush_thread1 = Thread(name="test-autologging-flush-1", target=flush_metrics_queue)
|
|
flush_thread1.start()
|
|
flush_thread1.join()
|
|
assert len(_metrics_queue) == 1
|
|
assert _metrics_queue[0] == metric_queue_item
|
|
_metrics_queue_lock.release()
|
|
|
|
# Verify that, if no other thread holds a lock on the metric queue leveraged by
|
|
# _flush_queue, _flush_queue flushes the queue as expected
|
|
flush_thread2 = Thread(name="test-autologging-flush-2", target=flush_metrics_queue)
|
|
flush_thread2.start()
|
|
flush_thread2.join()
|
|
assert len(_metrics_queue) == 0
|
|
|
|
|
|
def test_double_patch_does_not_overwrite(monkeypatch):
|
|
original_showwarning = warnings.showwarning
|
|
monkeypatch.setattr(warnings, "showwarning", original_showwarning)
|
|
|
|
controller = _WarningsController()
|
|
|
|
assert warnings.showwarning == original_showwarning
|
|
assert not controller._did_patch_showwarning
|
|
|
|
controller.set_non_mlflow_warnings_disablement_state_for_current_thread(True)
|
|
|
|
assert controller._did_patch_showwarning
|
|
assert warnings.showwarning == controller._patched_showwarning
|
|
|
|
patched_func = warnings.showwarning
|
|
|
|
controller._modify_patch_state_if_necessary()
|
|
|
|
assert warnings.showwarning == patched_func
|
|
|
|
controller.set_non_mlflow_warnings_disablement_state_for_current_thread(False)
|
|
|
|
assert warnings.showwarning == original_showwarning
|
|
|
|
|
|
def test_showwarning_captures_user_handler():
|
|
"""Verify that a user-set warnings.showwarning handler survives the patch/unpatch cycle.
|
|
|
|
Regression test for https://github.com/mlflow/mlflow/issues/21689
|
|
"""
|
|
original_showwarning = warnings.showwarning
|
|
|
|
# Simulate a user setting a custom handler after import (e.g. logging.captureWarnings)
|
|
custom_handler = lambda *args, **kwargs: None # noqa: E731
|
|
warnings.showwarning = custom_handler
|
|
|
|
try:
|
|
controller = _WarningsController()
|
|
|
|
# Patch: controller should capture the user's custom handler
|
|
controller.set_non_mlflow_warnings_disablement_state_for_current_thread(True)
|
|
assert warnings.showwarning == controller._patched_showwarning
|
|
assert controller._original_showwarning is custom_handler
|
|
|
|
# Unpatch: controller should restore the user's custom handler, not the default
|
|
controller.set_non_mlflow_warnings_disablement_state_for_current_thread(False)
|
|
assert warnings.showwarning is custom_handler
|
|
finally:
|
|
warnings.showwarning = original_showwarning
|