wehub-resource-sync
268 lines
9.3 KiB
Python
268 lines
9.3 KiB
Python
import logging
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow.entities import SpanLogLevel
|
|
from mlflow.entities.span import Span, SpanType
|
|
from mlflow.entities.span_event import SpanEvent
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.tracing.constant import SpanAttributeKey
|
|
from mlflow.tracing.utils.default_log_level import default_log_level_for_span_type
|
|
|
|
from tests.tracing.helper import get_traces
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("value", "expected"),
|
|
[
|
|
(SpanLogLevel.DEBUG, SpanLogLevel.DEBUG),
|
|
(SpanLogLevel.CRITICAL, SpanLogLevel.CRITICAL),
|
|
("DEBUG", SpanLogLevel.DEBUG),
|
|
("info", SpanLogLevel.INFO),
|
|
("Warning", SpanLogLevel.WARNING),
|
|
(" ERROR ", SpanLogLevel.ERROR),
|
|
],
|
|
)
|
|
def test_from_value_accepts_enum_and_string_forms(value, expected):
|
|
assert SpanLogLevel.from_value(value) is expected
|
|
|
|
|
|
@pytest.mark.parametrize("value", ["NOPE", "TRACE", "FATAL", "INFOO", "", "WARN", "warn"])
|
|
def test_from_value_rejects_invalid_string(value):
|
|
# "WARN" is not a valid alias; only the full names are accepted.
|
|
with pytest.raises(MlflowException, match="Invalid SpanLogLevel"):
|
|
SpanLogLevel.from_value(value)
|
|
|
|
|
|
@pytest.mark.parametrize("value", [0, 7, 10, 20, 30, 40, 50, 100, -1, logging.INFO])
|
|
def test_from_value_rejects_int(value):
|
|
# Raw integers — including the canonical 10/20/30/40/50 and `logging.*` —
|
|
# are rejected: the API surface is `SpanLogLevel | str` only. Use the enum
|
|
# member or its name string instead.
|
|
with pytest.raises(MlflowException, match="must be"):
|
|
SpanLogLevel.from_value(value)
|
|
|
|
|
|
@pytest.mark.parametrize("value", [None, 1.5, ["INFO"], object(), True, False])
|
|
def test_from_value_rejects_invalid_type(value):
|
|
with pytest.raises(MlflowException, match="must be"):
|
|
SpanLogLevel.from_value(value)
|
|
|
|
|
|
def test_log_level_constructor_default_for_unknown_span():
|
|
# No span_type provided -> defaults to UNKNOWN -> DEBUG via the constructor.
|
|
with mlflow.start_span("s"):
|
|
pass
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.DEBUG
|
|
assert persisted.attributes[SpanAttributeKey.LOG_LEVEL] == int(SpanLogLevel.DEBUG)
|
|
|
|
|
|
def test_log_level_constructor_default_for_info_span_type():
|
|
# CHAT_MODEL is in the INFO set -> constructor stamps INFO automatically.
|
|
with mlflow.start_span("s", span_type=SpanType.CHAT_MODEL):
|
|
pass
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.INFO
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("set_value", "expected"),
|
|
[
|
|
(SpanLogLevel.WARNING, SpanLogLevel.WARNING),
|
|
(SpanLogLevel.ERROR, SpanLogLevel.ERROR),
|
|
("info", SpanLogLevel.INFO),
|
|
("CRITICAL", SpanLogLevel.CRITICAL),
|
|
],
|
|
)
|
|
def test_set_log_level_normalizes_input(set_value, expected):
|
|
with mlflow.start_span("s") as span:
|
|
span.set_log_level(set_value)
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is expected
|
|
# Stored as the raw int under the reserved attribute key for portability.
|
|
assert persisted.attributes[SpanAttributeKey.LOG_LEVEL] == int(expected)
|
|
|
|
|
|
def test_set_log_level_rejects_invalid_input():
|
|
with (
|
|
mlflow.start_span("s") as span,
|
|
pytest.raises(MlflowException, match="Invalid SpanLogLevel"),
|
|
):
|
|
span.set_log_level("NOPE")
|
|
|
|
|
|
def test_start_span_kwarg_overrides_constructor_default():
|
|
# CHAT_MODEL would default to INFO; explicit kwarg should win.
|
|
with mlflow.start_span("s", span_type=SpanType.CHAT_MODEL, log_level="WARNING"):
|
|
pass
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.WARNING
|
|
|
|
|
|
def test_start_span_no_context_kwarg():
|
|
span = mlflow.start_span_no_context("s", log_level=SpanLogLevel.ERROR)
|
|
span.end()
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.ERROR
|
|
|
|
|
|
def test_trace_decorator_kwarg_sync():
|
|
@mlflow.trace(log_level="DEBUG")
|
|
def fn(x):
|
|
return x + 1
|
|
|
|
fn(1)
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.DEBUG
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_decorator_kwarg_async():
|
|
@mlflow.trace(log_level=SpanLogLevel.INFO)
|
|
async def fn(x):
|
|
return x + 1
|
|
|
|
await fn(1)
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.INFO
|
|
|
|
|
|
def test_trace_decorator_kwarg_generator():
|
|
@mlflow.trace(log_level="WARNING")
|
|
def gen():
|
|
yield 1
|
|
yield 2
|
|
|
|
list(gen())
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.WARNING
|
|
|
|
|
|
def test_log_level_round_trips_through_to_dict_from_dict():
|
|
with mlflow.start_span("s", log_level=SpanLogLevel.ERROR):
|
|
pass
|
|
persisted = get_traces()[0].data.spans[0]
|
|
|
|
rebuilt = Span.from_dict(persisted.to_dict())
|
|
|
|
assert rebuilt.log_level is SpanLogLevel.ERROR
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("span_type", "expected"),
|
|
[
|
|
# INFO set: user-visible semantic operations.
|
|
(SpanType.LLM, SpanLogLevel.INFO),
|
|
(SpanType.CHAT_MODEL, SpanLogLevel.INFO),
|
|
(SpanType.AGENT, SpanLogLevel.INFO),
|
|
(SpanType.TOOL, SpanLogLevel.INFO),
|
|
(SpanType.RETRIEVER, SpanLogLevel.INFO),
|
|
(SpanType.EMBEDDING, SpanLogLevel.INFO),
|
|
# DEBUG set: internal/glue work and unclassified types.
|
|
(SpanType.CHAIN, SpanLogLevel.DEBUG),
|
|
(SpanType.PARSER, SpanLogLevel.DEBUG),
|
|
(SpanType.RERANKER, SpanLogLevel.DEBUG),
|
|
(SpanType.MEMORY, SpanLogLevel.DEBUG),
|
|
(SpanType.WORKFLOW, SpanLogLevel.DEBUG),
|
|
(SpanType.TASK, SpanLogLevel.DEBUG),
|
|
(SpanType.GUARDRAIL, SpanLogLevel.DEBUG),
|
|
(SpanType.EVALUATOR, SpanLogLevel.DEBUG),
|
|
(SpanType.UNKNOWN, SpanLogLevel.DEBUG),
|
|
# Custom (non-built-in) span types fall through to DEBUG.
|
|
("MY_CUSTOM_TYPE", SpanLogLevel.DEBUG),
|
|
(None, SpanLogLevel.DEBUG),
|
|
],
|
|
)
|
|
def test_default_log_level_for_span_type_mapping(span_type, expected):
|
|
assert default_log_level_for_span_type(span_type) is expected
|
|
|
|
|
|
def test_constructor_stamps_default_for_info_span_type():
|
|
span = mlflow.start_span_no_context("s", span_type=SpanType.CHAT_MODEL)
|
|
span.end()
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.INFO
|
|
|
|
|
|
def test_constructor_stamps_default_for_debug_span_type():
|
|
span = mlflow.start_span_no_context("s", span_type=SpanType.PARSER)
|
|
span.end()
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.DEBUG
|
|
|
|
|
|
def test_multi_span_trace_carries_per_span_levels():
|
|
@mlflow.trace(log_level="DEBUG", name="inner")
|
|
def inner():
|
|
return 1
|
|
|
|
@mlflow.trace(log_level="WARNING", name="root")
|
|
def root():
|
|
return inner()
|
|
|
|
root()
|
|
spans = {s.name: s for s in get_traces()[0].data.spans}
|
|
assert spans["root"].log_level is SpanLogLevel.WARNING
|
|
assert spans["inner"].log_level is SpanLogLevel.DEBUG
|
|
|
|
|
|
# ---- Exception → ERROR bump --------------------------------------------------
|
|
|
|
|
|
def test_exception_event_bumps_debug_span_to_error():
|
|
# PARSER defaults to DEBUG via the constructor; an exception event should
|
|
# promote it to ERROR so users with the filter at INFO/WARNING still see it.
|
|
with mlflow.start_span("s", span_type=SpanType.PARSER) as span:
|
|
span.add_event(SpanEvent.from_exception(ValueError("boom")))
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.ERROR
|
|
|
|
|
|
def test_exception_event_bumps_info_span_to_error():
|
|
with mlflow.start_span("s", span_type=SpanType.CHAT_MODEL) as span:
|
|
span.add_event(SpanEvent.from_exception(RuntimeError("boom")))
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.ERROR
|
|
|
|
|
|
def test_exception_event_does_not_demote_critical():
|
|
# User-set CRITICAL must be preserved when an exception fires.
|
|
with mlflow.start_span("s", log_level=SpanLogLevel.CRITICAL) as span:
|
|
span.add_event(SpanEvent.from_exception(RuntimeError("boom")))
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.CRITICAL
|
|
|
|
|
|
def test_record_exception_bumps_to_error():
|
|
# record_exception() goes through add_event under the hood, so the bump
|
|
# should fire here too.
|
|
span = mlflow.start_span_no_context("s", span_type=SpanType.PARSER)
|
|
span.record_exception(ValueError("boom"))
|
|
span.end()
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.ERROR
|
|
|
|
|
|
def test_traced_function_that_raises_is_promoted_to_error():
|
|
# A plain @mlflow.trace function that throws records an exception event via
|
|
# the decorator's error-handling path, which should promote the span.
|
|
@mlflow.trace
|
|
def fn():
|
|
raise ValueError("boom")
|
|
|
|
with pytest.raises(ValueError, match="boom"):
|
|
fn()
|
|
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.ERROR
|
|
|
|
|
|
def test_non_exception_event_does_not_bump_log_level():
|
|
# Plain (non-exception) events must not move the level.
|
|
with mlflow.start_span("s", span_type=SpanType.CHAT_MODEL) as span:
|
|
span.add_event(SpanEvent(name="my_event", attributes={"k": "v"}))
|
|
persisted = get_traces()[0].data.spans[0]
|
|
assert persisted.log_level is SpanLogLevel.INFO
|