mlflow--mlflow/tests/tracing/processor/test_inference_table_processor.py

import json
from unittest import mock

import pytest

from mlflow.entities.span import LiveSpan
from mlflow.entities.trace_state import TraceState
from mlflow.tracing.constant import (
    TRACE_SCHEMA_VERSION,
    TRACE_SCHEMA_VERSION_KEY,
    SpanAttributeKey,
    TraceMetadataKey,
)
from mlflow.tracing.processor.inference_table import (
    _HEADER_REQUEST_ID_KEY,
    InferenceTableSpanProcessor,
)
from mlflow.tracing.trace_manager import InMemoryTraceManager
from mlflow.tracing.utils import generate_trace_id_v3
from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_MODEL_SERVING_ENDPOINT_NAME

from tests.tracing.helper import (
    create_mock_otel_span,
    create_test_trace_info,
    skip_module_when_testing_trace_sdk,
)

skip_module_when_testing_trace_sdk()

from mlflow.pyfunc.context import Context, set_prediction_context
from mlflow.tracking.fluent import set_active_model

_OTEL_TRACE_ID = 12345
_DATABRICKS_REQUEST_ID = "databricks-request-id"


@pytest.mark.parametrize("context_type", ["mlflow", "flask"])
def test_on_start(context_type):
    # Root span should create a new trace on start
    span = create_mock_otel_span(
        trace_id=_OTEL_TRACE_ID, span_id=1, parent_id=None, start_time=5_000_000
    )
    trace_manager = InMemoryTraceManager.get_instance()
    processor = InferenceTableSpanProcessor(span_exporter=mock.MagicMock())
    model = set_active_model(name="test-model")

    if context_type == "mlflow":
        with set_prediction_context(
            Context(request_id=_DATABRICKS_REQUEST_ID, endpoint_name="test-endpoint")
        ):
            processor.on_start(span)
    else:
        with mock.patch(
            "mlflow.tracing.processor.inference_table._get_flask_request"
        ) as mock_get_flask_request:
            request = mock_get_flask_request.return_value
            request.headers = {_HEADER_REQUEST_ID_KEY: _DATABRICKS_REQUEST_ID}

            processor.on_start(span)

    expected_trace_id = generate_trace_id_v3(span)

    with trace_manager.get_trace(expected_trace_id) as trace:
        # Trace ID should be generated by MLflow
        assert trace.info.trace_id == expected_trace_id
        # Databricks request ID should be set to the client request ID
        assert trace.info.client_request_id == _DATABRICKS_REQUEST_ID
        assert trace.info.experiment_id is None
        assert trace.info.timestamp_ms == 5
        assert trace.info.execution_time_ms is None
        assert trace.info.state == TraceState.IN_PROGRESS

        if context_type == "mlflow":
            assert trace.info.request_metadata == {
                TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION),
                MLFLOW_DATABRICKS_MODEL_SERVING_ENDPOINT_NAME: "test-endpoint",
                TraceMetadataKey.MODEL_ID: model.model_id,
            }

    # Child span should not create a new trace
    child_span = create_mock_otel_span(
        trace_id=_OTEL_TRACE_ID, span_id=2, parent_id=1, start_time=8_000_000
    )

    with set_prediction_context(Context(request_id=_DATABRICKS_REQUEST_ID)):
        processor.on_start(child_span)

    assert child_span.attributes.get(SpanAttributeKey.REQUEST_ID) == json.dumps(expected_trace_id)

    # start time should not be overwritten
    with trace_manager.get_trace(expected_trace_id) as trace:
        assert trace.info.timestamp_ms == 5


def test_on_start_with_experiment_id_env_var(monkeypatch):
    # When the MLFLOW_EXPERIMENT_ID env var is set, it should be populated into the trace info
    monkeypatch.setenv("MLFLOW_EXPERIMENT_ID", "123")

    span = create_mock_otel_span(
        trace_id=_OTEL_TRACE_ID, span_id=1, parent_id=None, start_time=5_000_000
    )
    trace_manager = InMemoryTraceManager.get_instance()
    processor = InferenceTableSpanProcessor(span_exporter=mock.MagicMock())

    with set_prediction_context(Context(request_id=_DATABRICKS_REQUEST_ID)):
        processor.on_start(span)

    expected_trace_id = generate_trace_id_v3(span)
    with trace_manager.get_trace(expected_trace_id) as trace:
        assert trace.info.trace_id == expected_trace_id
        assert trace.info.client_request_id == _DATABRICKS_REQUEST_ID
        assert trace.info.experiment_id == "123"


def test_on_end():
    otel_span = create_mock_otel_span(
        name="foo",
        trace_id=_OTEL_TRACE_ID,
        span_id=1,
        parent_id=None,
        start_time=5_000_000,
        end_time=9_000_000,
    )

    trace_id = generate_trace_id_v3(otel_span)
    trace_info = create_test_trace_info(trace_id, 0)
    trace_manager = InMemoryTraceManager.get_instance()
    trace_manager.register_trace(_OTEL_TRACE_ID, trace_info)

    span = LiveSpan(otel_span, trace_id)
    span.set_status("OK")
    span.set_inputs({"input1": "very long input" * 100})
    span.set_outputs({"output": "very long output" * 100})

    mock_exporter = mock.MagicMock()
    processor = InferenceTableSpanProcessor(span_exporter=mock_exporter)

    processor.on_end(otel_span)

    mock_exporter.export.assert_called_once_with((otel_span,))
    # Trace info should be updated according to the span attributes
    assert trace_info.state == TraceState.OK
    assert trace_info.execution_duration == 4

    # Non-root span should not be exported
    mock_exporter.reset_mock()
    child_span = create_mock_otel_span(trace_id=_OTEL_TRACE_ID, span_id=2, parent_id=1)
    processor.on_end(child_span)
    mock_exporter.export.assert_not_called()


def test_on_end_preserves_user_set_trace_state():
    otel_span = create_mock_otel_span(
        name="foo",
        trace_id=_OTEL_TRACE_ID,
        span_id=1,
        parent_id=None,
        start_time=5_000_000,
        end_time=9_000_000,
    )

    trace_id = generate_trace_id_v3(otel_span)
    trace_info = create_test_trace_info(trace_id, 0)
    trace_manager = InMemoryTraceManager.get_instance()
    trace_manager.register_trace(_OTEL_TRACE_ID, trace_info)

    # Explicitly set trace state to ERROR (user action)
    with trace_manager.get_trace(trace_id) as trace:
        trace.info.state = TraceState.ERROR

    span = LiveSpan(otel_span, trace_id)
    span.set_status("OK")  # Span status is OK
    span.set_inputs({"input1": "test"})
    span.set_outputs({"output": "test"})

    mock_exporter = mock.MagicMock()
    processor = InferenceTableSpanProcessor(span_exporter=mock_exporter)

    processor.on_end(otel_span)

    # Trace state should remain ERROR (user-set), not be overwritten by span status (OK)
    with trace_manager.get_trace(trace_id) as trace:
        assert trace.info.state == TraceState.ERROR
    assert trace_info.execution_duration == 4


def test_on_end_updates_trace_state_when_in_progress():
    otel_span = create_mock_otel_span(
        name="foo",
        trace_id=_OTEL_TRACE_ID,
        span_id=1,
        parent_id=None,
        start_time=5_000_000,
        end_time=9_000_000,
    )

    trace_id = generate_trace_id_v3(otel_span)
    trace_info = create_test_trace_info(trace_id, 0, state=TraceState.IN_PROGRESS)
    trace_manager = InMemoryTraceManager.get_instance()
    trace_manager.register_trace(_OTEL_TRACE_ID, trace_info)

    # Trace state remains IN_PROGRESS (not explicitly set by user)
    with trace_manager.get_trace(trace_id) as trace:
        assert trace.info.state == TraceState.IN_PROGRESS

    span = LiveSpan(otel_span, trace_id)
    span.set_status("ERROR")  # Span status is ERROR
    span.set_inputs({"input1": "test"})
    span.set_outputs({"output": "test"})

    mock_exporter = mock.MagicMock()
    processor = InferenceTableSpanProcessor(span_exporter=mock_exporter)

    processor.on_end(otel_span)

    # Trace state should be updated to ERROR from span status
    with trace_manager.get_trace(trace_id) as trace:
        assert trace.info.state == TraceState.ERROR
    assert trace_info.execution_duration == 4