wehub-resource-sync
349 lines
12 KiB
Python
349 lines
12 KiB
Python
import os
|
|
import time
|
|
import uuid
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
from unittest import mock
|
|
|
|
import opentelemetry.trace as trace_api
|
|
import pytest
|
|
from opentelemetry.sdk.trace import Event, ReadableSpan
|
|
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
|
|
|
import mlflow
|
|
from mlflow.entities import Trace, TraceData, TraceInfo
|
|
from mlflow.entities.trace_location import TraceLocation
|
|
from mlflow.entities.trace_state import TraceState
|
|
from mlflow.ml_package_versions import FLAVOR_TO_MODULE_NAME
|
|
from mlflow.tracing.client import TracingClient
|
|
from mlflow.tracing.constant import TRACE_SCHEMA_VERSION, TRACE_SCHEMA_VERSION_KEY
|
|
from mlflow.tracing.export.inference_table import pop_trace
|
|
from mlflow.tracing.processor.mlflow_v3 import MlflowV3SpanProcessor
|
|
from mlflow.tracing.processor.otel import OtelSpanProcessor
|
|
from mlflow.tracing.provider import _get_tracer
|
|
from mlflow.tracking.fluent import _get_experiment_id
|
|
from mlflow.utils.autologging_utils import AUTOLOGGING_INTEGRATIONS, get_autolog_function
|
|
from mlflow.utils.autologging_utils.safety import revert_patches
|
|
from mlflow.version import IS_TRACING_SDK_ONLY
|
|
|
|
|
|
def create_mock_otel_span(
|
|
trace_id: int,
|
|
span_id: int,
|
|
name: str = "test_span",
|
|
parent_id: int | None = None,
|
|
start_time: int | None = None,
|
|
end_time: int | None = None,
|
|
):
|
|
"""
|
|
Create a mock OpenTelemetry span for testing purposes.
|
|
|
|
OpenTelemetry doesn't allow creating a span outside of a tracer. So here we create a mock span
|
|
that extends ReadableSpan (data object) and exposes the necessary attributes for testing.
|
|
"""
|
|
|
|
@dataclass
|
|
class _MockSpanContext:
|
|
trace_id: str
|
|
span_id: str
|
|
trace_flags: trace_api.TraceFlags = trace_api.TraceFlags(1)
|
|
trace_state: trace_api.TraceState = field(default_factory=trace_api.TraceState)
|
|
|
|
class _MockOTelSpan(trace_api.Span, ReadableSpan):
|
|
def __init__(
|
|
self,
|
|
name,
|
|
context,
|
|
parent,
|
|
start_time=None,
|
|
end_time=None,
|
|
status=trace_api.Status(trace_api.StatusCode.UNSET),
|
|
):
|
|
self._name = name
|
|
self._parent = parent
|
|
self._context = context
|
|
self._start_time = start_time if start_time is not None else int(time.time() * 1e9)
|
|
self._end_time = end_time
|
|
self._status = status
|
|
self._attributes = {}
|
|
self._events = []
|
|
|
|
# NB: The following methods are defined as abstract method in the Span class.
|
|
def set_attributes(self, attributes):
|
|
self._attributes.update(attributes)
|
|
|
|
def set_attribute(self, key, value):
|
|
self._attributes[key] = value
|
|
|
|
def set_status(self, status):
|
|
self._status = status
|
|
|
|
def add_event(self, name, attributes=None, timestamp=None):
|
|
self._events.append(Event(name, attributes, timestamp))
|
|
|
|
def get_span_context(self):
|
|
return self._context
|
|
|
|
def is_recording(self):
|
|
return self._end_time is None
|
|
|
|
def update_name(self, name):
|
|
self.name = name
|
|
|
|
def end(self, end_time_ns=None):
|
|
pass
|
|
|
|
def record_exception():
|
|
pass
|
|
|
|
return _MockOTelSpan(
|
|
name=name,
|
|
context=_MockSpanContext(trace_id, span_id),
|
|
parent=_MockSpanContext(trace_id, parent_id) if parent_id else None,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|
|
|
|
|
|
def create_trace(request_id) -> Trace:
|
|
return Trace(info=create_test_trace_info(request_id), data=TraceData())
|
|
|
|
|
|
def create_test_trace_info(
|
|
trace_id,
|
|
experiment_id="test",
|
|
request_time=0,
|
|
execution_duration=1,
|
|
state=TraceState.OK,
|
|
trace_metadata=None,
|
|
tags=None,
|
|
):
|
|
# Add schema version to metadata if not provided, to match real trace creation behavior
|
|
final_metadata = trace_metadata or {}
|
|
if TRACE_SCHEMA_VERSION_KEY not in final_metadata:
|
|
final_metadata = final_metadata.copy()
|
|
final_metadata[TRACE_SCHEMA_VERSION_KEY] = str(TRACE_SCHEMA_VERSION)
|
|
|
|
return TraceInfo(
|
|
trace_id=trace_id,
|
|
trace_location=TraceLocation.from_experiment_id(experiment_id),
|
|
request_time=request_time,
|
|
execution_duration=execution_duration,
|
|
state=state,
|
|
trace_metadata=final_metadata,
|
|
tags=tags or {},
|
|
)
|
|
|
|
|
|
def create_test_trace_info_with_uc_table(
|
|
trace_id: str, catalog_name: str, schema_name: str
|
|
) -> TraceInfo:
|
|
return TraceInfo(
|
|
trace_id=trace_id,
|
|
trace_location=TraceLocation.from_databricks_uc_schema(catalog_name, schema_name),
|
|
request_time=0,
|
|
execution_duration=1,
|
|
state=TraceState.OK,
|
|
trace_metadata={TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)},
|
|
tags={},
|
|
)
|
|
|
|
|
|
def get_traces(experiment_id=None) -> list[Trace]:
|
|
# Flush any pending async trace writes before querying so tests see complete results.
|
|
mlflow.flush_trace_async_logging()
|
|
# Get all traces from the backend
|
|
return TracingClient().search_traces(
|
|
locations=[experiment_id or _get_experiment_id()],
|
|
)
|
|
|
|
|
|
def purge_traces(experiment_id=None):
|
|
if len(get_traces(experiment_id)) == 0:
|
|
return
|
|
|
|
# Delete all traces from the backend
|
|
TracingClient().delete_traces(
|
|
experiment_id=experiment_id or _get_experiment_id(),
|
|
max_traces=1000,
|
|
max_timestamp_millis=int(time.time() * 1000),
|
|
)
|
|
|
|
|
|
def get_tracer_tracking_uri() -> str | None:
|
|
"""Get current tracking URI configured as the trace export destination."""
|
|
from opentelemetry import trace
|
|
|
|
tracer = _get_tracer(__name__)
|
|
if isinstance(tracer, trace.ProxyTracer):
|
|
tracer = tracer._tracer
|
|
span_processor = tracer.span_processor._span_processors[0]
|
|
|
|
if isinstance(span_processor, MlflowV3SpanProcessor):
|
|
return span_processor.span_exporter._client.tracking_uri
|
|
|
|
|
|
@pytest.fixture
|
|
def reset_autolog_state():
|
|
"""Reset autologging state to avoid interference between tests"""
|
|
yield
|
|
|
|
for flavor in FLAVOR_TO_MODULE_NAME:
|
|
# 1. Remove post-import hooks (registered by global mlflow.autolog() function)
|
|
mlflow.utils.import_hooks._post_import_hooks.pop(flavor, None)
|
|
|
|
for flavor in AUTOLOGGING_INTEGRATIONS.keys():
|
|
# 2. Disable autologging for the flavor. This is necessary because some autologging
|
|
# update global settings (e.g. callbacks) and we need to revert them.
|
|
try:
|
|
if autolog := get_autolog_function(flavor):
|
|
autolog(disable=True)
|
|
except ImportError:
|
|
pass
|
|
|
|
# 3. Revert any patches applied by autologging
|
|
revert_patches(flavor)
|
|
|
|
AUTOLOGGING_INTEGRATIONS.clear()
|
|
|
|
|
|
def score_in_model_serving(model_uri: str, model_input: dict[str, Any]):
|
|
"""
|
|
A helper function to emulate model prediction inside a Databricks model serving environment.
|
|
|
|
This is highly simplified version, but captures important aspects for testing tracing:
|
|
1. Setting env vars that users set for enable tracing in model serving
|
|
2. Load the model in a background thread
|
|
"""
|
|
from mlflow.pyfunc.context import Context, set_prediction_context
|
|
|
|
with mock.patch.dict(
|
|
"os.environ",
|
|
os.environ | {"IS_IN_DB_MODEL_SERVING_ENV": "true", "ENABLE_MLFLOW_TRACING": "true"},
|
|
clear=True,
|
|
):
|
|
# Reset tracing setup to start fresh w/ model serving environment
|
|
mlflow.tracing.reset()
|
|
|
|
def _load_model():
|
|
return mlflow.pyfunc.load_model(model_uri)
|
|
|
|
with ThreadPoolExecutor(
|
|
max_workers=1, thread_name_prefix="test-tracing-helper"
|
|
) as executor:
|
|
model = executor.submit(_load_model).result()
|
|
|
|
# Score the model
|
|
request_id = uuid.uuid4().hex
|
|
with set_prediction_context(Context(request_id=request_id)):
|
|
predictions = model.predict(model_input)
|
|
|
|
trace = pop_trace(request_id)
|
|
return (request_id, predictions, trace)
|
|
|
|
|
|
def skip_when_testing_trace_sdk(f):
|
|
# Decorator to Skip the test if only mlflow-tracing package is installed and
|
|
# not the full mlflow package.
|
|
msg = "Skipping test because it requires mlflow or mlflow-skinny to be installed."
|
|
skip_decorator = pytest.mark.skipif(IS_TRACING_SDK_ONLY, reason=msg)
|
|
return skip_decorator(f)
|
|
|
|
|
|
def skip_module_when_testing_trace_sdk():
|
|
"""Skip the entire module if only mlflow-tracing package is installed"""
|
|
if IS_TRACING_SDK_ONLY:
|
|
pytest.skip(
|
|
"Skipping test because it requires mlflow or mlflow-skinny to be installed.",
|
|
allow_module_level=True,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def capture_otel_export():
|
|
"""Capture traces in memory for testing otel export."""
|
|
from mlflow.tracing.provider import provider
|
|
|
|
exporter = InMemorySpanExporter()
|
|
provider.get_or_init_tracer("test")
|
|
tp = provider.get()
|
|
processor = OtelSpanProcessor(span_exporter=exporter, export_metrics=False)
|
|
processor._should_register_traces = False
|
|
tp.add_span_processor(processor)
|
|
yield exporter, processor
|
|
processor.force_flush(timeout_millis=5000)
|
|
processor.shutdown()
|
|
|
|
|
|
V2_TRACE_DICT = {
|
|
"info": {
|
|
"request_id": "58f4e27101304034b15c512b603bf1b2",
|
|
"experiment_id": "0",
|
|
"timestamp_ms": 100,
|
|
"execution_time_ms": 200,
|
|
"status": "OK",
|
|
"request_metadata": {
|
|
"mlflow.trace_schema.version": "2",
|
|
"mlflow.traceInputs": '{"x": 2, "y": 5}',
|
|
"mlflow.traceOutputs": "8",
|
|
},
|
|
"tags": {
|
|
"mlflow.source.name": "test",
|
|
"mlflow.source.type": "LOCAL",
|
|
"mlflow.traceName": "predict",
|
|
"mlflow.artifactLocation": "/path/to/artifact",
|
|
},
|
|
"assessments": [],
|
|
},
|
|
"data": {
|
|
"spans": [
|
|
{
|
|
"name": "predict",
|
|
"context": {
|
|
"span_id": "0d48a6670588966b",
|
|
"trace_id": "63076d0c1b90f1df0970f897dc428bd6",
|
|
},
|
|
"parent_id": None,
|
|
"start_time": 100,
|
|
"end_time": 200,
|
|
"status_code": "OK",
|
|
"status_message": "",
|
|
"attributes": {
|
|
"mlflow.traceRequestId": '"58f4e27101304034b15c512b603bf1b2"',
|
|
"mlflow.spanType": '"UNKNOWN"',
|
|
"mlflow.spanFunctionName": '"predict"',
|
|
"mlflow.spanInputs": '{"x": 2, "y": 5}',
|
|
"mlflow.spanOutputs": "8",
|
|
},
|
|
"events": [],
|
|
},
|
|
{
|
|
"name": "add_one_with_custom_name",
|
|
"context": {
|
|
"span_id": "6fc32f36ef591f60",
|
|
"trace_id": "63076d0c1b90f1df0970f897dc428bd6",
|
|
},
|
|
"parent_id": "0d48a6670588966b",
|
|
"start_time": 300,
|
|
"end_time": 400,
|
|
"status_code": "OK",
|
|
"status_message": "",
|
|
"attributes": {
|
|
"mlflow.traceRequestId": '"58f4e27101304034b15c512b603bf1b2"',
|
|
"mlflow.spanType": '"LLM"',
|
|
"delta": "1",
|
|
"metadata": '{"foo": "bar"}',
|
|
"datetime": '"2025-04-29 08:37:06.772253"',
|
|
"mlflow.spanFunctionName": '"add_one"',
|
|
"mlflow.spanInputs": '{"z": 7}',
|
|
"mlflow.spanOutputs": "8",
|
|
},
|
|
"events": [],
|
|
},
|
|
],
|
|
"request": '{"x": 2, "y": 5}',
|
|
"response": "8",
|
|
},
|
|
}
|