[tracing] sglang tracing v2: support exporting tracing data asynchronously (#30023)

python/sglang/srt/environ.py

@@ -364,6 +364,8 @@ class Envs:
     SGLANG_MEM_PROFILE_MAX_ENTRIES = EnvInt(100000)
     SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500)
     SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64)
+    SGLANG_TRACE_ASYNC = EnvBool(False)
+    SGLANG_TRACE_ASYNC_FLUSH_THRESHOLD = EnvInt(100)
     SGLANG_NATIVE_MOVE_KV_CACHE = EnvBool(False)
     # Disable lazy compaction in the unified memory pool allocator and
     # fall back to the per-free eager compaction. Used for production

python/sglang/srt/managers/scheduler.py

@@ -268,6 +268,7 @@ from sglang.srt.model_loader.utils import get_resolved_model_impl
 from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin
 from sglang.srt.observability.metrics_collector import SchedulerMetricsCollector
 from sglang.srt.observability.req_time_stats import (
+    flush_trace_batch,
     set_schedule_time_batch,
     set_time_batch,
 )
@@ -3889,6 +3890,9 @@ class Scheduler(
         batch: ScheduleBatch,
         result: Union[GenerationBatchResult, EmbeddingBatchResult],
     ):
+        # Flush async trace ops here: in overlap mode this CPU work runs while
+        # the next batch's GPU forward is in flight, giving free overlap.
+        flush_trace_batch(batch.reqs)
         self.publish_load_snapshot(force=batch.forward_mode.is_extend())
 
         if batch.forward_mode.is_decode():

python/sglang/srt/observability/req_time_stats.py

@@ -37,6 +37,10 @@ from sglang.srt.observability.trace import (
     TraceSliceContext,
     get_global_tracing_enabled,
 )
+from sglang.srt.observability.trace_async import (
+    TraceReqContextAsync,
+    is_async_tracing_available,
+)
 from sglang.srt.utils import get_bool_env_var
 
 if TYPE_CHECKING:
@@ -286,13 +290,22 @@ class ReqTimeStatsBase:
         bootstrap_room: Optional[int],
         external_trace_header: Optional[Dict[str, str]] = None,
     ):
-        self.trace_ctx = TraceReqContext(
-            rid=rid,
-            bootstrap_room=bootstrap_room,
-            role=self.disagg_mode_str(),
-            module_name="request",
-            external_trace_header=external_trace_header,
-        )
+        if is_async_tracing_available():
+            self.trace_ctx = TraceReqContextAsync(
+                rid=rid,
+                bootstrap_room=bootstrap_room,
+                role=self.disagg_mode_str(),
+                module_name="request",
+                external_trace_header=external_trace_header,
+            )
+        else:
+            self.trace_ctx = TraceReqContext(
+                rid=rid,
+                bootstrap_room=bootstrap_room,
+                role=self.disagg_mode_str(),
+                module_name="request",
+                external_trace_header=external_trace_header,
+            )
 
         if not self.trace_ctx.tracing_enable:
             self.trace_ctx = TraceNullContext()
@@ -339,8 +352,12 @@ class ReqTimeStatsBase:
         trace_ctx_state = state.get("trace_ctx")
         if isinstance(trace_ctx_state, dict):
             if trace_ctx_state.get("tracing_enable"):
-                trace_ctx = object.__new__(TraceReqContext)
-                trace_ctx.__setstate__(trace_ctx_state)
+                if trace_ctx_state.get("is_async"):
+                    trace_ctx = object.__new__(TraceReqContextAsync)
+                    trace_ctx.__setstate__(trace_ctx_state)
+                else:
+                    trace_ctx = object.__new__(TraceReqContext)
+                    trace_ctx.__setstate__(trace_ctx_state)
                 state["trace_ctx"] = trace_ctx
             else:
                 state["trace_ctx"] = TraceNullContext()
@@ -1247,3 +1264,19 @@ def set_time_batch(
             method(ts)
         else:
             method(ts, attrs)
+
+
+def flush_trace_batch(reqs: List[Any]):
+    """Proactively flush buffered trace ops for a batch of requests.
+
+    Call at natural CPU/GPU overlap points (e.g., right before run_batch)
+    so the ZMQ send overlaps with GPU forward compute.
+    """
+    if reqs is None or not get_global_tracing_enabled():
+        return
+    for req in reqs:
+        time_stats = getattr(req, "time_stats", None)
+        if time_stats is not None:
+            trace_ctx = getattr(time_stats, "trace_ctx", None)
+            if trace_ctx is not None:
+                trace_ctx.flush()

python/sglang/srt/observability/trace.py

@@ -24,6 +24,7 @@ import uuid
 from dataclasses import dataclass
 from typing import Any, Dict, List, Mapping, Optional
 
+from sglang.srt.environ import envs
 from sglang.srt.utils import get_int_env_var
 
 logger = logging.getLogger(__name__)
@@ -124,10 +125,29 @@ class TraceThreadContext:
 
 
 class TraceCustomIdGenerator(id_generator.IdGenerator):
+    """Custom ID generator with support for pre-setting the next span ID.
+
+    Why custom IDs are needed:
+      The default IdGenerator may produce duplicate trace IDs across
+      multiple TP scheduler processes.
+
+    Preset mechanism (used by async tracing):
+      When SGLANG_TRACE_ASYNC=1, span creation is deferred to an exporter
+      process while the caller process needs to know span IDs in advance
+      for cross-process span linking.  The caller pre-generates a span ID
+      and sends it to the exporter.  Before calling tracer.start_span(),
+      the exporter calls preset_next_span_id(id) — the next
+      generate_span_id() call consumes it, then falls back to random
+      generation.  This avoids modifying the standard OTel start_span()
+      API while giving the caller full control over span IDs.
+
+      Thread-safety: _preset_local is a threading.local(), so concurrent
+      callers in different threads cannot interfere.  The exporter process
+      is single-threaded, so no additional locking is needed.
     """
-    The default IdGenerator may produce duplicate trace IDs across multiple TP scheduler processes,
-    hence a custom IdGenerator is implemented.
-    """
+
+    # Thread-local storage for the next span ID to use preset.
+    _preset_local = threading.local()
 
     def __init__(self):
         super().__init__()
@@ -138,12 +158,30 @@ class TraceCustomIdGenerator(id_generator.IdGenerator):
         return self.local_random.getrandbits(64)
 
     def generate_span_id(self) -> int:
+        # If a preset span ID was injected, consume it (one-shot).
+        preset = getattr(self._preset_local, "span_id", None)
+        if preset is not None:
+            self._preset_local.span_id = None
+            return preset
         return self.local_random.getrandbits(64)
 
+    @classmethod
+    def preset_next_span_id(cls, span_id: int):
+        """Inject a pre-generated span ID for the next start_span() call.
+
+        The ID is consumed exactly once by generate_span_id() and then
+        cleared.  Call this immediately before tracer.start_span().
+        """
+        cls._preset_local.span_id = span_id
+
 
 # global variables
 threads_info: Dict[int, TraceThreadInfo] = {}
 
+# Optional callback invoked when a new thread registers its trace info.
+# Used by trace_async to forward thread info to the exporter process.
+_on_thread_info_set = None
+
 get_cur_time_ns = lambda: int(time.time() * 1e9)
 if hasattr(time, "time_ns"):
     get_cur_time_ns = lambda: int(time.time_ns())
@@ -199,12 +237,8 @@ def process_tracing_init(
             resource=resource, id_generator=TraceCustomIdGenerator()
         )
 
-        schedule_delay_millis = get_int_env_var(
-            "SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS", 500
-        )
-        max_export_batch_size = get_int_env_var(
-            "SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE", 64
-        )
+        schedule_delay_millis = envs.SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS.get()
+        max_export_batch_size = envs.SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE.get()
 
         processor = BatchSpanProcessor(
             span_exporter=get_otlp_span_exporter(otlp_endpoint),
@@ -222,6 +256,12 @@ def process_tracing_init(
     opentelemetry_initialized = True
     tracer = trace.get_tracer("sglang server")
 
+    # Auto-start async trace exporter when SGLANG_TRACE_ASYNC=1
+    if envs.SGLANG_TRACE_ASYNC.get():
+        from sglang.srt.observability.trace_async import start_trace_exporter
+
+        start_trace_exporter(otlp_endpoint, server_name, trace_modules=trace_modules)
+
 
 def get_global_tracing_enabled():
     return opentelemetry_initialized
@@ -266,6 +306,9 @@ def trace_set_thread_info(
         pp_rank=pp_rank,
     )
 
+    if _on_thread_info_set is not None:
+        _on_thread_info_set(threads_info[pid])
+
 
 class TraceReqContext:
     def __init__(
@@ -275,9 +318,12 @@ class TraceReqContext:
         role="unified",
         module_name="",
         external_trace_header: Optional[Dict[str, str]] = None,
+        trace_level: Optional[int] = None,
     ):
         self.rid: str = str(rid)
-        self.trace_level = get_global_trace_level()
+        self.trace_level = (
+            trace_level if trace_level is not None else get_global_trace_level()
+        )
         self.tracing_enable: bool = opentelemetry_initialized and self.trace_level > 0
 
         # Filter by --trace-modules only for explicitly named modules; contexts
@@ -453,7 +499,7 @@ class TraceReqContext:
         copied.trace_level = self.trace_level
         copied.module_name = self.module_name
         copied.is_copy = True  # Mark as copy
-        copied.pid = self.pid
+        copied.pid = None
 
         # thread_context is None, will be rebuilt via rebuild_thread_context()
         copied.thread_context = None
@@ -476,11 +522,14 @@ class TraceReqContext:
 
         return copied
 
-    def rebuild_thread_context(self, ts: Optional[int] = None):
+    def rebuild_thread_context(
+        self, ts: Optional[int] = None, pid: Optional[int] = None
+    ):
         if not self.tracing_enable:
             return
 
         ts = ts or get_cur_time_ns()
+        self.pid = pid if pid is not None else threading.get_native_id()
         self.thread_context = self.__create_thread_context(ts)
 
     def trace_req_start(
@@ -770,6 +819,9 @@ class TraceReqContext:
             self.thread_context.thread_span.end(end_time=ts)
         self.thread_context = None
 
+    def flush(self):
+        pass
+
     def __del__(self):
         self.abort(abort_info={"reason": "have unclosed span, auto closed"})