wehub-resource-sync
1255 lines
45 KiB
Python
1255 lines
45 KiB
Python
import json
|
|
from datetime import datetime
|
|
from unittest import mock
|
|
|
|
import opentelemetry.trace as trace_api
|
|
import pytest
|
|
from opentelemetry.proto.resource.v1.resource_pb2 import Resource as OTelProtoResource
|
|
from opentelemetry.proto.trace.v1.trace_pb2 import Span as OTelProtoSpan
|
|
from opentelemetry.proto.trace.v1.trace_pb2 import Status as OTelProtoStatus
|
|
from opentelemetry.sdk.resources import Resource as OTelResource
|
|
from opentelemetry.sdk.trace import Event as OTelEvent
|
|
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
|
|
from opentelemetry.trace import Status as OTelStatus
|
|
from opentelemetry.trace import StatusCode as OTelStatusCode
|
|
|
|
import mlflow
|
|
from mlflow.entities import LiveSpan, Span, SpanEvent, SpanStatus, SpanStatusCode, SpanType
|
|
from mlflow.entities.span import (
|
|
NoOpSpan,
|
|
create_mlflow_span,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.tracing.constant import TRACE_ID_V4_PREFIX
|
|
from mlflow.tracing.provider import _get_tracer, trace_disabled
|
|
from mlflow.tracing.utils import build_otel_context, encode_span_id, encode_trace_id
|
|
from mlflow.tracing.utils.otlp import (
|
|
_decode_otel_proto_anyvalue,
|
|
_set_otel_proto_anyvalue,
|
|
resource_to_otel_proto,
|
|
)
|
|
|
|
|
|
def test_create_live_span():
|
|
trace_id = "tr-12345"
|
|
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("parent") as parent_span:
|
|
span = create_mlflow_span(parent_span, trace_id=trace_id, span_type=SpanType.LLM)
|
|
assert isinstance(span, LiveSpan)
|
|
assert span.trace_id == trace_id
|
|
assert span._trace_id == encode_trace_id(parent_span.context.trace_id)
|
|
assert span.span_id == encode_span_id(parent_span.context.span_id)
|
|
assert span.name == "parent"
|
|
assert span.start_time_ns == parent_span.start_time
|
|
assert span.end_time_ns is None
|
|
assert span.parent_id is None
|
|
|
|
span.set_inputs({"input": 1})
|
|
span.set_outputs(2)
|
|
assert span.inputs == {"input": 1}
|
|
assert span.outputs == 2
|
|
|
|
span.set_attribute("key", 3)
|
|
assert span.get_attribute("key") == 3
|
|
|
|
# non-serializable value should be stored as string
|
|
non_serializable = datetime.now()
|
|
span.set_attribute("non_serializable", non_serializable)
|
|
assert span.get_attribute("non_serializable") == str(non_serializable)
|
|
# NB: `mlflow.spanLogLevel` is intentionally absent here — the live
|
|
# span hasn't ended yet, and log-level resolution happens in `end()`.
|
|
# Persisted spans always carry the attribute (covered elsewhere).
|
|
assert parent_span._attributes == {
|
|
"mlflow.traceRequestId": json.dumps(trace_id),
|
|
"mlflow.spanInputs": '{"input": 1}',
|
|
"mlflow.spanOutputs": "2",
|
|
"mlflow.spanType": '"LLM"',
|
|
"key": "3",
|
|
"non_serializable": json.dumps(str(non_serializable)),
|
|
}
|
|
|
|
span.set_status("OK")
|
|
assert span.status == SpanStatus(SpanStatusCode.OK)
|
|
|
|
span.add_event(SpanEvent("test_event", timestamp=99999, attributes={"foo": "bar"}))
|
|
assert len(span.events) == 1
|
|
assert span.events[0].name == "test_event"
|
|
assert span.events[0].timestamp == 99999
|
|
assert span.events[0].attributes == {"foo": "bar"}
|
|
|
|
# Test child span
|
|
with tracer.start_as_current_span("child") as child_span:
|
|
span = create_mlflow_span(child_span, trace_id=trace_id)
|
|
assert isinstance(span, LiveSpan)
|
|
assert span.name == "child"
|
|
assert span.parent_id == encode_span_id(parent_span.context.span_id)
|
|
|
|
|
|
def test_create_non_live_span():
|
|
trace_id = "tr-12345"
|
|
parent_span_context = trace_api.SpanContext(
|
|
trace_id=12345, span_id=111, is_remote=False, trace_flags=trace_api.TraceFlags(1)
|
|
)
|
|
readable_span = OTelReadableSpan(
|
|
name="test",
|
|
context=trace_api.SpanContext(
|
|
trace_id=12345, span_id=222, is_remote=False, trace_flags=trace_api.TraceFlags(1)
|
|
),
|
|
parent=parent_span_context,
|
|
attributes={
|
|
"mlflow.traceRequestId": json.dumps(trace_id),
|
|
"mlflow.spanInputs": '{"input": 1, "nested": {"foo": "bar"}}',
|
|
"mlflow.spanOutputs": "2",
|
|
"key": "3",
|
|
},
|
|
start_time=99999,
|
|
end_time=100000,
|
|
)
|
|
span = create_mlflow_span(readable_span, trace_id)
|
|
|
|
assert isinstance(span, Span)
|
|
assert not isinstance(span, LiveSpan)
|
|
assert not isinstance(span, NoOpSpan)
|
|
assert span.trace_id == trace_id
|
|
assert span._trace_id == encode_trace_id(12345)
|
|
assert span.span_id == encode_span_id(222)
|
|
assert span.name == "test"
|
|
assert span.start_time_ns == 99999
|
|
assert span.end_time_ns == 100000
|
|
assert span.parent_id == encode_span_id(111)
|
|
assert span.inputs == {"input": 1, "nested": {"foo": "bar"}}
|
|
assert span.outputs == 2
|
|
assert span.status == SpanStatus(SpanStatusCode.UNSET, description="")
|
|
assert span.get_attribute("key") == 3
|
|
|
|
# Non-live span should not implement setter methods
|
|
with pytest.raises(AttributeError, match="set_inputs"):
|
|
span.set_inputs({"input": 1})
|
|
|
|
|
|
def test_create_noop_span():
|
|
trace_id = "tr-12345"
|
|
|
|
@trace_disabled
|
|
def f():
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("span") as otel_span:
|
|
span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
assert isinstance(span, NoOpSpan)
|
|
|
|
# create from None
|
|
span = create_mlflow_span(None, trace_id=trace_id)
|
|
assert isinstance(span, NoOpSpan)
|
|
|
|
|
|
def test_create_raise_for_invalid_otel_span():
|
|
with pytest.raises(MlflowException, match=r"The `otel_span` argument must be"):
|
|
create_mlflow_span(otel_span=123, trace_id="tr-12345")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"status",
|
|
[SpanStatus("OK"), SpanStatus(SpanStatusCode.ERROR, "Error!"), "OK", "ERROR"],
|
|
)
|
|
def test_set_status(status):
|
|
with mlflow.start_span("test_span") as span:
|
|
span.set_status(status)
|
|
|
|
assert isinstance(span.status, SpanStatus)
|
|
|
|
|
|
def test_set_status_raise_for_invalid_value():
|
|
with mlflow.start_span("test_span") as span:
|
|
with pytest.raises(MlflowException, match=r"INVALID is not a valid SpanStatusCode value."):
|
|
span.set_status("INVALID")
|
|
|
|
|
|
def test_dict_conversion():
|
|
with mlflow.start_span("parent"):
|
|
with mlflow.start_span("child", span_type=SpanType.LLM) as span:
|
|
span.set_inputs({"input": 1})
|
|
span.set_outputs(2)
|
|
span.set_attribute("key", 3)
|
|
span.set_status("OK")
|
|
span.add_event(SpanEvent("test_event", timestamp=0, attributes={"foo": "bar"}))
|
|
|
|
span_dict = span.to_dict()
|
|
recovered_span = Span.from_dict(span_dict)
|
|
|
|
assert span.trace_id == recovered_span.trace_id
|
|
assert span._trace_id == recovered_span._trace_id
|
|
assert span.span_id == recovered_span.span_id
|
|
assert span.name == recovered_span.name
|
|
assert span.span_type == recovered_span.span_type
|
|
assert span.start_time_ns == recovered_span.start_time_ns
|
|
assert span.end_time_ns == recovered_span.end_time_ns
|
|
assert span.parent_id == recovered_span.parent_id
|
|
assert span.status == recovered_span.status
|
|
assert span.inputs == recovered_span.inputs
|
|
assert span.outputs == recovered_span.outputs
|
|
assert span.attributes == recovered_span.attributes
|
|
assert span.events == recovered_span.events
|
|
|
|
# Loaded span should not implement setter methods
|
|
with pytest.raises(AttributeError, match="set_status"):
|
|
recovered_span.set_status("OK")
|
|
|
|
|
|
def test_dict_conversion_with_exception_event():
|
|
with pytest.raises(ValueError, match="Test exception"):
|
|
with mlflow.start_span("test") as span:
|
|
raise ValueError("Test exception")
|
|
|
|
span_dict = span.to_dict()
|
|
recovered_span = Span.from_dict(span_dict)
|
|
|
|
assert span.request_id == recovered_span.request_id
|
|
assert span._trace_id == recovered_span._trace_id
|
|
assert span.span_id == recovered_span.span_id
|
|
assert span.name == recovered_span.name
|
|
assert span.span_type == recovered_span.span_type
|
|
assert span.start_time_ns == recovered_span.start_time_ns
|
|
assert span.end_time_ns == recovered_span.end_time_ns
|
|
assert span.parent_id == recovered_span.parent_id
|
|
assert span.status == recovered_span.status
|
|
assert span.inputs == recovered_span.inputs
|
|
assert span.outputs == recovered_span.outputs
|
|
assert span.attributes == recovered_span.attributes
|
|
assert span.events == recovered_span.events
|
|
|
|
# Loaded span should not implement setter methods
|
|
with pytest.raises(AttributeError, match="set_status"):
|
|
recovered_span.set_status("OK")
|
|
|
|
|
|
def test_from_v2_dict():
|
|
span = Span.from_dict({
|
|
"name": "test",
|
|
"context": {
|
|
"span_id": "8a90fc46e65ea5a4",
|
|
"trace_id": "0125978dc5c5a9456d7ca9ef1f7cf4af",
|
|
},
|
|
"parent_id": None,
|
|
"start_time": 1738662897576578992,
|
|
"end_time": 1738662899068969049,
|
|
"status_code": "OK",
|
|
"status_message": "",
|
|
"attributes": {
|
|
"mlflow.traceRequestId": '"tr-123"',
|
|
"mlflow.spanType": '"LLM"',
|
|
"mlflow.spanInputs": '{"input": 1}',
|
|
"mlflow.spanOutputs": "2",
|
|
"key": "3",
|
|
},
|
|
"events": [],
|
|
})
|
|
|
|
assert span.request_id == "tr-123"
|
|
assert span.name == "test"
|
|
assert span.span_type == SpanType.LLM
|
|
assert span.parent_id is None
|
|
assert span.status == SpanStatus(SpanStatusCode.OK, description="")
|
|
assert span.inputs == {"input": 1}
|
|
assert span.outputs == 2
|
|
assert span.events == []
|
|
|
|
|
|
def test_to_immutable_span():
|
|
trace_id = "tr-12345"
|
|
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("parent") as parent_span:
|
|
live_span = LiveSpan(parent_span, trace_id=trace_id, span_type=SpanType.LLM)
|
|
live_span.set_inputs({"input": 1})
|
|
live_span.set_outputs(2)
|
|
live_span.set_attribute("key", 3)
|
|
live_span.set_status("OK")
|
|
live_span.add_event(SpanEvent("test_event", timestamp=0, attributes={"foo": "bar"}))
|
|
|
|
span = live_span.to_immutable_span()
|
|
|
|
assert isinstance(span, Span)
|
|
assert span.trace_id == trace_id
|
|
assert span._trace_id == encode_trace_id(parent_span.context.trace_id)
|
|
assert span.span_id == encode_span_id(parent_span.context.span_id)
|
|
assert span.name == "parent"
|
|
assert span.start_time_ns == parent_span.start_time
|
|
assert span.end_time_ns is not None
|
|
assert span.parent_id is None
|
|
assert span.inputs == {"input": 1}
|
|
assert span.outputs == 2
|
|
assert span.get_attribute("key") == 3
|
|
assert span.status == SpanStatus(SpanStatusCode.OK, description="")
|
|
assert span.events == [SpanEvent("test_event", timestamp=0, attributes={"foo": "bar"})]
|
|
|
|
with pytest.raises(AttributeError, match="set_attribute"):
|
|
span.set_attribute("OK")
|
|
|
|
|
|
def test_from_dict_raises_when_trace_id_is_empty():
|
|
with pytest.raises(MlflowException, match=r"Failed to create a Span object from "):
|
|
Span.from_dict({
|
|
"name": "predict",
|
|
"context": {
|
|
"trace_id": "12345",
|
|
"span_id": "12345",
|
|
},
|
|
"parent_id": None,
|
|
"start_time": 0,
|
|
"end_time": 1,
|
|
"status_code": "OK",
|
|
"status_message": "",
|
|
"attributes": {
|
|
"mlflow.traceRequestId": None,
|
|
},
|
|
"events": [],
|
|
})
|
|
|
|
|
|
def test_set_attribute_directly_to_otel_span():
|
|
with mlflow.start_span("test") as span:
|
|
span._span.set_attribute("int", 1)
|
|
span._span.set_attribute("str", "a")
|
|
|
|
assert span.get_attribute("int") == 1
|
|
assert span.get_attribute("str") == "a"
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_otel_span_for_conversion():
|
|
"""Create a sample OTelReadableSpan for testing."""
|
|
return OTelReadableSpan(
|
|
name="test_span",
|
|
context=build_otel_context(
|
|
trace_id=0x12345678901234567890123456789012,
|
|
span_id=0x1234567890123456,
|
|
),
|
|
parent=build_otel_context(
|
|
trace_id=0x12345678901234567890123456789012,
|
|
span_id=0x0987654321098765,
|
|
),
|
|
start_time=1000000000,
|
|
end_time=2000000000,
|
|
attributes={
|
|
"mlflow.traceRequestId": "tr-12345678901234567890123456789012",
|
|
"mlflow.spanType": "LLM",
|
|
"mlflow.spanInputs": '{"prompt": "Hello"}',
|
|
"mlflow.spanOutputs": '{"response": "Hi"}',
|
|
"custom_attr": '{"key": "value"}',
|
|
},
|
|
status=OTelStatus(OTelStatusCode.OK, "Success"),
|
|
events=[
|
|
OTelEvent(
|
|
name="event1",
|
|
timestamp=1500000000,
|
|
attributes={"event_key": "event_value"},
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"attributes",
|
|
[
|
|
# Empty attributes
|
|
{},
|
|
# String attributes
|
|
{"str_key": "str_value"},
|
|
# Numeric attributes
|
|
{"int_key": 42, "float_key": 3.14},
|
|
# Boolean attributes
|
|
{"bool_key": True, "false_key": False},
|
|
# Bytes attributes
|
|
{"bytes_key": b"binary_data"},
|
|
# Empty list and dict attributes
|
|
{"empty_list": [], "empty_dict": {}},
|
|
# List attributes
|
|
{"list_str": ["a", "b", "c"], "list_int": [1, 2, 3], "list_float": [1.1, 2.2]},
|
|
# Dict attributes
|
|
{"dict_key": {"nested": "value", "number": 123}},
|
|
# Mixed complex attributes
|
|
{
|
|
"complex": {
|
|
"nested_list": [1, "two", 3.0],
|
|
"nested_dict": {"deep": {"deeper": "value"}},
|
|
},
|
|
"simple": "string",
|
|
},
|
|
],
|
|
)
|
|
def test_otel_attribute_conversion(attributes):
|
|
from opentelemetry.proto.common.v1.common_pb2 import KeyValue
|
|
|
|
from mlflow.tracing.utils.otlp import _decode_otel_proto_anyvalue, _set_otel_proto_anyvalue
|
|
|
|
# Convert attributes to proto format
|
|
proto_attrs = []
|
|
for key, value in attributes.items():
|
|
kv = KeyValue()
|
|
kv.key = key
|
|
_set_otel_proto_anyvalue(kv.value, value)
|
|
proto_attrs.append(kv)
|
|
|
|
# Decode back and verify
|
|
decoded = {}
|
|
for kv in proto_attrs:
|
|
decoded[kv.key] = _decode_otel_proto_anyvalue(kv.value)
|
|
|
|
assert decoded == attributes
|
|
|
|
|
|
def test_span_to_otel_proto_conversion(sample_otel_span_for_conversion):
|
|
# Create MLflow span from OTel span
|
|
mlflow_span = Span(sample_otel_span_for_conversion)
|
|
|
|
# Convert to OTel proto
|
|
otel_proto = mlflow_span.to_otel_proto()
|
|
|
|
# Verify basic fields
|
|
assert otel_proto.name == "test_span"
|
|
assert otel_proto.start_time_unix_nano == 1000000000
|
|
assert otel_proto.end_time_unix_nano == 2000000000
|
|
|
|
# Verify IDs (should be in bytes format)
|
|
assert len(otel_proto.trace_id) == 16 # 128-bit trace ID
|
|
assert len(otel_proto.span_id) == 8 # 64-bit span ID
|
|
assert len(otel_proto.parent_span_id) == 8
|
|
|
|
# Verify status
|
|
assert otel_proto.status.code == OTelProtoStatus.STATUS_CODE_OK
|
|
# OTel SDK clears description for non-ERROR statuses
|
|
assert otel_proto.status.message == ""
|
|
|
|
# Verify attributes exist
|
|
assert len(otel_proto.attributes) == 5
|
|
attr_keys = {attr.key for attr in otel_proto.attributes}
|
|
assert "mlflow.spanType" in attr_keys
|
|
assert "custom_attr" in attr_keys
|
|
|
|
# Verify events
|
|
assert len(otel_proto.events) == 1
|
|
assert otel_proto.events[0].name == "event1"
|
|
assert otel_proto.events[0].time_unix_nano == 1500000000
|
|
|
|
|
|
def test_span_from_otel_proto_conversion():
|
|
# Create OTel proto span
|
|
otel_proto = OTelProtoSpan()
|
|
otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012")
|
|
otel_proto.span_id = bytes.fromhex("1234567890123456")
|
|
otel_proto.parent_span_id = bytes.fromhex("0987654321098765")
|
|
otel_proto.name = "proto_span"
|
|
otel_proto.start_time_unix_nano = 1000000000
|
|
otel_proto.end_time_unix_nano = 2000000000
|
|
|
|
# Add status
|
|
otel_proto.status.code = OTelProtoStatus.STATUS_CODE_ERROR
|
|
otel_proto.status.message = "Error occurred"
|
|
|
|
# Add attributes
|
|
from mlflow.tracing.utils.otlp import _set_otel_proto_anyvalue
|
|
|
|
attr2 = otel_proto.attributes.add()
|
|
attr2.key = "mlflow.spanType"
|
|
_set_otel_proto_anyvalue(attr2.value, "CHAIN")
|
|
|
|
attr3 = otel_proto.attributes.add()
|
|
attr3.key = "custom"
|
|
_set_otel_proto_anyvalue(attr3.value, {"nested": {"value": 123}})
|
|
|
|
# Add event
|
|
event = otel_proto.events.add()
|
|
event.name = "test_event"
|
|
event.time_unix_nano = 1500000000
|
|
event_attr = event.attributes.add()
|
|
event_attr.key = "event_data"
|
|
_set_otel_proto_anyvalue(event_attr.value, "event_value")
|
|
|
|
# Add link
|
|
proto_link = otel_proto.links.add()
|
|
proto_link.trace_id = bytes.fromhex("aabbccddeeff00112233445566778899")
|
|
proto_link.span_id = bytes.fromhex("1122334455667788")
|
|
link_attr = proto_link.attributes.add()
|
|
link_attr.key = "rel"
|
|
_set_otel_proto_anyvalue(link_attr.value, "causal")
|
|
|
|
# Convert to MLflow span
|
|
mlflow_span = Span.from_otel_proto(otel_proto)
|
|
|
|
# Verify basic fields
|
|
assert mlflow_span.name == "proto_span"
|
|
assert mlflow_span.start_time_ns == 1000000000
|
|
assert mlflow_span.end_time_ns == 2000000000
|
|
|
|
# Verify IDs
|
|
assert mlflow_span.trace_id == "tr-12345678901234567890123456789012"
|
|
assert mlflow_span.span_id == "1234567890123456"
|
|
assert mlflow_span.parent_id == "0987654321098765"
|
|
|
|
# Verify status
|
|
assert mlflow_span.status.status_code == SpanStatusCode.ERROR
|
|
assert mlflow_span.status.description == "Error occurred"
|
|
|
|
# Verify attributes
|
|
assert mlflow_span.span_type == "CHAIN"
|
|
assert mlflow_span.get_attribute("custom") == {"nested": {"value": 123}}
|
|
|
|
# Verify events
|
|
assert len(mlflow_span.events) == 1
|
|
assert mlflow_span.events[0].name == "test_event"
|
|
assert mlflow_span.events[0].timestamp == 1500000000
|
|
assert mlflow_span.events[0].attributes["event_data"] == "event_value"
|
|
|
|
# Verify resource is empty when none is provided
|
|
assert len(mlflow_span._span.resource.attributes) == 0
|
|
|
|
# Verify links use v3 format when no location_id
|
|
assert len(mlflow_span.links) == 1
|
|
assert mlflow_span.links[0].trace_id == "tr-aabbccddeeff00112233445566778899"
|
|
assert mlflow_span.links[0].span_id == "1122334455667788"
|
|
assert mlflow_span.links[0].attributes == {"rel": "causal"}
|
|
|
|
|
|
def test_span_from_otel_proto_with_resource():
|
|
otel_proto = OTelProtoSpan()
|
|
otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012")
|
|
otel_proto.span_id = bytes.fromhex("1234567890123456")
|
|
otel_proto.name = "span_with_resource"
|
|
otel_proto.start_time_unix_nano = 1000000000
|
|
otel_proto.end_time_unix_nano = 2000000000
|
|
|
|
# Build a proto Resource with typical OTel resource attributes
|
|
resource = OTelProtoResource()
|
|
for key, value in [
|
|
("service.name", "my-app"),
|
|
("telemetry.sdk.language", "java"),
|
|
("telemetry.sdk.name", "opentelemetry"),
|
|
("telemetry.sdk.version", "1.40.0"),
|
|
]:
|
|
attr = resource.attributes.add()
|
|
attr.key = key
|
|
_set_otel_proto_anyvalue(attr.value, value)
|
|
|
|
mlflow_span = Span.from_otel_proto(otel_proto, resource=resource)
|
|
|
|
assert mlflow_span.name == "span_with_resource"
|
|
res_attrs = dict(mlflow_span._span.resource.attributes)
|
|
assert res_attrs["service.name"] == "my-app"
|
|
assert res_attrs["telemetry.sdk.language"] == "java"
|
|
assert res_attrs["telemetry.sdk.name"] == "opentelemetry"
|
|
assert res_attrs["telemetry.sdk.version"] == "1.40.0"
|
|
|
|
|
|
def test_span_from_otel_proto_with_empty_resource():
|
|
otel_proto = OTelProtoSpan()
|
|
otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012")
|
|
otel_proto.span_id = bytes.fromhex("1234567890123456")
|
|
otel_proto.name = "span_empty_resource"
|
|
otel_proto.start_time_unix_nano = 1000000000
|
|
otel_proto.end_time_unix_nano = 2000000000
|
|
|
|
# An empty proto Resource (no attributes)
|
|
resource = OTelProtoResource()
|
|
mlflow_span = Span.from_otel_proto(otel_proto, resource=resource)
|
|
|
|
assert len(mlflow_span._span.resource.attributes) == 0
|
|
|
|
|
|
def test_span_from_otel_proto_with_location():
|
|
# Create OTel proto span
|
|
otel_proto = OTelProtoSpan()
|
|
otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012")
|
|
otel_proto.span_id = bytes.fromhex("1234567890123456")
|
|
otel_proto.parent_span_id = bytes.fromhex("0987654321098765")
|
|
otel_proto.name = "proto_span_v4"
|
|
otel_proto.start_time_unix_nano = 1000000000
|
|
otel_proto.end_time_unix_nano = 2000000000
|
|
|
|
# Add status
|
|
otel_proto.status.code = OTelProtoStatus.STATUS_CODE_OK
|
|
otel_proto.status.message = ""
|
|
|
|
# Add attributes
|
|
attr1 = otel_proto.attributes.add()
|
|
attr1.key = "mlflow.spanType"
|
|
_set_otel_proto_anyvalue(attr1.value, "LLM")
|
|
|
|
# Add link
|
|
proto_link = otel_proto.links.add()
|
|
proto_link.trace_id = bytes.fromhex("aabbccddeeff00112233445566778899")
|
|
proto_link.span_id = bytes.fromhex("1122334455667788")
|
|
|
|
# Convert to MLflow span with location
|
|
location = "catalog.schema"
|
|
mlflow_span = Span.from_otel_proto(otel_proto, location_id=location)
|
|
|
|
# Verify basic fields
|
|
assert mlflow_span.name == "proto_span_v4"
|
|
assert mlflow_span.start_time_ns == 1000000000
|
|
assert mlflow_span.end_time_ns == 2000000000
|
|
|
|
# Verify IDs
|
|
assert mlflow_span.span_id == "1234567890123456"
|
|
assert mlflow_span.parent_id == "0987654321098765"
|
|
|
|
# Verify trace_id is in v4 format with location
|
|
expected_trace_id = f"{TRACE_ID_V4_PREFIX}{location}/12345678901234567890123456789012"
|
|
assert mlflow_span.trace_id == expected_trace_id
|
|
|
|
# Verify the REQUEST_ID attribute also uses v4 format
|
|
request_id = mlflow_span.get_attribute("mlflow.traceRequestId")
|
|
assert request_id == expected_trace_id
|
|
|
|
# Links are skipped for v4 UC traces
|
|
assert len(mlflow_span.links) == 0
|
|
|
|
|
|
def test_span_from_otel_proto_with_pre_encoded_request_id():
|
|
"""
|
|
Regression test: when the MLflow SDK (OtelSpanProcessor.on_start) runs before OTLP export,
|
|
it JSON-encodes mlflow.traceRequestId into the live OTel span via LiveSpan.__init__.
|
|
The OTLP proto then carries the already-encoded value (e.g. '"tr-abc"' with inner quotes).
|
|
from_otel_proto() must not double-encode it — the server-computed ID must always win.
|
|
"""
|
|
otel_proto = OTelProtoSpan()
|
|
otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012")
|
|
otel_proto.span_id = bytes.fromhex("1234567890123456")
|
|
otel_proto.name = "langchain_span"
|
|
otel_proto.start_time_unix_nano = 1000000000
|
|
otel_proto.end_time_unix_nano = 2000000000
|
|
otel_proto.status.code = OTelProtoStatus.STATUS_CODE_OK
|
|
|
|
# Simulate what OtelSpanProcessor.on_start + LiveSpan.__init__ produce:
|
|
# _SpanAttributesRegistry.set() calls dump_span_attribute_value("tr-abc") = '"tr-abc"'
|
|
# so the OTel span attribute stored in memory is the JSON-encoded string '"tr-abc"'.
|
|
expected_trace_id = "tr-12345678901234567890123456789012"
|
|
already_encoded_id = json.dumps(expected_trace_id) # '"tr-12345678901234567890123456789012"'
|
|
|
|
attr = otel_proto.attributes.add()
|
|
attr.key = "mlflow.traceRequestId"
|
|
_set_otel_proto_anyvalue(attr.value, already_encoded_id)
|
|
|
|
mlflow_span = Span.from_otel_proto(otel_proto)
|
|
|
|
# Server-computed ID must win — must start with 'tr-', not with '"'
|
|
assert mlflow_span.trace_id == expected_trace_id
|
|
assert mlflow_span.trace_id.startswith("tr-")
|
|
|
|
|
|
def test_span_from_otel_proto_can_preserve_request_id_for_round_trip():
|
|
otel_proto = OTelProtoSpan()
|
|
otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012")
|
|
otel_proto.span_id = bytes.fromhex("1234567890123456")
|
|
otel_proto.name = "archived_span"
|
|
otel_proto.start_time_unix_nano = 1000000000
|
|
otel_proto.end_time_unix_nano = 2000000000
|
|
otel_proto.status.code = OTelProtoStatus.STATUS_CODE_OK
|
|
|
|
attr = otel_proto.attributes.add()
|
|
attr.key = "mlflow.traceRequestId"
|
|
_set_otel_proto_anyvalue(attr.value, "tr-abc123")
|
|
|
|
mlflow_span = Span.from_otel_proto(otel_proto, preserve_request_id=True)
|
|
|
|
assert mlflow_span.trace_id == "tr-abc123"
|
|
|
|
|
|
def test_otel_roundtrip_conversion(sample_otel_span_for_conversion):
|
|
# Start with OTel span -> MLflow span
|
|
mlflow_span = Span(sample_otel_span_for_conversion)
|
|
|
|
# Convert to OTel proto
|
|
otel_proto = mlflow_span.to_otel_proto()
|
|
|
|
# Convert back to MLflow span
|
|
roundtrip_span = Span.from_otel_proto(otel_proto)
|
|
|
|
# Verify key fields are preserved
|
|
assert roundtrip_span.name == mlflow_span.name
|
|
assert roundtrip_span.span_id == mlflow_span.span_id
|
|
assert roundtrip_span.parent_id == mlflow_span.parent_id
|
|
assert roundtrip_span.start_time_ns == mlflow_span.start_time_ns
|
|
assert roundtrip_span.end_time_ns == mlflow_span.end_time_ns
|
|
assert roundtrip_span.status.status_code == mlflow_span.status.status_code
|
|
assert roundtrip_span.status.description == mlflow_span.status.description
|
|
|
|
# Verify span attributes are preserved
|
|
assert roundtrip_span.span_type == mlflow_span.span_type
|
|
assert roundtrip_span.inputs == mlflow_span.inputs
|
|
assert roundtrip_span.outputs == mlflow_span.outputs
|
|
|
|
# Verify ALL attributes are preserved by iterating through them
|
|
# Get all attribute keys from both spans
|
|
original_attributes = mlflow_span.attributes
|
|
roundtrip_attributes = roundtrip_span.attributes
|
|
|
|
# Check we have the same number of attributes
|
|
assert len(original_attributes) == len(roundtrip_attributes)
|
|
|
|
# Check each attribute is preserved correctly
|
|
for attr_key in original_attributes:
|
|
assert attr_key in roundtrip_attributes, f"Attribute {attr_key} missing after roundtrip"
|
|
original_value = original_attributes[attr_key]
|
|
roundtrip_value = roundtrip_attributes[attr_key]
|
|
assert original_value == roundtrip_value, (
|
|
f"Attribute {attr_key} changed: {original_value} != {roundtrip_value}"
|
|
)
|
|
|
|
# Also explicitly verify specific important attributes
|
|
# The original span has a custom_attr that should be preserved
|
|
original_custom_attr = mlflow_span.get_attribute("custom_attr")
|
|
roundtrip_custom_attr = roundtrip_span.get_attribute("custom_attr")
|
|
assert original_custom_attr == roundtrip_custom_attr
|
|
assert original_custom_attr == {"key": "value"}
|
|
|
|
# Verify the trace request ID is preserved
|
|
assert roundtrip_span.request_id == mlflow_span.request_id
|
|
assert roundtrip_span.request_id == "tr-12345678901234567890123456789012"
|
|
|
|
# Verify events
|
|
assert len(roundtrip_span.events) == len(mlflow_span.events)
|
|
for orig_event, rt_event in zip(mlflow_span.events, roundtrip_span.events):
|
|
assert rt_event.name == orig_event.name
|
|
assert rt_event.timestamp == orig_event.timestamp
|
|
assert rt_event.attributes == orig_event.attributes
|
|
|
|
|
|
def test_resource_to_otel_proto():
|
|
resource = OTelResource.create({
|
|
"service.name": "test-service",
|
|
"service.version": "1.0.0",
|
|
"custom.int": 42,
|
|
"custom.bool": True,
|
|
})
|
|
resource_proto = resource_to_otel_proto(resource)
|
|
|
|
# Convert proto attributes to dict for easier verification
|
|
attrs = {}
|
|
for attr in resource_proto.attributes:
|
|
attrs[attr.key] = _decode_otel_proto_anyvalue(attr.value)
|
|
|
|
assert attrs["service.name"] == "test-service"
|
|
assert attrs["service.version"] == "1.0.0"
|
|
assert attrs["custom.int"] == 42
|
|
assert attrs["custom.bool"] is True
|
|
|
|
|
|
def test_resource_to_otel_proto_none():
|
|
resource_proto = resource_to_otel_proto(None)
|
|
assert len(resource_proto.attributes) == 0
|
|
|
|
|
|
def test_span_from_dict_old_format():
|
|
span_dict = {
|
|
"trace_id": "6ST7JNq8BC4JRp0HA/vD6Q==",
|
|
"span_id": "Sd/l0Zs4M3g=",
|
|
"parent_span_id": None,
|
|
"name": "test_span",
|
|
"start_time_unix_nano": 1000000000,
|
|
"end_time_unix_nano": 2000000000,
|
|
"status": {"code": "STATUS_CODE_ERROR", "message": "Error occurred"},
|
|
"attributes": {
|
|
"mlflow.spanInputs": '{"query": "test"}',
|
|
"mlflow.spanOutputs": '{"result": "success"}',
|
|
"custom": "value",
|
|
"mlflow.traceRequestId": '"tr-e924fb24dabc042e09469d0703fbc3e9"',
|
|
},
|
|
"events": [],
|
|
}
|
|
|
|
# Deserialize it
|
|
recovered_span = Span.from_dict(span_dict)
|
|
|
|
# Verify all fields are recovered correctly
|
|
assert recovered_span.trace_id == "tr-e924fb24dabc042e09469d0703fbc3e9"
|
|
assert recovered_span.span_id == "49dfe5d19b383378"
|
|
assert recovered_span.parent_id is None
|
|
assert recovered_span.name == span_dict["name"]
|
|
assert recovered_span.start_time_ns == span_dict["start_time_unix_nano"]
|
|
assert recovered_span.end_time_ns == span_dict["end_time_unix_nano"]
|
|
assert recovered_span.status.status_code.value == "ERROR"
|
|
assert recovered_span.inputs == {"query": "test"}
|
|
assert recovered_span.outputs == {"result": "success"}
|
|
assert recovered_span.get_attribute("custom") == "value"
|
|
|
|
|
|
def test_span_dict_v4_with_no_parent():
|
|
with mlflow.start_span("root_span") as span:
|
|
span.set_inputs({"x": 1})
|
|
span.set_outputs({"y": 2})
|
|
|
|
span_dict = span.to_dict()
|
|
|
|
# Root span should have None for parent_span_id
|
|
assert span_dict["parent_span_id"] is None
|
|
|
|
# Deserialize and verify
|
|
recovered = Span.from_dict(span_dict)
|
|
assert recovered.parent_id is None
|
|
assert recovered.name == "root_span"
|
|
assert recovered.inputs == {"x": 1}
|
|
assert recovered.outputs == {"y": 2}
|
|
|
|
|
|
def test_span_from_dict_supports_both_status_code_formats():
|
|
with mlflow.start_span("test") as span:
|
|
span.set_status("OK")
|
|
|
|
span_dict = span.to_dict()
|
|
|
|
# Current code serializes as protobuf enum name
|
|
assert span_dict["status"]["code"] == "STATUS_CODE_OK"
|
|
|
|
# Verify we can deserialize protobuf enum name format (backward compatibility)
|
|
span_dict["status"]["code"] = "STATUS_CODE_ERROR"
|
|
recovered = Span.from_dict(span_dict)
|
|
assert recovered.status.status_code == SpanStatusCode.ERROR
|
|
|
|
# Verify we can also deserialize enum value format
|
|
# (forward compatibility with older serialized data)
|
|
span_dict["status"]["code"] = "OK"
|
|
recovered = Span.from_dict(span_dict)
|
|
assert recovered.status.status_code == SpanStatusCode.OK
|
|
|
|
span_dict["status"]["code"] = "UNSET"
|
|
recovered = Span.from_dict(span_dict)
|
|
assert recovered.status.status_code == SpanStatusCode.UNSET
|
|
|
|
span_dict["status"]["code"] = "ERROR"
|
|
recovered = Span.from_dict(span_dict)
|
|
assert recovered.status.status_code == SpanStatusCode.ERROR
|
|
|
|
|
|
def test_load_from_old_span_dict():
|
|
span_dict = {
|
|
"trace_id": "ZqqBulxlq2cwRCfKHxmDVA==",
|
|
"span_id": "/mCYZRxbTqw=",
|
|
"trace_state": "",
|
|
"parent_span_id": "f6qlKYqTw2E=",
|
|
"name": "custom",
|
|
"start_time_unix_nano": 1761103703884225000,
|
|
"end_time_unix_nano": 1761103703884454000,
|
|
"attributes": {
|
|
"mlflow.spanOutputs": "4",
|
|
"mlflow.spanType": '"LLM"',
|
|
"mlflow.spanInputs": '{"z": 3}',
|
|
"mlflow.traceRequestId": '"tr-66aa81ba5c65ab67304427ca1f198354"',
|
|
"delta": "1",
|
|
"mlflow.spanFunctionName": '"add_one"',
|
|
},
|
|
"status": {"message": "", "code": "STATUS_CODE_OK"},
|
|
"events": [
|
|
{
|
|
"time_unix_nano": 1761105506649041,
|
|
"name": "agent_action",
|
|
"attributes": {
|
|
"tool": "search_web",
|
|
"tool_input": '"What is MLflow?"',
|
|
"log": "test",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
span = Span.from_dict(span_dict)
|
|
assert span.trace_id == "tr-66aa81ba5c65ab67304427ca1f198354"
|
|
assert span.span_id == "fe6098651c5b4eac"
|
|
assert span.parent_id == "7faaa5298a93c361"
|
|
assert span.name == "custom"
|
|
assert span.start_time_ns == 1761103703884225000
|
|
assert span.end_time_ns == 1761103703884454000
|
|
assert span.status == SpanStatus(SpanStatusCode.OK, description="")
|
|
assert span.inputs == {"z": 3}
|
|
assert span.outputs == 4
|
|
assert len(span.events) == 1
|
|
assert span.events[0].name == "agent_action"
|
|
assert span.events[0].timestamp == 1761105506649041
|
|
assert span.events[0].attributes == {
|
|
"tool": "search_web",
|
|
"tool_input": '"What is MLflow?"',
|
|
"log": "test",
|
|
}
|
|
|
|
|
|
def test_load_from_3_5_0_span_dict():
|
|
span_dict = {
|
|
"trace_id": "tr-66aa81ba5c65ab67304427ca1f198354",
|
|
"span_id": "fe6098651c5b4eac",
|
|
"trace_state": "",
|
|
"parent_span_id": "7faaa5298a93c361",
|
|
"name": "custom",
|
|
"start_time_unix_nano": 1761103703884225000,
|
|
"end_time_unix_nano": 1761103703884454000,
|
|
"attributes": {
|
|
"mlflow.spanOutputs": "4",
|
|
"mlflow.spanType": '"LLM"',
|
|
"mlflow.spanInputs": '{"z": 3}',
|
|
"mlflow.traceRequestId": '"tr-66aa81ba5c65ab67304427ca1f198354"',
|
|
"delta": "1",
|
|
"mlflow.spanFunctionName": '"add_one"',
|
|
},
|
|
"status": {"message": "", "code": "OK"},
|
|
"events": [
|
|
{
|
|
"time_unix_nano": 1761105506649041,
|
|
"name": "agent_action",
|
|
"attributes": {
|
|
"tool": "search_web",
|
|
"tool_input": '"What is MLflow?"',
|
|
"log": "test",
|
|
},
|
|
}
|
|
],
|
|
}
|
|
span = Span.from_dict(span_dict)
|
|
assert span.trace_id == "tr-66aa81ba5c65ab67304427ca1f198354"
|
|
assert span.span_id == "fe6098651c5b4eac"
|
|
assert span.parent_id == "7faaa5298a93c361"
|
|
assert span.name == "custom"
|
|
assert span.start_time_ns == 1761103703884225000
|
|
assert span.end_time_ns == 1761103703884454000
|
|
assert span.status == SpanStatus(SpanStatusCode.OK, description="")
|
|
assert span.inputs == {"z": 3}
|
|
assert span.outputs == 4
|
|
assert len(span.events) == 1
|
|
assert span.events[0].name == "agent_action"
|
|
assert span.events[0].timestamp == 1761105506649041
|
|
assert span.events[0].attributes == {
|
|
"tool": "search_web",
|
|
"tool_input": '"What is MLflow?"',
|
|
"log": "test",
|
|
}
|
|
|
|
|
|
def test_live_span_add_link():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
|
|
link = Link(
|
|
trace_id="tr-abc123",
|
|
span_id="aabbccddeeff0011",
|
|
attributes={"relationship": "triggered_by"},
|
|
)
|
|
span.add_link(link)
|
|
|
|
assert len(span.links) == 1
|
|
assert span.links[0].trace_id == "tr-abc123"
|
|
assert span.links[0].span_id == "aabbccddeeff0011"
|
|
assert span.links[0].attributes == {"relationship": "triggered_by"}
|
|
|
|
# Verify the link was forwarded to the underlying OTel span
|
|
assert len(otel_span.links) == 1
|
|
assert otel_span.links[0].context.trace_id == 0xABC123
|
|
assert otel_span.links[0].context.span_id == 0xAABBCCDDEEFF0011
|
|
|
|
|
|
def test_add_link_after_end_does_not_store_link():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
otel_span = tracer.start_span("test_span")
|
|
span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
span.end()
|
|
|
|
with mock.patch("mlflow.entities.span._logger.debug") as mock_debug:
|
|
span.add_link(
|
|
Link(
|
|
trace_id="tr-abc123",
|
|
span_id="aabbccddeeff0011",
|
|
attributes={"relationship": "triggered_by"},
|
|
)
|
|
)
|
|
|
|
assert len(span.links) == 0
|
|
assert len(otel_span.links) == 0
|
|
mock_debug.assert_called_once_with(
|
|
"Skipping link addition because the span is no longer recording."
|
|
)
|
|
|
|
|
|
def test_add_link_rejects_invalid_ids():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
|
|
with pytest.raises(MlflowException, match="Invalid link"):
|
|
span.add_link(Link(trace_id="tr-abc123", span_id="not-hex"))
|
|
|
|
with pytest.raises(MlflowException, match="Invalid link"):
|
|
span.add_link(Link(trace_id="bad-format", span_id="aabbccddeeff0011"))
|
|
|
|
with pytest.raises(MlflowException, match="Invalid link"):
|
|
span.add_link(Link(trace_id="tr-abc123", span_id="aabbccddeeff0011aabbccddeeff0011"))
|
|
|
|
assert len(span.links) == 0
|
|
|
|
|
|
def test_span_seeds_links_from_otel_span():
|
|
trace_id = "tr-12345"
|
|
otel_link = trace_api.Link(
|
|
context=trace_api.SpanContext(
|
|
trace_id=0xAABBCCDDEEFF00112233445566778899,
|
|
span_id=0xAABBCCDDEEFF0011,
|
|
is_remote=False,
|
|
trace_flags=trace_api.TraceFlags(1),
|
|
),
|
|
attributes={"type": "causal"},
|
|
)
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span", links=[otel_link]) as otel_span:
|
|
span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
|
|
immutable_span = span.to_immutable_span()
|
|
assert len(immutable_span.links) == 1
|
|
assert immutable_span.links[0].trace_id == "tr-aabbccddeeff00112233445566778899"
|
|
assert immutable_span.links[0].span_id == "aabbccddeeff0011"
|
|
assert immutable_span.links[0].attributes == {"type": "causal"}
|
|
|
|
|
|
def test_span_to_dict_with_links():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
live_span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
link = Link(trace_id="tr-abc123", span_id="aabbccddeeff0011", attributes={"type": "test"})
|
|
live_span.add_link(link)
|
|
|
|
immutable_span = live_span.to_immutable_span()
|
|
span_dict = immutable_span.to_dict()
|
|
|
|
assert "links" in span_dict
|
|
assert len(span_dict["links"]) == 1
|
|
assert span_dict["links"][0]["trace_id"] == "tr-abc123"
|
|
assert span_dict["links"][0]["span_id"] == "aabbccddeeff0011"
|
|
assert span_dict["links"][0]["attributes"] == {"type": "test"}
|
|
|
|
|
|
def test_span_from_dict_with_links():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
live_span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
link = Link(trace_id="tr-abc123", span_id="aabbccddeeff0011", attributes={"type": "test"})
|
|
live_span.add_link(link)
|
|
|
|
immutable_span = live_span.to_immutable_span()
|
|
span_dict = immutable_span.to_dict()
|
|
|
|
reconstructed_span = Span.from_dict(span_dict)
|
|
assert len(reconstructed_span.links) == 1
|
|
assert reconstructed_span.links[0].trace_id == "tr-abc123"
|
|
assert reconstructed_span.links[0].span_id == "aabbccddeeff0011"
|
|
assert reconstructed_span.links[0].attributes == {"type": "test"}
|
|
|
|
|
|
def test_span_from_dict_without_links():
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
|
|
immutable_span = span.to_immutable_span()
|
|
span_dict = immutable_span.to_dict()
|
|
|
|
assert span_dict["links"] == []
|
|
|
|
reconstructed_span = Span.from_dict(span_dict)
|
|
assert len(reconstructed_span.links) == 0
|
|
|
|
|
|
def test_span_links_otel_proto_roundtrip():
|
|
otel_proto = OTelProtoSpan()
|
|
otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012")
|
|
otel_proto.span_id = bytes.fromhex("1234567890123456")
|
|
otel_proto.name = "test"
|
|
|
|
link_trace_bytes = bytes.fromhex("aabbccddeeff00112233445566778899")
|
|
link_span_bytes = bytes.fromhex("aabbccddeeff0011")
|
|
proto_link = otel_proto.links.add()
|
|
proto_link.trace_id = link_trace_bytes
|
|
proto_link.span_id = link_span_bytes
|
|
link_attr = proto_link.attributes.add()
|
|
link_attr.key = "type"
|
|
_set_otel_proto_anyvalue(link_attr.value, "causal")
|
|
|
|
mlflow_span = Span.from_otel_proto(otel_proto)
|
|
assert len(mlflow_span.links) == 1
|
|
assert mlflow_span.links[0].attributes == {"type": "causal"}
|
|
|
|
rt = mlflow_span.to_otel_proto()
|
|
assert rt.links[0].trace_id == link_trace_bytes
|
|
assert rt.links[0].span_id == link_span_bytes
|
|
|
|
|
|
def test_span_skips_links_for_v4_traces():
|
|
otel_link = trace_api.Link(
|
|
context=trace_api.SpanContext(
|
|
trace_id=0xAABBCCDDEEFF00112233445566778899,
|
|
span_id=0xAABBCCDDEEFF0011,
|
|
is_remote=False,
|
|
trace_flags=trace_api.TraceFlags(1),
|
|
),
|
|
attributes={"type": "causal"},
|
|
)
|
|
v4_trace_id = f"{TRACE_ID_V4_PREFIX}catalog.schema/12345678901234567890123456789012"
|
|
otel_span = OTelReadableSpan(
|
|
name="test_v4",
|
|
context=build_otel_context(0x12345678901234567890123456789012, 0x1234567890123456),
|
|
attributes={"mlflow.traceRequestId": json.dumps(v4_trace_id)},
|
|
links=[otel_link],
|
|
resource=OTelResource.get_empty(),
|
|
)
|
|
|
|
span = Span(otel_span)
|
|
assert span.trace_id == v4_trace_id
|
|
assert len(span.links) == 0
|
|
|
|
|
|
def test_span_links_property_returns_deep_copy():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
live_span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
live_span.add_link(
|
|
Link(trace_id="tr-abc", span_id="aabbccddeeff0011", attributes={"k": "v"})
|
|
)
|
|
|
|
immutable_span = live_span.to_immutable_span()
|
|
links = immutable_span.links
|
|
links[0].attributes["k"] = "mutated"
|
|
|
|
assert immutable_span.links[0].attributes["k"] == "v"
|
|
|
|
|
|
def test_add_link_clones_on_insert():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
live_span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
link = Link(trace_id="tr-abc", span_id="aabbccddeeff0011", attributes={"k": "v"})
|
|
live_span.add_link(link)
|
|
link.attributes["k"] = "mutated"
|
|
|
|
assert live_span.links[0].attributes["k"] == "v"
|
|
|
|
|
|
def test_to_immutable_span_deep_copies_links():
|
|
from mlflow.entities.link import Link
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("test_span") as otel_span:
|
|
live_span = create_mlflow_span(otel_span, trace_id=trace_id)
|
|
live_span.add_link(
|
|
Link(trace_id="tr-abc", span_id="aabbccddeeff0011", attributes={"k": "v"})
|
|
)
|
|
|
|
immutable_span = live_span.to_immutable_span()
|
|
live_span._links[0].attributes["k"] = "mutated"
|
|
|
|
assert immutable_span.links[0].attributes["k"] == "v"
|
|
|
|
|
|
def test_lazy_span_to_dict_does_not_materialize():
|
|
from mlflow.entities.span import LazySpan
|
|
|
|
with mlflow.start_span("parent"):
|
|
with mlflow.start_span("child", span_type=SpanType.LLM) as span:
|
|
span.set_inputs({"input": 1})
|
|
span.set_outputs(2)
|
|
|
|
span_dict = span.to_dict()
|
|
lazy = LazySpan(span_dict)
|
|
|
|
assert isinstance(lazy, Span)
|
|
assert not isinstance(lazy, LiveSpan)
|
|
assert lazy.to_dict() is span_dict
|
|
assert lazy.__dict__["_materialized"] is False
|
|
|
|
|
|
def test_lazy_span_materializes_on_property_access():
|
|
from mlflow.entities.span import LazySpan
|
|
|
|
with mlflow.start_span("parent"):
|
|
with mlflow.start_span("child", span_type=SpanType.LLM) as span:
|
|
span.set_inputs({"input": 1})
|
|
span.set_outputs(2)
|
|
span.set_status("OK")
|
|
|
|
span_dict = span.to_dict()
|
|
lazy = LazySpan(span_dict)
|
|
|
|
assert lazy.name == "child"
|
|
assert lazy.__dict__["_materialized"] is True
|
|
assert lazy.inputs == {"input": 1}
|
|
assert lazy.outputs == 2
|
|
assert lazy.span_type == SpanType.LLM
|
|
assert lazy.to_dict() == span_dict
|
|
|
|
|
|
def test_lazy_span_round_trips_through_trace_data_to_dict():
|
|
from mlflow.entities.span import LazySpan
|
|
from mlflow.entities.trace_data import TraceData
|
|
|
|
with mlflow.start_span("parent"):
|
|
with mlflow.start_span("child", span_type=SpanType.LLM) as span:
|
|
span.set_inputs({"input": 1})
|
|
|
|
lazy = LazySpan(span.to_dict())
|
|
trace_data = TraceData(spans=[lazy])
|
|
|
|
dumped = trace_data.to_dict()
|
|
assert dumped["spans"][0]["name"] == "child"
|
|
assert lazy.__dict__["_materialized"] is False
|
|
|
|
|
|
def test_lazy_span_matches_eager_span_after_materialization():
|
|
from mlflow.entities.link import Link
|
|
from mlflow.entities.span import LazySpan
|
|
from mlflow.tracing.otel.translation import translate_loaded_span
|
|
|
|
trace_id = "tr-12345"
|
|
tracer = _get_tracer("test")
|
|
with tracer.start_as_current_span("parent"):
|
|
with tracer.start_as_current_span("child") as otel_span:
|
|
span = create_mlflow_span(otel_span, trace_id=trace_id, span_type=SpanType.LLM)
|
|
span.set_inputs({"input": 1})
|
|
span.set_outputs({"output": 2})
|
|
span.set_status("OK")
|
|
span.add_event(SpanEvent("test_event", timestamp=0, attributes={"foo": "bar"}))
|
|
span.add_link(
|
|
Link(trace_id="tr-abc123", span_id="aabbccddeeff0011", attributes={"type": "test"})
|
|
)
|
|
|
|
span_dict = translate_loaded_span(span.to_dict())
|
|
eager = Span.from_dict(span_dict)
|
|
lazy = LazySpan(span_dict)
|
|
|
|
assert lazy.to_dict() == eager.to_dict()
|
|
assert lazy.to_otel_proto().SerializeToString() == eager.to_otel_proto().SerializeToString()
|
|
|
|
_ = lazy.name
|
|
assert lazy.__dict__["_materialized"] is True
|
|
assert lazy.to_dict() == eager.to_dict()
|
|
assert lazy.attributes == eager.attributes
|
|
assert lazy.links == eager.links
|
|
assert lazy.status == eager.status
|
|
assert lazy.events == eager.events
|
|
assert lazy.to_otel_proto().SerializeToString() == eager.to_otel_proto().SerializeToString()
|