feat(observability): add OpenTelemetry tracing for pipeline parallelism (#23169)

Signed-off-by: Yinzuo Jiang <jiangyinzuo@foxmail.com>

python/sglang/srt/managers/scheduler.py

@@ -3795,7 +3795,7 @@ def run_scheduler_process(
             thread_label = "Prefill Scheduler"
         elif server_args.disaggregation_mode == "decode":
             thread_label = "Decode Scheduler"
-        trace_set_thread_info(thread_label, tp_rank, dp_rank)
+        trace_set_thread_info(thread_label, tp_rank, dp_rank, pp_rank)
 
     # Create a scheduler and run the event loop
     try:

python/sglang/srt/managers/scheduler_pp_mixin.py

@@ -29,6 +29,7 @@ from sglang.srt.managers.utils import (
     get_logprob_from_pp_outputs,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
+from sglang.srt.observability.req_time_stats import set_time_batch
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.utils import DynamicGradMode, broadcast_pyobj, point_to_point_pyobj
 from sglang.srt.utils.common import get_device_module, is_xpu
@@ -1162,7 +1163,18 @@ class SchedulerPPMixin:
         with torch.profiler.record_function("run_batch"):
             with self.forward_stream_ctx:
                 self.forward_stream.wait_stream(self.schedule_stream)
+                set_time_batch(
+                    self.cur_batch.reqs,
+                    "set_run_batch_cpu_start_time",
+                    trace_only=True,
+                )
                 result = self.run_batch(self.cur_batch, pp_proxy_tensors)
+                set_time_batch(
+                    self.cur_batch.reqs,
+                    "set_run_batch_cpu_end_time",
+                    trace_only=True,
+                    attrs={"pp_mb_id": mb_id},
+                )
                 mb_metadata[mb_id] = PPBatchMetadata(
                     can_run_cuda_graph=result.can_run_cuda_graph,
                 )

python/sglang/srt/observability/req_time_stats.py

@@ -205,6 +205,13 @@ class RequestStage:
         "spec_draft_extend",
         level=3,
     )
+
+    # CPU-side run batch
+    RUN_BATCH_CPU = RequestStageConfig(
+        "run_batch_cpu",
+        level=4,
+    )
+
     # other
     ANONYMOUS = RequestStageConfig("")
 
@@ -565,6 +572,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase):
     last_decode_scheduled_time: float = 0.0
     last_forward_entry_time: float = 0.0
     last_prefill_finished_time: float = 0.0
+    run_batch_cpu_start_time: float = 0.0
 
     # speculative decoding
     spec_draft_start_time: float = 0.0
@@ -633,6 +641,18 @@ class SchedulerReqTimeStats(ReqTimeStatsBase):
         stage = RequestStage.SPEC_DRAFT_EXTEND
         self.trace_slice(stage, self.spec_draft_extend_start_time, ts)
 
+    def set_run_batch_cpu_start_time(self, ts=None, attrs=None):
+        ts = ts or time.perf_counter()
+        self.run_batch_cpu_start_time = ts
+
+    def set_run_batch_cpu_end_time(self, ts=None, attrs=None):
+        ts = ts or time.perf_counter()
+        if self.run_batch_cpu_start_time > 0.0:
+            self.trace_slice(
+                RequestStage.RUN_BATCH_CPU, self.run_batch_cpu_start_time, ts, attrs
+            )
+            self.run_batch_cpu_start_time = 0.0
+
     def set_retract_time(self, ts=None):
         ts = ts or time.perf_counter()
         # retract
@@ -1105,7 +1125,12 @@ def set_schedule_time_batch(batch: ScheduleBatch):
         req.time_stats.set_last_scheduled_time(batch.forward_mode, ts, _attrs)
 
 
-def set_time_batch(reqs: List[Any], set_func: str, trace_only: bool = False):
+def set_time_batch(
+    reqs: List[Any],
+    set_func: str,
+    trace_only: bool = False,
+    attrs: Optional[Dict[str, Any]] = None,
+):
     if reqs is None or len(reqs) == 0:
         return
     if trace_only and not get_global_tracing_enabled():
@@ -1114,4 +1139,7 @@ def set_time_batch(reqs: List[Any], set_func: str, trace_only: bool = False):
     ts = time.perf_counter()
     for req in reqs:
         method = getattr(req.time_stats, set_func)
-        method(ts)
+        if attrs is None:
+            method(ts)
+        else:
+            method(ts, attrs)

python/sglang/srt/observability/trace.py

@@ -83,6 +83,7 @@ class TraceThreadInfo:
     thread_label: str
     tp_rank: int
     dp_rank: int
+    pp_rank: int
 
 
 @dataclass
@@ -223,7 +224,10 @@ def get_otlp_span_exporter(endpoint):
 
 # Should be called by each tracked thread.
 def trace_set_thread_info(
-    thread_label: str, tp_rank: Optional[int] = None, dp_rank: Optional[int] = None
+    thread_label: str,
+    tp_rank: Optional[int] = None,
+    dp_rank: Optional[int] = None,
+    pp_rank: Optional[int] = None,
 ):
     if not opentelemetry_initialized:
         return
@@ -238,6 +242,7 @@ def trace_set_thread_info(
         thread_label=thread_label,
         tp_rank=tp_rank,
         dp_rank=dp_rank,
+        pp_rank=pp_rank,
     )
 
 
@@ -292,6 +297,10 @@ class TraceReqContext:
         thread_name = f"{thread_info.thread_label}"
         if thread_info.tp_rank is not None:
             thread_name += f" [TP {thread_info.tp_rank}] "
+        if thread_info.pp_rank is not None:
+            thread_name += f" [PP {thread_info.pp_rank}] "
+        if thread_info.dp_rank is not None:
+            thread_name += f" [DP {thread_info.dp_rank}] "
         thread_name += f"(host:{thread_info.host_id[:8]} | pid:{self.pid})"
         thread_context.thread_span = tracer.start_span(
             name=thread_name,
@@ -299,8 +308,15 @@ class TraceReqContext:
             context=self.root_span_context,
         )
 
+        rank_attrs = {}
         if thread_info.tp_rank is not None:
-            thread_context.thread_span.set_attributes({"tp_rank": thread_info.tp_rank})
+            rank_attrs["tp_rank"] = thread_info.tp_rank
+        if thread_info.pp_rank is not None:
+            rank_attrs["pp_rank"] = thread_info.pp_rank
+        if thread_info.dp_rank is not None:
+            rank_attrs["dp_rank"] = thread_info.dp_rank
+        if rank_attrs:
+            thread_context.thread_span.set_attributes(rank_attrs)
 
         thread_context.thread_span.set_attributes(
             {

scripts/convert_otel_2_perfetto.py

@@ -237,6 +237,10 @@ def generate_perfetto_span(engine_root_spans, smg_otel_spans, thread_meta_data):
             pid = int(thread_span["attributes"]["pid"])
             host_id = thread_span["attributes"]["host_id"]
             thread_name = f'{thread_span["attributes"]["host_id"][:8]}:{thread_span["attributes"]["thread_label"]}'
+            if "pp_rank" in thread_span["attributes"]:
+                thread_name += f"-PP{thread_span['attributes']['pp_rank']}"
+            if "dp_rank" in thread_span["attributes"]:
+                thread_name += f"-DP{thread_span['attributes']['dp_rank']}"
             if "tp_rank" in thread_span["attributes"]:
                 thread_name += f"-TP{thread_span['attributes']['tp_rank']}"
 

test/registered/unit/observability/test_trace.py

@@ -67,7 +67,7 @@ class TestTraceFunctions(unittest.TestCase):
 
 class TestDataclasses(unittest.TestCase):
     def test_trace_thread_info(self):
-        info = TraceThreadInfo("host", 123, "label", 0, 1)
+        info = TraceThreadInfo("host", 123, "label", 0, 1, 0)
         self.assertEqual(info.thread_label, "label")
 
     def test_trace_event(self):
@@ -79,7 +79,7 @@ class TestDataclasses(unittest.TestCase):
         self.assertEqual(s.slice_name, "slice")
 
     def test_trace_thread_context(self):
-        info = TraceThreadInfo("h", 1, "l", 0, 0)
+        info = TraceThreadInfo("h", 1, "l", 0, 0, 0)
         ctx = TraceThreadContext(thread_info=info, cur_slice_stack=[])
         self.assertEqual(len(ctx.cur_slice_stack), 0)
 
@@ -514,7 +514,7 @@ class TestTraceReqContextEnabled(unittest.TestCase):
 
         pid = threading.get_native_id()
         mod.threads_info[pid] = TraceThreadInfo(
-            "host", pid, "sched", tp_rank=0, dp_rank=0
+            "host", pid, "sched", tp_rank=0, dp_rank=0, pp_rank=0
         )
         ctx = TraceReqContext(rid="req-1")
         ctx.trace_req_start(ts=1000)