wehub-resource-sync
1085 lines
37 KiB
Python
1085 lines
37 KiB
Python
import gzip
|
|
import time
|
|
import zlib
|
|
from pathlib import Path
|
|
from typing import Iterator
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
import requests
|
|
from opentelemetry import trace as otel_trace
|
|
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
|
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
|
|
ExportTraceServiceRequest,
|
|
ExportTraceServiceResponse,
|
|
)
|
|
from opentelemetry.proto.common.v1.common_pb2 import AnyValue, InstrumentationScope, KeyValue
|
|
from opentelemetry.proto.resource.v1.resource_pb2 import Resource
|
|
from opentelemetry.proto.trace.v1.trace_pb2 import ResourceSpans, ScopeSpans
|
|
from opentelemetry.proto.trace.v1.trace_pb2 import Span as OTelProtoSpan
|
|
from opentelemetry.sdk.resources import Resource as OTelSDKResource
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
|
|
from opentelemetry.util._once import Once
|
|
|
|
import mlflow
|
|
from mlflow.server import handlers
|
|
from mlflow.server.fastapi_app import app as mlflow_app
|
|
from mlflow.server.handlers import initialize_backend_stores
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
from mlflow.telemetry.client import TelemetryClient
|
|
from mlflow.telemetry.events import TraceSource, TracesReceivedByServerEvent
|
|
from mlflow.tracing.utils import encode_trace_id
|
|
from mlflow.tracing.utils.otlp import MLFLOW_EXPERIMENT_ID_HEADER
|
|
from mlflow.version import IS_TRACING_SDK_ONLY
|
|
|
|
from tests.helper_functions import get_safe_port
|
|
from tests.tracking.integration_test_utils import ServerThread
|
|
|
|
if IS_TRACING_SDK_ONLY:
|
|
pytest.skip("OTel endpoint tests require full MLflow server", allow_module_level=True)
|
|
|
|
|
|
@pytest.fixture
|
|
def mlflow_server(tmp_path: Path, db_uri: str) -> Iterator[str]:
|
|
artifact_root = tmp_path.as_uri()
|
|
|
|
handlers._tracking_store = None
|
|
handlers._model_registry_store = None
|
|
initialize_backend_stores(db_uri, default_artifact_root=artifact_root)
|
|
|
|
# Start the FastAPI app in a background thread and yield its URL.
|
|
with ServerThread(mlflow_app, get_safe_port()) as url:
|
|
yield url
|
|
# Drain any pending async trace exports while the server is still up.
|
|
# The autouse `enable_async_trace_logging` fixture in tests/tracing/conftest.py
|
|
# also calls flush on teardown, but by that time this ServerThread has
|
|
# already exited, causing the worker to retry against a dead server.
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
|
|
def test_otel_client_sends_spans_to_mlflow_database(mlflow_server: str, monkeypatch):
|
|
"""
|
|
Test end-to-end: OpenTelemetry client sends spans via experiment ID header to MLflow.
|
|
|
|
Note: This test verifies that spans are successfully accepted by the server.
|
|
Without artifact upload, traces won't be retrievable via search_traces.
|
|
"""
|
|
# Enable synchronous trace logging to ensure traces are immediately available
|
|
monkeypatch.setenv("MLFLOW_ASYNC_TRACE_LOGGING", "false")
|
|
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
|
|
experiment = mlflow.set_experiment("otel-test-experiment")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
resource = OTelSDKResource.create({"service.name": "test-service-e2e"})
|
|
tracer_provider = TracerProvider(resource=resource)
|
|
|
|
# First, verify the endpoint is reachable
|
|
test_response = requests.get(f"{mlflow_server}/health", timeout=5)
|
|
assert test_response.status_code == 200, (
|
|
f"Server health check failed: {test_response.status_code}"
|
|
)
|
|
|
|
exporter = OTLPSpanExporter(
|
|
endpoint=f"{mlflow_server}/v1/traces",
|
|
headers={MLFLOW_EXPERIMENT_ID_HEADER: experiment_id},
|
|
timeout=10, # Explicit timeout
|
|
)
|
|
|
|
# Use SimpleSpanProcessor for immediate span export in tests
|
|
# This ensures spans are sent immediately rather than batched
|
|
span_processor = SimpleSpanProcessor(exporter)
|
|
tracer_provider.add_span_processor(span_processor)
|
|
|
|
# Reset the global tracer provider to avoid conflicts with other tests.
|
|
# This is necessary because OpenTelemetry doesn't allow overriding an already-set provider.
|
|
#
|
|
# NOTE: We're using internal APIs here (_TRACER_PROVIDER_SET_ONCE and _TRACER_PROVIDER)
|
|
# because OpenTelemetry doesn't provide a public API to reset the global tracer provider.
|
|
# The library is designed to set the provider once at application startup, which doesn't
|
|
# work well for testing scenarios where different tests need different configurations.
|
|
# This pattern is also used in tests/semantic_kernel/conftest.py for the same reason.
|
|
otel_trace._TRACER_PROVIDER_SET_ONCE = Once()
|
|
otel_trace._TRACER_PROVIDER = None
|
|
|
|
# Set the tracer provider
|
|
otel_trace.set_tracer_provider(tracer_provider)
|
|
tracer = otel_trace.get_tracer(__name__)
|
|
|
|
with tracer.start_as_current_span("otel-e2e-test-span") as span:
|
|
span.set_attribute("test.e2e.attribute", "e2e-test-value")
|
|
# Capture the OTel trace ID to verify it matches the MLflow trace ID
|
|
otel_trace_id = span.get_span_context().trace_id
|
|
# Verify the span was actually created and has valid context
|
|
assert span.get_span_context().is_valid, "Span context is not valid"
|
|
assert otel_trace_id != 0, "Trace ID is zero"
|
|
|
|
# Add a small delay to ensure the server has processed the spans
|
|
time.sleep(0.5)
|
|
|
|
# Wait up to 30 seconds for search_traces() to return a trace
|
|
traces = []
|
|
for _ in range(30):
|
|
traces = mlflow.search_traces(
|
|
locations=[experiment_id], include_spans=False, return_type="list"
|
|
)
|
|
if traces:
|
|
break
|
|
time.sleep(1)
|
|
|
|
assert len(traces) > 0, "No traces found in the database after sending spans"
|
|
|
|
# Verify the trace ID matches the expected format based on the OTel span
|
|
expected_trace_id = f"tr-{encode_trace_id(otel_trace_id)}"
|
|
actual_trace_id = traces[0].info.trace_id
|
|
assert actual_trace_id == expected_trace_id, (
|
|
f"Trace ID mismatch: expected {expected_trace_id}, got {actual_trace_id}"
|
|
)
|
|
|
|
|
|
def test_otel_endpoint_requires_experiment_id_header(mlflow_server: str):
|
|
"""
|
|
Test that the OTel endpoint requires experiment ID header.
|
|
"""
|
|
# Create protobuf request
|
|
span = OTelProtoSpan()
|
|
span.trace_id = bytes.fromhex("0000000000000002" + "0" * 16)
|
|
span.span_id = bytes.fromhex("00000002" + "0" * 8)
|
|
span.name = "test-span-no-header"
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={"Content-Type": "application/x-protobuf"},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_invalid_otel_span_format_returns_400(mlflow_server: str):
|
|
"""
|
|
Test that invalid OpenTelemetry protobuf format returns HTTP 400.
|
|
"""
|
|
# Send completely invalid protobuf data
|
|
invalid_protobuf_data = b"this is not valid protobuf data at all"
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=invalid_protobuf_data,
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: "test-experiment",
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 400, f"Expected 400, got {response.status_code}"
|
|
|
|
|
|
def test_missing_required_span_fields_returns_422(mlflow_server: str):
|
|
"""
|
|
Test that spans that fail MLflow conversion return HTTP 422.
|
|
"""
|
|
# Create protobuf request with missing span_id (this should cause MLflow conversion to fail)
|
|
span = OTelProtoSpan()
|
|
span.trace_id = bytes.fromhex("0000000000000001" + "0" * 16)
|
|
# Don't set span_id - this should cause from_otel_proto to fail
|
|
span.name = "incomplete-span"
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: "0",
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_missing_experiment_id_header_returns_422(mlflow_server: str):
|
|
"""
|
|
Test that missing experiment ID header returns HTTP 422 (FastAPI validation error).
|
|
"""
|
|
# Create valid protobuf request
|
|
span = OTelProtoSpan()
|
|
span.trace_id = bytes.fromhex("0000000000000003" + "0" * 16)
|
|
span.span_id = bytes.fromhex("00000003" + "0" * 8)
|
|
span.name = "test-span"
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={"Content-Type": "application/x-protobuf"},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
|
|
|
|
def test_invalid_content_type_returns_400(mlflow_server: str):
|
|
"""
|
|
Test that invalid Content-Type header returns HTTP 400.
|
|
"""
|
|
# Create a valid OTLP request
|
|
span = OTelProtoSpan()
|
|
span.trace_id = b"1234567890123456"
|
|
span.span_id = b"12345678"
|
|
span.name = "test-span"
|
|
span.start_time_unix_nano = 1000000000
|
|
span.end_time_unix_nano = 2000000000
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
# Send request with unsupported Content-Type
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "text/plain",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: "test-experiment",
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
assert "Invalid Content-Type" in response.text
|
|
|
|
|
|
def test_empty_resource_spans_returns_400(mlflow_server: str):
|
|
request = ExportTraceServiceRequest()
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: "test-experiment",
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
assert "no spans found" in response.text
|
|
|
|
|
|
def test_batch_span_processor_with_multiple_traces(mlflow_server: str):
|
|
"""
|
|
Test that BatchSpanProcessor can send spans from multiple traces in a single request.
|
|
This verifies the server-side grouping by trace_id functionality.
|
|
"""
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
|
|
experiment = mlflow.set_experiment("otel-batch-test-experiment")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
resource = OTelSDKResource.create({"service.name": "test-batch-service"})
|
|
tracer_provider = TracerProvider(resource=resource)
|
|
|
|
exporter = OTLPSpanExporter(
|
|
endpoint=f"{mlflow_server}/v1/traces",
|
|
headers={MLFLOW_EXPERIMENT_ID_HEADER: experiment_id},
|
|
timeout=10,
|
|
)
|
|
|
|
# Use BatchSpanProcessor to batch spans from multiple traces
|
|
span_processor = BatchSpanProcessor(exporter)
|
|
tracer_provider.add_span_processor(span_processor)
|
|
|
|
# Reset the global tracer provider
|
|
otel_trace._TRACER_PROVIDER_SET_ONCE = Once()
|
|
otel_trace._TRACER_PROVIDER = None
|
|
otel_trace.set_tracer_provider(tracer_provider)
|
|
|
|
tracer = otel_trace.get_tracer(__name__)
|
|
|
|
# Create multiple traces with spans
|
|
trace_ids = []
|
|
for i in range(3):
|
|
with tracer.start_as_current_span(f"batch-test-span-{i}") as span:
|
|
span.set_attribute("test.batch.index", i)
|
|
otel_trace_id = span.get_span_context().trace_id
|
|
trace_ids.append(otel_trace_id)
|
|
assert otel_trace_id != 0
|
|
|
|
# Force flush to send all batched spans
|
|
span_processor.force_flush()
|
|
|
|
traces = mlflow.search_traces(
|
|
locations=[experiment_id], include_spans=False, return_type="list"
|
|
)
|
|
|
|
assert len(traces) == 3
|
|
|
|
# Verify all expected trace IDs are present
|
|
expected_trace_ids = {f"tr-{encode_trace_id(tid)}" for tid in trace_ids}
|
|
actual_trace_ids = {trace.info.trace_id for trace in traces}
|
|
|
|
assert expected_trace_ids == actual_trace_ids
|
|
|
|
|
|
def test_multiple_traces_in_single_request(mlflow_server: str):
|
|
"""
|
|
Test that a single request containing spans from multiple traces is handled correctly.
|
|
This simulates what BatchSpanProcessor does internally.
|
|
"""
|
|
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-multi-trace-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
# Create protobuf request with spans from 3 different traces
|
|
request = ExportTraceServiceRequest()
|
|
|
|
for trace_num in range(3):
|
|
# Create a span with unique trace_id
|
|
span = OTelProtoSpan()
|
|
trace_id_hex = f"{trace_num:016x}" + "0" * 16
|
|
span.trace_id = bytes.fromhex(trace_id_hex)
|
|
span.span_id = bytes.fromhex(f"{trace_num:08x}" + "0" * 8)
|
|
span.name = f"multi-trace-span-{trace_num}"
|
|
span.start_time_unix_nano = 1000000000 + trace_num * 1000
|
|
span.end_time_unix_nano = 2000000000 + trace_num * 1000
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
# Send the request with multiple traces
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
traces = mlflow.search_traces(
|
|
locations=[experiment_id], include_spans=False, return_type="list"
|
|
)
|
|
|
|
assert len(traces) == 3
|
|
|
|
|
|
def test_logging_many_traces_in_single_request(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-many-traces-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
# Create a request with 15 different traces (exceeds the 10 thread pool limit)
|
|
request = ExportTraceServiceRequest()
|
|
num_traces = 15
|
|
|
|
for trace_num in range(num_traces):
|
|
span = OTelProtoSpan()
|
|
trace_id_hex = f"{trace_num + 1000:016x}" + "0" * 16
|
|
span.trace_id = bytes.fromhex(trace_id_hex)
|
|
span.span_id = bytes.fromhex(f"{trace_num + 1000:08x}" + "0" * 8)
|
|
span.name = f"many-traces-test-span-{trace_num}"
|
|
span.start_time_unix_nano = 1000000000 + trace_num * 1000
|
|
span.end_time_unix_nano = 2000000000 + trace_num * 1000
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "many-traces-test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
# Send the request and measure response time
|
|
requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
traces = mlflow.search_traces(
|
|
locations=[experiment_id], include_spans=False, return_type="list"
|
|
)
|
|
|
|
assert len(traces) == num_traces
|
|
|
|
|
|
def test_mixed_trace_spans_in_single_request(mlflow_server: str):
|
|
"""
|
|
Test that multiple spans from the same trace, mixed with spans from other traces,
|
|
are grouped and logged correctly.
|
|
"""
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-mixed-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
request = ExportTraceServiceRequest()
|
|
|
|
# Create 2 spans for trace A, 1 span for trace B, 2 spans for trace C
|
|
trace_configs = [
|
|
("A", 0, 2), # trace A with 2 spans
|
|
("B", 1, 1), # trace B with 1 span
|
|
("C", 2, 2), # trace C with 2 spans
|
|
]
|
|
|
|
for trace_name, trace_id_num, span_count in trace_configs:
|
|
trace_id_hex = f"{trace_id_num + 2000:016x}" + "0" * 16
|
|
|
|
for span_num in range(span_count):
|
|
span = OTelProtoSpan(
|
|
trace_id=bytes.fromhex(trace_id_hex),
|
|
span_id=bytes.fromhex(f"{trace_id_num * 100 + span_num:08x}" + "0" * 8),
|
|
name=f"mixed-span-{trace_name}-{span_num}",
|
|
start_time_unix_nano=1000000000 + span_num * 1000,
|
|
end_time_unix_nano=2000000000 + span_num * 1000,
|
|
)
|
|
|
|
scope = InstrumentationScope(
|
|
name="mixed-test-scope",
|
|
)
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
traces = mlflow.search_traces(locations=[experiment_id], include_spans=True, return_type="list")
|
|
|
|
assert len(traces) == 3
|
|
span_counts = [len(trace.data.spans) for trace in traces]
|
|
assert span_counts == [2, 1, 2]
|
|
|
|
|
|
def test_error_logging_spans(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-error-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
resource = OTelSDKResource.create({"service.name": "test-batch-service"})
|
|
tracer_provider = TracerProvider(resource=resource)
|
|
|
|
exporter = OTLPSpanExporter(
|
|
endpoint=f"{mlflow_server}/v1/traces",
|
|
headers={MLFLOW_EXPERIMENT_ID_HEADER: experiment_id},
|
|
timeout=10,
|
|
)
|
|
|
|
# Use BatchSpanProcessor to batch spans from multiple traces
|
|
span_processor = BatchSpanProcessor(exporter)
|
|
tracer_provider.add_span_processor(span_processor)
|
|
|
|
# Reset the global tracer provider
|
|
otel_trace._TRACER_PROVIDER_SET_ONCE = Once()
|
|
otel_trace._TRACER_PROVIDER = None
|
|
otel_trace.set_tracer_provider(tracer_provider)
|
|
|
|
tracer = otel_trace.get_tracer(__name__)
|
|
|
|
original_log_spans = SqlAlchemyStore.log_spans
|
|
call_count = {"count": 0}
|
|
|
|
def mock_log_spans(self, *args, **kwargs):
|
|
if call_count["count"] == 0:
|
|
call_count["count"] += 1
|
|
raise Exception("test_error")
|
|
else:
|
|
return original_log_spans(self, *args, **kwargs)
|
|
|
|
with (
|
|
mock.patch.object(SqlAlchemyStore, "log_spans", mock_log_spans),
|
|
mock.patch(
|
|
"opentelemetry.exporter.otlp.proto.http.trace_exporter._logger.error"
|
|
) as mock_error,
|
|
):
|
|
for _ in range(2):
|
|
with tracer.start_as_current_span("batch-test-span-0"):
|
|
pass
|
|
|
|
span_processor.force_flush()
|
|
|
|
assert any("Failed to export" in error[0][0] for error in mock_error.call_args_list)
|
|
|
|
traces = mlflow.search_traces(
|
|
locations=[experiment_id], include_spans=False, return_type="list"
|
|
)
|
|
|
|
# The OTLP endpoint now calls log_spans once for all spans in the batch.
|
|
# If that call fails, all spans in the batch are dropped together (HTTP 422 is
|
|
# non-retryable for the OTel OTLP exporter). Previously, per-trace calls meant
|
|
# the second trace could still succeed. With the unified log_spans call, the
|
|
# result is 0 stored traces.
|
|
assert len(traces) == 0
|
|
|
|
|
|
def test_otel_trace_received_telemetry_from_mlflow_client(mlflow_server: str):
|
|
"""
|
|
Test TraceReceivedByServerEvent telemetry shows source=MLFLOW_PYTHON_CLIENT for standard client.
|
|
|
|
Uses @mlflow.trace with standard MLflow client configuration, which automatically sends
|
|
User-Agent and X-MLflow-Client-Version headers to identify traces from MLflow client.
|
|
"""
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
mlflow.set_experiment("otel-telemetry-mlflow-client-test")
|
|
|
|
with mock.patch("mlflow.telemetry.track.get_telemetry_client") as mock_get_client:
|
|
mock_client = mock.MagicMock(spec=TelemetryClient)
|
|
mock_get_client.return_value = mock_client
|
|
|
|
@mlflow.trace
|
|
def test_function():
|
|
return "test result"
|
|
|
|
result = test_function()
|
|
assert result == "test result"
|
|
|
|
time.sleep(1)
|
|
|
|
if mock_client.add_record.called:
|
|
record = mock_client.add_record.call_args[0][0]
|
|
assert record.event_name == TracesReceivedByServerEvent.name
|
|
assert record.params["source"] == TraceSource.MLFLOW_PYTHON_CLIENT.value
|
|
assert record.params["count"] == 1
|
|
|
|
|
|
def test_otel_trace_received_telemetry_from_external_client(mlflow_server: str):
|
|
"""
|
|
Test TracesReceivedByServerEvent telemetry shows source=UNKNOWN for external clients.
|
|
|
|
Sends a direct protobuf request without MLflow client headers to simulate an external
|
|
OpenTelemetry client (not MLflow client). Tests with 2 traces to verify count field.
|
|
"""
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-telemetry-external-client-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
trace_id_1 = bytes.fromhex("0000000000000100" + "0" * 16)
|
|
trace_id_2 = bytes.fromhex("0000000000000200" + "0" * 16)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
|
|
# First trace with root span and child spans
|
|
request.resource_spans.append(
|
|
ResourceSpans(
|
|
scope_spans=[
|
|
ScopeSpans(
|
|
scope=InstrumentationScope(name="telemetry-test-scope"),
|
|
spans=[
|
|
OTelProtoSpan(
|
|
trace_id=trace_id_1,
|
|
span_id=bytes.fromhex("00000001" + "0" * 8),
|
|
name="root-span-1",
|
|
start_time_unix_nano=1000000000,
|
|
end_time_unix_nano=2000000000,
|
|
),
|
|
OTelProtoSpan(
|
|
trace_id=trace_id_1,
|
|
span_id=bytes.fromhex("00000002" + "0" * 8),
|
|
parent_span_id=bytes.fromhex("00000001" + "0" * 8),
|
|
name="child-span-1",
|
|
start_time_unix_nano=1100000000,
|
|
end_time_unix_nano=1500000000,
|
|
),
|
|
],
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
# Second trace with root span
|
|
request.resource_spans.append(
|
|
ResourceSpans(
|
|
scope_spans=[
|
|
ScopeSpans(
|
|
scope=InstrumentationScope(name="telemetry-test-scope"),
|
|
spans=[
|
|
OTelProtoSpan(
|
|
trace_id=trace_id_2,
|
|
span_id=bytes.fromhex("00000003" + "0" * 8),
|
|
name="root-span-2",
|
|
start_time_unix_nano=1600000000,
|
|
end_time_unix_nano=1900000000,
|
|
),
|
|
],
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
with mock.patch("mlflow.telemetry.track.get_telemetry_client") as mock_get_client:
|
|
mock_client = mock.MagicMock(spec=TelemetryClient)
|
|
mock_client.config = None # Ensure telemetry is not disabled for any event
|
|
mock_get_client.return_value = mock_client
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
mock_client.add_record.assert_called_once()
|
|
record = mock_client.add_record.call_args[0][0]
|
|
|
|
assert record.event_name == TracesReceivedByServerEvent.name
|
|
assert record.status.value == "success"
|
|
assert record.params["source"] == TraceSource.UNKNOWN.value
|
|
assert record.params["count"] == 2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("service_name", "expected_source", "expected_service_names"),
|
|
[
|
|
("codex_cli_rs", TraceSource.EXTERNAL_OTEL_CLIENT, ["codex_cli_rs"]),
|
|
("gemini-cli", TraceSource.EXTERNAL_OTEL_CLIENT, ["gemini-cli"]),
|
|
("qwen-code", TraceSource.EXTERNAL_OTEL_CLIENT, ["qwen-code"]),
|
|
# Unknown service names are not on the allowlist — source falls back to UNKNOWN
|
|
("my-custom-app", TraceSource.UNKNOWN, None),
|
|
],
|
|
)
|
|
def test_otel_trace_received_telemetry_from_external_otel_client_with_service_name(
|
|
mlflow_server: str,
|
|
service_name: str,
|
|
expected_source: TraceSource,
|
|
expected_service_names: list[str] | None,
|
|
):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-telemetry-service-name-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
trace_id = bytes.fromhex("0000000000000500" + "0" * 16)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(
|
|
ResourceSpans(
|
|
resource=Resource(
|
|
attributes=[
|
|
KeyValue(key="service.name", value=AnyValue(string_value=service_name)),
|
|
]
|
|
),
|
|
scope_spans=[
|
|
ScopeSpans(
|
|
scope=InstrumentationScope(name="test-scope"),
|
|
spans=[
|
|
OTelProtoSpan(
|
|
trace_id=trace_id,
|
|
span_id=bytes.fromhex("00000005" + "0" * 8),
|
|
name="root-span",
|
|
start_time_unix_nano=1000000000,
|
|
end_time_unix_nano=2000000000,
|
|
),
|
|
],
|
|
)
|
|
],
|
|
)
|
|
)
|
|
|
|
with mock.patch("mlflow.telemetry.track.get_telemetry_client") as mock_get_client:
|
|
mock_client = mock.MagicMock(spec=TelemetryClient)
|
|
mock_client.config = None
|
|
mock_get_client.return_value = mock_client
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
mock_client.add_record.assert_called_once()
|
|
record = mock_client.add_record.call_args[0][0]
|
|
|
|
assert record.event_name == TracesReceivedByServerEvent.name
|
|
assert record.params["source"] == expected_source.value
|
|
assert record.params["count"] == 1
|
|
if expected_service_names is not None:
|
|
assert record.params["service_names"] == expected_service_names
|
|
else:
|
|
assert "service_names" not in record.params
|
|
|
|
|
|
def test_service_name_propagated_to_root_span(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-service-name-on-span-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
trace_id = bytes.fromhex("0000000000000600" + "0" * 16)
|
|
root_span_id = bytes.fromhex("00000006" + "0" * 8)
|
|
child_span_id = bytes.fromhex("00000007" + "0" * 8)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(
|
|
ResourceSpans(
|
|
resource=Resource(
|
|
attributes=[
|
|
KeyValue(key="service.name", value=AnyValue(string_value="gemini-cli")),
|
|
]
|
|
),
|
|
scope_spans=[
|
|
ScopeSpans(
|
|
scope=InstrumentationScope(name="test-scope"),
|
|
spans=[
|
|
OTelProtoSpan(
|
|
trace_id=trace_id,
|
|
span_id=root_span_id,
|
|
name="root-span",
|
|
start_time_unix_nano=1000000000,
|
|
end_time_unix_nano=2000000000,
|
|
),
|
|
OTelProtoSpan(
|
|
trace_id=trace_id,
|
|
span_id=child_span_id,
|
|
parent_span_id=root_span_id,
|
|
name="child-span",
|
|
start_time_unix_nano=1100000000,
|
|
end_time_unix_nano=1500000000,
|
|
),
|
|
],
|
|
)
|
|
],
|
|
)
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
traces = mlflow.search_traces(locations=[experiment_id], include_spans=True, return_type="list")
|
|
assert len(traces) == 1
|
|
|
|
root_span = next(s for s in traces[0].data.spans if s.parent_id is None)
|
|
child_span = next(s for s in traces[0].data.spans if s.parent_id is not None)
|
|
|
|
# service.name should be on the root span only
|
|
assert root_span.get_attribute("service.name") == "gemini-cli"
|
|
assert child_span.get_attribute("service.name") is None
|
|
|
|
|
|
def test_otlp_json_encoding(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-json-encoding-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
trace_id_hex = "61181d5c16b3b25d966891c3fc595af9"
|
|
span_id_hex = "508755a859b294e3"
|
|
|
|
payload = {
|
|
"resourceSpans": [
|
|
{
|
|
"resource": {
|
|
"attributes": [
|
|
{"key": "service.name", "value": {"stringValue": "gemini-cli"}},
|
|
]
|
|
},
|
|
"scopeSpans": [
|
|
{
|
|
"scope": {"name": "gemini-cli"},
|
|
"spans": [
|
|
{
|
|
"traceId": trace_id_hex,
|
|
"spanId": span_id_hex,
|
|
"name": "user_prompt",
|
|
"startTimeUnixNano": str(int(time.time() * 1e9)),
|
|
"endTimeUnixNano": str(int(time.time() * 1e9) + 1000000000),
|
|
"attributes": [
|
|
{
|
|
"key": "gen_ai.operation.name",
|
|
"value": {"stringValue": "user_prompt"},
|
|
},
|
|
],
|
|
}
|
|
],
|
|
}
|
|
],
|
|
}
|
|
]
|
|
}
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
json=payload,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
traces = mlflow.search_traces(locations=[experiment_id], include_spans=True, return_type="list")
|
|
assert len(traces) == 1
|
|
|
|
root_span = traces[0].data.spans[0]
|
|
assert root_span.name == "user_prompt"
|
|
assert root_span.get_attribute("service.name") == "gemini-cli"
|
|
|
|
|
|
def test_response_is_protobuf_format(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-protobuf-response-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
# Create a valid OTLP request
|
|
span = OTelProtoSpan()
|
|
span.trace_id = bytes.fromhex("0000000000000400" + "0" * 16)
|
|
span.span_id = bytes.fromhex("00000004" + "0" * 8)
|
|
span.name = "protobuf-test-span"
|
|
span.start_time_unix_nano = 1000000000
|
|
span.end_time_unix_nano = 2000000000
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=request.SerializeToString(),
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
# Verify response status and content-type
|
|
assert response.status_code == 200
|
|
assert response.headers.get("content-type") == "application/x-protobuf"
|
|
|
|
# Verify the response can be parsed as a valid ExportTraceServiceResponse
|
|
response_message = ExportTraceServiceResponse()
|
|
response_message.ParseFromString(response.content)
|
|
|
|
|
|
def _build_valid_otlp_request() -> str:
|
|
"""Helper: Build a valid OTLP ExportTraceServiceRequest protobuf."""
|
|
|
|
span = OTelProtoSpan()
|
|
span.trace_id = bytes.fromhex("0000000000000400" + "0" * 16)
|
|
span.span_id = bytes.fromhex("00000004" + "0" * 8)
|
|
span.name = "test-span"
|
|
span.start_time_unix_nano = 1000000000
|
|
span.end_time_unix_nano = 2000000000
|
|
|
|
scope = InstrumentationScope()
|
|
scope.name = "test-scope"
|
|
|
|
scope_spans = ScopeSpans()
|
|
scope_spans.scope.CopyFrom(scope)
|
|
scope_spans.spans.append(span)
|
|
|
|
resource = Resource()
|
|
resource_spans = ResourceSpans()
|
|
resource_spans.resource.CopyFrom(resource)
|
|
resource_spans.scope_spans.append(scope_spans)
|
|
|
|
request = ExportTraceServiceRequest()
|
|
request.resource_spans.append(resource_spans)
|
|
|
|
return request.SerializeToString()
|
|
|
|
|
|
def test_otlp_traces_no_compression(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-identity-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
data = _build_valid_otlp_request()
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=data,
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
# No Content-Encoding -> no compression
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.headers.get("content-type") == "application/x-protobuf"
|
|
|
|
msg = ExportTraceServiceResponse()
|
|
msg.ParseFromString(response.content)
|
|
|
|
|
|
def test_otlp_traces_gzip_compression(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-gzip-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
raw = _build_valid_otlp_request()
|
|
compressed = gzip.compress(raw)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=compressed,
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
"Content-Encoding": "gzip",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.headers.get("content-type") == "application/x-protobuf"
|
|
|
|
msg = ExportTraceServiceResponse()
|
|
msg.ParseFromString(response.content)
|
|
|
|
|
|
def test_otlp_traces_deflate_compression(mlflow_server: str):
|
|
mlflow.set_tracking_uri(mlflow_server)
|
|
experiment = mlflow.set_experiment("otel-deflate-test")
|
|
experiment_id = experiment.experiment_id
|
|
|
|
raw = _build_valid_otlp_request()
|
|
compressed = zlib.compress(raw)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_server}/v1/traces",
|
|
data=compressed,
|
|
headers={
|
|
"Content-Type": "application/x-protobuf",
|
|
"Content-Encoding": "deflate",
|
|
MLFLOW_EXPERIMENT_ID_HEADER: experiment_id,
|
|
},
|
|
timeout=10,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.headers.get("content-type") == "application/x-protobuf"
|
|
|
|
msg = ExportTraceServiceResponse()
|
|
msg.ParseFromString(response.content)
|