wehub-resource-sync
146 lines
4.6 KiB
Python
146 lines
4.6 KiB
Python
from unittest import mock
|
|
|
|
from mlflow.entities.span import LiveSpan, Span
|
|
from mlflow.tracing.attachments import Attachment
|
|
|
|
|
|
def _make_live_span(trace_id="tr-test123"):
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
|
|
tracer = TracerProvider().get_tracer("test")
|
|
otel_span = tracer.start_span("test_span")
|
|
return LiveSpan(otel_span, trace_id=trace_id)
|
|
|
|
|
|
def test_set_inputs_extracts_attachment():
|
|
span = _make_live_span()
|
|
att = Attachment(content_type="image/png", content_bytes=b"img")
|
|
span.set_inputs({"image": att})
|
|
|
|
assert att.id in span._attachments
|
|
assert span._attachments[att.id] is att
|
|
inputs = span.inputs
|
|
assert isinstance(inputs["image"], str)
|
|
assert inputs["image"].startswith("mlflow-attachment://")
|
|
parsed = Attachment.parse_ref(inputs["image"])
|
|
assert parsed["trace_id"] == "tr-test123"
|
|
assert parsed["attachment_id"] == att.id
|
|
|
|
|
|
def test_set_outputs_extracts_attachment():
|
|
span = _make_live_span()
|
|
att = Attachment(content_type="audio/wav", content_bytes=b"audio")
|
|
span.set_outputs({"sound": att})
|
|
|
|
assert att.id in span._attachments
|
|
outputs = span.outputs
|
|
assert isinstance(outputs["sound"], str)
|
|
parsed = Attachment.parse_ref(outputs["sound"])
|
|
assert parsed["content_type"] == "audio/wav"
|
|
|
|
|
|
def test_nested_dict_extraction():
|
|
span = _make_live_span()
|
|
att = Attachment(content_type="image/jpeg", content_bytes=b"jpg")
|
|
span.set_inputs({"nested": {"deep": att}})
|
|
|
|
assert att.id in span._attachments
|
|
inputs = span.inputs
|
|
assert isinstance(inputs["nested"]["deep"], str)
|
|
assert att.id in inputs["nested"]["deep"]
|
|
|
|
|
|
def test_list_extraction():
|
|
span = _make_live_span()
|
|
att1 = Attachment(content_type="image/png", content_bytes=b"a")
|
|
att2 = Attachment(content_type="image/png", content_bytes=b"b")
|
|
span.set_inputs({"images": [att1, att2]})
|
|
|
|
assert att1.id in span._attachments
|
|
assert att2.id in span._attachments
|
|
inputs = span.inputs
|
|
assert len(inputs["images"]) == 2
|
|
assert all(isinstance(v, str) for v in inputs["images"])
|
|
|
|
|
|
def test_mixed_values_only_replaces_attachments():
|
|
span = _make_live_span()
|
|
att = Attachment(content_type="image/png", content_bytes=b"img")
|
|
span.set_inputs({"image": att, "text": "hello", "number": 42})
|
|
|
|
assert att.id in span._attachments
|
|
inputs = span.inputs
|
|
assert inputs["text"] == "hello"
|
|
assert inputs["number"] == 42
|
|
assert isinstance(inputs["image"], str)
|
|
|
|
|
|
def test_non_dict_input_without_attachment():
|
|
span = _make_live_span()
|
|
span.set_inputs("plain string")
|
|
assert span._attachments == {}
|
|
assert span.inputs == "plain string"
|
|
|
|
|
|
def test_tuple_extraction():
|
|
span = _make_live_span()
|
|
att = Attachment(content_type="image/png", content_bytes=b"img")
|
|
span.set_inputs({"images": (att, "text")})
|
|
|
|
assert att.id in span._attachments
|
|
inputs = span.inputs
|
|
assert isinstance(inputs["images"], list)
|
|
assert isinstance(inputs["images"][0], str)
|
|
assert inputs["images"][0].startswith("mlflow-attachment://")
|
|
assert inputs["images"][1] == "text"
|
|
|
|
|
|
def test_set_inputs_twice_accumulates_attachments():
|
|
span = _make_live_span()
|
|
att1 = Attachment(content_type="image/png", content_bytes=b"first")
|
|
att2 = Attachment(content_type="image/jpeg", content_bytes=b"second")
|
|
|
|
span.set_inputs({"img": att1})
|
|
span.set_outputs({"img": att2})
|
|
|
|
assert att1.id in span._attachments
|
|
assert att2.id in span._attachments
|
|
assert len(span._attachments) == 2
|
|
|
|
|
|
def test_to_immutable_span_propagates_attachments():
|
|
span = _make_live_span()
|
|
att = Attachment(content_type="image/png", content_bytes=b"img")
|
|
span.set_inputs({"image": att})
|
|
|
|
immutable = span.to_immutable_span()
|
|
assert isinstance(immutable, Span)
|
|
assert att.id in immutable._attachments
|
|
assert immutable._attachments[att.id] is att
|
|
|
|
|
|
def test_set_inputs_after_end_does_not_store_attachment():
|
|
span = _make_live_span()
|
|
span.end()
|
|
|
|
att = Attachment(content_type="image/png", content_bytes=b"img")
|
|
with mock.patch("mlflow.entities.span._logger.debug") as mock_debug:
|
|
span.set_inputs({"image": att})
|
|
|
|
assert att.id not in span._attachments
|
|
assert span.inputs is None
|
|
mock_debug.assert_called_once_with(
|
|
"Skipping attachment storage because the span is no longer recording."
|
|
)
|
|
|
|
|
|
def test_set_outputs_after_end_does_not_store_attachment():
|
|
span = _make_live_span()
|
|
span.end()
|
|
|
|
att = Attachment(content_type="audio/wav", content_bytes=b"audio")
|
|
span.set_outputs({"sound": att})
|
|
|
|
assert att.id not in span._attachments
|
|
assert span.outputs is None
|