wehub-resource-sync
247 lines
9.3 KiB
Python
247 lines
9.3 KiB
Python
import json
|
|
import logging
|
|
import threading
|
|
from typing import Any
|
|
|
|
from haystack.tracing import OpenTelemetryTracer, enable_tracing
|
|
from opentelemetry import trace
|
|
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 import TracerProvider as SDKTracerProvider
|
|
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
|
|
from opentelemetry.trace import (
|
|
NoOpTracerProvider,
|
|
ProxyTracerProvider,
|
|
get_tracer_provider,
|
|
set_tracer_provider,
|
|
)
|
|
|
|
from mlflow.entities import LiveSpan, SpanType
|
|
from mlflow.entities.span import create_mlflow_span
|
|
from mlflow.tracing.constant import SpanAttributeKey, TokenUsageKey
|
|
from mlflow.tracing.provider import _get_tracer
|
|
from mlflow.tracing.trace_manager import InMemoryTraceManager
|
|
from mlflow.tracing.utils import (
|
|
_bypass_attribute_guard,
|
|
generate_trace_id_v3,
|
|
get_mlflow_span_for_otel_span,
|
|
set_span_cost_attribute,
|
|
set_span_model_attribute,
|
|
should_compute_cost_client_side,
|
|
)
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def setup_haystack_tracing():
|
|
from haystack import tracing as hs_tracing
|
|
|
|
hs_tracing.tracer.is_content_tracing_enabled = True
|
|
|
|
provider = get_tracer_provider()
|
|
hs_processor = HaystackSpanProcessor()
|
|
if isinstance(provider, (NoOpTracerProvider, ProxyTracerProvider)):
|
|
new_provider = SDKTracerProvider()
|
|
new_provider.add_span_processor(hs_processor)
|
|
set_tracer_provider(new_provider)
|
|
else:
|
|
if not any(
|
|
isinstance(p, HaystackSpanProcessor)
|
|
for p in provider._active_span_processor._span_processors
|
|
):
|
|
provider.add_span_processor(hs_processor)
|
|
|
|
tracer = trace.get_tracer(__name__)
|
|
enable_tracing(OpenTelemetryTracer(tracer))
|
|
|
|
|
|
def _infer_span_type_from_haystack(
|
|
comp_type: str | None,
|
|
comp_alias: str | None,
|
|
span: OTelReadableSpan,
|
|
) -> SpanType:
|
|
s = (comp_type or comp_alias or span.name or "").lower()
|
|
|
|
if any(
|
|
k in s
|
|
for k in (
|
|
"llm",
|
|
"chat",
|
|
"generator",
|
|
"completion",
|
|
"textgen",
|
|
"chatgenerator",
|
|
"openai",
|
|
"anthropic",
|
|
"mistral",
|
|
"cohere",
|
|
"gemini",
|
|
)
|
|
):
|
|
return SpanType.LLM
|
|
|
|
if "embedder" in s:
|
|
return SpanType.EMBEDDING
|
|
|
|
if "retriever" in s:
|
|
return SpanType.RETRIEVER
|
|
|
|
if "ranker" in s:
|
|
return SpanType.RERANKER
|
|
|
|
if "agent" in s:
|
|
return SpanType.AGENT
|
|
|
|
return SpanType.TOOL
|
|
|
|
|
|
class HaystackSpanProcessor(SimpleSpanProcessor):
|
|
def __init__(self):
|
|
self.span_exporter = SpanExporter()
|
|
self._pipeline_io: dict[str, tuple[dict[str, Any], dict[str, Any]]] = {}
|
|
self._processing_local = threading.local()
|
|
|
|
def on_start(self, span: OTelSpan, parent_context: Context | None = None):
|
|
# Recursion guard: with MLFLOW_USE_DEFAULT_TRACER_PROVIDER=false (shared provider),
|
|
# tracer.span_processor.on_start() routes back through the same composite processor,
|
|
# re-entering this method and causing infinite recursion.
|
|
if getattr(self._processing_local, "in_on_start", False):
|
|
return
|
|
self._processing_local.in_on_start = True
|
|
try:
|
|
tracer = _get_tracer(__name__)
|
|
tracer.span_processor.on_start(span, parent_context)
|
|
|
|
trace_id = generate_trace_id_v3(span)
|
|
mlflow_span = create_mlflow_span(span, trace_id)
|
|
InMemoryTraceManager.get_instance().register_span(mlflow_span)
|
|
finally:
|
|
self._processing_local.in_on_start = False
|
|
|
|
def on_end(self, span: OTelReadableSpan) -> None:
|
|
# Recursion guard: with MLFLOW_USE_DEFAULT_TRACER_PROVIDER=false (shared provider),
|
|
# tracer.span_processor.on_end() routes back through the same composite processor,
|
|
# re-entering this method and causing infinite recursion.
|
|
if getattr(self._processing_local, "in_on_end", False):
|
|
return
|
|
self._processing_local.in_on_end = True
|
|
try:
|
|
mlflow_span = get_mlflow_span_for_otel_span(span)
|
|
if mlflow_span is None:
|
|
_logger.debug("Span not found in the map. Skipping end.")
|
|
return
|
|
|
|
with _bypass_attribute_guard(mlflow_span._span):
|
|
if span.name in ("haystack.pipeline.run", "haystack.async_pipeline.run"):
|
|
self.set_pipeline_info(mlflow_span, span)
|
|
elif span.name in ("haystack.component.run"):
|
|
self.set_component_info(mlflow_span, span)
|
|
|
|
tracer = _get_tracer(__name__)
|
|
tracer.span_processor.on_end(span)
|
|
finally:
|
|
self._processing_local.in_on_end = False
|
|
|
|
def set_component_info(self, mlflow_span: LiveSpan, span: OTelReadableSpan) -> None:
|
|
comp_alias = span.attributes.get("haystack.component.name")
|
|
comp_type = span.attributes.get("haystack.component.type")
|
|
mlflow_span.set_span_type(_infer_span_type_from_haystack(comp_type, comp_alias, span))
|
|
|
|
# Haystack spans originally have name='haystack.component.run'. We need to update both the
|
|
# _name field of the Otel span and the _original_name field of the MLflow span to
|
|
# customize the span name here, as otherwise it would be overwritten in the
|
|
# deduplication process
|
|
span_name = comp_type or comp_alias or span.name
|
|
mlflow_span._span._name = span_name
|
|
mlflow_span._original_name = span_name
|
|
|
|
if (inputs := span.attributes.get("haystack.component.input")) is not None:
|
|
try:
|
|
mlflow_span.set_inputs(json.loads(inputs))
|
|
except Exception:
|
|
mlflow_span.set_inputs(inputs)
|
|
if (outputs := span.attributes.get("haystack.component.output")) is not None:
|
|
try:
|
|
mlflow_span.set_outputs(json.loads(outputs))
|
|
except Exception:
|
|
mlflow_span.set_outputs(outputs)
|
|
|
|
if isinstance(mlflow_span.inputs, dict):
|
|
set_span_model_attribute(mlflow_span, mlflow_span.inputs)
|
|
|
|
if usage := _parse_token_usage(mlflow_span.outputs):
|
|
mlflow_span.set_attribute(SpanAttributeKey.CHAT_USAGE, usage)
|
|
if should_compute_cost_client_side():
|
|
set_span_cost_attribute(mlflow_span)
|
|
|
|
if parent_id := mlflow_span.parent_id:
|
|
key = comp_alias or comp_type or mlflow_span.name
|
|
inputs_agg, outputs_agg = self._pipeline_io.setdefault(parent_id, ({}, {}))
|
|
if mlflow_span.inputs is not None:
|
|
inputs_agg[key] = mlflow_span.inputs
|
|
if mlflow_span.outputs is not None:
|
|
outputs_agg[key] = mlflow_span.outputs
|
|
|
|
def set_pipeline_info(self, mlflow_span: LiveSpan, span: OTelReadableSpan) -> None:
|
|
# Pipelines are CHAINs
|
|
mlflow_span.set_span_type(SpanType.CHAIN)
|
|
|
|
if pipe_name := span.attributes.get("haystack.pipeline.name"):
|
|
mlflow_span._span._name = pipe_name
|
|
|
|
if (inputs := span.attributes.get("haystack.pipeline.input")) is not None:
|
|
try:
|
|
mlflow_span.set_inputs(json.loads(inputs))
|
|
except Exception:
|
|
mlflow_span.set_inputs(inputs)
|
|
if (outputs := span.attributes.get("haystack.pipeline.output")) is not None:
|
|
try:
|
|
mlflow_span.set_outputs(json.loads(outputs))
|
|
except Exception:
|
|
mlflow_span.set_outputs(outputs)
|
|
|
|
if mlflow_span.span_id in self._pipeline_io:
|
|
inputs_agg, outputs_agg = self._pipeline_io.pop(mlflow_span.span_id)
|
|
if mlflow_span.inputs is None and inputs_agg:
|
|
mlflow_span.set_inputs(inputs_agg)
|
|
if mlflow_span.outputs is None and outputs_agg:
|
|
mlflow_span.set_outputs(outputs_agg)
|
|
|
|
|
|
def _parse_token_usage(outputs: Any) -> dict[str, int] | None:
|
|
try:
|
|
if not isinstance(outputs, dict):
|
|
return None
|
|
|
|
replies = outputs.get("replies")
|
|
if isinstance(replies, list) and len(replies) > 0:
|
|
usage = (
|
|
replies[0].get("meta", {}).get("usage", {}) if isinstance(replies[0], dict) else {}
|
|
)
|
|
|
|
meta = outputs.get("meta")
|
|
if isinstance(meta, list) and len(meta) > 0:
|
|
usage = meta[0].get("usage", {}) if isinstance(meta[0], dict) else {}
|
|
|
|
if isinstance(usage, dict):
|
|
in_tok = usage.get("prompt_tokens", 0)
|
|
out_tok = usage.get("completion_tokens", 0)
|
|
tot_tok = usage.get("total_tokens", 0)
|
|
return {
|
|
TokenUsageKey.INPUT_TOKENS: in_tok,
|
|
TokenUsageKey.OUTPUT_TOKENS: out_tok,
|
|
TokenUsageKey.TOTAL_TOKENS: tot_tok,
|
|
}
|
|
except Exception:
|
|
_logger.debug("Failed to parse token usage from outputs.", exc_info=True)
|
|
|
|
|
|
def teardown_haystack_tracing():
|
|
provider = get_tracer_provider()
|
|
if isinstance(provider, SDKTracerProvider):
|
|
span_processors = getattr(provider._active_span_processor, "_span_processors", ())
|
|
provider._active_span_processor._span_processors = tuple(
|
|
p for p in span_processors if not isinstance(p, HaystackSpanProcessor)
|
|
)
|