wehub-resource-sync
131 lines
3.7 KiB
Python
131 lines
3.7 KiB
Python
import mlflow
|
|
from mlflow.entities import Link
|
|
from mlflow.entities.span import NoOpSpan
|
|
from mlflow.tracing.fluent import start_span_no_context
|
|
|
|
from tests.tracing.helper import get_traces
|
|
|
|
# --- Integration tests for start_span with links ---
|
|
|
|
|
|
def test_start_span_with_links():
|
|
links = [
|
|
Link(trace_id="tr-0123456789abcdef0123456789abcdef", span_id="0123456789abcdef"),
|
|
Link(
|
|
trace_id="tr-abcdef0123456789abcdef0123456789",
|
|
span_id="abcdef0123456789",
|
|
attributes={"kind": "follows_from"},
|
|
),
|
|
]
|
|
|
|
with mlflow.start_span(name="linked_span", links=links) as span:
|
|
span.set_outputs("done")
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
|
|
linked_span = traces[0].data.spans[0]
|
|
assert len(linked_span.links) == 2
|
|
assert linked_span.links[0].trace_id == "tr-0123456789abcdef0123456789abcdef"
|
|
assert linked_span.links[0].span_id == "0123456789abcdef"
|
|
assert linked_span.links[1].attributes == {"kind": "follows_from"}
|
|
|
|
|
|
def test_start_span_no_context_with_links():
|
|
links = [
|
|
Link(trace_id="tr-0123456789abcdef0123456789abcdef", span_id="0123456789abcdef"),
|
|
]
|
|
|
|
root = start_span_no_context("root_span", links=links)
|
|
root.end()
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert len(traces[0].data.spans[0].links) == 1
|
|
assert traces[0].data.spans[0].links[0].span_id == "0123456789abcdef"
|
|
|
|
|
|
def test_start_span_without_links():
|
|
with mlflow.start_span(name="no_links_span") as span:
|
|
span.set_outputs("done")
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert traces[0].data.spans[0].links == []
|
|
|
|
|
|
# --- Decorator tests ---
|
|
|
|
|
|
def test_trace_decorator_with_static_links():
|
|
links = [
|
|
Link(trace_id="tr-0123456789abcdef0123456789abcdef", span_id="0123456789abcdef"),
|
|
]
|
|
|
|
@mlflow.trace(links=links)
|
|
def my_func(x):
|
|
return x * 2
|
|
|
|
my_func(5)
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert len(traces[0].data.spans[0].links) == 1
|
|
assert traces[0].data.spans[0].links[0].span_id == "0123456789abcdef"
|
|
|
|
|
|
def test_trace_decorator_with_generator_and_links():
|
|
links = [
|
|
Link(trace_id="tr-0123456789abcdef0123456789abcdef", span_id="0123456789abcdef"),
|
|
]
|
|
|
|
@mlflow.trace(links=links)
|
|
def my_generator():
|
|
yield 1
|
|
yield 2
|
|
yield 3
|
|
|
|
result = list(my_generator())
|
|
assert result == [1, 2, 3]
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
assert len(traces[0].data.spans[0].links) == 1
|
|
|
|
|
|
# --- NoOpSpan safety ---
|
|
|
|
|
|
def test_noop_span_add_link_does_not_raise():
|
|
noop = NoOpSpan()
|
|
link = Link(trace_id="tr-0123456789abcdef0123456789abcdef", span_id="0123456789abcdef")
|
|
noop.add_link(link)
|
|
|
|
|
|
# --- End-to-end: links persist through store round-trip ---
|
|
|
|
|
|
def test_links_persist_through_store_roundtrip():
|
|
links = [
|
|
Link(
|
|
trace_id="tr-0123456789abcdef0123456789abcdef",
|
|
span_id="0123456789abcdef",
|
|
attributes={"type": "caused_by"},
|
|
),
|
|
Link(trace_id="tr-fedcba9876543210fedcba9876543210", span_id="fedcba9876543210"),
|
|
]
|
|
|
|
with mlflow.start_span(name="roundtrip_span", links=links) as span:
|
|
span.set_outputs("test")
|
|
|
|
traces = get_traces()
|
|
assert len(traces) == 1
|
|
stored_links = traces[0].data.spans[0].links
|
|
assert len(stored_links) == 2
|
|
assert stored_links[0].trace_id == "tr-0123456789abcdef0123456789abcdef"
|
|
assert stored_links[0].span_id == "0123456789abcdef"
|
|
assert stored_links[0].attributes == {"type": "caused_by"}
|
|
assert stored_links[1].trace_id == "tr-fedcba9876543210fedcba9876543210"
|
|
assert stored_links[1].span_id == "fedcba9876543210"
|
|
assert stored_links[1].attributes is None
|