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(
             {