wehub-resource-sync
190 lines
8.1 KiB
Python
190 lines
8.1 KiB
Python
import json
|
|
import logging
|
|
|
|
from opentelemetry.context import Context
|
|
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
|
|
from opentelemetry.sdk.trace import Span as OTelSpan
|
|
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
|
|
|
|
from mlflow.entities.span import create_mlflow_span
|
|
from mlflow.entities.trace_info import TraceInfo
|
|
from mlflow.entities.trace_location import TraceLocation
|
|
from mlflow.entities.trace_state import TraceState
|
|
from mlflow.environment_variables import MLFLOW_EXPERIMENT_ID
|
|
from mlflow.tracing.constant import (
|
|
TRACE_SCHEMA_VERSION,
|
|
TRACE_SCHEMA_VERSION_KEY,
|
|
SpanAttributeKey,
|
|
TraceMetadataKey,
|
|
)
|
|
from mlflow.tracing.trace_manager import InMemoryTraceManager
|
|
from mlflow.tracing.utils import (
|
|
_try_get_prediction_context,
|
|
aggregate_cost_from_spans,
|
|
aggregate_usage_from_spans,
|
|
generate_trace_id_v3,
|
|
get_otel_attribute,
|
|
maybe_get_dependencies_schemas,
|
|
maybe_get_request_id,
|
|
should_compute_cost_client_side,
|
|
update_trace_state_from_span_conditionally,
|
|
)
|
|
from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_MODEL_SERVING_ENDPOINT_NAME
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
_HEADER_REQUEST_ID_KEY = "X-Request-Id"
|
|
|
|
|
|
# Extracting for testing purposes
|
|
def _get_flask_request():
|
|
import flask
|
|
|
|
if flask.has_request_context():
|
|
return flask.request
|
|
|
|
|
|
class InferenceTableSpanProcessor(SimpleSpanProcessor):
|
|
"""
|
|
Defines custom hooks to be executed when a span is started or ended (before exporting).
|
|
|
|
This processor is used when the tracing destination is Databricks Inference Table.
|
|
"""
|
|
|
|
def __init__(self, span_exporter: SpanExporter):
|
|
super().__init__(span_exporter)
|
|
self._trace_manager = InMemoryTraceManager.get_instance()
|
|
|
|
def on_start(self, span: OTelSpan, parent_context: Context | None = None):
|
|
"""
|
|
Handle the start of a span. This method is called when an OpenTelemetry span is started.
|
|
|
|
Args:
|
|
span: An OpenTelemetry Span object that is started.
|
|
parent_context: The context of the span. Note that this is only passed when the context
|
|
object is explicitly specified to OpenTelemetry start_span call. If the parent
|
|
span is obtained from the global context, it won't be passed here so we should not
|
|
rely on it.
|
|
"""
|
|
databricks_request_id = maybe_get_request_id()
|
|
if databricks_request_id is None:
|
|
# NB: This is currently used for streaming inference in Databricks Model Serving.
|
|
# In normal prediction, serving logic pass the request ID using the
|
|
# `with set_prediction_context` context manager that wraps `model.predict`
|
|
# call. However, in streaming case, the context manager is not applicable
|
|
# so we still need to rely on Flask request context (which is set to the
|
|
# stream response via flask.stream_with_context()
|
|
if flask_request := _get_flask_request():
|
|
databricks_request_id = flask_request.headers.get(_HEADER_REQUEST_ID_KEY)
|
|
if not databricks_request_id:
|
|
_logger.warning(
|
|
"Databricks request ID not found in the request headers. "
|
|
"Skipping trace processing."
|
|
)
|
|
return
|
|
else:
|
|
_logger.warning(
|
|
"Failed to get Databricks request ID from the request headers because "
|
|
"request context is not available. Skipping trace processing."
|
|
)
|
|
return
|
|
|
|
trace_id = generate_trace_id_v3(span)
|
|
span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(trace_id))
|
|
tags = {}
|
|
if dependencies_schema := maybe_get_dependencies_schemas():
|
|
tags.update(dependencies_schema)
|
|
|
|
if span._parent is None:
|
|
trace_info = TraceInfo(
|
|
trace_id=trace_id,
|
|
client_request_id=databricks_request_id,
|
|
# NB: Agent framework populate the MLFLOW_EXPERIMENT_ID environment variable
|
|
# with the experiment ID to which the model is logged. We don't use the
|
|
# _get_experiment_id() method because it will fallback to the default
|
|
# experiment if the MLFLOW_EXPERIMENT_ID is not set.
|
|
trace_location=TraceLocation.from_experiment_id(MLFLOW_EXPERIMENT_ID.get()),
|
|
request_time=span.start_time // 1_000_000, # nanosecond to millisecond
|
|
execution_duration=None,
|
|
state=TraceState.IN_PROGRESS,
|
|
trace_metadata=self._get_trace_metadata(),
|
|
tags=tags,
|
|
)
|
|
self._trace_manager.register_trace(span.context.trace_id, trace_info)
|
|
|
|
self._trace_manager.register_span(create_mlflow_span(span, trace_id))
|
|
|
|
def on_end(self, span: OTelReadableSpan) -> None:
|
|
"""
|
|
Handle the end of a span. This method is called when an OpenTelemetry span is ended.
|
|
|
|
Args:
|
|
span: An OpenTelemetry ReadableSpan object that is ended.
|
|
"""
|
|
# Processing the trace only when the root span is found.
|
|
if span._parent is not None:
|
|
return
|
|
|
|
trace_id = get_otel_attribute(span, SpanAttributeKey.REQUEST_ID)
|
|
with self._trace_manager.get_trace(trace_id) as trace:
|
|
if trace is None:
|
|
_logger.debug(f"Trace data with trace ID {trace_id} not found.")
|
|
return
|
|
|
|
trace.info.execution_duration = (span.end_time - span.start_time) // 1_000_000
|
|
|
|
# Update trace state from span status, but only if the user hasn't explicitly set
|
|
# a different trace status
|
|
update_trace_state_from_span_conditionally(trace, span)
|
|
|
|
spans = list(trace.span_dict.values())
|
|
|
|
# Aggregate token usage and cost as best-effort: this metadata is optional, and
|
|
# a failure here must never abort root-span export / trace finalization (#24344).
|
|
try:
|
|
if usage := aggregate_usage_from_spans(spans):
|
|
trace.info.request_metadata[TraceMetadataKey.TOKEN_USAGE] = json.dumps(usage)
|
|
|
|
if should_compute_cost_client_side() and (cost := aggregate_cost_from_spans(spans)):
|
|
trace.info.request_metadata[TraceMetadataKey.COST] = json.dumps(cost)
|
|
except Exception as e:
|
|
_logger.warning(
|
|
f"Failed to aggregate token usage/cost for trace {trace_id}: {e}. "
|
|
"Continuing finalization without it. For full traceback, set logging "
|
|
"level to debug.",
|
|
exc_info=_logger.isEnabledFor(logging.DEBUG),
|
|
)
|
|
|
|
super().on_end(span)
|
|
|
|
def _get_trace_metadata(self) -> dict[str, str]:
|
|
metadata = {TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)}
|
|
|
|
context = _try_get_prediction_context()
|
|
if context:
|
|
metadata[MLFLOW_DATABRICKS_MODEL_SERVING_ENDPOINT_NAME] = context.endpoint_name or ""
|
|
|
|
# The model ID fetch order is
|
|
# 1. from the context, this could be set by databricks serving
|
|
# 2. from the active model, this could be set by model loading or with environment variable
|
|
if context and context.model_id:
|
|
metadata[TraceMetadataKey.MODEL_ID] = context.model_id
|
|
_logger.debug(f"Model id fetched from the context: {context.model_id}")
|
|
else:
|
|
try:
|
|
from mlflow.tracking.fluent import _get_active_model_id_global
|
|
|
|
if active_model_id := _get_active_model_id_global():
|
|
metadata[TraceMetadataKey.MODEL_ID] = active_model_id
|
|
_logger.debug(
|
|
f"Adding active model ID {active_model_id} to the trace metadata."
|
|
)
|
|
except Exception as e:
|
|
_logger.debug(
|
|
f"Failed to get model ID from the active model: {e}. "
|
|
"Skipping adding model ID to trace metadata.",
|
|
exc_info=True,
|
|
)
|
|
return metadata
|