wehub-resource-sync
331 lines
11 KiB
Python
331 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
from mlflow.entities._mlflow_object import _MlflowObject
|
|
from mlflow.entities.span import Span, SpanType
|
|
from mlflow.entities.trace_data import TraceData
|
|
from mlflow.entities.trace_info import TraceInfo
|
|
from mlflow.entities.trace_info_v2 import TraceInfoV2
|
|
from mlflow.environment_variables import MLFLOW_TRACING_SQL_WAREHOUSE_ID
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
from mlflow.protos.service_pb2 import Trace as ProtoTrace
|
|
|
|
if TYPE_CHECKING:
|
|
from mlflow.entities.assessment import Assessment
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class Trace(_MlflowObject):
|
|
"""A trace object.
|
|
|
|
Args:
|
|
info: A lightweight object that contains the metadata of a trace.
|
|
data: A container object that holds the spans data of a trace.
|
|
"""
|
|
|
|
info: TraceInfo
|
|
data: TraceData
|
|
|
|
def __post_init__(self):
|
|
if isinstance(self.info, TraceInfoV2):
|
|
self.info = self.info.to_v3(request=self.data.request, response=self.data.response)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"Trace(trace_id={self.info.trace_id})"
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"info": self.info.to_dict(), "data": self.data.to_dict()}
|
|
|
|
def to_json(self, pretty=False) -> str:
|
|
from mlflow.tracing.utils import TraceJSONEncoder
|
|
|
|
return json.dumps(self.to_dict(), cls=TraceJSONEncoder, indent=2 if pretty else None)
|
|
|
|
@classmethod
|
|
def from_dict(cls, trace_dict: dict[str, Any]) -> Trace:
|
|
info = trace_dict.get("info")
|
|
data = trace_dict.get("data")
|
|
if info is None or data is None:
|
|
raise MlflowException(
|
|
"Unable to parse Trace from dictionary. Expected keys: 'info' and 'data'. "
|
|
f"Received keys: {list(trace_dict.keys())}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
return cls(
|
|
info=TraceInfo.from_dict(info),
|
|
data=TraceData.from_dict(data),
|
|
)
|
|
|
|
@classmethod
|
|
def from_json(cls, trace_json: str) -> Trace:
|
|
try:
|
|
trace_dict = json.loads(trace_json)
|
|
except json.JSONDecodeError as e:
|
|
raise MlflowException(
|
|
f"Unable to parse trace JSON: {trace_json}. Error: {e}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
return cls.from_dict(trace_dict)
|
|
|
|
def _serialize_for_mimebundle(self):
|
|
# databricks notebooks will use the trace ID to
|
|
# fetch the trace from the backend. including the
|
|
# full JSON can cause notebooks to exceed size limits
|
|
return json.dumps({
|
|
"trace_id": self.info.trace_id,
|
|
# TODO: remove this once sql_warehouse_id
|
|
# is optional in the v4 tracing APIs
|
|
"sql_warehouse_id": MLFLOW_TRACING_SQL_WAREHOUSE_ID.get(),
|
|
})
|
|
|
|
def _repr_mimebundle_(self, include=None, exclude=None):
|
|
"""
|
|
This method is used to trigger custom display logic in IPython notebooks.
|
|
See https://ipython.readthedocs.io/en/stable/config/integrating.html#MyObject
|
|
for more details.
|
|
|
|
At the moment, the only supported MIME type is "application/databricks.mlflow.trace",
|
|
which contains a JSON representation of the Trace object. This object is deserialized
|
|
in Databricks notebooks to display the Trace object in a nicer UI.
|
|
"""
|
|
from mlflow.tracing.display import (
|
|
get_display_handler,
|
|
get_notebook_iframe_html,
|
|
is_using_tracking_server,
|
|
)
|
|
from mlflow.utils.databricks_utils import is_in_databricks_runtime
|
|
|
|
bundle = {"text/plain": repr(self)}
|
|
|
|
if not get_display_handler().disabled:
|
|
if is_in_databricks_runtime():
|
|
bundle["application/databricks.mlflow.trace"] = self._serialize_for_mimebundle()
|
|
elif is_using_tracking_server():
|
|
bundle["text/html"] = get_notebook_iframe_html([self])
|
|
|
|
return bundle
|
|
|
|
def to_pandas_dataframe_row(self) -> dict[str, Any]:
|
|
return {
|
|
"trace_id": self.info.trace_id,
|
|
"trace": self.to_json(), # json string to be compatible with Spark DataFrame
|
|
"client_request_id": self.info.client_request_id,
|
|
"state": self.info.state,
|
|
"request_time": self.info.request_time,
|
|
"execution_duration": self.info.execution_duration,
|
|
"request": self._deserialize_json_attr(self.data.request),
|
|
"response": self._deserialize_json_attr(self.data.response),
|
|
"trace_metadata": self.info.trace_metadata,
|
|
"tags": self.info.tags,
|
|
"spans": [span.to_dict() for span in self.data.spans],
|
|
"assessments": [assessment.to_dictionary() for assessment in self.info.assessments],
|
|
}
|
|
|
|
def _deserialize_json_attr(self, value: str):
|
|
try:
|
|
return json.loads(value)
|
|
except Exception:
|
|
_logger.debug(f"Failed to deserialize JSON attribute: {value}", exc_info=True)
|
|
return value
|
|
|
|
def search_spans(
|
|
self,
|
|
span_type: SpanType | None = None,
|
|
name: str | re.Pattern | None = None,
|
|
span_id: str | None = None,
|
|
) -> list[Span]:
|
|
"""
|
|
Search for spans that match the given criteria within the trace.
|
|
|
|
Args:
|
|
span_type: The type of the span to search for.
|
|
name: The name of the span to search for. This can be a string or a regular expression.
|
|
span_id: The ID of the span to search for.
|
|
|
|
Returns:
|
|
A list of spans that match the given criteria.
|
|
If there is no match, an empty list is returned.
|
|
|
|
.. code-block:: python
|
|
|
|
import mlflow
|
|
import re
|
|
from mlflow.entities import SpanType
|
|
|
|
|
|
@mlflow.trace(span_type=SpanType.CHAIN)
|
|
def run(x: int) -> int:
|
|
x = add_one(x)
|
|
x = add_two(x)
|
|
x = multiply_by_two(x)
|
|
return x
|
|
|
|
|
|
@mlflow.trace(span_type=SpanType.TOOL)
|
|
def add_one(x: int) -> int:
|
|
return x + 1
|
|
|
|
|
|
@mlflow.trace(span_type=SpanType.TOOL)
|
|
def add_two(x: int) -> int:
|
|
return x + 2
|
|
|
|
|
|
@mlflow.trace(span_type=SpanType.TOOL)
|
|
def multiply_by_two(x: int) -> int:
|
|
return x * 2
|
|
|
|
|
|
# Run the function and get the trace
|
|
y = run(2)
|
|
trace_id = mlflow.get_last_active_trace_id()
|
|
trace = mlflow.get_trace(trace_id)
|
|
|
|
# 1. Search spans by name (exact match)
|
|
spans = trace.search_spans(name="add_one")
|
|
print(spans)
|
|
# Output: [Span(name='add_one', ...)]
|
|
|
|
# 2. Search spans by name (regular expression)
|
|
pattern = re.compile(r"add.*")
|
|
spans = trace.search_spans(name=pattern)
|
|
print(spans)
|
|
# Output: [Span(name='add_one', ...), Span(name='add_two', ...)]
|
|
|
|
# 3. Search spans by type
|
|
spans = trace.search_spans(span_type=SpanType.LLM)
|
|
print(spans)
|
|
# Output: [Span(name='run', ...)]
|
|
|
|
# 4. Search spans by name and type
|
|
spans = trace.search_spans(name="add_one", span_type=SpanType.TOOL)
|
|
print(spans)
|
|
# Output: [Span(name='add_one', ...)]
|
|
"""
|
|
|
|
def _match_name(span: Span) -> bool:
|
|
if isinstance(name, str):
|
|
return span.name == name
|
|
elif isinstance(name, re.Pattern):
|
|
return name.search(span.name) is not None
|
|
elif name is None:
|
|
return True
|
|
else:
|
|
raise MlflowException(
|
|
f"Invalid type for 'name'. Expected str or re.Pattern. Got: {type(name)}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
def _match_type(span: Span) -> bool:
|
|
if isinstance(span_type, str):
|
|
return span.span_type == span_type
|
|
elif span_type is None:
|
|
return True
|
|
else:
|
|
raise MlflowException(
|
|
"Invalid type for 'span_type'. Expected str or mlflow.entities.SpanType. "
|
|
f"Got: {type(span_type)}",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
def _match_id(span: Span) -> bool:
|
|
if span_id is None:
|
|
return True
|
|
else:
|
|
return span.span_id == span_id
|
|
|
|
return [
|
|
span
|
|
for span in self.data.spans
|
|
if _match_name(span) and _match_type(span) and _match_id(span)
|
|
]
|
|
|
|
def search_assessments(
|
|
self,
|
|
name: str | None = None,
|
|
*,
|
|
span_id: str | None = None,
|
|
all: bool = False,
|
|
type: Literal["expectation", "feedback"] | None = None,
|
|
) -> list["Assessment"]:
|
|
"""
|
|
Get assessments for a given name / span ID. By default, this only returns assessments
|
|
that are valid (i.e. have not been overridden by another assessment). To return all
|
|
assessments, specify `all=True`.
|
|
|
|
Args:
|
|
name: The name of the assessment to get. If not provided, this will match
|
|
all assessment names.
|
|
span_id: The span ID to get assessments for.
|
|
If not provided, this will match all spans.
|
|
all: If True, return all assessments regardless of validity.
|
|
type: The type of assessment to get (one of "feedback" or "expectation").
|
|
If not provided, this will match all assessment types.
|
|
|
|
Returns:
|
|
A list of assessments that meet the given conditions.
|
|
"""
|
|
|
|
def validate_type(assessment: Assessment) -> bool:
|
|
from mlflow.entities.assessment import Expectation, Feedback
|
|
|
|
if type == "expectation":
|
|
return isinstance(assessment, Expectation)
|
|
elif type == "feedback":
|
|
return isinstance(assessment, Feedback)
|
|
|
|
return True
|
|
|
|
return [
|
|
assessment
|
|
for assessment in self.info.assessments
|
|
if (name is None or assessment.name == name)
|
|
and (span_id is None or assessment.span_id == span_id)
|
|
# valid defaults to true, so Nones are valid
|
|
and (all or assessment.valid in (True, None))
|
|
and (type is None or validate_type(assessment))
|
|
]
|
|
|
|
@staticmethod
|
|
def pandas_dataframe_columns() -> list[str]:
|
|
return [
|
|
"trace_id",
|
|
"trace",
|
|
"client_request_id",
|
|
"state",
|
|
"request_time",
|
|
"execution_duration",
|
|
"request",
|
|
"response",
|
|
"trace_metadata",
|
|
"tags",
|
|
"spans",
|
|
"assessments",
|
|
]
|
|
|
|
def to_proto(self):
|
|
"""
|
|
Convert into a proto object to sent to the MLflow backend.
|
|
"""
|
|
|
|
return ProtoTrace(
|
|
trace_info=self.info.to_proto(),
|
|
spans=[span.to_otel_proto() for span in self.data.spans],
|
|
)
|
|
|
|
@classmethod
|
|
def from_proto(cls, proto: ProtoTrace) -> "Trace":
|
|
return cls(
|
|
info=TraceInfo.from_proto(proto.trace_info),
|
|
data=TraceData(spans=[Span.from_otel_proto(span) for span in proto.spans]),
|
|
)
|