test_span.py
1 import json 2 from datetime import datetime 3 4 import opentelemetry.trace as trace_api 5 import pytest 6 from opentelemetry.proto.trace.v1.trace_pb2 import Span as OTelProtoSpan 7 from opentelemetry.proto.trace.v1.trace_pb2 import Status as OTelProtoStatus 8 from opentelemetry.sdk.resources import Resource as OTelResource 9 from opentelemetry.sdk.trace import Event as OTelEvent 10 from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan 11 from opentelemetry.trace import Status as OTelStatus 12 from opentelemetry.trace import StatusCode as OTelStatusCode 13 14 import mlflow 15 from mlflow.entities import LiveSpan, Span, SpanEvent, SpanStatus, SpanStatusCode, SpanType 16 from mlflow.entities.span import ( 17 NoOpSpan, 18 create_mlflow_span, 19 ) 20 from mlflow.exceptions import MlflowException 21 from mlflow.tracing.constant import TRACE_ID_V4_PREFIX 22 from mlflow.tracing.provider import _get_tracer, trace_disabled 23 from mlflow.tracing.utils import build_otel_context, encode_span_id, encode_trace_id 24 from mlflow.tracing.utils.otlp import ( 25 _decode_otel_proto_anyvalue, 26 _set_otel_proto_anyvalue, 27 resource_to_otel_proto, 28 ) 29 30 31 def test_create_live_span(): 32 trace_id = "tr-12345" 33 34 tracer = _get_tracer("test") 35 with tracer.start_as_current_span("parent") as parent_span: 36 span = create_mlflow_span(parent_span, trace_id=trace_id, span_type=SpanType.LLM) 37 assert isinstance(span, LiveSpan) 38 assert span.trace_id == trace_id 39 assert span._trace_id == encode_trace_id(parent_span.context.trace_id) 40 assert span.span_id == encode_span_id(parent_span.context.span_id) 41 assert span.name == "parent" 42 assert span.start_time_ns == parent_span.start_time 43 assert span.end_time_ns is None 44 assert span.parent_id is None 45 46 span.set_inputs({"input": 1}) 47 span.set_outputs(2) 48 assert span.inputs == {"input": 1} 49 assert span.outputs == 2 50 51 span.set_attribute("key", 3) 52 assert span.get_attribute("key") == 3 53 54 # non-serializable value should be stored as string 55 non_serializable = datetime.now() 56 span.set_attribute("non_serializable", non_serializable) 57 assert span.get_attribute("non_serializable") == str(non_serializable) 58 assert parent_span._attributes == { 59 "mlflow.traceRequestId": json.dumps(trace_id), 60 "mlflow.spanInputs": '{"input": 1}', 61 "mlflow.spanOutputs": "2", 62 "mlflow.spanType": '"LLM"', 63 "key": "3", 64 "non_serializable": json.dumps(str(non_serializable)), 65 } 66 67 span.set_status("OK") 68 assert span.status == SpanStatus(SpanStatusCode.OK) 69 70 span.add_event(SpanEvent("test_event", timestamp=99999, attributes={"foo": "bar"})) 71 assert len(span.events) == 1 72 assert span.events[0].name == "test_event" 73 assert span.events[0].timestamp == 99999 74 assert span.events[0].attributes == {"foo": "bar"} 75 76 # Test child span 77 with tracer.start_as_current_span("child") as child_span: 78 span = create_mlflow_span(child_span, trace_id=trace_id) 79 assert isinstance(span, LiveSpan) 80 assert span.name == "child" 81 assert span.parent_id == encode_span_id(parent_span.context.span_id) 82 83 84 def test_create_non_live_span(): 85 trace_id = "tr-12345" 86 parent_span_context = trace_api.SpanContext( 87 trace_id=12345, span_id=111, is_remote=False, trace_flags=trace_api.TraceFlags(1) 88 ) 89 readable_span = OTelReadableSpan( 90 name="test", 91 context=trace_api.SpanContext( 92 trace_id=12345, span_id=222, is_remote=False, trace_flags=trace_api.TraceFlags(1) 93 ), 94 parent=parent_span_context, 95 attributes={ 96 "mlflow.traceRequestId": json.dumps(trace_id), 97 "mlflow.spanInputs": '{"input": 1, "nested": {"foo": "bar"}}', 98 "mlflow.spanOutputs": "2", 99 "key": "3", 100 }, 101 start_time=99999, 102 end_time=100000, 103 ) 104 span = create_mlflow_span(readable_span, trace_id) 105 106 assert isinstance(span, Span) 107 assert not isinstance(span, LiveSpan) 108 assert not isinstance(span, NoOpSpan) 109 assert span.trace_id == trace_id 110 assert span._trace_id == encode_trace_id(12345) 111 assert span.span_id == encode_span_id(222) 112 assert span.name == "test" 113 assert span.start_time_ns == 99999 114 assert span.end_time_ns == 100000 115 assert span.parent_id == encode_span_id(111) 116 assert span.inputs == {"input": 1, "nested": {"foo": "bar"}} 117 assert span.outputs == 2 118 assert span.status == SpanStatus(SpanStatusCode.UNSET, description="") 119 assert span.get_attribute("key") == 3 120 121 # Non-live span should not implement setter methods 122 with pytest.raises(AttributeError, match="set_inputs"): 123 span.set_inputs({"input": 1}) 124 125 126 def test_create_noop_span(): 127 trace_id = "tr-12345" 128 129 @trace_disabled 130 def f(): 131 tracer = _get_tracer("test") 132 with tracer.start_as_current_span("span") as otel_span: 133 span = create_mlflow_span(otel_span, trace_id=trace_id) 134 assert isinstance(span, NoOpSpan) 135 136 # create from None 137 span = create_mlflow_span(None, trace_id=trace_id) 138 assert isinstance(span, NoOpSpan) 139 140 141 def test_create_raise_for_invalid_otel_span(): 142 with pytest.raises(MlflowException, match=r"The `otel_span` argument must be"): 143 create_mlflow_span(otel_span=123, trace_id="tr-12345") 144 145 146 @pytest.mark.parametrize( 147 "status", 148 [SpanStatus("OK"), SpanStatus(SpanStatusCode.ERROR, "Error!"), "OK", "ERROR"], 149 ) 150 def test_set_status(status): 151 with mlflow.start_span("test_span") as span: 152 span.set_status(status) 153 154 assert isinstance(span.status, SpanStatus) 155 156 157 def test_set_status_raise_for_invalid_value(): 158 with mlflow.start_span("test_span") as span: 159 with pytest.raises(MlflowException, match=r"INVALID is not a valid SpanStatusCode value."): 160 span.set_status("INVALID") 161 162 163 def test_dict_conversion(): 164 with mlflow.start_span("parent"): 165 with mlflow.start_span("child", span_type=SpanType.LLM) as span: 166 span.set_inputs({"input": 1}) 167 span.set_outputs(2) 168 span.set_attribute("key", 3) 169 span.set_status("OK") 170 span.add_event(SpanEvent("test_event", timestamp=0, attributes={"foo": "bar"})) 171 172 span_dict = span.to_dict() 173 recovered_span = Span.from_dict(span_dict) 174 175 assert span.trace_id == recovered_span.trace_id 176 assert span._trace_id == recovered_span._trace_id 177 assert span.span_id == recovered_span.span_id 178 assert span.name == recovered_span.name 179 assert span.span_type == recovered_span.span_type 180 assert span.start_time_ns == recovered_span.start_time_ns 181 assert span.end_time_ns == recovered_span.end_time_ns 182 assert span.parent_id == recovered_span.parent_id 183 assert span.status == recovered_span.status 184 assert span.inputs == recovered_span.inputs 185 assert span.outputs == recovered_span.outputs 186 assert span.attributes == recovered_span.attributes 187 assert span.events == recovered_span.events 188 189 # Loaded span should not implement setter methods 190 with pytest.raises(AttributeError, match="set_status"): 191 recovered_span.set_status("OK") 192 193 194 def test_dict_conversion_with_exception_event(): 195 with pytest.raises(ValueError, match="Test exception"): 196 with mlflow.start_span("test") as span: 197 raise ValueError("Test exception") 198 199 span_dict = span.to_dict() 200 recovered_span = Span.from_dict(span_dict) 201 202 assert span.request_id == recovered_span.request_id 203 assert span._trace_id == recovered_span._trace_id 204 assert span.span_id == recovered_span.span_id 205 assert span.name == recovered_span.name 206 assert span.span_type == recovered_span.span_type 207 assert span.start_time_ns == recovered_span.start_time_ns 208 assert span.end_time_ns == recovered_span.end_time_ns 209 assert span.parent_id == recovered_span.parent_id 210 assert span.status == recovered_span.status 211 assert span.inputs == recovered_span.inputs 212 assert span.outputs == recovered_span.outputs 213 assert span.attributes == recovered_span.attributes 214 assert span.events == recovered_span.events 215 216 # Loaded span should not implement setter methods 217 with pytest.raises(AttributeError, match="set_status"): 218 recovered_span.set_status("OK") 219 220 221 def test_from_v2_dict(): 222 span = Span.from_dict({ 223 "name": "test", 224 "context": { 225 "span_id": "8a90fc46e65ea5a4", 226 "trace_id": "0125978dc5c5a9456d7ca9ef1f7cf4af", 227 }, 228 "parent_id": None, 229 "start_time": 1738662897576578992, 230 "end_time": 1738662899068969049, 231 "status_code": "OK", 232 "status_message": "", 233 "attributes": { 234 "mlflow.traceRequestId": '"tr-123"', 235 "mlflow.spanType": '"LLM"', 236 "mlflow.spanInputs": '{"input": 1}', 237 "mlflow.spanOutputs": "2", 238 "key": "3", 239 }, 240 "events": [], 241 }) 242 243 assert span.request_id == "tr-123" 244 assert span.name == "test" 245 assert span.span_type == SpanType.LLM 246 assert span.parent_id is None 247 assert span.status == SpanStatus(SpanStatusCode.OK, description="") 248 assert span.inputs == {"input": 1} 249 assert span.outputs == 2 250 assert span.events == [] 251 252 253 def test_to_immutable_span(): 254 trace_id = "tr-12345" 255 256 tracer = _get_tracer("test") 257 with tracer.start_as_current_span("parent") as parent_span: 258 live_span = LiveSpan(parent_span, trace_id=trace_id, span_type=SpanType.LLM) 259 live_span.set_inputs({"input": 1}) 260 live_span.set_outputs(2) 261 live_span.set_attribute("key", 3) 262 live_span.set_status("OK") 263 live_span.add_event(SpanEvent("test_event", timestamp=0, attributes={"foo": "bar"})) 264 265 span = live_span.to_immutable_span() 266 267 assert isinstance(span, Span) 268 assert span.trace_id == trace_id 269 assert span._trace_id == encode_trace_id(parent_span.context.trace_id) 270 assert span.span_id == encode_span_id(parent_span.context.span_id) 271 assert span.name == "parent" 272 assert span.start_time_ns == parent_span.start_time 273 assert span.end_time_ns is not None 274 assert span.parent_id is None 275 assert span.inputs == {"input": 1} 276 assert span.outputs == 2 277 assert span.get_attribute("key") == 3 278 assert span.status == SpanStatus(SpanStatusCode.OK, description="") 279 assert span.events == [SpanEvent("test_event", timestamp=0, attributes={"foo": "bar"})] 280 281 with pytest.raises(AttributeError, match="set_attribute"): 282 span.set_attribute("OK") 283 284 285 def test_from_dict_raises_when_trace_id_is_empty(): 286 with pytest.raises(MlflowException, match=r"Failed to create a Span object from "): 287 Span.from_dict({ 288 "name": "predict", 289 "context": { 290 "trace_id": "12345", 291 "span_id": "12345", 292 }, 293 "parent_id": None, 294 "start_time": 0, 295 "end_time": 1, 296 "status_code": "OK", 297 "status_message": "", 298 "attributes": { 299 "mlflow.traceRequestId": None, 300 }, 301 "events": [], 302 }) 303 304 305 def test_set_attribute_directly_to_otel_span(): 306 with mlflow.start_span("test") as span: 307 span._span.set_attribute("int", 1) 308 span._span.set_attribute("str", "a") 309 310 assert span.get_attribute("int") == 1 311 assert span.get_attribute("str") == "a" 312 313 314 @pytest.fixture 315 def sample_otel_span_for_conversion(): 316 """Create a sample OTelReadableSpan for testing.""" 317 return OTelReadableSpan( 318 name="test_span", 319 context=build_otel_context( 320 trace_id=0x12345678901234567890123456789012, 321 span_id=0x1234567890123456, 322 ), 323 parent=build_otel_context( 324 trace_id=0x12345678901234567890123456789012, 325 span_id=0x0987654321098765, 326 ), 327 start_time=1000000000, 328 end_time=2000000000, 329 attributes={ 330 "mlflow.traceRequestId": "tr-12345678901234567890123456789012", 331 "mlflow.spanType": "LLM", 332 "mlflow.spanInputs": '{"prompt": "Hello"}', 333 "mlflow.spanOutputs": '{"response": "Hi"}', 334 "custom_attr": '{"key": "value"}', 335 }, 336 status=OTelStatus(OTelStatusCode.OK, "Success"), 337 events=[ 338 OTelEvent( 339 name="event1", 340 timestamp=1500000000, 341 attributes={"event_key": "event_value"}, 342 ) 343 ], 344 ) 345 346 347 @pytest.mark.parametrize( 348 "attributes", 349 [ 350 # Empty attributes 351 {}, 352 # String attributes 353 {"str_key": "str_value"}, 354 # Numeric attributes 355 {"int_key": 42, "float_key": 3.14}, 356 # Boolean attributes 357 {"bool_key": True, "false_key": False}, 358 # Bytes attributes 359 {"bytes_key": b"binary_data"}, 360 # Empty list and dict attributes 361 {"empty_list": [], "empty_dict": {}}, 362 # List attributes 363 {"list_str": ["a", "b", "c"], "list_int": [1, 2, 3], "list_float": [1.1, 2.2]}, 364 # Dict attributes 365 {"dict_key": {"nested": "value", "number": 123}}, 366 # Mixed complex attributes 367 { 368 "complex": { 369 "nested_list": [1, "two", 3.0], 370 "nested_dict": {"deep": {"deeper": "value"}}, 371 }, 372 "simple": "string", 373 }, 374 ], 375 ) 376 def test_otel_attribute_conversion(attributes): 377 from opentelemetry.proto.common.v1.common_pb2 import KeyValue 378 379 from mlflow.tracing.utils.otlp import _decode_otel_proto_anyvalue, _set_otel_proto_anyvalue 380 381 # Convert attributes to proto format 382 proto_attrs = [] 383 for key, value in attributes.items(): 384 kv = KeyValue() 385 kv.key = key 386 _set_otel_proto_anyvalue(kv.value, value) 387 proto_attrs.append(kv) 388 389 # Decode back and verify 390 decoded = {} 391 for kv in proto_attrs: 392 decoded[kv.key] = _decode_otel_proto_anyvalue(kv.value) 393 394 assert decoded == attributes 395 396 397 def test_span_to_otel_proto_conversion(sample_otel_span_for_conversion): 398 # Create MLflow span from OTel span 399 mlflow_span = Span(sample_otel_span_for_conversion) 400 401 # Convert to OTel proto 402 otel_proto = mlflow_span.to_otel_proto() 403 404 # Verify basic fields 405 assert otel_proto.name == "test_span" 406 assert otel_proto.start_time_unix_nano == 1000000000 407 assert otel_proto.end_time_unix_nano == 2000000000 408 409 # Verify IDs (should be in bytes format) 410 assert len(otel_proto.trace_id) == 16 # 128-bit trace ID 411 assert len(otel_proto.span_id) == 8 # 64-bit span ID 412 assert len(otel_proto.parent_span_id) == 8 413 414 # Verify status 415 assert otel_proto.status.code == OTelProtoStatus.STATUS_CODE_OK 416 # OTel SDK clears description for non-ERROR statuses 417 assert otel_proto.status.message == "" 418 419 # Verify attributes exist 420 assert len(otel_proto.attributes) == 5 421 attr_keys = {attr.key for attr in otel_proto.attributes} 422 assert "mlflow.spanType" in attr_keys 423 assert "custom_attr" in attr_keys 424 425 # Verify events 426 assert len(otel_proto.events) == 1 427 assert otel_proto.events[0].name == "event1" 428 assert otel_proto.events[0].time_unix_nano == 1500000000 429 430 431 def test_span_from_otel_proto_conversion(): 432 # Create OTel proto span 433 otel_proto = OTelProtoSpan() 434 otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012") 435 otel_proto.span_id = bytes.fromhex("1234567890123456") 436 otel_proto.parent_span_id = bytes.fromhex("0987654321098765") 437 otel_proto.name = "proto_span" 438 otel_proto.start_time_unix_nano = 1000000000 439 otel_proto.end_time_unix_nano = 2000000000 440 441 # Add status 442 otel_proto.status.code = OTelProtoStatus.STATUS_CODE_ERROR 443 otel_proto.status.message = "Error occurred" 444 445 # Add attributes 446 from mlflow.tracing.utils.otlp import _set_otel_proto_anyvalue 447 448 attr2 = otel_proto.attributes.add() 449 attr2.key = "mlflow.spanType" 450 _set_otel_proto_anyvalue(attr2.value, "CHAIN") 451 452 attr3 = otel_proto.attributes.add() 453 attr3.key = "custom" 454 _set_otel_proto_anyvalue(attr3.value, {"nested": {"value": 123}}) 455 456 # Add event 457 event = otel_proto.events.add() 458 event.name = "test_event" 459 event.time_unix_nano = 1500000000 460 event_attr = event.attributes.add() 461 event_attr.key = "event_data" 462 _set_otel_proto_anyvalue(event_attr.value, "event_value") 463 464 # Convert to MLflow span 465 mlflow_span = Span.from_otel_proto(otel_proto) 466 467 # Verify basic fields 468 assert mlflow_span.name == "proto_span" 469 assert mlflow_span.start_time_ns == 1000000000 470 assert mlflow_span.end_time_ns == 2000000000 471 472 # Verify IDs 473 assert mlflow_span.trace_id == "tr-12345678901234567890123456789012" 474 assert mlflow_span.span_id == "1234567890123456" 475 assert mlflow_span.parent_id == "0987654321098765" 476 477 # Verify status 478 assert mlflow_span.status.status_code == SpanStatusCode.ERROR 479 assert mlflow_span.status.description == "Error occurred" 480 481 # Verify attributes 482 assert mlflow_span.span_type == "CHAIN" 483 assert mlflow_span.get_attribute("custom") == {"nested": {"value": 123}} 484 485 # Verify events 486 assert len(mlflow_span.events) == 1 487 assert mlflow_span.events[0].name == "test_event" 488 assert mlflow_span.events[0].timestamp == 1500000000 489 assert mlflow_span.events[0].attributes["event_data"] == "event_value" 490 491 492 def test_span_from_otel_proto_with_location(): 493 # Create OTel proto span 494 otel_proto = OTelProtoSpan() 495 otel_proto.trace_id = bytes.fromhex("12345678901234567890123456789012") 496 otel_proto.span_id = bytes.fromhex("1234567890123456") 497 otel_proto.parent_span_id = bytes.fromhex("0987654321098765") 498 otel_proto.name = "proto_span_v4" 499 otel_proto.start_time_unix_nano = 1000000000 500 otel_proto.end_time_unix_nano = 2000000000 501 502 # Add status 503 otel_proto.status.code = OTelProtoStatus.STATUS_CODE_OK 504 otel_proto.status.message = "" 505 506 # Add attributes 507 attr1 = otel_proto.attributes.add() 508 attr1.key = "mlflow.spanType" 509 _set_otel_proto_anyvalue(attr1.value, "LLM") 510 511 # Convert to MLflow span with location 512 location = "catalog.schema" 513 mlflow_span = Span.from_otel_proto(otel_proto, location_id=location) 514 515 # Verify basic fields 516 assert mlflow_span.name == "proto_span_v4" 517 assert mlflow_span.start_time_ns == 1000000000 518 assert mlflow_span.end_time_ns == 2000000000 519 520 # Verify IDs 521 assert mlflow_span.span_id == "1234567890123456" 522 assert mlflow_span.parent_id == "0987654321098765" 523 524 # Verify trace_id is in v4 format with location 525 expected_trace_id = f"{TRACE_ID_V4_PREFIX}{location}/12345678901234567890123456789012" 526 assert mlflow_span.trace_id == expected_trace_id 527 528 # Verify the REQUEST_ID attribute also uses v4 format 529 request_id = mlflow_span.get_attribute("mlflow.traceRequestId") 530 assert request_id == expected_trace_id 531 532 533 def test_otel_roundtrip_conversion(sample_otel_span_for_conversion): 534 # Start with OTel span -> MLflow span 535 mlflow_span = Span(sample_otel_span_for_conversion) 536 537 # Convert to OTel proto 538 otel_proto = mlflow_span.to_otel_proto() 539 540 # Convert back to MLflow span 541 roundtrip_span = Span.from_otel_proto(otel_proto) 542 543 # Verify key fields are preserved 544 assert roundtrip_span.name == mlflow_span.name 545 assert roundtrip_span.span_id == mlflow_span.span_id 546 assert roundtrip_span.parent_id == mlflow_span.parent_id 547 assert roundtrip_span.start_time_ns == mlflow_span.start_time_ns 548 assert roundtrip_span.end_time_ns == mlflow_span.end_time_ns 549 assert roundtrip_span.status.status_code == mlflow_span.status.status_code 550 assert roundtrip_span.status.description == mlflow_span.status.description 551 552 # Verify span attributes are preserved 553 assert roundtrip_span.span_type == mlflow_span.span_type 554 assert roundtrip_span.inputs == mlflow_span.inputs 555 assert roundtrip_span.outputs == mlflow_span.outputs 556 557 # Verify ALL attributes are preserved by iterating through them 558 # Get all attribute keys from both spans 559 original_attributes = mlflow_span.attributes 560 roundtrip_attributes = roundtrip_span.attributes 561 562 # Check we have the same number of attributes 563 assert len(original_attributes) == len(roundtrip_attributes) 564 565 # Check each attribute is preserved correctly 566 for attr_key in original_attributes: 567 assert attr_key in roundtrip_attributes, f"Attribute {attr_key} missing after roundtrip" 568 original_value = original_attributes[attr_key] 569 roundtrip_value = roundtrip_attributes[attr_key] 570 assert original_value == roundtrip_value, ( 571 f"Attribute {attr_key} changed: {original_value} != {roundtrip_value}" 572 ) 573 574 # Also explicitly verify specific important attributes 575 # The original span has a custom_attr that should be preserved 576 original_custom_attr = mlflow_span.get_attribute("custom_attr") 577 roundtrip_custom_attr = roundtrip_span.get_attribute("custom_attr") 578 assert original_custom_attr == roundtrip_custom_attr 579 assert original_custom_attr == {"key": "value"} 580 581 # Verify the trace request ID is preserved 582 assert roundtrip_span.request_id == mlflow_span.request_id 583 assert roundtrip_span.request_id == "tr-12345678901234567890123456789012" 584 585 # Verify events 586 assert len(roundtrip_span.events) == len(mlflow_span.events) 587 for orig_event, rt_event in zip(mlflow_span.events, roundtrip_span.events): 588 assert rt_event.name == orig_event.name 589 assert rt_event.timestamp == orig_event.timestamp 590 assert rt_event.attributes == orig_event.attributes 591 592 593 def test_resource_to_otel_proto(): 594 resource = OTelResource.create({ 595 "service.name": "test-service", 596 "service.version": "1.0.0", 597 "custom.int": 42, 598 "custom.bool": True, 599 }) 600 resource_proto = resource_to_otel_proto(resource) 601 602 # Convert proto attributes to dict for easier verification 603 attrs = {} 604 for attr in resource_proto.attributes: 605 attrs[attr.key] = _decode_otel_proto_anyvalue(attr.value) 606 607 assert attrs["service.name"] == "test-service" 608 assert attrs["service.version"] == "1.0.0" 609 assert attrs["custom.int"] == 42 610 assert attrs["custom.bool"] is True 611 612 613 def test_resource_to_otel_proto_none(): 614 resource_proto = resource_to_otel_proto(None) 615 assert len(resource_proto.attributes) == 0 616 617 618 def test_span_from_dict_old_format(): 619 span_dict = { 620 "trace_id": "6ST7JNq8BC4JRp0HA/vD6Q==", 621 "span_id": "Sd/l0Zs4M3g=", 622 "parent_span_id": None, 623 "name": "test_span", 624 "start_time_unix_nano": 1000000000, 625 "end_time_unix_nano": 2000000000, 626 "status": {"code": "STATUS_CODE_ERROR", "message": "Error occurred"}, 627 "attributes": { 628 "mlflow.spanInputs": '{"query": "test"}', 629 "mlflow.spanOutputs": '{"result": "success"}', 630 "custom": "value", 631 "mlflow.traceRequestId": '"tr-e924fb24dabc042e09469d0703fbc3e9"', 632 }, 633 "events": [], 634 } 635 636 # Deserialize it 637 recovered_span = Span.from_dict(span_dict) 638 639 # Verify all fields are recovered correctly 640 assert recovered_span.trace_id == "tr-e924fb24dabc042e09469d0703fbc3e9" 641 assert recovered_span.span_id == "49dfe5d19b383378" 642 assert recovered_span.parent_id is None 643 assert recovered_span.name == span_dict["name"] 644 assert recovered_span.start_time_ns == span_dict["start_time_unix_nano"] 645 assert recovered_span.end_time_ns == span_dict["end_time_unix_nano"] 646 assert recovered_span.status.status_code.value == "ERROR" 647 assert recovered_span.inputs == {"query": "test"} 648 assert recovered_span.outputs == {"result": "success"} 649 assert recovered_span.get_attribute("custom") == "value" 650 651 652 def test_span_dict_v4_with_no_parent(): 653 with mlflow.start_span("root_span") as span: 654 span.set_inputs({"x": 1}) 655 span.set_outputs({"y": 2}) 656 657 span_dict = span.to_dict() 658 659 # Root span should have None for parent_span_id 660 assert span_dict["parent_span_id"] is None 661 662 # Deserialize and verify 663 recovered = Span.from_dict(span_dict) 664 assert recovered.parent_id is None 665 assert recovered.name == "root_span" 666 assert recovered.inputs == {"x": 1} 667 assert recovered.outputs == {"y": 2} 668 669 670 def test_span_from_dict_supports_both_status_code_formats(): 671 with mlflow.start_span("test") as span: 672 span.set_status("OK") 673 674 span_dict = span.to_dict() 675 676 # Current code serializes as protobuf enum name 677 assert span_dict["status"]["code"] == "STATUS_CODE_OK" 678 679 # Verify we can deserialize protobuf enum name format (backward compatibility) 680 span_dict["status"]["code"] = "STATUS_CODE_ERROR" 681 recovered = Span.from_dict(span_dict) 682 assert recovered.status.status_code == SpanStatusCode.ERROR 683 684 # Verify we can also deserialize enum value format 685 # (forward compatibility with older serialized data) 686 span_dict["status"]["code"] = "OK" 687 recovered = Span.from_dict(span_dict) 688 assert recovered.status.status_code == SpanStatusCode.OK 689 690 span_dict["status"]["code"] = "UNSET" 691 recovered = Span.from_dict(span_dict) 692 assert recovered.status.status_code == SpanStatusCode.UNSET 693 694 span_dict["status"]["code"] = "ERROR" 695 recovered = Span.from_dict(span_dict) 696 assert recovered.status.status_code == SpanStatusCode.ERROR 697 698 699 def test_load_from_old_span_dict(): 700 span_dict = { 701 "trace_id": "ZqqBulxlq2cwRCfKHxmDVA==", 702 "span_id": "/mCYZRxbTqw=", 703 "trace_state": "", 704 "parent_span_id": "f6qlKYqTw2E=", 705 "name": "custom", 706 "start_time_unix_nano": 1761103703884225000, 707 "end_time_unix_nano": 1761103703884454000, 708 "attributes": { 709 "mlflow.spanOutputs": "4", 710 "mlflow.spanType": '"LLM"', 711 "mlflow.spanInputs": '{"z": 3}', 712 "mlflow.traceRequestId": '"tr-66aa81ba5c65ab67304427ca1f198354"', 713 "delta": "1", 714 "mlflow.spanFunctionName": '"add_one"', 715 }, 716 "status": {"message": "", "code": "STATUS_CODE_OK"}, 717 "events": [ 718 { 719 "time_unix_nano": 1761105506649041, 720 "name": "agent_action", 721 "attributes": { 722 "tool": "search_web", 723 "tool_input": '"What is MLflow?"', 724 "log": "test", 725 }, 726 } 727 ], 728 } 729 span = Span.from_dict(span_dict) 730 assert span.trace_id == "tr-66aa81ba5c65ab67304427ca1f198354" 731 assert span.span_id == "fe6098651c5b4eac" 732 assert span.parent_id == "7faaa5298a93c361" 733 assert span.name == "custom" 734 assert span.start_time_ns == 1761103703884225000 735 assert span.end_time_ns == 1761103703884454000 736 assert span.status == SpanStatus(SpanStatusCode.OK, description="") 737 assert span.inputs == {"z": 3} 738 assert span.outputs == 4 739 assert len(span.events) == 1 740 assert span.events[0].name == "agent_action" 741 assert span.events[0].timestamp == 1761105506649041 742 assert span.events[0].attributes == { 743 "tool": "search_web", 744 "tool_input": '"What is MLflow?"', 745 "log": "test", 746 } 747 748 749 def test_load_from_3_5_0_span_dict(): 750 span_dict = { 751 "trace_id": "tr-66aa81ba5c65ab67304427ca1f198354", 752 "span_id": "fe6098651c5b4eac", 753 "trace_state": "", 754 "parent_span_id": "7faaa5298a93c361", 755 "name": "custom", 756 "start_time_unix_nano": 1761103703884225000, 757 "end_time_unix_nano": 1761103703884454000, 758 "attributes": { 759 "mlflow.spanOutputs": "4", 760 "mlflow.spanType": '"LLM"', 761 "mlflow.spanInputs": '{"z": 3}', 762 "mlflow.traceRequestId": '"tr-66aa81ba5c65ab67304427ca1f198354"', 763 "delta": "1", 764 "mlflow.spanFunctionName": '"add_one"', 765 }, 766 "status": {"message": "", "code": "OK"}, 767 "events": [ 768 { 769 "time_unix_nano": 1761105506649041, 770 "name": "agent_action", 771 "attributes": { 772 "tool": "search_web", 773 "tool_input": '"What is MLflow?"', 774 "log": "test", 775 }, 776 } 777 ], 778 } 779 span = Span.from_dict(span_dict) 780 assert span.trace_id == "tr-66aa81ba5c65ab67304427ca1f198354" 781 assert span.span_id == "fe6098651c5b4eac" 782 assert span.parent_id == "7faaa5298a93c361" 783 assert span.name == "custom" 784 assert span.start_time_ns == 1761103703884225000 785 assert span.end_time_ns == 1761103703884454000 786 assert span.status == SpanStatus(SpanStatusCode.OK, description="") 787 assert span.inputs == {"z": 3} 788 assert span.outputs == 4 789 assert len(span.events) == 1 790 assert span.events[0].name == "agent_action" 791 assert span.events[0].timestamp == 1761105506649041 792 assert span.events[0].attributes == { 793 "tool": "search_web", 794 "tool_input": '"What is MLflow?"', 795 "log": "test", 796 }