chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Feng Su <sufeng@linux.alibaba.com>
This commit is contained in:
Feng Su
2025-10-28 01:25:46 -07:00
committed by GitHub
parent b1e13e7cea
commit ea96106000
19 changed files with 824 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+25
View File
@@ -58,6 +58,11 @@ from sglang.srt.mem_cache.memory_pool import (
ReqToTokenPool,
SWAKVPool,
)
from sglang.srt.tracing.trace import (
trace_event_batch,
trace_slice_batch,
trace_slice_end,
)
from sglang.srt.utils import get_int_env_var, require_mlp_sync
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
@@ -313,6 +318,7 @@ class DecodePreallocQueue:
)
req.add_latency(RequestStage.DECODE_PREPARE)
trace_slice_end(RequestStage.DECODE_PREPARE, req.rid, auto_next_anon=True)
self.queue.append(
DecodeRequest(req=req, kv_receiver=kv_receiver, waiting_for_input=False)
)
@@ -527,6 +533,9 @@ class DecodePreallocQueue:
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
)
self.queue = [
entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
@@ -775,8 +784,19 @@ class DecodeTransferQueue:
[decode_req.req], decode_req.req.return_logprob
)
self.tree_cache.cache_finished_req(decode_req.req)
trace_slice_end(
RequestStage.DECODE_QUICK_FINISH,
decode_req.req.rid,
thread_finish_flag=True,
)
else:
transferred_reqs.append(decode_req.req)
trace_slice_end(
RequestStage.DECODE_TRANSFERRED,
decode_req.req.rid,
auto_next_anon=True,
)
elif poll in [
KVPoll.Bootstrapping,
KVPoll.WaitingForInput,
@@ -822,6 +842,7 @@ class SchedulerDisaggregationDecodeMixin:
self.stream_output(
batch.reqs, any(req.return_logprob for req in batch.reqs)
)
trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
if prepare_mlp_sync_flag:
self._prepare_idle_batch_and_run(None)
else:
@@ -871,6 +892,7 @@ class SchedulerDisaggregationDecodeMixin:
self.stream_output(
batch.reqs, any(req.return_logprob for req in batch.reqs)
)
trace_slice_batch(RequestStage.DECODE_FAKE_OUTPUT, batch.reqs)
if prepare_mlp_sync_flag:
batch_, batch_result = self._prepare_idle_batch_and_run(
None, delay_process=True
@@ -953,6 +975,9 @@ class SchedulerDisaggregationDecodeMixin:
self.running_batch = self.update_running_batch(self.running_batch)
ret = self.running_batch if not self.running_batch.is_empty() else None
if ret:
attrs = {"bid": hex(id(ret)), "batch_size": ret.batch_size()}
trace_event_batch("schedule", ret.reqs, attrs=attrs)
return ret
def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]:
python/sglang/srt/disaggregation/prefill.py
+19
View File
@@ -53,6 +53,7 @@ from sglang.srt.mem_cache.memory_pool import (
NSATokenToKVPool,
SWAKVPool,
)
from sglang.srt.tracing.trace import trace_event_batch, trace_slice, trace_slice_end
from sglang.srt.utils import broadcast_pyobj, point_to_point_pyobj, require_mlp_sync
if TYPE_CHECKING:
@@ -198,6 +199,7 @@ class PrefillBootstrapQueue:
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)
def extend(self, reqs: List[Req], num_kv_heads: int) -> None:
for req in reqs:
@@ -289,6 +291,10 @@ class PrefillBootstrapQueue:
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
)
self.queue = [
entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
]
@@ -316,6 +322,9 @@ class SchedulerDisaggregationPrefillMixin:
)
self.process_prefill_chunk()
batch = self.get_new_batch_prefill()
if batch:
attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
trace_event_batch("schedule", batch.reqs, attrs=attrs)
if require_mlp_sync(self.server_args):
batch = self.prepare_mlp_sync_batch(batch)
@@ -348,6 +357,9 @@ class SchedulerDisaggregationPrefillMixin:
)
self.process_prefill_chunk()
batch = self.get_new_batch_prefill()
if batch:
attrs = {"bid": hex(id(batch)), "batch_size": batch.batch_size()}
trace_event_batch("schedule", batch.reqs, attrs=attrs)
if require_mlp_sync(self.server_args):
batch = self.prepare_mlp_sync_batch(batch)
@@ -423,6 +435,7 @@ class SchedulerDisaggregationPrefillMixin:
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]
@@ -487,6 +500,9 @@ class SchedulerDisaggregationPrefillMixin:
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
)
self.maybe_send_health_check_signal()
@@ -558,6 +574,9 @@ class SchedulerDisaggregationPrefillMixin:
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
)
self.disagg_prefill_inflight_queue = undone_reqs
python/sglang/srt/entrypoints/engine.py
+7 -4
View File
@@ -143,10 +143,13 @@ class Engine(EngineBase):
# Enable tracing
if server_args.enable_trace:
process_tracing_init(server_args.oltp_traces_endpoint, "sglang")
if server_args.disaggregation_mode == "null":
thread_label = "Tokenizer"
trace_set_thread_info(thread_label)
process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
thread_label = "Tokenizer"
if server_args.disaggregation_mode == "prefill":
thread_label = "Prefill Tokenizer"
elif server_args.disaggregation_mode == "decode":
thread_label = "Decode Tokenizer"
trace_set_thread_info(thread_label)
try:
self.loop = asyncio.get_running_loop()
python/sglang/srt/entrypoints/http_server.py
+6 -3
View File
@@ -220,9 +220,12 @@ async def lifespan(fast_api_app: FastAPI):
# Init tracing
if server_args.enable_trace:
process_tracing_init(server_args.oltp_traces_endpoint, "sglang")
if server_args.disaggregation_mode == "null":
trace_set_thread_info(thread_label)
process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
if server_args.disaggregation_mode == "prefill":
thread_label = "Prefill" + thread_label
elif server_args.disaggregation_mode == "decode":
thread_label = "Decode" + thread_label
trace_set_thread_info(thread_label)
# Initialize OpenAI serving handlers
fast_api_app.state.openai_serving_completion = OpenAIServingCompletion(
python/sglang/srt/environ.py
+2
View File
@@ -129,6 +129,8 @@ class Envs:
SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)
SGLANG_SIMULATE_ACC_METHOD = EnvStr("multinomial")
SGLANG_TORCH_PROFILER_DIR = EnvStr("/tmp")
SGLANG_OTLP_EXPORTER_SCHEDULE_DELAY_MILLIS = EnvInt(500)
SGLANG_OTLP_EXPORTER_MAX_EXPORT_BATCH_SIZE = EnvInt(64)
# Scheduler: memory leak test
SGLANG_TEST_RETRACT = EnvBool(False)
python/sglang/srt/managers/data_parallel_controller.py
+30 -3
View File
@@ -34,13 +34,21 @@ from sglang.srt.managers.io_struct import (
TokenizedGenerateReqInput,
WatchLoadUpdateReq,
)
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_batch import Req, RequestStage
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.utils import (
bind_port,
configure_logger,
@@ -170,11 +178,22 @@ class DataParallelController:
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)
self.dispatching(req)
if self.server_args.enable_trace:
trace_slice_end(RequestStage.DC_DISPATCH, req.rid, thread_finish_flag=True)
def init_dispatcher(self):
self._request_dispatcher = TypeBasedDispatcher(
[
(TokenizedGenerateReqInput, self.dispatching),
(TokenizedEmbeddingReqInput, self.dispatching),
(TokenizedGenerateReqInput, self.dispatching_with_trace),
(TokenizedEmbeddingReqInput, self.dispatching_with_trace),
(BlockReqInput, self.send_to_all_workers),
(WatchLoadUpdateReq, self.handle_load_update_req),
]
@@ -487,6 +506,14 @@ def run_data_parallel_controller_process(
pipe_writer,
):
kill_itself_when_parent_died()
if server_args.enable_trace:
process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
thread_label = "DP Controller"
if server_args.disaggregation_mode == "prefill":
thread_label = "Prefill DP Controller"
elif server_args.disaggregation_mode == "decode":
thread_label = "Decode DP Controller"
trace_set_thread_info(thread_label)
setproctitle.setproctitle("sglang::data_parallel_controller")
faulthandler.enable()
configure_logger(server_args)
python/sglang/srt/managers/schedule_batch.py
+14 -2
View File
@@ -396,13 +396,23 @@ class MultimodalInputs:
class RequestStage(str, enum.Enum):
# prefill
# Tokenizer
TOKENIZE = "tokenize"
TOKENIZER_DISPATCH = "dispatch"
# DP controller
DC_DISPATCH = "dc_dispatch"
# common/non-disaggregation
PREFILL_WAITING = "prefill_waiting"
REQUEST_PROCESS = "request_process"
DECODE_LOOP = "decode_loop"
PREFILL_FORWARD = "prefill_forward"
PREFILL_CHUNKED_FORWARD = "chunked_prefill"
# disaggregation prefill
PREFILL_PREPARE = "prefill_prepare"
PREFILL_BOOTSTRAP = "prefill_bootstrap"
PREFILL_FORWARD = "prefill_forward"
PREFILL_TRANSFER_KV_CACHE = "prefill_transfer_kv_cache"
# disaggregation decode
@@ -410,6 +420,8 @@ class RequestStage(str, enum.Enum):
DECODE_BOOTSTRAP = "decode_bootstrap"
DECODE_WAITING = "decode_waiting"
DECODE_TRANSFERRED = "decode_transferred"
DECODE_FAKE_OUTPUT = "fake_output"
DECODE_QUICK_FINISH = "quick_finish"
class Req:
python/sglang/srt/managers/scheduler.py
+14 -9
View File
@@ -157,6 +157,7 @@ from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.tracing.trace import (
process_tracing_init,
trace_event_batch,
trace_set_proc_propagate_context,
trace_set_thread_info,
trace_slice_batch,
@@ -1354,7 +1355,7 @@ class Scheduler(
self._prefetch_kvcache(req)
self.waiting_queue.append(req)
req.time_stats.wait_queue_entry_time = time.perf_counter()
trace_slice_end("process req", req.rid, auto_next_anon=True)
trace_slice_end(RequestStage.REQUEST_PROCESS, req.rid, auto_next_anon=True)
elif self.disaggregation_mode == DisaggregationMode.PREFILL:
self._prefetch_kvcache(req)
self.disagg_prefill_bootstrap_queue.add(
@@ -1618,6 +1619,10 @@ class Scheduler(
if need_dp_attn_preparation:
ret = self.prepare_mlp_sync_batch(ret)
if ret:
attrs = {"bid": hex(id(ret)), "batch_size": ret.batch_size()}
trace_event_batch("schedule", ret.reqs, attrs=attrs)
return ret
def get_num_allocatable_reqs(self, running_bs):
@@ -1993,13 +1998,10 @@ class Scheduler(
):
if batch.forward_mode.is_decode():
self.process_batch_result_decode(batch, result)
if self.enable_trace:
trace_slice_batch("decode loop", batch.reqs)
trace_slice_batch(RequestStage.DECODE_LOOP, batch.reqs)
elif batch.forward_mode.is_extend():
self.process_batch_result_prefill(batch, result)
if self.enable_trace:
trace_slice_batch("prefill", batch.reqs)
elif batch.forward_mode.is_idle():
if self.enable_overlap:
@@ -2741,10 +2743,13 @@ def run_scheduler_process(
# Set up tracing
if server_args.enable_trace:
process_tracing_init(server_args.oltp_traces_endpoint, "sglang")
if server_args.disaggregation_mode == "null":
thread_label = "Scheduler"
trace_set_thread_info(thread_label, tp_rank, dp_rank)
process_tracing_init(server_args.otlp_traces_endpoint, "sglang")
thread_label = "Scheduler"
if server_args.disaggregation_mode == "prefill":
thread_label = "Prefill Scheduler"
elif server_args.disaggregation_mode == "decode":
thread_label = "Decode Scheduler"
trace_set_thread_info(thread_label, tp_rank, dp_rank)
# Create a scheduler and run the event loop
try:
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+28 -1
View File
@@ -14,7 +14,13 @@ from sglang.srt.managers.io_struct import (
BatchEmbeddingOutput,
BatchTokenIDOutput,
)
from sglang.srt.managers.schedule_batch import BaseFinishReason, Req, ScheduleBatch
from sglang.srt.managers.schedule_batch import (
BaseFinishReason,
Req,
RequestStage,
ScheduleBatch,
)
from sglang.srt.tracing.trace import trace_slice
from sglang.srt.utils.common import ceil_div
if TYPE_CHECKING:
@@ -160,6 +166,14 @@ class SchedulerOutputProcessorMixin:
)
self.abort_request(AbortReq(rid=req.rid))
req.grammar.finished = req.finished()
trace_slice(
RequestStage.PREFILL_FORWARD,
req.rid,
auto_next_anon=not req.finished(),
thread_finish_flag=req.finished(),
)
else:
# being chunked reqs' prefill is not finished
req.is_chunked -= 1
@@ -188,6 +202,12 @@ class SchedulerOutputProcessorMixin:
)
logprob_pt += num_input_logprobs
trace_slice(
RequestStage.PREFILL_CHUNKED_FORWARD,
req.rid,
auto_next_anon=True,
)
else: # embedding or reward model
is_sparse = envs.SGLANG_EMBEDDINGS_SPARSE_HEAD.is_set()
@@ -227,6 +247,13 @@ class SchedulerOutputProcessorMixin:
# being chunked reqs' prefill is not finished
req.is_chunked -= 1
trace_slice(
RequestStage.PREFILL_FORWARD,
req.rid,
auto_next_anon=not req.finished(),
thread_finish_flag=req.finished(),
)
self.stream_output(batch.reqs, batch.return_logprob, skip_stream_req)
def _resolve_spec_overlap_token_ids(
python/sglang/srt/managers/tokenizer_manager.py
+24 -6
View File
@@ -68,6 +68,7 @@ from sglang.srt.managers.io_struct import (
)
from sglang.srt.managers.mm_utils import TensorTransportMode
from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
from sglang.srt.managers.schedule_batch import RequestStage
from sglang.srt.managers.scheduler import is_health_check_generate_req
from sglang.srt.managers.scheduler_input_blocker import input_blocker_guard_region
from sglang.srt.managers.tokenizer_communicator_mixin import TokenizerCommunicatorMixin
@@ -79,6 +80,7 @@ from sglang.srt.tracing.trace import (
trace_get_proc_propagate_context,
trace_req_finish,
trace_req_start,
trace_set_remote_propagate_context,
trace_slice_end,
trace_slice_start,
)
@@ -383,6 +385,10 @@ class TokenizerManager(TokenizerCommunicatorMixin):
self.auto_create_handle_loop()
obj.normalize_batch_and_arguments()
if request:
if "trace_context" in request.headers:
trace_set_remote_propagate_context(request.headers["trace_context"])
if self.server_args.tokenizer_worker_num > 1:
self._attach_multi_http_worker_info(obj)
@@ -605,7 +611,7 @@ class TokenizerManager(TokenizerCommunicatorMixin):
mm_inputs = None
self._validate_one_request(obj, input_ids)
trace_slice_end("tokenize", obj.rid)
trace_slice_end(RequestStage.TOKENIZE, obj.rid)
return self._create_tokenized_object(
obj, input_text, input_ids, input_embeds, mm_inputs, token_type_ids
)
@@ -831,7 +837,7 @@ class TokenizerManager(TokenizerCommunicatorMixin):
req, req.text, input_ids_list[i], None, None, token_type_ids
)
)
trace_slice_end("tokenize", req.rid)
trace_slice_end(RequestStage.TOKENIZE, req.rid)
logger.debug(f"Completed batch processing for {batch_size} requests")
return tokenized_objs
@@ -883,12 +889,14 @@ class TokenizerManager(TokenizerCommunicatorMixin):
tokenized_obj: Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput],
created_time: Optional[float] = None,
):
trace_slice_start("dispatch", obj.rid)
trace_slice_start(RequestStage.TOKENIZER_DISPATCH, obj.rid)
tokenized_obj.trace_context = trace_get_proc_propagate_context(obj.rid)
self.send_to_scheduler.send_pyobj(tokenized_obj)
state = ReqState([], False, asyncio.Event(), obj, created_time=created_time)
self.rid_to_state[obj.rid] = state
trace_slice_end("dispatch", obj.rid, thread_finish_flag=True)
trace_slice_end(
RequestStage.TOKENIZER_DISPATCH, obj.rid, thread_finish_flag=True
)
return state
def _send_batch_request(
@@ -2131,7 +2139,12 @@ class TokenizerManager(TokenizerCommunicatorMixin):
bootstrap_room = (
obj.bootstrap_room if hasattr(obj, "bootstrap_room") else None
)
trace_req_start(obj.rid, bootstrap_room, ts=int(created_time * 1e9))
trace_req_start(
obj.rid,
bootstrap_room,
ts=int(created_time * 1e9),
role=self.server_args.disaggregation_mode,
)
trace_slice_start("", obj.rid, ts=int(created_time * 1e9), anonymous=True)
else:
for i in range(len(obj.rid)):
@@ -2140,7 +2153,12 @@ class TokenizerManager(TokenizerCommunicatorMixin):
if hasattr(obj, "bootstrap_room") and obj.bootstrap_room
else None
)
trace_req_start(obj.rid[i], bootstrap_room, ts=int(created_time * 1e9))
trace_req_start(
obj.rid[i],
bootstrap_room,
ts=int(created_time * 1e9),
role=self.server_args.disaggregation_mode,
)
trace_slice_start(
"", obj.rid[i], ts=int(created_time * 1e9), anonymous=True
)
python/sglang/srt/server_args.py
+2 -2
View File
@@ -299,7 +299,7 @@ class ServerArgs:
enable_request_time_stats_logging: bool = False
kv_events_config: Optional[str] = None
enable_trace: bool = False
oltp_traces_endpoint: str = "localhost:4317"
otlp_traces_endpoint: str = "localhost:4317"
# API related
api_key: Optional[str] = None
@@ -2340,7 +2340,7 @@ class ServerArgs:
help="Enable opentelemetry trace",
)
parser.add_argument(
"--oltp-traces-endpoint",
"--otlp-traces-endpoint",
type=str,
default="localhost:4317",
help="Config opentelemetry collector endpoint if --enable-trace is set. format: <ip>:<port>",
python/sglang/srt/tracing/trace.py
+121 -12
View File
@@ -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)