chore(telemetry): code cleanup (#3897)

# What does this PR do?
Clean up telemetry code since the telemetry API has been remove.
- moved telemetry files out of providers to core
- removed from Api

## Test Plan

❯ OTEL_SERVICE_NAME=llama_stack
OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 uv run llama stack run
starter
❯ curl http://localhost:8321/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "openai/gpt-4o-mini",
    "messages": [
      {
        "role": "user",
        "content": "Hello!"
      }
    ]
  }'

-> verify traces in Grafana

CI

llama_stack/core/telemetry/__init__.py

@@ -0,0 +1,32 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .telemetry import Telemetry
+from .trace_protocol import serialize_value, trace_protocol
+from .tracing import (
+    CURRENT_TRACE_CONTEXT,
+    ROOT_SPAN_MARKERS,
+    end_trace,
+    enqueue_event,
+    get_current_span,
+    setup_logger,
+    span,
+    start_trace,
+)
+
+__all__ = [
+    "Telemetry",
+    "trace_protocol",
+    "serialize_value",
+    "CURRENT_TRACE_CONTEXT",
+    "ROOT_SPAN_MARKERS",
+    "end_trace",
+    "enqueue_event",
+    "get_current_span",
+    "setup_logger",
+    "span",
+    "start_trace",
+]

llama_stack/core/telemetry/telemetry.py

@@ -0,0 +1,250 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import threading
+from typing import Any
+
+from opentelemetry import metrics, trace
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from llama_stack.apis.telemetry import (
+    Event,
+    MetricEvent,
+    SpanEndPayload,
+    SpanStartPayload,
+    SpanStatus,
+    StructuredLogEvent,
+    UnstructuredLogEvent,
+)
+from llama_stack.apis.telemetry import (
+    Telemetry as TelemetryBase,
+)
+from llama_stack.core.telemetry.tracing import ROOT_SPAN_MARKERS
+from llama_stack.log import get_logger
+
+_GLOBAL_STORAGE: dict[str, dict[str | int, Any]] = {
+    "active_spans": {},
+    "counters": {},
+    "gauges": {},
+    "up_down_counters": {},
+}
+_global_lock = threading.Lock()
+_TRACER_PROVIDER = None
+
+logger = get_logger(name=__name__, category="telemetry")
+
+
+def is_tracing_enabled(tracer):
+    with tracer.start_as_current_span("check_tracing") as span:
+        return span.is_recording()
+
+
+class Telemetry(TelemetryBase):
+    def __init__(self) -> None:
+        self.meter = None
+
+        global _TRACER_PROVIDER
+        # Initialize the correct span processor based on the provider state.
+        # This is needed since once the span processor is set, it cannot be unset.
+        # Recreating the telemetry adapter multiple times will result in duplicate span processors.
+        # Since the library client can be recreated multiple times in a notebook,
+        # the kernel will hold on to the span processor and cause duplicate spans to be written.
+        if os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"):
+            if _TRACER_PROVIDER is None:
+                provider = TracerProvider()
+                trace.set_tracer_provider(provider)
+                _TRACER_PROVIDER = provider
+
+                # Use single OTLP endpoint for all telemetry signals
+
+                # Let OpenTelemetry SDK handle endpoint construction automatically
+                # The SDK will read OTEL_EXPORTER_OTLP_ENDPOINT and construct appropriate URLs
+                # https://opentelemetry.io/docs/languages/sdk-configuration/otlp-exporter
+                span_exporter = OTLPSpanExporter()
+                span_processor = BatchSpanProcessor(span_exporter)
+                trace.get_tracer_provider().add_span_processor(span_processor)
+
+                metric_reader = PeriodicExportingMetricReader(OTLPMetricExporter())
+                metric_provider = MeterProvider(metric_readers=[metric_reader])
+                metrics.set_meter_provider(metric_provider)
+            self.is_otel_endpoint_set = True
+        else:
+            logger.warning("OTEL_EXPORTER_OTLP_ENDPOINT is not set, skipping telemetry")
+            self.is_otel_endpoint_set = False
+
+        self.meter = metrics.get_meter(__name__)
+        self._lock = _global_lock
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        if self.is_otel_endpoint_set:
+            trace.get_tracer_provider().force_flush()
+
+    async def log_event(self, event: Event, ttl_seconds: int = 604800) -> None:
+        if isinstance(event, UnstructuredLogEvent):
+            self._log_unstructured(event, ttl_seconds)
+        elif isinstance(event, MetricEvent):
+            self._log_metric(event)
+        elif isinstance(event, StructuredLogEvent):
+            self._log_structured(event, ttl_seconds)
+        else:
+            raise ValueError(f"Unknown event type: {event}")
+
+    def _log_unstructured(self, event: UnstructuredLogEvent, ttl_seconds: int) -> None:
+        with self._lock:
+            # Use global storage instead of instance storage
+            span_id = int(event.span_id, 16)
+            span = _GLOBAL_STORAGE["active_spans"].get(span_id)
+
+            if span:
+                timestamp_ns = int(event.timestamp.timestamp() * 1e9)
+                span.add_event(
+                    name=event.type.value,
+                    attributes={
+                        "message": event.message,
+                        "severity": event.severity.value,
+                        "__ttl__": ttl_seconds,
+                        **(event.attributes or {}),
+                    },
+                    timestamp=timestamp_ns,
+                )
+            else:
+                print(f"Warning: No active span found for span_id {span_id}. Dropping event: {event}")
+
+    def _get_or_create_counter(self, name: str, unit: str) -> metrics.Counter:
+        assert self.meter is not None
+        if name not in _GLOBAL_STORAGE["counters"]:
+            _GLOBAL_STORAGE["counters"][name] = self.meter.create_counter(
+                name=name,
+                unit=unit,
+                description=f"Counter for {name}",
+            )
+        return _GLOBAL_STORAGE["counters"][name]
+
+    def _get_or_create_gauge(self, name: str, unit: str) -> metrics.ObservableGauge:
+        assert self.meter is not None
+        if name not in _GLOBAL_STORAGE["gauges"]:
+            _GLOBAL_STORAGE["gauges"][name] = self.meter.create_gauge(
+                name=name,
+                unit=unit,
+                description=f"Gauge for {name}",
+            )
+        return _GLOBAL_STORAGE["gauges"][name]
+
+    def _log_metric(self, event: MetricEvent) -> None:
+        # Add metric as an event to the current span
+        try:
+            with self._lock:
+                # Only try to add to span if we have a valid span_id
+                if event.span_id:
+                    try:
+                        span_id = int(event.span_id, 16)
+                        span = _GLOBAL_STORAGE["active_spans"].get(span_id)
+
+                        if span:
+                            timestamp_ns = int(event.timestamp.timestamp() * 1e9)
+                            span.add_event(
+                                name=f"metric.{event.metric}",
+                                attributes={
+                                    "value": event.value,
+                                    "unit": event.unit,
+                                    **(event.attributes or {}),
+                                },
+                                timestamp=timestamp_ns,
+                            )
+                    except (ValueError, KeyError):
+                        # Invalid span_id or span not found, but we already logged to console above
+                        pass
+        except Exception:
+            # Lock acquisition failed
+            logger.debug("Failed to acquire lock to add metric to span")
+
+        # Log to OpenTelemetry meter if available
+        if self.meter is None:
+            return
+        if isinstance(event.value, int):
+            counter = self._get_or_create_counter(event.metric, event.unit)
+            counter.add(event.value, attributes=event.attributes)
+        elif isinstance(event.value, float):
+            up_down_counter = self._get_or_create_up_down_counter(event.metric, event.unit)
+            up_down_counter.add(event.value, attributes=event.attributes)
+
+    def _get_or_create_up_down_counter(self, name: str, unit: str) -> metrics.UpDownCounter:
+        assert self.meter is not None
+        if name not in _GLOBAL_STORAGE["up_down_counters"]:
+            _GLOBAL_STORAGE["up_down_counters"][name] = self.meter.create_up_down_counter(
+                name=name,
+                unit=unit,
+                description=f"UpDownCounter for {name}",
+            )
+        return _GLOBAL_STORAGE["up_down_counters"][name]
+
+    def _log_structured(self, event: StructuredLogEvent, ttl_seconds: int) -> None:
+        with self._lock:
+            span_id = int(event.span_id, 16)
+            tracer = trace.get_tracer(__name__)
+            if event.attributes is None:
+                event.attributes = {}
+            event.attributes["__ttl__"] = ttl_seconds
+
+            # Extract these W3C trace context attributes so they are not written to
+            # underlying storage, as we just need them to propagate the trace context.
+            traceparent = event.attributes.pop("traceparent", None)
+            tracestate = event.attributes.pop("tracestate", None)
+            if traceparent:
+                # If we have a traceparent header value, we're not the root span.
+                for root_attribute in ROOT_SPAN_MARKERS:
+                    event.attributes.pop(root_attribute, None)
+
+            if isinstance(event.payload, SpanStartPayload):
+                # Check if span already exists to prevent duplicates
+                if span_id in _GLOBAL_STORAGE["active_spans"]:
+                    return
+
+                context = None
+                if event.payload.parent_span_id:
+                    parent_span_id = int(event.payload.parent_span_id, 16)
+                    parent_span = _GLOBAL_STORAGE["active_spans"].get(parent_span_id)
+                    context = trace.set_span_in_context(parent_span)
+                elif traceparent:
+                    carrier = {
+                        "traceparent": traceparent,
+                        "tracestate": tracestate,
+                    }
+                    context = TraceContextTextMapPropagator().extract(carrier=carrier)
+
+                span = tracer.start_span(
+                    name=event.payload.name,
+                    context=context,
+                    attributes=event.attributes or {},
+                )
+                _GLOBAL_STORAGE["active_spans"][span_id] = span
+
+            elif isinstance(event.payload, SpanEndPayload):
+                span = _GLOBAL_STORAGE["active_spans"].get(span_id)
+                if span:
+                    if event.attributes:
+                        span.set_attributes(event.attributes)
+
+                    status = (
+                        trace.Status(status_code=trace.StatusCode.OK)
+                        if event.payload.status == SpanStatus.OK
+                        else trace.Status(status_code=trace.StatusCode.ERROR)
+                    )
+                    span.set_status(status)
+                    span.end()
+                    _GLOBAL_STORAGE["active_spans"].pop(span_id, None)
+            else:
+                raise ValueError(f"Unknown structured log event: {event}")

llama_stack/core/telemetry/trace_protocol.py

@@ -0,0 +1,145 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+import inspect
+import json
+from collections.abc import AsyncGenerator, Callable
+from functools import wraps
+from typing import Any, cast
+
+from pydantic import BaseModel
+
+from llama_stack.models.llama.datatypes import Primitive
+
+type JSONValue = Primitive | list["JSONValue"] | dict[str, "JSONValue"]
+
+
+def serialize_value(value: Any) -> str:
+    return str(_prepare_for_json(value))
+
+
+def _prepare_for_json(value: Any) -> JSONValue:
+    """Serialize a single value into JSON-compatible format."""
+    if value is None:
+        return ""
+    elif isinstance(value, str | int | float | bool):
+        return value
+    elif hasattr(value, "_name_"):
+        return cast(str, value._name_)
+    elif isinstance(value, BaseModel):
+        return cast(JSONValue, json.loads(value.model_dump_json()))
+    elif isinstance(value, list | tuple | set):
+        return [_prepare_for_json(item) for item in value]
+    elif isinstance(value, dict):
+        return {str(k): _prepare_for_json(v) for k, v in value.items()}
+    else:
+        try:
+            json.dumps(value)
+            return cast(JSONValue, value)
+        except Exception:
+            return str(value)
+
+
+def trace_protocol[T: type[Any]](cls: T) -> T:
+    """
+    A class decorator that automatically traces all methods in a protocol/base class
+    and its inheriting classes.
+    """
+
+    def trace_method(method: Callable[..., Any]) -> Callable[..., Any]:
+        is_async = asyncio.iscoroutinefunction(method)
+        is_async_gen = inspect.isasyncgenfunction(method)
+
+        def create_span_context(self: Any, *args: Any, **kwargs: Any) -> tuple[str, str, dict[str, Primitive]]:
+            class_name = self.__class__.__name__
+            method_name = method.__name__
+            span_type = "async_generator" if is_async_gen else "async" if is_async else "sync"
+            sig = inspect.signature(method)
+            param_names = list(sig.parameters.keys())[1:]  # Skip 'self'
+            combined_args: dict[str, str] = {}
+            for i, arg in enumerate(args):
+                param_name = param_names[i] if i < len(param_names) else f"position_{i + 1}"
+                combined_args[param_name] = serialize_value(arg)
+            for k, v in kwargs.items():
+                combined_args[str(k)] = serialize_value(v)
+
+            span_attributes: dict[str, Primitive] = {
+                "__autotraced__": True,
+                "__class__": class_name,
+                "__method__": method_name,
+                "__type__": span_type,
+                "__args__": json.dumps(combined_args),
+            }
+
+            return class_name, method_name, span_attributes
+
+        @wraps(method)
+        async def async_gen_wrapper(self: Any, *args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]:
+            from llama_stack.core.telemetry import tracing
+
+            class_name, method_name, span_attributes = create_span_context(self, *args, **kwargs)
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                count = 0
+                try:
+                    async for item in method(self, *args, **kwargs):
+                        yield item
+                        count += 1
+                finally:
+                    span.set_attribute("chunk_count", count)
+
+        @wraps(method)
+        async def async_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            from llama_stack.core.telemetry import tracing
+
+            class_name, method_name, span_attributes = create_span_context(self, *args, **kwargs)
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = await method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as e:
+                    span.set_attribute("error", str(e))
+                    raise
+
+        @wraps(method)
+        def sync_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+            from llama_stack.core.telemetry import tracing
+
+            class_name, method_name, span_attributes = create_span_context(self, *args, **kwargs)
+
+            with tracing.span(f"{class_name}.{method_name}", span_attributes) as span:
+                try:
+                    result = method(self, *args, **kwargs)
+                    span.set_attribute("output", serialize_value(result))
+                    return result
+                except Exception as e:
+                    span.set_attribute("error", str(e))
+                    raise
+
+        if is_async_gen:
+            return async_gen_wrapper
+        elif is_async:
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    original_init_subclass = cast(Callable[..., Any] | None, getattr(cls, "__init_subclass__", None))
+
+    def __init_subclass__(cls_child: type[Any], **kwargs: Any) -> None:  # noqa: N807
+        if original_init_subclass:
+            cast(Callable[..., None], original_init_subclass)(**kwargs)
+
+        for name, method in vars(cls_child).items():
+            if inspect.isfunction(method) and not name.startswith("_"):
+                setattr(cls_child, name, trace_method(method))  # noqa: B010
+
+    cls_any = cast(Any, cls)
+    cls_any.__init_subclass__ = classmethod(__init_subclass__)
+
+    return cls

llama_stack/core/telemetry/tracing.py

@@ -0,0 +1,388 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import asyncio
+import contextvars
+import logging  # allow-direct-logging
+import queue
+import secrets
+import sys
+import threading
+import time
+from collections.abc import Callable
+from datetime import UTC, datetime
+from functools import wraps
+from typing import Any, Self
+
+from llama_stack.apis.telemetry import (
+    Event,
+    LogSeverity,
+    Span,
+    SpanEndPayload,
+    SpanStartPayload,
+    SpanStatus,
+    StructuredLogEvent,
+    Telemetry,
+    UnstructuredLogEvent,
+)
+from llama_stack.core.telemetry.trace_protocol import serialize_value
+from llama_stack.log import get_logger
+
+logger = get_logger(__name__, category="core")
+
+# Fallback logger that does NOT propagate to TelemetryHandler to avoid recursion
+_fallback_logger = logging.getLogger("llama_stack.telemetry.background")
+if not _fallback_logger.handlers:
+    _fallback_logger.propagate = False
+    _fallback_logger.setLevel(logging.ERROR)
+    _fallback_handler = logging.StreamHandler(sys.stderr)
+    _fallback_handler.setLevel(logging.ERROR)
+    _fallback_handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(name)s: %(message)s"))
+    _fallback_logger.addHandler(_fallback_handler)
+
+
+INVALID_SPAN_ID = 0x0000000000000000
+INVALID_TRACE_ID = 0x00000000000000000000000000000000
+
+ROOT_SPAN_MARKERS = ["__root__", "__root_span__"]
+# The logical root span may not be visible to this process if a parent context
+# is passed in. The local root span is the first local span in a trace.
+LOCAL_ROOT_SPAN_MARKER = "__local_root_span__"
+
+
+def trace_id_to_str(trace_id: int) -> str:
+    """Convenience trace ID formatting method
+    Args:
+        trace_id: Trace ID int
+
+    Returns:
+        The trace ID as 32-byte hexadecimal string
+    """
+    return format(trace_id, "032x")
+
+
+def span_id_to_str(span_id: int) -> str:
+    """Convenience span ID formatting method
+    Args:
+        span_id: Span ID int
+
+    Returns:
+        The span ID as 16-byte hexadecimal string
+    """
+    return format(span_id, "016x")
+
+
+def generate_span_id() -> str:
+    span_id = secrets.randbits(64)
+    while span_id == INVALID_SPAN_ID:
+        span_id = secrets.randbits(64)
+    return span_id_to_str(span_id)
+
+
+def generate_trace_id() -> str:
+    trace_id = secrets.randbits(128)
+    while trace_id == INVALID_TRACE_ID:
+        trace_id = secrets.randbits(128)
+    return trace_id_to_str(trace_id)
+
+
+LOG_QUEUE_FULL_LOG_INTERVAL_SECONDS = 60.0
+
+
+class BackgroundLogger:
+    def __init__(self, api: Telemetry, capacity: int = 100000):
+        self.api = api
+        self.log_queue: queue.Queue[Any] = queue.Queue(maxsize=capacity)
+        self.worker_thread = threading.Thread(target=self._worker, daemon=True)
+        self.worker_thread.start()
+        self._last_queue_full_log_time: float = 0.0
+        self._dropped_since_last_notice: int = 0
+
+    def log_event(self, event: Event) -> None:
+        try:
+            self.log_queue.put_nowait(event)
+        except queue.Full:
+            # Aggregate drops and emit at most once per interval via fallback logger
+            self._dropped_since_last_notice += 1
+            current_time = time.time()
+            if current_time - self._last_queue_full_log_time >= LOG_QUEUE_FULL_LOG_INTERVAL_SECONDS:
+                _fallback_logger.error(
+                    "Log queue is full; dropped %d events since last notice",
+                    self._dropped_since_last_notice,
+                )
+                self._last_queue_full_log_time = current_time
+                self._dropped_since_last_notice = 0
+
+    def _worker(self):
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(self._process_logs())
+
+    async def _process_logs(self):
+        while True:
+            try:
+                event = self.log_queue.get()
+                await self.api.log_event(event)
+            except Exception:
+                import traceback
+
+                traceback.print_exc()
+                print("Error processing log event")
+            finally:
+                self.log_queue.task_done()
+
+    def __del__(self) -> None:
+        self.log_queue.join()
+
+
+BACKGROUND_LOGGER: BackgroundLogger | None = None
+
+
+def enqueue_event(event: Event) -> None:
+    """Enqueue a telemetry event to the background logger if available.
+
+    This provides a non-blocking path for routers and other hot paths to
+    submit telemetry without awaiting the Telemetry API, reducing contention
+    with the main event loop.
+    """
+    global BACKGROUND_LOGGER
+    if BACKGROUND_LOGGER is None:
+        raise RuntimeError("Telemetry API not initialized")
+    BACKGROUND_LOGGER.log_event(event)
+
+
+class TraceContext:
+    def __init__(self, logger: BackgroundLogger, trace_id: str):
+        self.logger = logger
+        self.trace_id = trace_id
+        self.spans: list[Span] = []
+
+    def push_span(self, name: str, attributes: dict[str, Any] | None = None) -> Span:
+        current_span = self.get_current_span()
+        span = Span(
+            span_id=generate_span_id(),
+            trace_id=self.trace_id,
+            name=name,
+            start_time=datetime.now(UTC),
+            parent_span_id=current_span.span_id if current_span else None,
+            attributes=attributes,
+        )
+
+        self.logger.log_event(
+            StructuredLogEvent(
+                trace_id=span.trace_id,
+                span_id=span.span_id,
+                timestamp=span.start_time,
+                attributes=span.attributes,
+                payload=SpanStartPayload(
+                    name=span.name,
+                    parent_span_id=span.parent_span_id,
+                ),
+            )
+        )
+
+        self.spans.append(span)
+        return span
+
+    def pop_span(self, status: SpanStatus = SpanStatus.OK) -> None:
+        span = self.spans.pop()
+        if span is not None:
+            self.logger.log_event(
+                StructuredLogEvent(
+                    trace_id=span.trace_id,
+                    span_id=span.span_id,
+                    timestamp=span.start_time,
+                    attributes=span.attributes,
+                    payload=SpanEndPayload(
+                        status=status,
+                    ),
+                )
+            )
+
+    def get_current_span(self) -> Span | None:
+        return self.spans[-1] if self.spans else None
+
+
+CURRENT_TRACE_CONTEXT: contextvars.ContextVar[TraceContext | None] = contextvars.ContextVar(
+    "trace_context", default=None
+)
+
+
+def setup_logger(api: Telemetry, level: int = logging.INFO):
+    global BACKGROUND_LOGGER
+
+    if BACKGROUND_LOGGER is None:
+        BACKGROUND_LOGGER = BackgroundLogger(api)
+    root_logger = logging.getLogger()
+    root_logger.setLevel(level)
+    root_logger.addHandler(TelemetryHandler())
+
+
+async def start_trace(name: str, attributes: dict[str, Any] | None = None) -> TraceContext | None:
+    global CURRENT_TRACE_CONTEXT, BACKGROUND_LOGGER
+
+    if BACKGROUND_LOGGER is None:
+        logger.debug("No Telemetry implementation set. Skipping trace initialization...")
+        return None
+
+    trace_id = generate_trace_id()
+    context = TraceContext(BACKGROUND_LOGGER, trace_id)
+    # Mark this span as the root for the trace for now. The processing of
+    # traceparent context if supplied comes later and will result in the
+    # ROOT_SPAN_MARKERS being removed. Also mark this is the 'local' root,
+    # i.e. the root of the spans originating in this process as this is
+    # needed to ensure that we insert this 'local' root span's id into
+    # the trace record in sqlite store.
+    attributes = dict.fromkeys(ROOT_SPAN_MARKERS, True) | {LOCAL_ROOT_SPAN_MARKER: True} | (attributes or {})
+    context.push_span(name, attributes)
+
+    CURRENT_TRACE_CONTEXT.set(context)
+    return context
+
+
+async def end_trace(status: SpanStatus = SpanStatus.OK):
+    global CURRENT_TRACE_CONTEXT
+
+    context = CURRENT_TRACE_CONTEXT.get()
+    if context is None:
+        logger.debug("No trace context to end")
+        return
+
+    context.pop_span(status)
+    CURRENT_TRACE_CONTEXT.set(None)
+
+
+def severity(levelname: str) -> LogSeverity:
+    if levelname == "DEBUG":
+        return LogSeverity.DEBUG
+    elif levelname == "INFO":
+        return LogSeverity.INFO
+    elif levelname == "WARNING":
+        return LogSeverity.WARN
+    elif levelname == "ERROR":
+        return LogSeverity.ERROR
+    elif levelname == "CRITICAL":
+        return LogSeverity.CRITICAL
+    else:
+        raise ValueError(f"Unknown log level: {levelname}")
+
+
+# TODO: ideally, the actual emitting should be done inside a separate daemon
+# process completely isolated from the server
+class TelemetryHandler(logging.Handler):
+    def emit(self, record: logging.LogRecord) -> None:
+        # horrendous hack to avoid logging from asyncio and getting into an infinite loop
+        if record.module in ("asyncio", "selector_events"):
+            return
+
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT.get()
+        if context is None:
+            return
+
+        span = context.get_current_span()
+        if span is None:
+            return
+
+        enqueue_event(
+            UnstructuredLogEvent(
+                trace_id=span.trace_id,
+                span_id=span.span_id,
+                timestamp=datetime.now(UTC),
+                message=self.format(record),
+                severity=severity(record.levelname),
+            )
+        )
+
+    def close(self) -> None:
+        pass
+
+
+class SpanContextManager:
+    def __init__(self, name: str, attributes: dict[str, Any] | None = None):
+        self.name = name
+        self.attributes = attributes
+        self.span: Span | None = None
+
+    def __enter__(self) -> Self:
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to push span")
+            return self
+
+        self.span = context.push_span(self.name, self.attributes)
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to pop span")
+            return
+
+        context.pop_span()
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        if self.span:
+            if self.span.attributes is None:
+                self.span.attributes = {}
+            self.span.attributes[key] = serialize_value(value)
+
+    async def __aenter__(self) -> Self:
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to push span")
+            return self
+
+        self.span = context.push_span(self.name, self.attributes)
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+        global CURRENT_TRACE_CONTEXT
+        context = CURRENT_TRACE_CONTEXT.get()
+        if not context:
+            logger.debug("No trace context to pop span")
+            return
+
+        context.pop_span()
+
+    def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
+        @wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            with self:
+                return func(*args, **kwargs)
+
+        @wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            async with self:
+                return await func(*args, **kwargs)
+
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            if asyncio.iscoroutinefunction(func):
+                return async_wrapper(*args, **kwargs)
+            else:
+                return sync_wrapper(*args, **kwargs)
+
+        return wrapper
+
+
+def span(name: str, attributes: dict[str, Any] | None = None) -> SpanContextManager:
+    return SpanContextManager(name, attributes)
+
+
+def get_current_span() -> Span | None:
+    global CURRENT_TRACE_CONTEXT
+    if CURRENT_TRACE_CONTEXT is None:
+        logger.debug("No trace context to get current span")
+        return None
+
+    context = CURRENT_TRACE_CONTEXT.get()
+    if context:
+        return context.get_current_span()
+    return None