wehub-resource-sync
304 lines
9.5 KiB
Python
304 lines
9.5 KiB
Python
import json
|
|
from collections import defaultdict
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.tracing.display import (
|
|
IPythonTraceDisplayHandler,
|
|
get_display_handler,
|
|
get_notebook_iframe_html,
|
|
)
|
|
|
|
from tests.tracing.helper import create_trace, skip_module_when_testing_trace_sdk
|
|
|
|
skip_module_when_testing_trace_sdk()
|
|
|
|
|
|
class MockEventRegistry:
|
|
def __init__(self):
|
|
self.events = defaultdict(list)
|
|
|
|
def register(self, event, callback):
|
|
self.events[event].append(callback)
|
|
|
|
def trigger(self, event):
|
|
for callback in self.events[event]:
|
|
callback(None)
|
|
|
|
|
|
class MockIPython:
|
|
def __init__(self):
|
|
self.events = MockEventRegistry()
|
|
|
|
def mock_run_cell(self):
|
|
self.events.trigger("post_run_cell")
|
|
|
|
|
|
@pytest.fixture
|
|
def _in_databricks(monkeypatch):
|
|
monkeypatch.setenv("DATABRICKS_RUNTIME_VERSION", "15.x")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_singleton():
|
|
IPythonTraceDisplayHandler._instance = None
|
|
IPythonTraceDisplayHandler.disabled = False
|
|
|
|
|
|
in_databricks = pytest.mark.usefixtures(_in_databricks.__name__)
|
|
|
|
|
|
@in_databricks
|
|
def test_display_is_not_called_without_ipython(monkeypatch):
|
|
# in an IPython environment, the interactive shell will
|
|
# be returned. however, for test purposes, just mock that
|
|
# the value is not None.
|
|
mock_display = Mock()
|
|
monkeypatch.setattr("IPython.display.display", mock_display)
|
|
handler = get_display_handler()
|
|
|
|
handler.display_traces([create_trace("a")])
|
|
assert mock_display.call_count == 0
|
|
|
|
mock_ipython = MockIPython()
|
|
monkeypatch.setattr("IPython.get_ipython", lambda: mock_ipython)
|
|
|
|
# reset the singleton so the handler
|
|
# can register the post-display hook
|
|
IPythonTraceDisplayHandler._instance = None
|
|
handler = get_display_handler()
|
|
handler.display_traces([create_trace("b")])
|
|
|
|
# simulate cell execution
|
|
mock_ipython.mock_run_cell()
|
|
|
|
assert mock_display.call_count == 1
|
|
|
|
|
|
@in_databricks
|
|
def test_ipython_client_clears_display_after_execution(monkeypatch):
|
|
mock_ipython = MockIPython()
|
|
monkeypatch.setattr("IPython.get_ipython", lambda: mock_ipython)
|
|
handler = get_display_handler()
|
|
|
|
mock_display_handle = Mock()
|
|
mock_display = Mock(return_value=mock_display_handle)
|
|
monkeypatch.setattr("IPython.display.display", mock_display)
|
|
handler.display_traces([create_trace("a")])
|
|
handler.display_traces([create_trace("b")])
|
|
handler.display_traces([create_trace("c")])
|
|
|
|
mock_ipython.mock_run_cell()
|
|
# despite many calls to `display_traces`,
|
|
# there should only be one call to `display`
|
|
assert mock_display.call_count == 1
|
|
|
|
mock_ipython.mock_run_cell()
|
|
# expect that display is not called,
|
|
# since no traces should be present
|
|
assert mock_display.call_count == 1
|
|
|
|
|
|
@in_databricks
|
|
def test_display_is_called_in_correct_functions(monkeypatch):
|
|
mock_ipython = MockIPython()
|
|
monkeypatch.setattr("IPython.get_ipython", lambda: mock_ipython)
|
|
mock_display_handle = Mock()
|
|
mock_display = Mock(return_value=mock_display_handle)
|
|
monkeypatch.setattr("IPython.display.display", mock_display)
|
|
|
|
@mlflow.trace
|
|
def foo():
|
|
return 3
|
|
|
|
# display should be called after trace creation
|
|
foo()
|
|
mlflow.flush_trace_async_logging()
|
|
mock_ipython.mock_run_cell()
|
|
assert mock_display.call_count == 1
|
|
|
|
|
|
@in_databricks
|
|
def test_display_deduplicates_traces(monkeypatch):
|
|
mock_ipython = MockIPython()
|
|
monkeypatch.setattr("IPython.get_ipython", lambda: mock_ipython)
|
|
handler = get_display_handler()
|
|
|
|
mock_display = Mock()
|
|
monkeypatch.setattr("IPython.display.display", mock_display)
|
|
|
|
trace_a = create_trace("a")
|
|
trace_b = create_trace("b")
|
|
trace_c = create_trace("c")
|
|
|
|
# The display client should dedupe traces to display and only display 3 (not 6).
|
|
handler.display_traces([trace_a])
|
|
handler.display_traces([trace_b])
|
|
handler.display_traces([trace_c])
|
|
handler.display_traces([trace_a, trace_b, trace_c])
|
|
mock_ipython.mock_run_cell()
|
|
|
|
expected = [trace_a, trace_b, trace_c]
|
|
|
|
assert mock_display.call_count == 1
|
|
assert mock_display.call_args[0][0] == {
|
|
"application/databricks.mlflow.trace": json.dumps([
|
|
json.loads(t._serialize_for_mimebundle()) for t in expected
|
|
]),
|
|
"text/plain": repr(expected),
|
|
}
|
|
|
|
|
|
@in_databricks
|
|
def test_display_respects_max_limit(monkeypatch):
|
|
mock_ipython = MockIPython()
|
|
monkeypatch.setattr("IPython.get_ipython", lambda: mock_ipython)
|
|
handler = get_display_handler()
|
|
|
|
mock_display = Mock()
|
|
monkeypatch.setattr("IPython.display.display", mock_display)
|
|
|
|
monkeypatch.setenv("MLFLOW_MAX_TRACES_TO_DISPLAY_IN_NOTEBOOK", "1")
|
|
|
|
trace_a = create_trace("a")
|
|
trace_b = create_trace("b")
|
|
trace_c = create_trace("c")
|
|
handler.display_traces([trace_a, trace_b, trace_c])
|
|
mock_ipython.mock_run_cell()
|
|
|
|
assert mock_display.call_count == 1
|
|
assert mock_display.call_args[0][0] == {
|
|
"application/databricks.mlflow.trace": trace_a._serialize_for_mimebundle(),
|
|
"text/plain": repr(trace_a),
|
|
}
|
|
|
|
|
|
@in_databricks
|
|
def test_enable_and_disable_display(monkeypatch):
|
|
mock_ipython = MockIPython()
|
|
monkeypatch.setattr("IPython.get_ipython", lambda: mock_ipython)
|
|
mock_display_handle = Mock()
|
|
mock_display = Mock(return_value=mock_display_handle)
|
|
monkeypatch.setattr("IPython.display.display", mock_display)
|
|
trace_a = create_trace("a")
|
|
|
|
# test that disabling the display handler prevents display() from being called
|
|
mlflow.tracing.disable_notebook_display()
|
|
handler = get_display_handler()
|
|
handler.display_traces([trace_a])
|
|
mock_ipython.mock_run_cell()
|
|
|
|
mock_display.assert_not_called()
|
|
|
|
# test that re-enabling it will make things display again
|
|
mlflow.tracing.enable_notebook_display()
|
|
handler = get_display_handler()
|
|
handler.display_traces([trace_a])
|
|
mock_ipython.mock_run_cell()
|
|
|
|
assert mock_display.call_count == 1
|
|
assert mock_display.call_args[0][0] == {
|
|
"application/databricks.mlflow.trace": trace_a._serialize_for_mimebundle(),
|
|
"text/plain": repr(trace_a),
|
|
}
|
|
|
|
|
|
@in_databricks
|
|
def test_mimebundle_in_databricks():
|
|
# by default, it should contain the metadata
|
|
# necessary for rendering the trace UI
|
|
trace = create_trace("a")
|
|
assert trace._repr_mimebundle_() == {
|
|
"application/databricks.mlflow.trace": trace._serialize_for_mimebundle(),
|
|
"text/plain": repr(trace),
|
|
}
|
|
|
|
# if trace display is disabled, only "text/plain" should exist
|
|
mlflow.tracing.disable_notebook_display()
|
|
assert trace._repr_mimebundle_() == {
|
|
"text/plain": repr(trace),
|
|
}
|
|
|
|
# re-enabling should bring the metadata back
|
|
mlflow.tracing.enable_notebook_display()
|
|
assert trace._repr_mimebundle_() == {
|
|
"application/databricks.mlflow.trace": trace._serialize_for_mimebundle(),
|
|
"text/plain": repr(trace),
|
|
}
|
|
|
|
|
|
def test_mimebundle_in_oss():
|
|
# if the user is not using a tracking server, it should only contain text/plain
|
|
trace = create_trace("a")
|
|
assert trace._repr_mimebundle_() == {
|
|
"text/plain": repr(trace),
|
|
}
|
|
|
|
# if the user is using a tracking server, it
|
|
# should contain an iframe in the text/html key
|
|
mlflow.set_tracking_uri("http://localhost:5000")
|
|
assert trace._repr_mimebundle_() == {
|
|
"text/plain": repr(trace),
|
|
"text/html": get_notebook_iframe_html([trace]),
|
|
}
|
|
|
|
# disabling should remove this key, even if tracking server is used
|
|
mlflow.tracing.disable_notebook_display()
|
|
assert trace._repr_mimebundle_() == {
|
|
"text/plain": repr(trace),
|
|
}
|
|
|
|
|
|
def test_notebook_trace_renderer_base_url_override(monkeypatch):
|
|
trace = create_trace("a")
|
|
mlflow.set_tracking_uri("http://mlflow:5000")
|
|
monkeypatch.setenv("MLFLOW_NOTEBOOK_TRACE_RENDERER_BASE_URL", "http://localhost:5000")
|
|
|
|
html = get_notebook_iframe_html([trace])
|
|
assert "http://localhost:5000/static-files/lib/notebook-trace-renderer/index.html" in html
|
|
assert "http://mlflow:5000/static-files/lib/notebook-trace-renderer/index.html" not in html
|
|
|
|
|
|
def test_notebook_iframe_includes_workspace_query_param(monkeypatch):
|
|
trace = create_trace("a")
|
|
mlflow.set_tracking_uri("http://localhost:5000")
|
|
|
|
# Without workspace set, the query string should not contain workspace
|
|
html = get_notebook_iframe_html([trace])
|
|
assert "workspace=" not in html
|
|
|
|
# With workspace set, the query string should contain workspace
|
|
monkeypatch.setenv("MLFLOW_WORKSPACE", "my-workspace")
|
|
html = get_notebook_iframe_html([trace])
|
|
assert "workspace=my-workspace" in html
|
|
|
|
|
|
def test_display_in_oss(monkeypatch):
|
|
mock_ipython = MockIPython()
|
|
monkeypatch.setattr("IPython.get_ipython", lambda: mock_ipython)
|
|
mock_display_handle = Mock()
|
|
mock_display = Mock(return_value=mock_display_handle)
|
|
monkeypatch.setattr("IPython.display.display", mock_display)
|
|
monkeypatch.setattr("IPython.display.HTML", Mock(side_effect=lambda html: html))
|
|
|
|
handler = get_display_handler()
|
|
handler.display_traces([create_trace("a")])
|
|
|
|
mock_ipython.mock_run_cell()
|
|
|
|
# default tracking uri is sqlite, so no display call should be made
|
|
assert mock_display.call_count == 0
|
|
|
|
# after setting an HTTP tracking URI, it should work
|
|
mlflow.set_tracking_uri("http://localhost:5000")
|
|
|
|
handler = get_display_handler()
|
|
handler.display_traces([create_trace("a")])
|
|
|
|
mock_ipython.mock_run_cell()
|
|
|
|
assert mock_display.call_count == 1
|
|
assert "<iframe" in mock_display.call_args[0][0]["text/html"]
|