| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import functools |
| 5 | from collections.abc import Callable |
| 6 | from typing import Any, TypeAlias |
| 7 | |
| 8 | # Import the implementation details |
| 9 | from .otel import ( |
| 10 | SpanKind, |
| 11 | extract_trace_context, |
| 12 | init_otel_tracer, |
| 13 | init_otel_worker_tracer, |
| 14 | instrument_otel, |
| 15 | is_otel_available, |
| 16 | manual_instrument_otel, |
| 17 | otel_import_error_traceback, |
| 18 | ) |
| 19 | from .utils import ( |
| 20 | SpanAttributes, |
| 21 | contains_trace_headers, |
| 22 | extract_trace_headers, |
| 23 | log_tracing_disabled_warning, |
| 24 | ) |
| 25 | |
| 26 | __all__ = [ |
| 27 | "instrument", |
| 28 | "instrument_manual", |
| 29 | "init_tracer", |
| 30 | "maybe_init_worker_tracer", |
| 31 | "is_tracing_available", |
| 32 | "SpanAttributes", |
| 33 | "SpanKind", |
| 34 | "extract_trace_context", |
| 35 | "extract_trace_headers", |
| 36 | "log_tracing_disabled_warning", |
| 37 | "contains_trace_headers", |
| 38 | "otel_import_error_traceback", |
| 39 | ] |
| 40 | |
| 41 | BackendAvailableFunc: TypeAlias = Callable[[], bool] |
| 42 | InstrumentFunc: TypeAlias = Callable[..., Any] |
| 43 | InstrumentManualFunc: TypeAlias = Callable[..., Any] |
| 44 | InitTracerFunc: TypeAlias = Callable[..., Any] |
| 45 | InitWorkerTracerFunc: TypeAlias = Callable[..., Any] |
| 46 | _REGISTERED_TRACING_BACKENDS: dict[ |
| 47 | str, |
| 48 | tuple[ |
| 49 | BackendAvailableFunc, |
| 50 | InitTracerFunc, |
| 51 | InitWorkerTracerFunc, |
| 52 | InstrumentFunc, |
| 53 | InstrumentManualFunc, |
| 54 | ], |
| 55 | ] = { |
| 56 | "otel": ( |
| 57 | is_otel_available, |
| 58 | init_otel_tracer, |
| 59 | init_otel_worker_tracer, |
| 60 | instrument_otel, |
| 61 | manual_instrument_otel, |
| 62 | ), |
| 63 | } |
| 64 | |
| 65 | |
| 66 | def init_tracer( |
| 67 | instrumenting_module_name: str, |
| 68 | otlp_traces_endpoint: str, |
| 69 | extra_attributes: dict[str, str] | None = None, |
| 70 | ): |
| 71 | is_available, init_tracer_fn, _, _, _ = _REGISTERED_TRACING_BACKENDS["otel"] |
| 72 | if is_available(): |
| 73 | return init_tracer_fn( |
| 74 | instrumenting_module_name, otlp_traces_endpoint, extra_attributes |
| 75 | ) |
| 76 | |
| 77 | |
| 78 | def maybe_init_worker_tracer( |
| 79 | instrumenting_module_name: str, |
| 80 | process_kind: str, |
| 81 | process_name: str, |
| 82 | ): |
| 83 | is_available, _, init_worker_tracer_fn, _, _ = _REGISTERED_TRACING_BACKENDS["otel"] |
| 84 | if is_available(): |
| 85 | return init_worker_tracer_fn( |
| 86 | instrumenting_module_name, process_kind, process_name |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | def instrument( |
| 91 | obj: Callable | None = None, |
| 92 | *, |
| 93 | span_name: str = "", |
| 94 | attributes: dict[str, str] | None = None, |
| 95 | record_exception: bool = True, |
| 96 | ): |
| 97 | """Generic decorator to instrument functions.""" |
| 98 | if obj is None: |
| 99 | return functools.partial( |
| 100 | instrument, |
| 101 | span_name=span_name, |
| 102 | attributes=attributes, |
| 103 | record_exception=record_exception, |
| 104 | ) |
| 105 | |
| 106 | # Dispatch to OTel (and potentially others later) |
| 107 | is_available, _, _, otel_instrument, _ = _REGISTERED_TRACING_BACKENDS["otel"] |
| 108 | if is_available(): |
| 109 | return otel_instrument( |
| 110 | func=obj, |
| 111 | span_name=span_name, |
| 112 | attributes=attributes, |
| 113 | record_exception=record_exception, |
| 114 | ) |
| 115 | else: |
| 116 | return obj |
| 117 | |
| 118 | |
| 119 | def instrument_manual( |
| 120 | span_name: str, |
| 121 | start_time: int, |
| 122 | end_time: int | None = None, |
| 123 | attributes: dict[str, Any] | None = None, |
| 124 | context: Any = None, |
| 125 | kind: Any = None, |
| 126 | ): |
| 127 | """Manually create a span with explicit timestamps. |
| 128 | |
| 129 | Args: |
| 130 | span_name: Name of the span to create. |
| 131 | start_time: Start time in nanoseconds since epoch. |
| 132 | end_time: Optional end time in nanoseconds. If None, ends immediately. |
| 133 | attributes: Optional dict of span attributes. |
| 134 | context: Optional trace context (e.g., from extract_trace_context). |
| 135 | kind: Optional SpanKind (e.g., SpanKind.SERVER). |
| 136 | |
| 137 | """ |
| 138 | is_available, _, _, _, manual_instrument_fn = _REGISTERED_TRACING_BACKENDS["otel"] |
| 139 | if is_available(): |
| 140 | return manual_instrument_fn( |
| 141 | span_name, start_time, end_time, attributes, context, kind |
| 142 | ) |
| 143 | else: |
| 144 | return None |
| 145 | |
| 146 | |
| 147 | def is_tracing_available() -> bool: |
| 148 | """Returns True if any tracing backend (OTel, Profiler, etc.) is available. |
| 149 | Use this to guard expensive tracing logic in the main code. |
| 150 | """ |
| 151 | check_available = [ |
| 152 | is_available |
| 153 | for is_available, _, _, _, _ in _REGISTERED_TRACING_BACKENDS.values() |
| 154 | ] |
| 155 | return any(check_available) |