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/datatypes.py

@@ -31,6 +31,7 @@ from llama_stack.core.storage.datatypes import (
     StorageBackendType,
     StorageConfig,
 )
+from llama_stack.log import LoggingConfig
 from llama_stack.providers.datatypes import Api, ProviderSpec
 
 LLAMA_STACK_BUILD_CONFIG_VERSION = 2
@@ -195,14 +196,6 @@ class TelemetryConfig(BaseModel):
     enabled: bool = Field(default=False, description="enable or disable telemetry")
 
 
-class LoggingConfig(BaseModel):
-    category_levels: dict[str, str] = Field(
-        default_factory=dict,
-        description="""
- Dictionary of different logging configurations for different portions (ex: core, server) of llama stack""",
-    )
-
-
 class OAuth2JWKSConfig(BaseModel):
     # The JWKS URI for collecting public keys
     uri: str

llama_stack/core/distribution.py

@@ -25,7 +25,7 @@ from llama_stack.providers.datatypes import (
 logger = get_logger(name=__name__, category="core")
 
 
-INTERNAL_APIS = {Api.inspect, Api.providers, Api.prompts, Api.conversations, Api.telemetry}
+INTERNAL_APIS = {Api.inspect, Api.providers, Api.prompts, Api.conversations}
 
 
 def stack_apis() -> list[Api]:

llama_stack/core/library_client.py

@@ -32,7 +32,7 @@ from termcolor import cprint
 
 from llama_stack.core.build import print_pip_install_help
 from llama_stack.core.configure import parse_and_maybe_upgrade_config
-from llama_stack.core.datatypes import Api, BuildConfig, BuildProvider, DistributionSpec
+from llama_stack.core.datatypes import BuildConfig, BuildProvider, DistributionSpec
 from llama_stack.core.request_headers import (
     PROVIDER_DATA_VAR,
     request_provider_data_context,
@@ -44,11 +44,12 @@ from llama_stack.core.stack import (
     get_stack_run_config_from_distro,
     replace_env_vars,
 )
+from llama_stack.core.telemetry import Telemetry
+from llama_stack.core.telemetry.tracing import CURRENT_TRACE_CONTEXT, end_trace, setup_logger, start_trace
 from llama_stack.core.utils.config import redact_sensitive_fields
 from llama_stack.core.utils.context import preserve_contexts_async_generator
 from llama_stack.core.utils.exec import in_notebook
 from llama_stack.log import get_logger, setup_logging
-from llama_stack.providers.utils.telemetry.tracing import CURRENT_TRACE_CONTEXT, end_trace, setup_logger, start_trace
 from llama_stack.strong_typing.inspection import is_unwrapped_body_param
 
 logger = get_logger(name=__name__, category="core")
@@ -293,8 +294,8 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient):
             raise _e
 
         assert self.impls is not None
-        if Api.telemetry in self.impls:
-            setup_logger(self.impls[Api.telemetry])
+        if self.config.telemetry.enabled:
+            setup_logger(Telemetry())
 
         if not os.environ.get("PYTEST_CURRENT_TEST"):
             console = Console()

llama_stack/core/resolver.py

@@ -27,7 +27,6 @@ from llama_stack.apis.safety import Safety
 from llama_stack.apis.scoring import Scoring
 from llama_stack.apis.scoring_functions import ScoringFunctions
 from llama_stack.apis.shields import Shields
-from llama_stack.apis.telemetry import Telemetry
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
 from llama_stack.apis.vector_stores import VectorStore
@@ -49,7 +48,6 @@ from llama_stack.providers.datatypes import (
     Api,
     BenchmarksProtocolPrivate,
     DatasetsProtocolPrivate,
-    InlineProviderSpec,
     ModelsProtocolPrivate,
     ProviderSpec,
     RemoteProviderConfig,
@@ -98,7 +96,6 @@ def api_protocol_map(external_apis: dict[Api, ExternalApiSpec] | None = None) ->
         Api.files: Files,
         Api.prompts: Prompts,
         Api.conversations: Conversations,
-        Api.telemetry: Telemetry,
     }
 
     if external_apis:
@@ -241,24 +238,6 @@ def validate_and_prepare_providers(
         key = api_str if api not in router_apis else f"inner-{api_str}"
         providers_with_specs[key] = specs
 
-    # TODO: remove this logic, telemetry should not have providers.
-    # if telemetry has been enabled in the config initialize our internal impl
-    # telemetry is not an external API so it SHOULD NOT be auto-routed.
-    if run_config.telemetry.enabled:
-        specs = {}
-        p = InlineProviderSpec(
-            api=Api.telemetry,
-            provider_type="inline::meta-reference",
-            pip_packages=[],
-            optional_api_dependencies=[Api.datasetio],
-            module="llama_stack.providers.inline.telemetry.meta_reference",
-            config_class="llama_stack.providers.inline.telemetry.meta_reference.config.TelemetryConfig",
-            description="Meta's reference implementation of telemetry and observability using OpenTelemetry.",
-        )
-        spec = ProviderWithSpec(spec=p, provider_type="inline::meta-reference", provider_id="meta-reference")
-        specs["meta-reference"] = spec
-        providers_with_specs["telemetry"] = specs
-
     return providers_with_specs
 
 

llama_stack/core/routers/__init__.py

@@ -72,14 +72,6 @@ async def get_auto_router_impl(
         raise ValueError(f"API {api.value} not found in router map")
 
     api_to_dep_impl = {}
-    if run_config.telemetry.enabled:
-        api_to_deps = {
-            "inference": {"telemetry": Api.telemetry},
-        }
-        for dep_name, dep_api in api_to_deps.get(api.value, {}).items():
-            if dep_api in deps:
-                api_to_dep_impl[dep_name] = deps[dep_api]
-
     # TODO: move pass configs to routers instead
     if api == Api.inference:
         inference_ref = run_config.storage.stores.inference
@@ -92,6 +84,7 @@ async def get_auto_router_impl(
         )
         await inference_store.initialize()
         api_to_dep_impl["store"] = inference_store
+        api_to_dep_impl["telemetry_enabled"] = run_config.telemetry.enabled
 
     elif api == Api.vector_io:
         api_to_dep_impl["vector_stores_config"] = run_config.vector_stores

llama_stack/core/routers/inference.py

@@ -53,13 +53,13 @@ from llama_stack.apis.inference.inference import (
     OpenAIChatCompletionContentPartTextParam,
 )
 from llama_stack.apis.models import Model, ModelType
-from llama_stack.apis.telemetry import MetricEvent, MetricInResponse, Telemetry
+from llama_stack.apis.telemetry import MetricEvent, MetricInResponse
+from llama_stack.core.telemetry.tracing import enqueue_event, get_current_span
 from llama_stack.log import get_logger
 from llama_stack.models.llama.llama3.chat_format import ChatFormat
 from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.providers.datatypes import HealthResponse, HealthStatus, RoutingTable
 from llama_stack.providers.utils.inference.inference_store import InferenceStore
-from llama_stack.providers.utils.telemetry.tracing import enqueue_event, get_current_span
 
 logger = get_logger(name=__name__, category="core::routers")
 
@@ -70,14 +70,14 @@ class InferenceRouter(Inference):
     def __init__(
         self,
         routing_table: RoutingTable,
-        telemetry: Telemetry | None = None,
         store: InferenceStore | None = None,
+        telemetry_enabled: bool = False,
     ) -> None:
         logger.debug("Initializing InferenceRouter")
         self.routing_table = routing_table
-        self.telemetry = telemetry
+        self.telemetry_enabled = telemetry_enabled
         self.store = store
-        if self.telemetry:
+        if self.telemetry_enabled:
             self.tokenizer = Tokenizer.get_instance()
             self.formatter = ChatFormat(self.tokenizer)
 
@@ -159,7 +159,7 @@ class InferenceRouter(Inference):
         model: Model,
     ) -> list[MetricInResponse]:
         metrics = self._construct_metrics(prompt_tokens, completion_tokens, total_tokens, model)
-        if self.telemetry:
+        if self.telemetry_enabled:
             for metric in metrics:
                 enqueue_event(metric)
         return [MetricInResponse(metric=metric.metric, value=metric.value) for metric in metrics]
@@ -223,7 +223,7 @@ class InferenceRouter(Inference):
             # that we do not return an AsyncIterator, our tests expect a stream of chunks we cannot intercept currently.
 
         response = await provider.openai_completion(params)
-        if self.telemetry:
+        if self.telemetry_enabled:
             metrics = self._construct_metrics(
                 prompt_tokens=response.usage.prompt_tokens,
                 completion_tokens=response.usage.completion_tokens,
@@ -285,7 +285,7 @@ class InferenceRouter(Inference):
         if self.store:
             asyncio.create_task(self.store.store_chat_completion(response, params.messages))
 
-        if self.telemetry:
+        if self.telemetry_enabled:
             metrics = self._construct_metrics(
                 prompt_tokens=response.usage.prompt_tokens,
                 completion_tokens=response.usage.completion_tokens,
@@ -393,7 +393,7 @@ class InferenceRouter(Inference):
             else:
                 if hasattr(chunk, "delta"):
                     completion_text += chunk.delta
-                if hasattr(chunk, "stop_reason") and chunk.stop_reason and self.telemetry:
+                if hasattr(chunk, "stop_reason") and chunk.stop_reason and self.telemetry_enabled:
                     complete = True
                     completion_tokens = await self._count_tokens(completion_text)
             # if we are done receiving tokens
@@ -401,7 +401,7 @@ class InferenceRouter(Inference):
                 total_tokens = (prompt_tokens or 0) + (completion_tokens or 0)
 
                 # Create a separate span for streaming completion metrics
-                if self.telemetry:
+                if self.telemetry_enabled:
                     # Log metrics in the new span context
                     completion_metrics = self._construct_metrics(
                         prompt_tokens=prompt_tokens,
@@ -450,7 +450,7 @@ class InferenceRouter(Inference):
         total_tokens = (prompt_tokens or 0) + (completion_tokens or 0)
 
         # Create a separate span for completion metrics
-        if self.telemetry:
+        if self.telemetry_enabled:
             # Log metrics in the new span context
             completion_metrics = self._construct_metrics(
                 prompt_tokens=prompt_tokens,
@@ -548,7 +548,7 @@ class InferenceRouter(Inference):
                         completion_text += "".join(choice_data["content_parts"])
 
                     # Add metrics to the chunk
-                    if self.telemetry and hasattr(chunk, "usage") and chunk.usage:
+                    if self.telemetry_enabled and hasattr(chunk, "usage") and chunk.usage:
                         metrics = self._construct_metrics(
                             prompt_tokens=chunk.usage.prompt_tokens,
                             completion_tokens=chunk.usage.completion_tokens,

llama_stack/core/server/server.py

@@ -36,7 +36,6 @@ from llama_stack.apis.common.responses import PaginatedResponse
 from llama_stack.core.access_control.access_control import AccessDeniedError
 from llama_stack.core.datatypes import (
     AuthenticationRequiredError,
-    LoggingConfig,
     StackRunConfig,
     process_cors_config,
 )
@@ -53,19 +52,13 @@ from llama_stack.core.stack import (
     cast_image_name_to_string,
     replace_env_vars,
 )
+from llama_stack.core.telemetry import Telemetry
+from llama_stack.core.telemetry.tracing import CURRENT_TRACE_CONTEXT, setup_logger
 from llama_stack.core.utils.config import redact_sensitive_fields
 from llama_stack.core.utils.config_resolution import Mode, resolve_config_or_distro
 from llama_stack.core.utils.context import preserve_contexts_async_generator
-from llama_stack.log import get_logger, setup_logging
+from llama_stack.log import LoggingConfig, get_logger, setup_logging
 from llama_stack.providers.datatypes import Api
-from llama_stack.providers.inline.telemetry.meta_reference.config import TelemetryConfig
-from llama_stack.providers.inline.telemetry.meta_reference.telemetry import (
-    TelemetryAdapter,
-)
-from llama_stack.providers.utils.telemetry.tracing import (
-    CURRENT_TRACE_CONTEXT,
-    setup_logger,
-)
 
 from .auth import AuthenticationMiddleware
 from .quota import QuotaMiddleware
@@ -451,9 +444,7 @@ def create_app() -> StackApp:
             app.add_middleware(CORSMiddleware, **cors_config.model_dump())
 
     if config.telemetry.enabled:
-        setup_logger(impls[Api.telemetry])
-    else:
-        setup_logger(TelemetryAdapter(TelemetryConfig(), {}))
+        setup_logger(Telemetry())
 
     # Load external APIs if configured
     external_apis = load_external_apis(config)
@@ -511,7 +502,8 @@ def create_app() -> StackApp:
     app.exception_handler(RequestValidationError)(global_exception_handler)
     app.exception_handler(Exception)(global_exception_handler)
 
-    app.add_middleware(TracingMiddleware, impls=impls, external_apis=external_apis)
+    if config.telemetry.enabled:
+        app.add_middleware(TracingMiddleware, impls=impls, external_apis=external_apis)
 
     return app
 

llama_stack/core/server/tracing.py

@@ -7,8 +7,8 @@ from aiohttp import hdrs
 
 from llama_stack.core.external import ExternalApiSpec
 from llama_stack.core.server.routes import find_matching_route, initialize_route_impls
+from llama_stack.core.telemetry.tracing import end_trace, start_trace
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.telemetry.tracing import end_trace, start_trace
 
 logger = get_logger(name=__name__, category="core::server")
 

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