| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import atexit |
| 5 | import functools |
| 6 | import inspect |
| 7 | import os |
| 8 | import traceback |
| 9 | from collections.abc import Mapping |
| 10 | from contextlib import contextmanager |
| 11 | from typing import Any |
| 12 | |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.tracing.utils import TRACE_HEADERS, LoadingSpanAttributes |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | try: |
| 19 | from opentelemetry import trace |
| 20 | from opentelemetry.context.context import Context |
| 21 | from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( |
| 22 | OTLPSpanExporter as OTLPGrpcExporter, |
| 23 | ) |
| 24 | from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( |
| 25 | OTLPSpanExporter as OTLPHttpExporter, |
| 26 | ) |
| 27 | from opentelemetry.propagate import inject |
| 28 | from opentelemetry.sdk.environment_variables import ( |
| 29 | OTEL_EXPORTER_OTLP_TRACES_PROTOCOL, |
| 30 | ) |
| 31 | from opentelemetry.sdk.resources import Resource |
| 32 | from opentelemetry.sdk.trace import TracerProvider |
| 33 | from opentelemetry.sdk.trace.export import BatchSpanProcessor |
| 34 | from opentelemetry.trace import ( |
| 35 | SpanKind, # noqa: F401 |
| 36 | Tracer, |
| 37 | set_tracer_provider, |
| 38 | ) |
| 39 | from opentelemetry.trace.propagation.tracecontext import ( |
| 40 | TraceContextTextMapPropagator, |
| 41 | ) |
| 42 | |
| 43 | _IS_OTEL_AVAILABLE = True |
| 44 | otel_import_error_traceback = None |
| 45 | except ImportError: |
| 46 | _IS_OTEL_AVAILABLE = False |
| 47 | otel_import_error_traceback = traceback.format_exc() |
| 48 | trace = None # type: ignore |
| 49 | Context = Any # type: ignore |
| 50 | Tracer = Any # type: ignore |
| 51 | inject = None # type: ignore |
| 52 | Resource = None # type: ignore |
| 53 | SpanKind = Any # type: ignore |
| 54 | |
| 55 | |
| 56 | def is_otel_available() -> bool: |
| 57 | return _IS_OTEL_AVAILABLE |
| 58 | |
| 59 | |
| 60 | def init_otel_tracer( |
| 61 | instrumenting_module_name: str, |
| 62 | otlp_traces_endpoint: str, |
| 63 | extra_attributes: dict[str, str] | None = None, |
| 64 | ) -> Tracer: |
| 65 | """Initializes the OpenTelemetry tracer provider.""" |
| 66 | if not _IS_OTEL_AVAILABLE: |
| 67 | raise ValueError( |
| 68 | "OpenTelemetry is not available. Unable to initialize " |
| 69 | "a tracer. Ensure OpenTelemetry packages are installed. " |
| 70 | f"Original error:\n{otel_import_error_traceback}" |
| 71 | ) |
| 72 | |
| 73 | # Store the endpoint in environment so child processes can inherit it |
| 74 | os.environ["OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"] = otlp_traces_endpoint |
| 75 | |
| 76 | resource_attrs = {} |
| 77 | resource_attrs["vllm.instrumenting_module_name"] = instrumenting_module_name |
| 78 | resource_attrs["vllm.process_id"] = str(os.getpid()) |
| 79 | if extra_attributes: |
| 80 | resource_attrs.update(extra_attributes) |
| 81 | resource = Resource.create(resource_attrs) |
| 82 | |
| 83 | trace_provider = TracerProvider(resource=resource) |
| 84 | span_exporter = get_span_exporter(otlp_traces_endpoint) |
| 85 | trace_provider.add_span_processor(BatchSpanProcessor(span_exporter)) |
| 86 | set_tracer_provider(trace_provider) |
| 87 | |
| 88 | atexit.register(trace_provider.shutdown) |
| 89 | |
| 90 | tracer = trace_provider.get_tracer(instrumenting_module_name) |
| 91 | return tracer |
| 92 | |
| 93 | |
| 94 | def get_span_exporter(endpoint): |
| 95 | protocol = os.environ.get(OTEL_EXPORTER_OTLP_TRACES_PROTOCOL, "grpc") |
| 96 | if protocol == "grpc": |
| 97 | exporter = OTLPGrpcExporter(endpoint=endpoint, insecure=True) |
| 98 | elif protocol == "http/protobuf": |
| 99 | exporter = OTLPHttpExporter(endpoint=endpoint) |
| 100 | else: |
| 101 | raise ValueError(f"Unsupported OTLP protocol '{protocol}' is configured") |
| 102 | return exporter |
| 103 | |
| 104 | |
| 105 | def init_otel_worker_tracer( |
| 106 | instrumenting_module_name: str, |
| 107 | process_kind: str, |
| 108 | process_name: str, |
| 109 | ) -> Tracer: |
| 110 | """Backend-specific initialization for OpenTelemetry in a worker process.""" |
| 111 | # Initialize the tracer if an OTLP endpoint is configured. |
| 112 | # The endpoint is propagated via environment variable from the main process. |
| 113 | otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") |
| 114 | if not otlp_endpoint: |
| 115 | return None |
| 116 | |
| 117 | extra_attrs = { |
| 118 | "vllm.process_kind": process_kind, |
| 119 | "vllm.process_name": process_name, |
| 120 | } |
| 121 | |
| 122 | return init_otel_tracer(instrumenting_module_name, otlp_endpoint, extra_attrs) |
| 123 | |
| 124 | |
| 125 | def extract_trace_context(headers: Mapping[str, str] | None) -> Context | None: |
| 126 | """Extracts context from HTTP headers.""" |
| 127 | if _IS_OTEL_AVAILABLE and headers: |
| 128 | return TraceContextTextMapPropagator().extract(headers) |
| 129 | return None |
| 130 | |
| 131 | |
| 132 | def instrument_otel(func, span_name, attributes, record_exception): |
| 133 | """Internal wrapper logic for sync and async functions.""" |
| 134 | # Pre-calculate static code attributes once (these don't change) |
| 135 | code_attrs = { |
| 136 | LoadingSpanAttributes.CODE_FUNCTION: func.__qualname__, |
| 137 | LoadingSpanAttributes.CODE_NAMESPACE: func.__module__, |
| 138 | LoadingSpanAttributes.CODE_FILEPATH: func.__code__.co_filename, |
| 139 | LoadingSpanAttributes.CODE_LINENO: str(func.__code__.co_firstlineno), |
| 140 | } |
| 141 | if attributes: |
| 142 | code_attrs.update(attributes) |
| 143 | |
| 144 | final_span_name = span_name or func.__qualname__ |
| 145 | module_name = func.__module__ |
| 146 | |
| 147 | @functools.wraps(func) |
| 148 | async def async_wrapper(*args, **kwargs): |
| 149 | tracer = trace.get_tracer(module_name) |
| 150 | ctx = _get_smart_context() |
| 151 | with ( |
| 152 | tracer.start_as_current_span( |
| 153 | final_span_name, |
| 154 | context=ctx, |
| 155 | attributes=code_attrs, |
| 156 | record_exception=record_exception, |
| 157 | ), |
| 158 | propagate_trace_to_env(), |
| 159 | ): |
| 160 | return await func(*args, **kwargs) |
| 161 | |
| 162 | @functools.wraps(func) |
| 163 | def sync_wrapper(*args, **kwargs): |
| 164 | tracer = trace.get_tracer(module_name) |
| 165 | ctx = _get_smart_context() |
| 166 | with ( |
| 167 | tracer.start_as_current_span( |
| 168 | final_span_name, |
| 169 | context=ctx, |
| 170 | attributes=code_attrs, |
| 171 | record_exception=record_exception, |
| 172 | ), |
| 173 | propagate_trace_to_env(), |
| 174 | ): |
| 175 | return func(*args, **kwargs) |
| 176 | |
| 177 | return async_wrapper if inspect.iscoroutinefunction(func) else sync_wrapper |
| 178 | |
| 179 | |
| 180 | def manual_instrument_otel( |
| 181 | span_name: str, |
| 182 | start_time: int, |
| 183 | end_time: int | None = None, |
| 184 | attributes: dict[str, Any] | None = None, |
| 185 | context: Context | None = None, |
| 186 | kind: Any = None, # SpanKind, but typed as Any for when OTEL unavailable |
| 187 | ): |
| 188 | """Manually create and end a span with explicit timestamps.""" |
| 189 | if not _IS_OTEL_AVAILABLE: |
| 190 | return |
| 191 | |
| 192 | tracer = trace.get_tracer(__name__) |
| 193 | # Use provided context, or fall back to smart context detection |
| 194 | ctx = context if context is not None else _get_smart_context() |
| 195 | |
| 196 | span_kwargs: dict[str, Any] = { |
| 197 | "name": span_name, |
| 198 | "context": ctx, |
| 199 | "start_time": start_time, |
| 200 | } |
| 201 | if kind is not None: |
| 202 | span_kwargs["kind"] = kind |
| 203 | |
| 204 | span = tracer.start_span(**span_kwargs) |
| 205 | if attributes: |
| 206 | span.set_attributes(attributes) |
| 207 | if end_time is not None: |
| 208 | span.end(end_time=end_time) |
| 209 | else: |
| 210 | span.end() |
| 211 | |
| 212 | |
| 213 | def _get_smart_context() -> Context | None: |
| 214 | """Determines the parent context. |
| 215 | 1. If a Span is already active in this process, use it. |
| 216 | 2. If not, extract from os.environ, handling the case-sensitivity mismatch. |
| 217 | """ |
| 218 | current_span = trace.get_current_span() |
| 219 | if current_span.get_span_context().is_valid: |
| 220 | return None |
| 221 | |
| 222 | carrier = {} |
| 223 | |
| 224 | if tp := os.environ.get("traceparent", os.environ.get("TRACEPARENT")): # noqa: SIM112 |
| 225 | carrier["traceparent"] = tp |
| 226 | |
| 227 | if ts := os.environ.get("tracestate", os.environ.get("TRACESTATE")): # noqa: SIM112 |
| 228 | carrier["tracestate"] = ts |
| 229 | |
| 230 | if not carrier: |
| 231 | carrier = dict(os.environ) |
| 232 | |
| 233 | return TraceContextTextMapPropagator().extract(carrier) |
| 234 | |
| 235 | |
| 236 | @contextmanager |
| 237 | def propagate_trace_to_env(): |
| 238 | """Temporarily injects the current OTel context into os.environ. |
| 239 | This ensures that any subprocesses (like vLLM workers) spawned |
| 240 | within this context inherit the correct traceparent. |
| 241 | """ |
| 242 | if not _IS_OTEL_AVAILABLE: |
| 243 | yield |
| 244 | return |
| 245 | |
| 246 | # Capture original state of relevant keys |
| 247 | original_state = {k: os.environ.get(k) for k in TRACE_HEADERS} |
| 248 | |
| 249 | try: |
| 250 | # inject() writes 'traceparent' and 'tracestate' to os.environ |
| 251 | inject(os.environ) |
| 252 | yield |
| 253 | |
| 254 | finally: |
| 255 | # Restore original environment |
| 256 | for key, original_value in original_state.items(): |
| 257 | if original_value is None: |
| 258 | os.environ.pop(key, None) |
| 259 | else: |
| 260 | os.environ[key] = original_value |