wehub-resource-sync
1078 lines
43 KiB
Python
1078 lines
43 KiB
Python
import json
|
|
import logging
|
|
import time
|
|
from collections import defaultdict
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import nullcontext
|
|
from typing import TYPE_CHECKING, Sequence
|
|
|
|
import mlflow
|
|
|
|
if TYPE_CHECKING:
|
|
from mlflow.genai.label_schemas.label_schemas import (
|
|
InputCategorical,
|
|
InputNumeric,
|
|
InputPassFail,
|
|
InputText,
|
|
LabelSchema,
|
|
LabelSchemaType,
|
|
)
|
|
from mlflow.genai.review_queues import (
|
|
ReviewItemType,
|
|
ReviewQueue,
|
|
ReviewQueueItem,
|
|
ReviewQueueType,
|
|
ReviewStatus,
|
|
)
|
|
from mlflow.entities.assessment import Assessment
|
|
from mlflow.entities.issue import Issue, IssueSeverity, IssueStatus
|
|
from mlflow.entities.model_registry import PromptVersion
|
|
from mlflow.entities.span import NO_OP_SPAN_TRACE_ID, Span
|
|
from mlflow.entities.trace import Trace
|
|
from mlflow.entities.trace_data import TraceData
|
|
from mlflow.entities.trace_info import TraceInfo
|
|
from mlflow.entities.trace_location import UCSchemaLocation, UnityCatalog
|
|
from mlflow.environment_variables import (
|
|
_MLFLOW_SEARCH_TRACES_MAX_BATCH_SIZE,
|
|
MLFLOW_GET_TRACE_OTEL_INITIAL_RETRY_INTERVAL_SECONDS,
|
|
MLFLOW_GET_TRACE_OTEL_MAX_RETRY_INTERVAL_SECONDS,
|
|
MLFLOW_GET_TRACE_OTEL_RETRY_TIMEOUT_SECONDS,
|
|
MLFLOW_SEARCH_TRACES_MAX_THREADS,
|
|
MLFLOW_TRACING_SQL_WAREHOUSE_ID,
|
|
)
|
|
from mlflow.exceptions import (
|
|
MlflowException,
|
|
MlflowNotImplementedException,
|
|
MlflowTraceDataCorrupted,
|
|
MlflowTraceDataException,
|
|
MlflowTraceDataNotFound,
|
|
)
|
|
from mlflow.protos.databricks_pb2 import (
|
|
BAD_REQUEST,
|
|
INVALID_PARAMETER_VALUE,
|
|
NOT_FOUND,
|
|
RESOURCE_DOES_NOT_EXIST,
|
|
)
|
|
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
|
|
from mlflow.store.entities.paged_list import PagedList
|
|
from mlflow.store.tracking import SEARCH_TRACES_DEFAULT_MAX_RESULTS
|
|
from mlflow.telemetry.events import LogAssessmentEvent, StartTraceEvent, TraceAttachmentsEvent
|
|
from mlflow.telemetry.track import record_usage_event
|
|
from mlflow.tracing.attachments import Attachment
|
|
from mlflow.tracing.constant import (
|
|
SpansLocation,
|
|
TraceMetadataKey,
|
|
TraceTagKey,
|
|
)
|
|
from mlflow.tracing.trace_manager import InMemoryTraceManager
|
|
from mlflow.tracing.utils import TraceJSONEncoder, exclude_immutable_tags, parse_trace_id_v4
|
|
from mlflow.tracing.utils.artifact_utils import get_artifact_uri_for_trace
|
|
from mlflow.tracking._tracking_service.utils import _get_store, _resolve_tracking_uri
|
|
from mlflow.utils import is_uuid
|
|
from mlflow.utils.mlflow_tags import IMMUTABLE_TAGS
|
|
from mlflow.utils.uri import add_databricks_profile_info_to_artifact_uri, is_databricks_uri
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TracingClient:
|
|
"""
|
|
Client of an MLflow Tracking Server that creates and manages experiments and runs.
|
|
"""
|
|
|
|
def __init__(self, tracking_uri: str | None = None):
|
|
"""
|
|
Args:
|
|
tracking_uri: Address of local or remote tracking server.
|
|
"""
|
|
self.tracking_uri = _resolve_tracking_uri(tracking_uri)
|
|
# NB: Fetch the tracking store (`self.store`) upon client initialization to ensure that
|
|
# the tracking URI is valid and the store can be properly resolved. We define `store` as a
|
|
# property method to ensure that the client is serializable, even if the store is not
|
|
# self.store
|
|
self.store
|
|
|
|
@property
|
|
def store(self):
|
|
return _get_store(self.tracking_uri)
|
|
|
|
@record_usage_event(StartTraceEvent)
|
|
def start_trace(self, trace_info: TraceInfo) -> TraceInfo:
|
|
"""
|
|
Create a new trace in the backend.
|
|
|
|
Args:
|
|
trace_info: The TraceInfo object to record in the backend.
|
|
|
|
Returns:
|
|
The returned TraceInfoV3 object from the backend.
|
|
"""
|
|
return self.store.start_trace(trace_info=trace_info)
|
|
|
|
def log_spans(self, location: str, spans: list[Span]) -> list[Span]:
|
|
"""
|
|
Log spans to the backend.
|
|
|
|
Args:
|
|
location: The location to log spans to. It should either be an experiment ID or a
|
|
Unity Catalog table name.
|
|
spans: List of Span objects to log.
|
|
|
|
Returns:
|
|
List of logged Span objects from the backend.
|
|
"""
|
|
return self.store.log_spans(
|
|
location=location,
|
|
spans=spans,
|
|
tracking_uri=self.tracking_uri if is_databricks_uri(self.tracking_uri) else None,
|
|
)
|
|
|
|
def delete_traces(
|
|
self,
|
|
experiment_id: str,
|
|
max_timestamp_millis: int | None = None,
|
|
max_traces: int | None = None,
|
|
trace_ids: list[str] | None = None,
|
|
) -> int:
|
|
return self.store.delete_traces(
|
|
experiment_id=experiment_id,
|
|
max_timestamp_millis=max_timestamp_millis,
|
|
max_traces=max_traces,
|
|
trace_ids=trace_ids,
|
|
)
|
|
|
|
def get_trace_info(self, trace_id: str) -> TraceInfo:
|
|
"""
|
|
Get the trace info matching the ``trace_id``.
|
|
|
|
Args:
|
|
trace_id: String id of the trace to fetch.
|
|
|
|
Returns:
|
|
TraceInfo object, of type ``mlflow.entities.trace_info.TraceInfo``.
|
|
"""
|
|
with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
|
|
if trace is not None:
|
|
return trace.info
|
|
|
|
return self.store.get_trace_info(trace_id)
|
|
|
|
def get_trace(self, trace_id: str) -> Trace:
|
|
"""
|
|
Get the trace matching the ``trace_id``.
|
|
|
|
Args:
|
|
trace_id: String id of the trace to fetch.
|
|
|
|
Returns:
|
|
The fetched Trace object, of type ``mlflow.entities.Trace``.
|
|
"""
|
|
location, _ = parse_trace_id_v4(trace_id)
|
|
if location is not None:
|
|
# For a V4 trace, load spans from the v4 BatchGetTraces endpoint.
|
|
# BatchGetTraces returns an empty list if the trace is not found, which is
|
|
# retried with exponential backoff (capped per-interval) up to
|
|
# MLFLOW_GET_TRACE_OTEL_RETRY_TIMEOUT_SECONDS in total.
|
|
deadline = time.monotonic() + MLFLOW_GET_TRACE_OTEL_RETRY_TIMEOUT_SECONDS.get()
|
|
initial_interval = max(0.0, MLFLOW_GET_TRACE_OTEL_INITIAL_RETRY_INTERVAL_SECONDS.get())
|
|
max_interval = max(0.0, MLFLOW_GET_TRACE_OTEL_MAX_RETRY_INTERVAL_SECONDS.get())
|
|
attempt = 0
|
|
while True:
|
|
if traces := self.store.batch_get_traces([trace_id], location):
|
|
return traces[0]
|
|
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
break
|
|
|
|
interval = min(initial_interval * 2**attempt, max_interval, remaining)
|
|
attempt += 1
|
|
_logger.debug(
|
|
f"Trace not found, retrying in {interval:.2f} seconds (attempt {attempt})"
|
|
)
|
|
time.sleep(interval)
|
|
|
|
raise MlflowException(
|
|
message=f"Trace with ID {trace_id} is not found.",
|
|
error_code=NOT_FOUND,
|
|
)
|
|
else:
|
|
try:
|
|
trace_info = self.get_trace_info(trace_id)
|
|
# if the trace is stored in the tracking store or archive repo, load spans via the
|
|
# store/server path; otherwise, load spans from the artifact repository
|
|
if trace_info.tags.get(TraceTagKey.SPANS_LOCATION) in (
|
|
SpansLocation.TRACKING_STORE,
|
|
SpansLocation.ARCHIVE_REPO,
|
|
):
|
|
try:
|
|
return self.store.get_trace(trace_id)
|
|
except MlflowNotImplementedException:
|
|
pass
|
|
if traces := self.store.batch_get_traces([trace_info.trace_id]):
|
|
return traces[0]
|
|
else:
|
|
raise MlflowException(
|
|
f"Trace with ID {trace_id} is not found.",
|
|
error_code=NOT_FOUND,
|
|
)
|
|
else:
|
|
trace_data = self._download_trace_data(trace_info)
|
|
except MlflowTraceDataNotFound:
|
|
raise MlflowException(
|
|
message=(
|
|
f"Trace with ID {trace_id} cannot be loaded because it is missing span "
|
|
"data. Please try creating or loading another trace."
|
|
),
|
|
error_code=BAD_REQUEST,
|
|
) from None # Ensure the original spammy exception is not included in the traceback
|
|
except MlflowTraceDataCorrupted:
|
|
raise MlflowException(
|
|
message=(
|
|
f"Trace with ID {trace_id} cannot be loaded because its span data"
|
|
" is corrupted. Please try creating or loading another trace."
|
|
),
|
|
error_code=BAD_REQUEST,
|
|
) from None # Ensure the original spammy exception is not included in the traceback
|
|
return Trace(trace_info, trace_data)
|
|
|
|
def get_online_trace_details(
|
|
self,
|
|
trace_id: str,
|
|
source_inference_table: str,
|
|
source_databricks_request_id: str,
|
|
) -> str:
|
|
return self.store.get_online_trace_details(
|
|
trace_id=trace_id,
|
|
source_inference_table=source_inference_table,
|
|
source_databricks_request_id=source_databricks_request_id,
|
|
)
|
|
|
|
def _search_traces(
|
|
self,
|
|
experiment_ids: list[str] | None = None,
|
|
filter_string: str | None = None,
|
|
max_results: int = SEARCH_TRACES_DEFAULT_MAX_RESULTS,
|
|
order_by: list[str] | None = None,
|
|
page_token: str | None = None,
|
|
model_id: str | None = None,
|
|
locations: list[str] | None = None,
|
|
):
|
|
return self.store.search_traces(
|
|
experiment_ids=experiment_ids,
|
|
filter_string=filter_string,
|
|
max_results=max_results,
|
|
order_by=order_by,
|
|
page_token=page_token,
|
|
model_id=model_id,
|
|
locations=locations,
|
|
)
|
|
|
|
def search_traces(
|
|
self,
|
|
experiment_ids: list[str] | None = None,
|
|
filter_string: str | None = None,
|
|
max_results: int = SEARCH_TRACES_DEFAULT_MAX_RESULTS,
|
|
order_by: list[str] | None = None,
|
|
page_token: str | None = None,
|
|
run_id: str | None = None,
|
|
include_spans: bool = True,
|
|
model_id: str | None = None,
|
|
locations: list[str] | None = None,
|
|
) -> PagedList[Trace]:
|
|
"""
|
|
Return traces that match the given list of search expressions within the experiments.
|
|
|
|
Args:
|
|
experiment_ids: List of experiment ids to scope the search. Deprecated,
|
|
use `locations` instead.
|
|
filter_string: A search filter string.
|
|
max_results: Maximum number of traces desired.
|
|
order_by: List of order_by clauses.
|
|
page_token: Token specifying the next page of results. It should be obtained from
|
|
a ``search_traces`` call.
|
|
run_id: A run id to scope the search. When a trace is created under an active run,
|
|
it will be associated with the run and you can filter on the run id to retrieve
|
|
the trace.
|
|
include_spans: If ``True``, include spans in the returned traces. Otherwise, only
|
|
the trace metadata is returned, e.g., trace ID, start time, end time, etc,
|
|
without any spans.
|
|
model_id: If specified, return traces associated with the model ID.
|
|
locations: A list of locations to search over. To search over experiments, provide
|
|
a list of experiment IDs. To search over UC tables on databricks, provide
|
|
a list of locations in the format
|
|
`<catalog_name>.<schema_name>[.<table_prefix>]`.
|
|
|
|
Returns:
|
|
A :py:class:`PagedList <mlflow.store.entities.PagedList>` of
|
|
:py:class:`Trace <mlflow.entities.Trace>` objects that satisfy the search
|
|
expressions. If the underlying tracking store supports pagination, the token for the
|
|
next page may be obtained via the ``token`` attribute of the returned object; however,
|
|
some store implementations may not support pagination and thus the returned token would
|
|
not be meaningful in such cases.
|
|
"""
|
|
if model_id is not None:
|
|
if filter_string:
|
|
raise MlflowException(
|
|
message=(
|
|
"Cannot specify both `model_id` or `filter_string` in the search_traces "
|
|
"call."
|
|
),
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
# if sql_warehouse_id is not set then we convert model_id to filter_string,
|
|
# because `_search_unified_traces` requires sql warehouse id existing.
|
|
if MLFLOW_TRACING_SQL_WAREHOUSE_ID.get() is None:
|
|
filter_string = f"request_metadata.`mlflow.modelId` = '{model_id}'"
|
|
model_id = None
|
|
|
|
if run_id:
|
|
run = self.store.get_run(run_id)
|
|
if run.info.experiment_id not in locations:
|
|
raise MlflowException(
|
|
f"Run {run_id} belongs to experiment {run.info.experiment_id}, which is not "
|
|
f"in the list of locations provided: {locations}. Please include "
|
|
f"experiment {run.info.experiment_id} in the `locations` parameter to "
|
|
"search for traces from this run.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
|
|
additional_filter = f"attribute.run_id = '{run_id}'"
|
|
if filter_string:
|
|
if TraceMetadataKey.SOURCE_RUN in filter_string:
|
|
raise MlflowException(
|
|
"You cannot filter by run_id when it is already part of the filter string."
|
|
f"Please remove the {TraceMetadataKey.SOURCE_RUN} filter from the filter "
|
|
"string and try again.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|
|
filter_string += f" AND {additional_filter}"
|
|
else:
|
|
filter_string = additional_filter
|
|
|
|
traces = []
|
|
next_max_results = max_results
|
|
next_token = page_token
|
|
|
|
max_workers = MLFLOW_SEARCH_TRACES_MAX_THREADS.get()
|
|
executor = (
|
|
ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="MlflowTracingSearch")
|
|
if include_spans
|
|
else nullcontext()
|
|
)
|
|
with executor:
|
|
while len(traces) < max_results:
|
|
trace_infos, next_token = self._search_traces(
|
|
experiment_ids=experiment_ids,
|
|
filter_string=filter_string,
|
|
max_results=next_max_results,
|
|
order_by=order_by,
|
|
page_token=next_token,
|
|
model_id=model_id,
|
|
locations=locations,
|
|
)
|
|
|
|
if include_spans:
|
|
trace_infos_by_location = self._group_trace_infos_by_location(trace_infos)
|
|
traces.extend(self._load_traces_by_location(trace_infos_by_location, executor))
|
|
|
|
else:
|
|
traces.extend(Trace(t, TraceData(spans=[])) for t in trace_infos)
|
|
|
|
if not next_token:
|
|
break
|
|
|
|
next_max_results = max_results - len(traces)
|
|
|
|
return PagedList(traces, next_token)
|
|
|
|
def batch_get_traces(self, trace_ids: list[str], location: str | None = None) -> list[Trace]:
|
|
"""
|
|
Retrieve multiple traces by their IDs.
|
|
|
|
Args:
|
|
trace_ids: List of trace IDs to retrieve.
|
|
location: Optional location (e.g., "catalog.schema" for UC schema) to search for traces.
|
|
|
|
Returns:
|
|
List of Trace objects.
|
|
"""
|
|
if not trace_ids:
|
|
return []
|
|
|
|
# If location is provided, this is a UC schema/v4 call - delegate directly
|
|
if location is not None:
|
|
return self.store.batch_get_traces(trace_ids, location)
|
|
|
|
# Get trace infos (metadata only) to determine where spans are stored.
|
|
# Fall back to store.batch_get_traces directly if the store doesn't
|
|
# implement batch_get_trace_infos (e.g. DatabricksRestStore).
|
|
try:
|
|
trace_infos = self.store.batch_get_trace_infos(trace_ids)
|
|
except MlflowNotImplementedException:
|
|
return self.store.batch_get_traces(trace_ids, location)
|
|
|
|
trace_infos_by_location = self._group_trace_infos_by_location(trace_infos)
|
|
|
|
max_workers = min(len(trace_ids), MLFLOW_SEARCH_TRACES_MAX_THREADS.get())
|
|
with ThreadPoolExecutor(
|
|
max_workers=max_workers, thread_name_prefix="MlflowTracingBatchGet"
|
|
) as executor:
|
|
return self._load_traces_by_location(trace_infos_by_location, executor)
|
|
|
|
def _download_spans_from_batch_get_traces(
|
|
self, trace_ids: list[str], location: str | None, executor: ThreadPoolExecutor
|
|
) -> list[Trace]:
|
|
"""
|
|
Fetch full traces including spans from the store batch-get path.
|
|
BatchGetTrace endpoint only support up to 10 traces in a single call.
|
|
"""
|
|
traces = []
|
|
|
|
def _fetch_minibatch(ids: list[str]) -> list[Trace]:
|
|
return self.store.batch_get_traces(ids, location) or []
|
|
|
|
batch_size = _MLFLOW_SEARCH_TRACES_MAX_BATCH_SIZE.get()
|
|
batches = [trace_ids[i : i + batch_size] for i in range(0, len(trace_ids), batch_size)]
|
|
for minibatch_traces in executor.map(_fetch_minibatch, batches):
|
|
traces.extend(minibatch_traces)
|
|
return traces
|
|
|
|
def _load_traces_by_location(
|
|
self,
|
|
trace_infos_by_location: dict[str, list[TraceInfo]],
|
|
executor: ThreadPoolExecutor,
|
|
) -> list[Trace]:
|
|
traces = []
|
|
for location, location_trace_infos in trace_infos_by_location.items():
|
|
if location == SpansLocation.ARTIFACT_REPO:
|
|
traces.extend(
|
|
tr
|
|
for tr in executor.map(
|
|
self._download_spans_from_artifact_repo,
|
|
location_trace_infos,
|
|
)
|
|
if tr
|
|
)
|
|
else:
|
|
batch_get_location = None if location == SpansLocation.ARCHIVE_REPO else location
|
|
traces.extend(
|
|
self._download_spans_from_batch_get_traces(
|
|
[t.trace_id for t in location_trace_infos], batch_get_location, executor
|
|
)
|
|
)
|
|
return traces
|
|
|
|
def _download_spans_from_artifact_repo(self, trace_info: TraceInfo) -> Trace | None:
|
|
"""
|
|
Download trace data for the given trace_info and returns a Trace object.
|
|
If the download fails (e.g., the trace data is missing or corrupted), returns None.
|
|
|
|
This is used for traces whose spans are fetched directly from artifact storage, including
|
|
the existing artifact-backed path.
|
|
"""
|
|
is_online_trace = is_uuid(trace_info.trace_id)
|
|
is_databricks = is_databricks_uri(self.tracking_uri)
|
|
|
|
# For online traces in Databricks, we need to get trace data from a different endpoint
|
|
try:
|
|
if is_databricks and is_online_trace:
|
|
# For online traces, get data from the online API
|
|
trace_data = self.get_online_trace_details(
|
|
trace_id=trace_info.trace_id,
|
|
source_inference_table=trace_info.request_metadata.get("mlflow.sourceTable"),
|
|
source_databricks_request_id=trace_info.request_metadata.get(
|
|
"mlflow.databricksRequestId"
|
|
),
|
|
)
|
|
trace_data = TraceData.from_dict(json.loads(trace_data))
|
|
else:
|
|
# For offline traces, download data from artifact storage
|
|
trace_data = self._download_trace_data(trace_info)
|
|
except MlflowTraceDataException as e:
|
|
_logger.warning(
|
|
(
|
|
f"Failed to download trace data for trace {trace_info.trace_id!r} "
|
|
f"with {e.ctx}. For full traceback, set logging level to DEBUG."
|
|
),
|
|
exc_info=_logger.isEnabledFor(logging.DEBUG),
|
|
)
|
|
return None
|
|
else:
|
|
return Trace(trace_info, trace_data)
|
|
|
|
def _group_trace_infos_by_location(
|
|
self, trace_infos: list[TraceInfo]
|
|
) -> dict[str, list[TraceInfo]]:
|
|
"""
|
|
Group the trace infos based on where the trace data is stored.
|
|
|
|
Returns:
|
|
A dictionary mapping location to a list of trace infos.
|
|
"""
|
|
trace_infos_by_location = defaultdict(list)
|
|
for trace_info in trace_infos:
|
|
if uc_schema := trace_info.trace_location.uc_schema:
|
|
location = f"{uc_schema.catalog_name}.{uc_schema.schema_name}"
|
|
trace_infos_by_location[location].append(trace_info)
|
|
elif uc_tp := trace_info.trace_location.uc_table_prefix:
|
|
location = f"{uc_tp.catalog_name}.{uc_tp.schema_name}.{uc_tp.table_prefix}"
|
|
trace_infos_by_location[location].append(trace_info)
|
|
elif trace_info.trace_location.mlflow_experiment:
|
|
# DB-backed experiment traces use the tracking store. Legacy artifact-backed traces
|
|
# and archived traces are grouped by their spans-location tag so they can be
|
|
# fetched through the correct non-DB retrieval path.
|
|
spans_location = trace_info.tags.get(TraceTagKey.SPANS_LOCATION)
|
|
if spans_location == SpansLocation.TRACKING_STORE:
|
|
# location is not used for traces with mlflow experiment location in tracking
|
|
# store, so we use None as the location
|
|
trace_infos_by_location[None].append(trace_info)
|
|
elif spans_location in (SpansLocation.ARTIFACT_REPO, SpansLocation.ARCHIVE_REPO):
|
|
trace_infos_by_location[spans_location].append(trace_info)
|
|
else:
|
|
# Older traces may not set spansLocation and should continue to use the
|
|
# artifact-backed trace-data path.
|
|
trace_infos_by_location[SpansLocation.ARTIFACT_REPO].append(trace_info)
|
|
else:
|
|
_logger.warning(f"Unsupported location: {trace_info.trace_location}. Skipping.")
|
|
return trace_infos_by_location
|
|
|
|
def calculate_trace_filter_correlation(
|
|
self,
|
|
experiment_ids: list[str],
|
|
filter_string1: str,
|
|
filter_string2: str,
|
|
base_filter: str | None = None,
|
|
):
|
|
"""
|
|
Calculate the correlation (NPMI) between two trace filter conditions.
|
|
|
|
This method computes the Normalized Pointwise Mutual Information (NPMI)
|
|
between traces matching two different filter conditions, which measures
|
|
how much more (or less) likely traces are to satisfy both conditions
|
|
compared to if the conditions were independent.
|
|
|
|
Args:
|
|
experiment_ids: List of experiment IDs to search within.
|
|
filter_string1: First filter condition (e.g., "span.type = 'LLM'").
|
|
filter_string2: Second filter condition (e.g., "feedback.quality > 0.8").
|
|
base_filter: Optional base filter that both filter1 and filter2 are tested on top of
|
|
(e.g., 'request_time > ... and request_time < ...' for time windows).
|
|
|
|
Returns:
|
|
TraceFilterCorrelationResult containing:
|
|
- npmi: NPMI score from -1 (never co-occur) to 1 (always co-occur)
|
|
- npmi_smoothed: Smoothed NPMI value with Jeffreys prior for robustness
|
|
- filter1_count: Number of traces matching filter_string1
|
|
- filter2_count: Number of traces matching filter_string2
|
|
- joint_count: Number of traces matching both filters
|
|
- total_count: Total number of traces in the experiments
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.tracing.client import TracingClient
|
|
|
|
client = TracingClient()
|
|
result = client.calculate_trace_filter_correlation(
|
|
experiment_ids=["123"],
|
|
filter_string1="span.type = 'LLM'",
|
|
filter_string2="feedback.quality > 0.8",
|
|
)
|
|
print(f"NPMI: {result.npmi:.3f}")
|
|
# Output: NPMI: 0.456
|
|
"""
|
|
return self.store.calculate_trace_filter_correlation(
|
|
experiment_ids=experiment_ids,
|
|
filter_string1=filter_string1,
|
|
filter_string2=filter_string2,
|
|
base_filter=base_filter,
|
|
)
|
|
|
|
def set_trace_tags(self, trace_id: str, tags: dict[str, str]):
|
|
"""
|
|
Set tags on the trace with the given trace_id.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace.
|
|
tags: A dictionary of key-value pairs.
|
|
"""
|
|
tags = exclude_immutable_tags(tags)
|
|
for k, v in tags.items():
|
|
self.set_trace_tag(trace_id, k, v)
|
|
|
|
def set_trace_tag(self, trace_id: str, key: str, value: str):
|
|
"""
|
|
Set a tag on the trace with the given trace ID.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace to set the tag on.
|
|
key: The string key of the tag. Must be at most 250 characters long, otherwise
|
|
it will be truncated when stored.
|
|
value: The string value of the tag. Must be at most 250 characters long, otherwise
|
|
it will be truncated when stored.
|
|
"""
|
|
if not isinstance(value, str):
|
|
_logger.warning(
|
|
"Received non-string value for trace tag. Please note that non-string tag values"
|
|
"will automatically be stringified when the trace is logged."
|
|
)
|
|
|
|
if key in IMMUTABLE_TAGS:
|
|
_logger.warning(f"Tag '{key}' is immutable and cannot be set on a trace.")
|
|
return
|
|
|
|
# Trying to set the tag on the active trace first
|
|
with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
|
|
if trace:
|
|
trace.info.tags[key] = str(value)
|
|
return
|
|
|
|
self.store.set_trace_tag(trace_id, key, str(value))
|
|
|
|
def delete_trace_tag(self, trace_id: str, key: str):
|
|
"""
|
|
Delete a tag on the trace with the given trace ID.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace to delete the tag from.
|
|
key: The string key of the tag. Must be at most 250 characters long, otherwise
|
|
it will be truncated when stored.
|
|
"""
|
|
# Allow users to clear archival-failure markers so the scheduler can retry after
|
|
# manual intervention, while keeping other internal storage tags immutable.
|
|
if key in IMMUTABLE_TAGS and key != TraceTagKey.ARCHIVAL_FAILURE:
|
|
_logger.warning(f"Tag '{key}' is immutable and cannot be deleted on a trace.")
|
|
return
|
|
|
|
# Trying to delete the tag on the active trace first
|
|
with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
|
|
if trace:
|
|
if key in trace.info.tags:
|
|
trace.info.tags.pop(key)
|
|
return
|
|
else:
|
|
raise MlflowException(
|
|
f"Tag with key {key} not found in trace with ID {trace_id}.",
|
|
error_code=RESOURCE_DOES_NOT_EXIST,
|
|
)
|
|
|
|
self.store.delete_trace_tag(trace_id, key)
|
|
|
|
def get_assessment(self, trace_id: str, assessment_id: str) -> Assessment:
|
|
"""
|
|
Get an assessment entity from the backend store.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace.
|
|
assessment_id: The ID of the assessment to get.
|
|
|
|
Returns:
|
|
The Assessment object.
|
|
"""
|
|
|
|
return self.store.get_assessment(trace_id, assessment_id)
|
|
|
|
@record_usage_event(LogAssessmentEvent)
|
|
def log_assessment(self, trace_id: str, assessment: Assessment) -> Assessment:
|
|
"""
|
|
Log an assessment to a trace.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace.
|
|
assessment: The assessment object to log.
|
|
|
|
Returns:
|
|
The logged Assessment object.
|
|
"""
|
|
assessment.trace_id = trace_id
|
|
|
|
if trace_id is None or trace_id == NO_OP_SPAN_TRACE_ID:
|
|
_logger.debug(
|
|
"Skipping assessment logging for NO_OP_SPAN_TRACE_ID. This is expected when "
|
|
"tracing is disabled."
|
|
)
|
|
return assessment
|
|
|
|
# If the trace is the active trace, add the assessment to it in-memory.
|
|
# Exception: remote (distributed) traces use a dummy in-memory entry that is discarded
|
|
# when the context exits, so assessments must be persisted directly to the backend.
|
|
if trace_id == mlflow.get_active_trace_id():
|
|
with InMemoryTraceManager.get_instance().get_trace(trace_id) as trace:
|
|
if trace is None:
|
|
_logger.debug(
|
|
f"Trace {trace_id} is active but not found in the in-memory buffer. "
|
|
"Something is wrong with trace handling. Skipping assessment logging."
|
|
)
|
|
return assessment
|
|
if trace.is_remote_trace:
|
|
return self.store.create_assessment(assessment)
|
|
trace.info.assessments.append(assessment)
|
|
return assessment
|
|
return self.store.create_assessment(assessment)
|
|
|
|
def update_assessment(
|
|
self,
|
|
trace_id: str,
|
|
assessment_id: str,
|
|
assessment: Assessment,
|
|
):
|
|
"""
|
|
Update an existing assessment entity in the backend store.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace.
|
|
assessment_id: The ID of the feedback assessment to update.
|
|
assessment: The updated assessment.
|
|
"""
|
|
|
|
return self.store.update_assessment(
|
|
trace_id=trace_id,
|
|
assessment_id=assessment_id,
|
|
name=assessment.name,
|
|
expectation=assessment.expectation,
|
|
feedback=assessment.feedback,
|
|
rationale=assessment.rationale,
|
|
metadata=assessment.metadata,
|
|
)
|
|
|
|
def delete_assessment(self, trace_id: str, assessment_id: str):
|
|
"""
|
|
Delete an assessment associated with a trace.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace.
|
|
assessment_id: The ID of the assessment to delete.
|
|
"""
|
|
|
|
self.store.delete_assessment(trace_id=trace_id, assessment_id=assessment_id)
|
|
|
|
def _get_artifact_repo_for_trace(self, trace_info: TraceInfo):
|
|
artifact_uri = get_artifact_uri_for_trace(trace_info)
|
|
artifact_uri = add_databricks_profile_info_to_artifact_uri(artifact_uri, self.tracking_uri)
|
|
return get_artifact_repository(artifact_uri)
|
|
|
|
def _download_trace_data(self, trace_info: TraceInfo) -> TraceData:
|
|
"""
|
|
Download trace data from artifact repository.
|
|
|
|
Args:
|
|
trace_info: Either a TraceInfo or TraceInfoV3 object containing trace metadata.
|
|
|
|
Returns:
|
|
TraceData object representing the downloaded trace data.
|
|
"""
|
|
artifact_repo = self._get_artifact_repo_for_trace(trace_info)
|
|
return TraceData.from_dict(artifact_repo.download_trace_data())
|
|
|
|
def _upload_trace_data(self, trace_info: TraceInfo, trace_data: TraceData) -> None:
|
|
artifact_repo = self._get_artifact_repo_for_trace(trace_info)
|
|
trace_data_json = json.dumps(trace_data.to_dict(), cls=TraceJSONEncoder, ensure_ascii=False)
|
|
return artifact_repo.upload_trace_data(trace_data_json)
|
|
|
|
@record_usage_event(TraceAttachmentsEvent)
|
|
def _upload_attachments(
|
|
self,
|
|
trace_info: TraceInfo,
|
|
attachments: dict[str, Attachment],
|
|
) -> None:
|
|
artifact_repo = self._get_artifact_repo_for_trace(trace_info)
|
|
for attachment_id, attachment in attachments.items():
|
|
try:
|
|
artifact_repo.upload_attachment(attachment_id, attachment.content_bytes)
|
|
except Exception as e:
|
|
_logger.warning(f"Failed to upload attachment {attachment_id}: {e}")
|
|
|
|
def link_prompt_versions_to_trace(
|
|
self, trace_id: str, prompts: Sequence[PromptVersion]
|
|
) -> None:
|
|
"""
|
|
Link multiple prompt versions to a trace.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace to link prompts to.
|
|
prompts: List of PromptVersion objects to link to the trace.
|
|
"""
|
|
from mlflow.tracking._model_registry.utils import _get_store as _get_model_registry_store
|
|
|
|
registry_store = _get_model_registry_store()
|
|
registry_store.link_prompts_to_trace(prompt_versions=prompts, trace_id=trace_id)
|
|
|
|
def _set_experiment_trace_location(
|
|
self,
|
|
location: UCSchemaLocation,
|
|
experiment_id: str,
|
|
sql_warehouse_id: str | None = None,
|
|
) -> UCSchemaLocation:
|
|
if is_databricks_uri(self.tracking_uri):
|
|
return self.store.set_experiment_trace_location(
|
|
experiment_id=str(experiment_id),
|
|
location=location,
|
|
sql_warehouse_id=sql_warehouse_id,
|
|
)
|
|
raise MlflowException(
|
|
"Setting storage location is not supported on non-Databricks backends."
|
|
)
|
|
|
|
def _get_trace_location(self, telemetry_profile_id: str) -> UnityCatalog:
|
|
if is_databricks_uri(self.tracking_uri) and hasattr(self.store, "get_trace_location"):
|
|
return self.store.get_trace_location(telemetry_profile_id)
|
|
raise MlflowException("Getting trace location by ID is not supported on this backend.")
|
|
|
|
def _create_or_get_trace_location(
|
|
self, location: UnityCatalog, sql_warehouse_id: str | None = None
|
|
) -> UnityCatalog:
|
|
if is_databricks_uri(self.tracking_uri) and hasattr(
|
|
self.store, "create_or_get_trace_location"
|
|
):
|
|
return self.store.create_or_get_trace_location(location, sql_warehouse_id)
|
|
raise MlflowException("Creating trace location is not supported on this backend.")
|
|
|
|
def _link_trace_location(self, experiment_id: str, location: UnityCatalog) -> None:
|
|
if is_databricks_uri(self.tracking_uri) and hasattr(self.store, "link_trace_location"):
|
|
self.store.link_trace_location(experiment_id, location)
|
|
return
|
|
raise MlflowException("Linking trace location is not supported on this backend.")
|
|
|
|
def _unset_experiment_trace_location(
|
|
self, experiment_id: str, location: UCSchemaLocation | UnityCatalog
|
|
) -> None:
|
|
if is_databricks_uri(self.tracking_uri):
|
|
self.store.unset_experiment_trace_location(str(experiment_id), location)
|
|
else:
|
|
raise MlflowException(
|
|
"Clearing storage location is not supported on non-Databricks backends."
|
|
)
|
|
|
|
def _create_issue(
|
|
self,
|
|
experiment_id: str,
|
|
name: str,
|
|
description: str,
|
|
status: IssueStatus = IssueStatus.PENDING,
|
|
severity: IssueSeverity | None = None,
|
|
root_causes: list[str] | None = None,
|
|
source_run_id: str | None = None,
|
|
categories: list[str] | None = None,
|
|
created_by: str | None = None,
|
|
) -> Issue:
|
|
"""
|
|
Create a new issue in the tracking store.
|
|
|
|
Args:
|
|
experiment_id: The experiment ID.
|
|
name: Short descriptive name for the issue.
|
|
description: Detailed description of the issue.
|
|
status: Issue status. Defaults to IssueStatus.PENDING if not provided.
|
|
severity: Optional severity level indicator.
|
|
root_causes: Optional list of root cause analyses.
|
|
source_run_id: Optional MLflow run ID that discovered this issue.
|
|
categories: Optional list of categories for the issue.
|
|
created_by: Optional identifier for who created this issue.
|
|
|
|
Returns:
|
|
The created Issue entity.
|
|
"""
|
|
return self.store.create_issue(
|
|
experiment_id=experiment_id,
|
|
name=name,
|
|
description=description,
|
|
status=status,
|
|
severity=severity,
|
|
root_causes=root_causes,
|
|
source_run_id=source_run_id,
|
|
categories=categories,
|
|
created_by=created_by,
|
|
)
|
|
|
|
def _get_issue(self, issue_id: str) -> Issue:
|
|
"""
|
|
Get an issue by ID.
|
|
|
|
Args:
|
|
issue_id: The ID of the issue to retrieve.
|
|
|
|
Returns:
|
|
The Issue entity.
|
|
"""
|
|
return self.store.get_issue(issue_id)
|
|
|
|
# ----- Label schemas (tracking-store CRUD) -----
|
|
|
|
def _create_label_schema(
|
|
self,
|
|
experiment_id: str,
|
|
*,
|
|
name: str,
|
|
type: "LabelSchemaType | str",
|
|
input: "InputPassFail | InputCategorical | InputNumeric | InputText",
|
|
instruction: str | None = None,
|
|
enable_comment: bool = False,
|
|
) -> "LabelSchema":
|
|
"""Create a new label schema.
|
|
|
|
Args:
|
|
experiment_id: Parent experiment ID.
|
|
name: Schema name. Free text shown to reviewers as the label
|
|
prompt and used as the assessment key; unique within the
|
|
experiment.
|
|
type: ``"feedback"`` or ``"expectation"``.
|
|
input: One of :py:class:`InputPassFail` / :py:class:`InputCategorical`
|
|
/ :py:class:`InputNumeric` / :py:class:`InputText`.
|
|
instruction: Optional supplementary guidance (<=1000 chars).
|
|
enable_comment: UI hint; persisted but not consumed server-side.
|
|
|
|
Returns:
|
|
The created :py:class:`LabelSchema` with backend-generated
|
|
``schema_id`` and audit fields populated.
|
|
"""
|
|
return self.store.create_label_schema(
|
|
experiment_id=experiment_id,
|
|
name=name,
|
|
type=type,
|
|
input=input,
|
|
instruction=instruction,
|
|
enable_comment=enable_comment,
|
|
)
|
|
|
|
def _get_label_schema(self, schema_id: str) -> "LabelSchema":
|
|
return self.store.get_label_schema(schema_id)
|
|
|
|
def _get_label_schema_by_name(self, experiment_id: str, name: str) -> "LabelSchema":
|
|
return self.store.get_label_schema_by_name(experiment_id, name)
|
|
|
|
def _list_label_schemas(
|
|
self,
|
|
experiment_id: str,
|
|
max_results: int = 100,
|
|
page_token: str | None = None,
|
|
) -> "PagedList[LabelSchema]":
|
|
return self.store.list_label_schemas(
|
|
experiment_id, max_results=max_results, page_token=page_token
|
|
)
|
|
|
|
def _update_label_schema(
|
|
self,
|
|
schema_id: str,
|
|
*,
|
|
name: str | None = None,
|
|
instruction: str | None = None,
|
|
enable_comment: bool | None = None,
|
|
input: "InputPassFail | InputCategorical | InputNumeric | InputText | None" = None,
|
|
) -> "LabelSchema":
|
|
"""Sparse-update a label schema.
|
|
|
|
``type`` is immutable post-create and is not accepted. Fields left as
|
|
``None`` are unchanged on the server. ``enable_comment=None`` is
|
|
treated as "unchanged"; pass ``True`` or ``False`` to set explicitly.
|
|
|
|
Returns:
|
|
The updated :py:class:`LabelSchema`.
|
|
"""
|
|
return self.store.update_label_schema(
|
|
schema_id,
|
|
name=name,
|
|
instruction=instruction,
|
|
enable_comment=enable_comment,
|
|
input=input,
|
|
)
|
|
|
|
def _delete_label_schema(self, schema_id: str) -> None:
|
|
"""Delete a label schema. No-op when the schema doesn't exist."""
|
|
return self.store.delete_label_schema(schema_id)
|
|
|
|
# ----- Review queues (tracking-store CRUD) -----
|
|
|
|
def _create_review_queue(
|
|
self,
|
|
experiment_id: str,
|
|
*,
|
|
name: str,
|
|
queue_type: "ReviewQueueType | str",
|
|
users: list[str] | None = None,
|
|
schema_ids: list[str] | None = None,
|
|
) -> "ReviewQueue":
|
|
# `created_by` (the owner) is stamped server-side, never by the client.
|
|
return self.store.create_review_queue(
|
|
experiment_id,
|
|
name=name,
|
|
queue_type=queue_type,
|
|
users=users,
|
|
schema_ids=schema_ids,
|
|
)
|
|
|
|
def _get_or_create_user_queue(self, experiment_id: str, *, user: str) -> "ReviewQueue":
|
|
return self.store.get_or_create_user_queue(experiment_id, user=user)
|
|
|
|
def _get_review_queue(self, queue_id: str) -> "ReviewQueue":
|
|
return self.store.get_review_queue(queue_id)
|
|
|
|
def _get_review_queue_by_name(self, experiment_id: str, name: str) -> "ReviewQueue":
|
|
return self.store.get_review_queue_by_name(experiment_id, name=name)
|
|
|
|
def _list_review_queues(
|
|
self,
|
|
experiment_id: str,
|
|
*,
|
|
user: str | None = None,
|
|
max_results: int | None = None,
|
|
page_token: str | None = None,
|
|
) -> "PagedList[ReviewQueue]":
|
|
return self.store.list_review_queues(
|
|
experiment_id, user=user, max_results=max_results, page_token=page_token
|
|
)
|
|
|
|
def _update_review_queue(
|
|
self,
|
|
queue_id: str,
|
|
*,
|
|
name: str | None = None,
|
|
new_owner: str | None = None,
|
|
users: list[str] | None = None,
|
|
schema_ids: list[str] | None = None,
|
|
) -> "ReviewQueue":
|
|
return self.store.update_review_queue(
|
|
queue_id, name=name, new_owner=new_owner, users=users, schema_ids=schema_ids
|
|
)
|
|
|
|
def _delete_review_queue(self, queue_id: str) -> None:
|
|
return self.store.delete_review_queue(queue_id)
|
|
|
|
def _add_items_to_review_queue(
|
|
self,
|
|
queue_id: str,
|
|
*,
|
|
item_ids: list[str],
|
|
item_type: "ReviewItemType | str" = "trace",
|
|
) -> "list[ReviewQueueItem]":
|
|
return self.store.add_items_to_review_queue(
|
|
queue_id, item_ids=item_ids, item_type=item_type
|
|
)
|
|
|
|
def _remove_items_from_review_queue(self, queue_id: str, *, item_ids: list[str]) -> None:
|
|
return self.store.remove_items_from_review_queue(queue_id, item_ids=item_ids)
|
|
|
|
def _list_review_queue_items(
|
|
self,
|
|
queue_id: str,
|
|
*,
|
|
status: "ReviewStatus | str | None" = None,
|
|
max_results: int | None = None,
|
|
page_token: str | None = None,
|
|
) -> "PagedList[ReviewQueueItem]":
|
|
return self.store.list_review_queue_items(
|
|
queue_id, status=status, max_results=max_results, page_token=page_token
|
|
)
|
|
|
|
def _set_review_queue_item_status(
|
|
self,
|
|
queue_id: str,
|
|
*,
|
|
item_id: str,
|
|
status: "ReviewStatus | str",
|
|
completed_by: str | None = None,
|
|
) -> "ReviewQueueItem":
|
|
return self.store.set_review_queue_item_status(
|
|
queue_id, item_id=item_id, status=status, completed_by=completed_by
|
|
)
|