From 71160e4ddb435f44ed90076a3629b26181f42e9a Mon Sep 17 00:00:00 2001 From: Yinzuo Jiang Date: Tue, 28 Apr 2026 17:05:23 +0800 Subject: [PATCH] feat(observability): add OpenTelemetry tracing for pipeline parallelism (#23169) Signed-off-by: Yinzuo Jiang --- python/sglang/srt/managers/scheduler.py | 2 +- .../sglang/srt/managers/scheduler_pp_mixin.py | 12 +++++++ .../srt/observability/req_time_stats.py | 32 +++++++++++++++++-- python/sglang/srt/observability/trace.py | 20 ++++++++++-- scripts/convert_otel_2_perfetto.py | 4 +++ .../unit/observability/test_trace.py | 6 ++-- 6 files changed, 68 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 50797f988..47e8f659b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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: diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 29a32db46..939c83f3c 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/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, ) diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 43f3356a8..f8b3f3233 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/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) diff --git a/python/sglang/srt/observability/trace.py b/python/sglang/srt/observability/trace.py index 5f7139c6e..be0c84191 100644 --- a/python/sglang/srt/observability/trace.py +++ b/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( { diff --git a/scripts/convert_otel_2_perfetto.py b/scripts/convert_otel_2_perfetto.py index 3a82969a4..89534a38f 100644 --- a/scripts/convert_otel_2_perfetto.py +++ b/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']}" diff --git a/test/registered/unit/observability/test_trace.py b/test/registered/unit/observability/test_trace.py index 719f31270..82d7a5005 100644 --- a/test/registered/unit/observability/test_trace.py +++ b/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)