SGLang Tracing: Add trace-level, trace-module, and unify tracing/request-stage-metrics

docs/references/production_request_trace.md

@@ -1,4 +1,4 @@
-SGlang exports request trace data based on the OpenTelemetry Collector. You can enable tracing by adding the `--enable-trace` and configure the OpenTelemetry Collector endpoint using `--otlp-traces-endpoint` when launching the server.
+SGlang exports request trace data based on the OpenTelemetry Collector. You can enable tracing by adding the `--trace-level` and configure the OpenTelemetry Collector endpoint using `--otlp-traces-endpoint` when launching the server. The `--trace-level` option accepts configurable values from `1` to `3`, with higher numbers indicating more detailed tracing. Additionally, you can use `--trace-module` to specify the module to trace; currently, only `request` is supported.
 
 You can find example screenshots of the visualization in https://github.com/sgl-project/sglang/issues/8965.
 
@@ -26,7 +26,7 @@ This section explains how to configure the request tracing and export the trace
     export SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS=500
     export SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE=64
     # start the prefill and decode server
-    python -m sglang.launch_server --enable-trace --otlp-traces-endpoint 0.0.0.0:4317 <other option>
+    python -m sglang.launch_server --trace-level 3 --otlp-traces-endpoint 0.0.0.0:4317 [--trace-module request] <other option>
     # start the mini lb
     python -m sglang_router.launch_router --enable-trace --otlp-traces-endpoint 0.0.0.0:4317 <other option>
     ```
@@ -61,62 +61,65 @@ We have already inserted instrumentation points in the tokenizer and scheduler m
     ```
     The "thread label" can be regarded as the name of the thread, used to distinguish different threads in the visualization view.
 
-2. Mark the beginning and end of a request
+2. create a time recorder for a request
+    Each request needs to call `SglangStageContext()` to initialize a time recorder, which is used to generate slice spans and request stage metrics. You can either store it within the request object or maintain it as a global variable. A set of APIs for managing the global time recorder is provided in `python/sglang/srt/tracing/trace_metric_warpper.py`.
+
+3. Mark the beginning and end of a request
     ```
-    trace_req_start(rid, bootstrap_room)
-    trace_req_finish(rid)
+    # The time recorder calls trace_req_start() by default when it is created.
+    stage_context.trace_req_finish()
     ```
-    These two APIs must be called within the same process, for example, in the tokenizer.
+    SglangStageContext() and trace_req_finish() must be called within the same process, for example, in the tokenizer.
 
-3. Add tracing for slice
+4. Add tracing for slice
 
     * Add slice tracing normally:
         ```python
-        trace_slice_start("slice A", rid)
-        trace_slice_end("slice A", rid)
+        stage_context.metric_trace_slice_start(RequestStage.TOKENIZER)
+        stage_context.metric_trace_slice_end(RequestStage.TOKENIZER)
         ```
 
-    - Use the "anonymous" flag to not specify a slice name at the start of the slice, allowing the slice name to be determined by trace_slice_end.
+    - Use the `ANONYMOUS` to not specify a slice name at the start of the slice, allowing the slice name to be determined by trace_slice_end.
     <br>Note: Anonymous slices must not be nested.
         ```python
-        trace_slice_start("", rid, anonymous = True)
-        trace_slice_end("slice A", rid)
+        stage_context.metric_trace_slice_start(RequestStage.ANONYMOUS)
+        stage_context.metric_trace_slice_end(RequestStage.TOKENIZER)
         ```
 
     - In trace_slice_end, use auto_next_anon to automatically create the next anonymous slice, which can reduce the number of instrumentation points needed.
         ```python
-        trace_slice_start("", rid, anonymous = True)
-        trace_slice_end("slice A", rid, auto_next_anon = True)
-        trace_slice_end("slice B", rid, auto_next_anon = True)
-        trace_slice_end("slice C", rid, auto_next_anon = True)
-        trace_slice_end("slice D", rid)
+        stage_context.metric_trace_slice_start(RequestStage.ANONYMOUS)
+        stage_context.metric_trace_slice_end(RequestStage.A, auto_next_anon = True)
+        stage_context.metric_trace_slice_end(RequestStage.B, auto_next_anon = True)
+        stage_context.metric_trace_slice_end(RequestStage.C, auto_next_anon = True)
+        stage_context.metric_trace_slice_end(RequestStage.D)
         ```
     - The end of the last slice in a thread must be marked with thread_finish_flag=True; otherwise, the thread's span will not be properly generated.
         ```python
-        trace_slice_end("slice D", rid, thread_finish_flag = True)
+        stage_context.metric_trace_slice_end(RequestStage.D, thread_finish_flag = True)
         ```
 
-4. When the request execution flow transfers to another thread, the trace context needs to be explicitly propagated.
+5. When the request execution flow transfers to another thread, the trace context needs to be explicitly propagated.
     - sender: Execute the following code before sending the request to another thread via ZMQ
         ```python
-        trace_context = trace_get_proc_propagate_context(rid)
-        req.trace_context = trace_context
+        trace_context = stage_context.trace_get_proc_propagate_context(rid)
+        req.stage_context = trace_context
         ```
     - receiver: Execute the following code after receiving the request via ZMQ
         ```python
-        trace_set_proc_propagate_context(rid, req.trace_context)
+        stage_context = SglangStageContext(......,propagation_context = req.stage_context)
         ```
 
-5. When the request execution flow transfers to another node(PD disaggregation), the trace context needs to be explicitly propagated.
+6. When the request execution flow transfers to another node(PD disaggregation), the trace context needs to be explicitly propagated.
     - sender: Execute the following code before sending the request to node thread via http
         ```python
-        trace_context = trace_get_remote_propagate_context(bootstrap_room_list)
+        trace_context = trace_get_remote_propagate_context_batch(bootstrap_room_list)
         headers = {"trace_context": trace_context}
         session.post(url, headers=headers)
         ```
     - receiver: Execute the following code after receiving the request via http
         ```python
-        trace_set_remote_propagate_context(request.headers['trace_context'])
+        trace_set_remote_propagate_context_batch(request.headers['trace_context'])
         ```
 
 ## How to Extend the Tracing Framework to Support Complex Tracing Scenarios

python/sglang/srt/disaggregation/decode.py

@@ -47,7 +47,7 @@
     prepare_abort,
 )
 from sglang.srt.layers.dp_attention import get_attention_tp_size
-from sglang.srt.managers.schedule_batch import FINISH_ABORT, RequestStage, ScheduleBatch
+from sglang.srt.managers.schedule_batch import FINISH_ABORT, ScheduleBatch
 from sglang.srt.managers.utils import GenerationBatchResult
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
@@ -60,7 +60,7 @@
     ReqToTokenPool,
     SWAKVPool,
 )
-from sglang.srt.tracing.trace import trace_event_batch, trace_slice_end
+from sglang.srt.tracing.trace_metric_warpper import RequestStage, trace_event_batch
 from sglang.srt.utils import get_int_env_var
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
 
@@ -316,8 +316,9 @@ def add(self, req: Req, is_retracted: bool = False) -> None:
                 prefill_dp_rank=req.data_parallel_rank,
             )
 
-            req.add_latency(RequestStage.DECODE_PREPARE)
-            trace_slice_end(RequestStage.DECODE_PREPARE, req.rid, auto_next_anon=True)
+            req.stage_context.metric_trace_slice_end(
+                RequestStage.DECODE_PREPARE, auto_next_anon=True
+            )
             self.queue.append(
                 DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False)
             )
@@ -531,9 +532,9 @@ def pop_preallocated(self) -> List[DecodeRequest]:
             decode_req.req.time_stats.decode_transfer_queue_entry_time = (
                 time.perf_counter()
             )
-            decode_req.req.add_latency(RequestStage.DECODE_BOOTSTRAP)
-            trace_slice_end(
-                RequestStage.DECODE_BOOTSTRAP, decode_req.req.rid, auto_next_anon=True
+
+            decode_req.req.stage_context.metric_trace_slice_end(
+                RequestStage.DECODE_BOOTSTRAP, auto_next_anon=True
             )
 
         self.queue = [
@@ -715,9 +716,8 @@ def _commit_transfer_to_req(self, decode_req: DecodeRequest) -> None:
 
         decode_req.kv_receiver.clear()
         decode_req.kv_receiver = None
-        trace_slice_end(
+        decode_req.req.stage_context.metric_trace_slice_end(
             RequestStage.DECODE_TRANSFERRED,
-            decode_req.req.rid,
             auto_next_anon=True,
         )
         decode_req.req.time_stats.wait_queue_entry_time = time.perf_counter()
@@ -757,6 +757,7 @@ def pop_transferred(self) -> List[Req]:
                 self._commit_transfer_to_req(decode_req)
                 indices_to_remove.add(i)
                 transferred_reqs.append(decode_req.req)
+
             elif poll in [
                 KVPoll.Bootstrapping,
                 KVPoll.WaitingForInput,
@@ -769,7 +770,6 @@ def pop_transferred(self) -> List[Req]:
         for i in indices_to_remove:
             idx = self.queue[i].metadata_buffer_index
             assert idx != -1
-            self.queue[i].req.add_latency(RequestStage.DECODE_TRANSFERRED)
             self.req_to_metadata_buffer_idx_allocator.free(idx)
 
         self.queue = [
@@ -905,7 +905,9 @@ def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]:
             # we can only add at least `num_not_used_batch` new batch to the running queue
             if i < num_not_used_batch:
                 can_run_list.append(req)
-                req.add_latency(RequestStage.DECODE_WAITING)
+                req.stage_context.metric_trace_slice_end(
+                    RequestStage.DECODE_WAITING, auto_next_anon=True
+                )
                 req.init_next_round_input(self.tree_cache)
             else:
                 waiting_queue.append(req)

python/sglang/srt/disaggregation/prefill.py

@@ -42,19 +42,14 @@
     poll_and_all_reduce,
     prepare_abort,
 )
-from sglang.srt.managers.schedule_batch import (
-    FINISH_LENGTH,
-    Req,
-    RequestStage,
-    ScheduleBatch,
-)
+from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch
 from sglang.srt.mem_cache.common import release_kv_cache
 from sglang.srt.mem_cache.memory_pool import (
     HybridLinearKVPool,
     NSATokenToKVPool,
     SWAKVPool,
 )
-from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end
+from sglang.srt.tracing.trace_metric_warpper import RequestStage, trace_event_batch
 from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj
 
 if TYPE_CHECKING:
@@ -198,9 +193,10 @@ def add(self, req: Req, num_kv_heads: int) -> None:
             pp_rank=self.pp_rank,
         )
         self._process_req(req)
-        req.add_latency(RequestStage.PREFILL_PREPARE)
         self.queue.append(req)
-        trace_slice_end(RequestStage.PREFILL_PREPARE, req.rid, auto_next_anon=True)
+        req.stage_context.metric_trace_slice_end(
+            RequestStage.PREFILL_PREPARE, auto_next_anon=True
+        )
 
     def extend(self, reqs: List[Req], num_kv_heads: int) -> None:
         for req in reqs:
@@ -290,10 +286,9 @@ def pop_bootstrapped(
             bootstrapped_reqs.append(req)
             indices_to_remove.add(i)
             req.time_stats.wait_queue_entry_time = time.perf_counter()
-            req.add_latency(RequestStage.PREFILL_BOOTSTRAP)
 
-            trace_slice_end(
-                RequestStage.PREFILL_BOOTSTRAP, req.rid, auto_next_anon=True
+            req.stage_context.metric_trace_slice_end(
+                RequestStage.PREFILL_BOOTSTRAP, auto_next_anon=True
             )
 
         self.queue = [
@@ -431,8 +426,6 @@ def process_batch_result_disagg_prefill(
                 # There is no output_ids for prefill
                 req.output_ids.append(next_token_id)
                 self.tree_cache.cache_unfinished_req(req)  # update the tree and lock
-                req.add_latency(RequestStage.PREFILL_FORWARD)
-                trace_slice(RequestStage.PREFILL_FORWARD, req.rid, auto_next_anon=True)
                 self.disagg_prefill_inflight_queue.append(req)
                 if self.spec_algorithm.is_eagle() and batch.spec_info is not None:
                     req.output_topk_p = batch.spec_info.topk_p[i]
@@ -459,6 +452,9 @@ def process_batch_result_disagg_prefill(
                     logprob_pt += num_input_logprobs
                 self.send_kv_chunk(req, last_chunk=True)
                 req.time_stats.prefill_transfer_queue_entry_time = time.perf_counter()
+                req.stage_context.metric_trace_slice(
+                    RequestStage.PREFILL_FORWARD, auto_next_anon=True
+                )
 
                 if req.grammar is not None:
                     # FIXME: this try-except block is for handling unexpected xgrammar issue.
@@ -497,8 +493,9 @@ def process_batch_result_disagg_prefill(
 
                 if self.enable_overlap:
                     self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx)
-                trace_slice(
-                    RequestStage.PREFILL_CHUNKED_FORWARD, req.rid, auto_next_anon=True
+                req.stage_context.metric_trace_slice(
+                    RequestStage.PREFILL_CHUNKED_FORWARD,
+                    auto_next_anon=(req.is_chunked != 0),
                 )
 
         self.maybe_send_health_check_signal()
@@ -568,11 +565,10 @@ def process_disagg_prefill_inflight_queue(
         )
         for req in done_reqs:
             req: Req
-            req.add_latency(RequestStage.PREFILL_TRANSFER_KV_CACHE)
             self.req_to_metadata_buffer_idx_allocator.free(req.metadata_buffer_index)
             req.metadata_buffer_index = -1
-            trace_slice(
-                RequestStage.PREFILL_TRANSFER_KV_CACHE, req.rid, thread_finish_flag=True
+            req.stage_context.metric_trace_slice(
+                RequestStage.PREFILL_TRANSFER_KV_CACHE, thread_finish_flag=True
             )
 
         self.disagg_prefill_inflight_queue = undone_reqs

python/sglang/srt/entrypoints/engine.py

@@ -141,7 +141,7 @@ def __init__(self, **kwargs):
         )
 
         # Enable tracing
-        if server_args.enable_trace:
+        if server_args.trace_level > 0:
             process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
             thread_label = "Tokenizer"
             if server_args.disaggregation_mode == "prefill":

python/sglang/srt/entrypoints/http_server.py

@@ -219,7 +219,7 @@ async def lifespan(fast_api_app: FastAPI):
         enable_func_timer()
 
     # Init tracing
-    if server_args.enable_trace:
+    if server_args.trace_level > 0:
         process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
         if server_args.disaggregation_mode == "prefill":
             thread_label = "Prefill" + thread_label

python/sglang/srt/managers/data_parallel_controller.py

@@ -34,21 +34,15 @@
     TokenizedGenerateReqInput,
     WatchLoadUpdateReq,
 )
-from sglang.srt.managers.schedule_batch import Req, RequestStage
+from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.managers.scheduler import run_scheduler_process
 from sglang.srt.server_args import (
     DP_ATTENTION_HANDSHAKE_PORT_DELTA,
     PortArgs,
     ServerArgs,
 )
-from sglang.srt.tracing.trace import (
-    process_tracing_init,
-    trace_get_proc_propagate_context,
-    trace_set_proc_propagate_context,
-    trace_set_thread_info,
-    trace_slice_end,
-    trace_slice_start,
-)
+from sglang.srt.tracing.trace import process_tracing_init, trace_set_thread_info
+from sglang.srt.tracing.trace_metric_warpper import RequestStage, SglangStageContext
 from sglang.srt.utils.common import (
     bind_port,
     configure_ipv6,
@@ -178,15 +172,26 @@ def handle_load_update_req(self, obj):
         self.dp_budget.update_budget(obj)
 
     def dispatching_with_trace(self, req: Req):
-        if self.server_args.enable_trace:
-            trace_set_proc_propagate_context(req.rid, req.trace_context)
-            trace_slice_start(RequestStage.DC_DISPATCH, req.rid)
-            req.trace_context = trace_get_proc_propagate_context(req.rid)
+        if self.server_args.trace_level > 0:
+            bootstrap_room = (
+                req.bootstrap_room if hasattr(req, "bootstrap_room") else None
+            )
+            stage_context = SglangStageContext(
+                req.rid,
+                bootstrap_room=bootstrap_room,
+                module_name="request",
+                server_args=self.server_args,
+                propagation_context=req.stage_context,
+            )
+            stage_context.metric_trace_slice_start(RequestStage.DC_DISPATCH)
+            req.stage_context = stage_context.trace_get_proc_propagate_context()
 
         self.dispatching(req)
 
-        if self.server_args.enable_trace:
-            trace_slice_end(RequestStage.DC_DISPATCH, req.rid, thread_finish_flag=True)
+        if self.server_args.trace_level > 0:
+            stage_context.metric_trace_slice_end(
+                RequestStage.DC_DISPATCH, thread_finish_flag=True
+            )
 
     def init_dispatcher(self):
         self._request_dispatcher = TypeBasedDispatcher(
@@ -516,7 +521,7 @@ def run_data_parallel_controller_process(
     parent_process = psutil.Process().parent()
 
     configure_logger(server_args)
-    if server_args.enable_trace:
+    if server_args.trace_level > 0:
         process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
         thread_label = "DP Controller"
         if server_args.disaggregation_mode == "prefill":