wehub-resource-sync
230 lines
8.7 KiB
Python
230 lines
8.7 KiB
Python
import contextlib
|
|
import logging
|
|
import threading
|
|
from dataclasses import dataclass, field
|
|
from typing import Generator, Sequence
|
|
|
|
from mlflow.entities import LiveSpan, Trace, TraceData, TraceInfo
|
|
from mlflow.entities.model_registry import PromptVersion
|
|
from mlflow.environment_variables import MLFLOW_TRACE_TIMEOUT_SECONDS
|
|
from mlflow.tracing.constant import TraceTagKey
|
|
from mlflow.tracing.utils.prompt import update_linked_prompts_tag
|
|
from mlflow.tracing.utils.timeout import get_trace_cache_with_timeout
|
|
from mlflow.tracing.utils.truncation import set_request_response_preview
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Internal representation to keep the state of a trace.
|
|
# Dict[str, Span] is used instead of TraceData to allow access by span_id.
|
|
@dataclass
|
|
class _Trace:
|
|
info: TraceInfo
|
|
span_dict: dict[str, LiveSpan] = field(default_factory=dict)
|
|
prompts: list[PromptVersion] = field(default_factory=list)
|
|
is_remote_trace: bool = False
|
|
|
|
def to_mlflow_trace(self) -> Trace:
|
|
trace_data = TraceData()
|
|
for span in self.span_dict.values():
|
|
# Convert LiveSpan, mutable objects, into immutable Span objects before persisting.
|
|
trace_data.spans.append(span.to_immutable_span())
|
|
|
|
set_request_response_preview(self.info, trace_data)
|
|
return Trace(self.info, trace_data)
|
|
|
|
def get_root_span(self) -> LiveSpan | None:
|
|
for span in self.span_dict.values():
|
|
if span.parent_id is None:
|
|
return span
|
|
return None
|
|
|
|
|
|
@dataclass
|
|
class ManagerTrace:
|
|
"""
|
|
Wrapper around a trace and its associated prompts.
|
|
"""
|
|
|
|
trace: Trace
|
|
prompts: Sequence[PromptVersion]
|
|
is_remote_trace: bool = False
|
|
|
|
|
|
class InMemoryTraceManager:
|
|
"""
|
|
Manage spans and traces created by the tracing system in memory.
|
|
"""
|
|
|
|
_instance_lock = threading.RLock()
|
|
_instance = None
|
|
|
|
@classmethod
|
|
def get_instance(cls):
|
|
if cls._instance is None:
|
|
with cls._instance_lock:
|
|
if cls._instance is None:
|
|
cls._instance = InMemoryTraceManager()
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
# In-memory cache to store trace_od -> _Trace mapping.
|
|
self._traces = get_trace_cache_with_timeout()
|
|
|
|
# Store mapping between OpenTelemetry trace ID and MLflow trace ID
|
|
self._otel_id_to_mlflow_trace_id: dict[int, str] = {}
|
|
self._lock = threading.RLock() # Lock for _traces
|
|
|
|
def register_trace(self, otel_trace_id: int, trace_info: TraceInfo, is_remote_trace=False):
|
|
"""
|
|
Register a new trace info object to the in-memory trace registry.
|
|
|
|
Args:
|
|
otel_trace_id: The OpenTelemetry trace ID for the new trace.
|
|
trace_info: The trace info object to be stored.
|
|
is_remote_trace: Whether the trace is a remote trace. For a distributed trace, it is
|
|
registered in both client side and remote server side, for remote server side
|
|
registration, the 'is_remote_trace' flag is set to True.
|
|
"""
|
|
# Check for a new timeout setting whenever a new trace is created.
|
|
self._check_timeout_update()
|
|
with self._lock:
|
|
self._traces[trace_info.trace_id] = _Trace(trace_info, is_remote_trace=is_remote_trace)
|
|
self._otel_id_to_mlflow_trace_id[otel_trace_id] = trace_info.trace_id
|
|
|
|
def register_span(self, span: LiveSpan):
|
|
"""
|
|
Store the given span in the in-memory trace data.
|
|
|
|
Args:
|
|
span: The span to be stored.
|
|
"""
|
|
if not isinstance(span, LiveSpan):
|
|
_logger.debug(f"Invalid span object {type(span)} is passed. Skipping.")
|
|
return
|
|
|
|
with self._lock:
|
|
trace_data_dict = self._traces[span.request_id].span_dict
|
|
trace_data_dict[span.span_id] = span
|
|
|
|
def register_prompt(self, trace_id: str, prompt: PromptVersion):
|
|
"""
|
|
Register a prompt to link to the trace with the given trace ID.
|
|
|
|
Args:
|
|
trace_id: The ID of the trace to which the prompt belongs.
|
|
prompt: The prompt version to be registered.
|
|
"""
|
|
with self._lock:
|
|
if prompt not in self._traces[trace_id].prompts:
|
|
self._traces[trace_id].prompts.append(prompt)
|
|
|
|
# NB: Set prompt URIs in trace tags for linking.
|
|
try:
|
|
current_tag = self._traces[trace_id].info.tags.get(TraceTagKey.LINKED_PROMPTS)
|
|
updated_tag = update_linked_prompts_tag(current_tag, [prompt])
|
|
self._traces[trace_id].info.tags[TraceTagKey.LINKED_PROMPTS] = updated_tag
|
|
except Exception:
|
|
_logger.debug(f"Failed to update prompts tag for trace {trace_id}", exc_info=True)
|
|
raise
|
|
|
|
@contextlib.contextmanager
|
|
def get_trace(self, trace_id: str) -> Generator[_Trace | None, None, None]:
|
|
"""
|
|
Yield the trace info for the given trace ID..
|
|
This is designed to be used as a context manager to ensure the trace info is accessed
|
|
with the lock held.
|
|
"""
|
|
with self._lock:
|
|
yield self._traces.get(trace_id)
|
|
|
|
def get_span_from_id(self, trace_id: str, span_id: str) -> LiveSpan | None:
|
|
"""
|
|
Get a span object for the given trace_id and span_id.
|
|
"""
|
|
with self._lock:
|
|
trace = self._traces.get(trace_id)
|
|
|
|
return trace.span_dict.get(span_id) if trace else None
|
|
|
|
def get_root_span_id(self, trace_id) -> str | None:
|
|
"""
|
|
Get the root span ID for the given trace ID.
|
|
"""
|
|
with self._lock:
|
|
trace = self._traces.get(trace_id)
|
|
|
|
if trace:
|
|
for span in trace.span_dict.values():
|
|
if span.parent_id is None:
|
|
return span.span_id
|
|
|
|
return None
|
|
|
|
def get_mlflow_trace_id_from_otel_id(self, otel_trace_id: int) -> str | None:
|
|
"""
|
|
Get the MLflow trace ID for the given OpenTelemetry trace ID.
|
|
"""
|
|
return self._otel_id_to_mlflow_trace_id.get(otel_trace_id)
|
|
|
|
def has_open_spans(self, otel_trace_id: int) -> bool:
|
|
"""Return True if any span in the trace has not yet ended."""
|
|
mlflow_trace_id = self.get_mlflow_trace_id_from_otel_id(otel_trace_id)
|
|
if mlflow_trace_id is None:
|
|
return False
|
|
with self.get_trace(mlflow_trace_id) as trace:
|
|
if trace is None:
|
|
return False
|
|
return any(span.end_time_ns is None for span in trace.span_dict.values())
|
|
|
|
def set_trace_metadata(self, trace_id: str, key: str, value: str):
|
|
"""
|
|
Set the trace metadata for the given request ID.
|
|
"""
|
|
with self.get_trace(trace_id) as trace:
|
|
if trace:
|
|
trace.info.trace_metadata[key] = value
|
|
|
|
def pop_trace(self, otel_trace_id: int) -> ManagerTrace | None:
|
|
"""
|
|
Pop trace data for the given OpenTelemetry trace ID and
|
|
return it as a ManagerTrace wrapper containing the trace and prompts.
|
|
"""
|
|
with self._lock:
|
|
mlflow_trace_id = self._otel_id_to_mlflow_trace_id.pop(otel_trace_id, None)
|
|
internal_trace = self._traces.pop(mlflow_trace_id, None) if mlflow_trace_id else None
|
|
if internal_trace is None:
|
|
return None
|
|
return ManagerTrace(
|
|
trace=internal_trace.to_mlflow_trace(),
|
|
prompts=internal_trace.prompts,
|
|
is_remote_trace=internal_trace.is_remote_trace,
|
|
)
|
|
|
|
def _check_timeout_update(self):
|
|
"""
|
|
TTL/Timeout may be updated by users after initial cache creation. This method checks
|
|
for the update and create a new cache instance with the updated timeout.
|
|
"""
|
|
new_timeout = MLFLOW_TRACE_TIMEOUT_SECONDS.get()
|
|
if new_timeout != getattr(self._traces, "timeout", None):
|
|
if len(self._traces) > 0:
|
|
_logger.warning(
|
|
f"The timeout of the trace buffer has been updated to {new_timeout} seconds. "
|
|
"This operation discards all in-progress traces at the moment. Please make "
|
|
"sure to update the timeout when there are no in-progress traces."
|
|
)
|
|
|
|
with self._lock:
|
|
# We need to check here again in case this method runs in parallel
|
|
if new_timeout != getattr(self._traces, "timeout", None):
|
|
self._traces = get_trace_cache_with_timeout()
|
|
|
|
@classmethod
|
|
def reset(cls):
|
|
"""Clear all the aggregated trace data. This should only be used for testing."""
|
|
if cls._instance:
|
|
with cls._instance._lock:
|
|
cls._instance._traces.clear()
|
|
cls._instance = None
|