[Feature] Sglang Tracing: Fine-Grained Tracking for Request Latency - Part 2 (#10804)

Signed-off-by: Feng Su <sufeng@linux.alibaba.com>

python/sglang/srt/tracing/trace.py

@@ -15,6 +15,8 @@
 
 from __future__ import annotations
 
+import base64
+import json
 import logging
 import os
 import random
@@ -24,6 +26,8 @@ import uuid
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
+from sglang.srt.utils import get_int_env_var
+
 if TYPE_CHECKING:
     from sglang.srt.managers.scheduler import Req
 
@@ -85,6 +89,8 @@ class SglangTraceReqContext:
     # Indicates whether this instance is a replica from the main process.
     # When True, root_span is None and only root_span_context is preserved.
     is_copy: bool = False
+    bootstrap_room_span: Optional[trace.span.Span] = None
+    bootstrap_room_span_context: Optional[context.Context] = None
     root_span: Optional[trace.span.Span] = None
     root_span_context: Optional[context.Context] = None
 
@@ -96,8 +102,7 @@ class SglangTracePropagateContext:
 
     def to_dict(self):
         carrier: dict[str, str] = {}
-        context.attach(self.root_span_context)
-        propagate.inject(carrier)
+        propagate.inject(carrier, self.root_span_context)
 
         if self.prev_span_context:
             return {
@@ -149,6 +154,7 @@ class SglangTraceCustomIdGenerator(id_generator.IdGenerator):
 
 
 # global variables
+remote_trace_contexts: Dict[str, SglangTracePropagateContext] = {}
 threads_info: Dict[int, SglangTraceThreadInfo] = {}
 reqs_context: Dict[str, SglangTraceReqContext] = {}
 
@@ -193,8 +199,17 @@ def process_tracing_init(otlp_endpoint, server_name):
             resource=resource, id_generator=SglangTraceCustomIdGenerator()
         )
 
+        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
+        )
+
         processor = BatchSpanProcessor(
-            OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True)
+            OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True),
+            schedule_delay_millis=schedule_delay_millis,
+            max_export_batch_size=max_export_batch_size,
         )
         tracer_provider.add_span_processor(processor)
         trace.set_tracer_provider(tracer_provider)
@@ -266,7 +281,9 @@ def __create_thread_context(pid, req_span_context, ts: Optional[int] = None):
     return thread_context
 
 
-def trace_get_proc_propagate_context(rid) -> Optional[Dict[str, Any]]:
+def trace_get_proc_propagate_context(
+    rid, remote_propagate=False
+) -> Optional[Dict[str, Any]]:
     if not tracing_enabled:
         return None
 
@@ -283,9 +300,11 @@ def trace_get_proc_propagate_context(rid) -> Optional[Dict[str, Any]]:
     elif thread_context.last_span_context:
         prev_span_context = thread_context.last_span_context
 
-    trace_context = SglangTracePropagateContext(
-        reqs_context[rid].root_span_context, prev_span_context
-    )
+    root_span_context = reqs_context[rid].root_span_context
+    if remote_propagate:
+        root_span_context = reqs_context[rid].bootstrap_room_span_context
+
+    trace_context = SglangTracePropagateContext(root_span_context, prev_span_context)
     return trace_context.to_dict()
 
 
@@ -327,10 +346,54 @@ def trace_set_proc_propagate_context(rid, trace_context: Optional[Dict[str, Any]
     ].last_span_context = trace_context.prev_span_context
 
 
+def trace_get_remote_propagate_context(bootstrap_room_list: List[str]):
+    if not tracing_enabled:
+        return ""
+
+    reqs_trace_contexts = {}
+    for bootstrap_room in bootstrap_room_list:
+        # In the router, rid is also the bootstrap room.
+        bootstrap_room = str(bootstrap_room)
+
+        if bootstrap_room not in reqs_context:
+            continue
+
+        _context = trace_get_proc_propagate_context(
+            bootstrap_room, remote_propagate=True
+        )
+        reqs_trace_contexts[bootstrap_room] = _context
+
+    json_str = json.dumps(reqs_trace_contexts, ensure_ascii=False)
+    return base64.b64encode(json_str.encode("utf-8")).decode("utf-8")
+
+
+def trace_set_remote_propagate_context(base64_str):
+    if not tracing_enabled:
+        return
+
+    if base64_str is None or base64_str == "" or base64_str == "None":
+        return
+
+    base64_bytes = base64.b64decode(base64_str)
+    json_str = base64_bytes.decode("utf-8")
+    remote_reqs_trace_contexts = json.loads(json_str)
+
+    for bootstrap_room in remote_reqs_trace_contexts:
+        if bootstrap_room in remote_trace_contexts:
+            continue
+
+        remote_trace_contexts[bootstrap_room] = (
+            SglangTracePropagateContext.instance_from_dict(
+                remote_reqs_trace_contexts[bootstrap_room]
+            )
+        )
+
+
 def trace_req_start(
     rid: str,
     bootstrap_room: Optional[int] = None,
     ts: Optional[int] = None,
+    role: Optional[str] = "null",
 ):
     if not tracing_enabled:
         return
@@ -344,6 +407,7 @@ def trace_req_start(
         return
 
     # create req context and root span
+    bootstrap_room = 0 if bootstrap_room is None else bootstrap_room
     reqs_context[rid] = SglangTraceReqContext(
         rid=rid,
         start_time_ns=ts,
@@ -352,23 +416,42 @@ def trace_req_start(
         is_copy=False,
     )
 
+    # create bootstrap room span
+    tracer = threads_info[pid].tracer
+    if str(bootstrap_room) not in remote_trace_contexts:
+        attrs = {"bootstrap_room": str(hex(bootstrap_room))}
+        bootstrap_room_span = tracer.start_span(
+            name=f"Bootstrap Room {hex(bootstrap_room)}",
+            start_time=ts,
+            attributes=attrs,
+        )
+        reqs_context[rid].bootstrap_room_span = bootstrap_room_span
+        bootstrap_room_span_context = trace.set_span_in_context(bootstrap_room_span)
+    else:
+        bootstrap_room_span_context = remote_trace_contexts[
+            str(bootstrap_room)
+        ].root_span_context
+
     # Drop the worker_id added by MultiTokenizer
     orig_rid = rid.split("_")[-1]
-    tracer = threads_info[pid].tracer
+    role = "" if role == "null" else role
+    attrs = {"rid": orig_rid}
     root_span = tracer.start_span(
-        name=f"Req {orig_rid[:8]}",
+        name=f"{role} Req {orig_rid[:8]}",
         start_time=ts,
+        context=bootstrap_room_span_context,
+        attributes=attrs,
     )
 
     root_span.set_attributes(
         {
             "rid": rid,
-            "bootstrap_room": bootstrap_room if bootstrap_room else "None",
         }
     )
 
     reqs_context[rid].root_span = root_span
     reqs_context[rid].root_span_context = trace.set_span_in_context(root_span)
+    reqs_context[rid].bootstrap_room_span_context = bootstrap_room_span_context
 
     # create thread context and thread span
     reqs_context[rid].threads_context[pid] = __create_thread_context(
@@ -376,6 +459,10 @@ def trace_req_start(
         reqs_context[rid].root_span_context,
         ts,
     )
+    if str(bootstrap_room) in remote_trace_contexts:
+        reqs_context[rid].threads_context[pid].last_span_context = (
+            remote_trace_contexts[str(bootstrap_room)].prev_span_context
+        )
 
 
 def trace_req_finish(
@@ -399,6 +486,10 @@ def trace_req_finish(
         req_context.root_span.set_attributes(attrs)
 
     req_context.root_span.end(end_time=ts)
+    if str(req_context.bootstrap_room) in remote_trace_contexts:
+        del remote_trace_contexts[str(req_context.bootstrap_room)]
+    else:
+        req_context.bootstrap_room_span.end(end_time=ts)
 
     del reqs_context[rid]
 
@@ -518,7 +609,9 @@ trace_slice = trace_slice_end
 
 
 # Add event to the current slice on the same thread with the same rid.
-def trace_event(name: str, rid: str, ts: Optional[int] = None):
+def trace_event(
+    name: str, rid: str, ts: Optional[int] = None, attrs: Dict[str, Any] = None
+):
     if not tracing_enabled:
         return
 
@@ -539,7 +632,7 @@ def trace_event(name: str, rid: str, ts: Optional[int] = None):
     ts = ts or __get_cur_time_ns()
 
     slice_info = thread_context.cur_slice_stack[-1]
-    slice_info.span.add_event(name=name, timestamp=ts)
+    slice_info.span.add_event(name=name, timestamp=ts, attributes=attrs)
 
 
 # Add attrs to the current slice on the same thread with the same rid.
@@ -569,6 +662,9 @@ def trace_slice_batch(
     name: str,
     reqs: List[Req],
 ):
+    if not tracing_enabled:
+        return
+
     for req in reqs:
         trace_slice(
             name,
@@ -576,3 +672,16 @@ def trace_slice_batch(
             auto_next_anon=not req.finished(),
             thread_finish_flag=req.finished(),
         )
+
+
+def trace_event_batch(
+    name: str,
+    reqs: List[Req],
+    ts: Optional[int] = None,
+    attrs: Dict[str, Any] = None,
+):
+    if not tracing_enabled:
+        return
+
+    for req in reqs:
+        trace_event(name, req.rid, ts=ts, attrs=attrs)