Verify exact production traces and record rollout
This commit is contained in:
@@ -0,0 +1,217 @@
|
||||
"""Request correlation, OpenTelemetry spans and stage timing.
|
||||
|
||||
This module deliberately has a no-op default. The RAG domain remains runnable
|
||||
without the OpenTelemetry packages or collector, while a configured runtime
|
||||
gets one trace context shared by FastAPI's async middleware and its sync
|
||||
thread-pool endpoint.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from functools import wraps
|
||||
from time import perf_counter
|
||||
from typing import Any, Iterator, Mapping
|
||||
|
||||
from .metrics import Metrics, NullMetrics, STAGE_DURATION
|
||||
|
||||
_CORRELATION_ID: ContextVar[str | None] = ContextVar(
|
||||
"duocthu_correlation_id", default=None
|
||||
)
|
||||
_STAGE: ContextVar[str | None] = ContextVar("duocthu_stage", default=None)
|
||||
_SAFE_CORRELATION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$")
|
||||
_metrics: Metrics = NullMetrics()
|
||||
|
||||
try:
|
||||
from opentelemetry import trace
|
||||
|
||||
_tracer = trace.get_tracer("duocthu.ai-service")
|
||||
except ImportError: # pragma: no cover - exercised only in minimal installs
|
||||
trace = None
|
||||
_tracer = None
|
||||
|
||||
|
||||
def configure_telemetry(settings: Any, metrics: Metrics | None = None) -> bool:
|
||||
"""Configure one OTLP tracer provider; return whether tracing is active.
|
||||
|
||||
Metrics are configured independently and continue to work when tracing is
|
||||
disabled or optional OpenTelemetry packages are absent.
|
||||
"""
|
||||
global _metrics, _tracer
|
||||
_metrics = metrics or NullMetrics()
|
||||
if not getattr(settings, "otel_enabled", False):
|
||||
return False
|
||||
|
||||
try:
|
||||
from opentelemetry import trace as otel_trace
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
||||
from opentelemetry.sdk.resources import Resource
|
||||
from opentelemetry.sdk.trace import TracerProvider
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
current = otel_trace.get_tracer_provider()
|
||||
# OpenTelemetry only permits setting the global provider once. Reuse a
|
||||
# provider installed by a host/test instead of replacing it and warning.
|
||||
if isinstance(current, TracerProvider):
|
||||
_tracer = current.get_tracer("duocthu.ai-service")
|
||||
return True
|
||||
|
||||
provider = TracerProvider(
|
||||
resource=Resource.create(
|
||||
{
|
||||
"service.name": getattr(settings, "otel_service_name", "ai-service"),
|
||||
"service.version": "0.1.0",
|
||||
"deployment.environment": getattr(settings, "environment", "local"),
|
||||
}
|
||||
),
|
||||
sampler=ParentBased(
|
||||
TraceIdRatioBased(float(getattr(settings, "otel_sample_ratio", 1.0)))
|
||||
),
|
||||
)
|
||||
exporter = OTLPSpanExporter(endpoint=settings.otel_exporter_otlp_endpoint)
|
||||
provider.add_span_processor(BatchSpanProcessor(exporter))
|
||||
otel_trace.set_tracer_provider(provider)
|
||||
_tracer = provider.get_tracer("duocthu.ai-service")
|
||||
return True
|
||||
|
||||
|
||||
def normalize_correlation_id(candidate: str | None) -> str:
|
||||
value = (candidate or "").strip()
|
||||
return value if _SAFE_CORRELATION.fullmatch(value) else str(uuid.uuid4())
|
||||
|
||||
|
||||
@contextmanager
|
||||
def correlation_context(candidate: str | None) -> Iterator[str]:
|
||||
correlation_id = normalize_correlation_id(candidate)
|
||||
token = _CORRELATION_ID.set(correlation_id)
|
||||
try:
|
||||
yield correlation_id
|
||||
finally:
|
||||
_CORRELATION_ID.reset(token)
|
||||
|
||||
|
||||
def current_correlation_id() -> str:
|
||||
value = _CORRELATION_ID.get()
|
||||
return value or str(uuid.uuid4())
|
||||
|
||||
|
||||
def current_trace_id() -> str | None:
|
||||
if trace is None:
|
||||
return None
|
||||
context = trace.get_current_span().get_span_context()
|
||||
if not context.is_valid:
|
||||
return None
|
||||
return f"{context.trace_id:032x}"
|
||||
|
||||
|
||||
def current_stage() -> str | None:
|
||||
return _STAGE.get()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def request_span(
|
||||
method: str, route: str, headers: Mapping[str, str]
|
||||
) -> Iterator[Any]:
|
||||
"""Start a server span and extract an incoming W3C ``traceparent``."""
|
||||
if _tracer is None:
|
||||
yield None
|
||||
return
|
||||
try:
|
||||
from opentelemetry import propagate
|
||||
from opentelemetry.trace import SpanKind
|
||||
except ImportError: # pragma: no cover - minimal install fallback
|
||||
yield None
|
||||
return
|
||||
|
||||
parent = propagate.extract(headers)
|
||||
with _tracer.start_as_current_span(
|
||||
f"{method} {route}",
|
||||
context=parent,
|
||||
kind=SpanKind.SERVER,
|
||||
attributes={
|
||||
"http.request.method": method,
|
||||
"http.route": route,
|
||||
"duocthu.correlation_id": current_correlation_id(),
|
||||
},
|
||||
) as span:
|
||||
yield span
|
||||
|
||||
|
||||
@contextmanager
|
||||
def dependency_span(provider: str, operation: str) -> Iterator[Any]:
|
||||
"""Trace one external dependency call without creating another histogram."""
|
||||
if _tracer is None:
|
||||
yield None
|
||||
return
|
||||
with _tracer.start_as_current_span(
|
||||
f"provider.{provider}.{operation}",
|
||||
attributes={
|
||||
"duocthu.provider": provider,
|
||||
"duocthu.operation": operation,
|
||||
},
|
||||
) as span:
|
||||
yield span
|
||||
|
||||
|
||||
@contextmanager
|
||||
def stage(name: str, **attributes: Any) -> Iterator[Any]:
|
||||
"""Create a child span and observe a low-cardinality stage histogram."""
|
||||
started = perf_counter()
|
||||
outcome = "ok"
|
||||
span = None
|
||||
token = _STAGE.set(name)
|
||||
try:
|
||||
if _tracer is None:
|
||||
yield None
|
||||
else:
|
||||
safe_attributes = {
|
||||
key: value
|
||||
for key, value in attributes.items()
|
||||
if value is not None and isinstance(value, (str, bool, int, float))
|
||||
}
|
||||
safe_attributes["duocthu.stage"] = name
|
||||
with _tracer.start_as_current_span(
|
||||
f"rag.stage.{name}", attributes=safe_attributes
|
||||
) as span:
|
||||
yield span
|
||||
except BaseException as exc:
|
||||
outcome = "cancelled" if exc.__class__.__name__ == "CancelledError" else "error"
|
||||
if span is not None:
|
||||
span.set_attribute("duocthu.outcome", outcome)
|
||||
raise
|
||||
finally:
|
||||
elapsed = perf_counter() - started
|
||||
if span is not None:
|
||||
span.set_attribute("duocthu.outcome", outcome)
|
||||
span.set_attribute("duocthu.duration_ms", elapsed * 1000)
|
||||
_metrics.observe(STAGE_DURATION, elapsed, stage=name, outcome=outcome)
|
||||
_STAGE.reset(token)
|
||||
|
||||
|
||||
def annotate_current_span(**attributes: Any) -> None:
|
||||
if trace is None:
|
||||
return
|
||||
span = trace.get_current_span()
|
||||
if not span.is_recording():
|
||||
return
|
||||
for key, value in attributes.items():
|
||||
if value is not None and isinstance(value, (str, bool, int, float)):
|
||||
span.set_attribute(key, value)
|
||||
|
||||
|
||||
def traced_stage(name: str):
|
||||
"""Decorator form for pure functions that should remain otherwise untouched."""
|
||||
def decorate(function):
|
||||
@wraps(function)
|
||||
def wrapped(*args, **kwargs):
|
||||
with stage(name):
|
||||
return function(*args, **kwargs)
|
||||
|
||||
return wrapped
|
||||
|
||||
return decorate
|
||||
Reference in New Issue
Block a user