From 694cea8656373bbef2d1073cab92a6a6857c3d63 Mon Sep 17 00:00:00 2001 From: Joectwm <158221336+Joectwm@users.noreply.github.com> Date: Fri, 12 Jun 2026 03:16:14 -0500 Subject: [PATCH] Add EPD disaggregated encode tracing (#25994) --- .../srt/disaggregation/encode_receiver.py | 11 ++- .../srt/disaggregation/encode_server.py | 43 +++++++++- .../sglang/srt/managers/tokenizer_manager.py | 10 ++- .../srt/observability/req_time_stats.py | 80 ++++++++++++++++++- 4 files changed, 135 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index a4aa338c8..b409638f6 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -1684,10 +1684,10 @@ class MMReceiverBase(ABC): finally: recv_socket.close() - def send_encode_request(self, obj): - self._send_encode_request(obj) + def send_encode_request(self, obj, time_stats_json=None): + self._send_encode_request(obj, time_stats_json=time_stats_json) - def _send_encode_request(self, obj): + def _send_encode_request(self, obj, time_stats_json=None): mm_data = self._extract_url_data(obj) if obj.rid is None: obj.rid = uuid.uuid4().hex @@ -1730,6 +1730,7 @@ class MMReceiverBase(ABC): num_items_assigned, None, encode_urls, + time_stats_json, ), daemon=True, ) @@ -1839,6 +1840,7 @@ class MMReceiverBase(ABC): num_items_assigned, embedding_port, encode_urls=None, + time_stats_json=None, ): try: asyncio.run( @@ -1850,6 +1852,7 @@ class MMReceiverBase(ABC): endpoint_send=None, num_items_assigned=num_items_assigned, encode_urls=encode_urls, + time_stats_json=time_stats_json, ) ) except Exception as e: @@ -2069,6 +2072,7 @@ class MMReceiverHTTP(MMReceiverBase): endpoint_send, num_items_assigned=None, encode_urls=None, + time_stats_json=None, ): if len(mm_data) == 0: return @@ -2121,6 +2125,7 @@ class MMReceiverHTTP(MMReceiverBase): "modality": modality.name, # convert enum to string for json serialization "prefill_host": self.host, "embedding_port": embedding_port, + "time_stats_json": time_stats_json, } ) cum_idx += 1 diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index ef9e6017a..6ea60c397 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -48,6 +48,11 @@ from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.model_loader import get_model from sglang.srt.multimodal.processors.qwen_vl import preprocess_video +from sglang.srt.observability.req_time_stats import EncoderReqTimeStats +from sglang.srt.observability.trace import ( + process_tracing_init, + trace_set_thread_info, +) from sglang.srt.server_args import ( PortArgs, ServerArgs, @@ -2447,6 +2452,10 @@ async def _dp_worker_encode_and_send( # Mooncake returns metadata for main to forward; zmq inlines the send. # Soft errors raise MMError so the dispatcher route maps them to HTTP. req_id = request["req_id"] + time_stats_json = request.pop("time_stats_json", None) + time_stats = EncoderReqTimeStats() + if time_stats_json: + time_stats.decode_json(time_stats_json) request["enter_time"] = time.time() modality = Modality.from_str(request["modality"]) backend = enc.server_args.encoder_transfer_backend @@ -2461,14 +2470,20 @@ async def _dp_worker_encode_and_send( code=HTTPStatus.BAD_REQUEST, ) + time_stats.set_mm_encode_start_time() encode_coro = ( sched.submit(request) if sched is not None and modality in _BATCHABLE_MODALITIES else enc.encode_request(request, modality) ) - nbytes, embedding_len, embedding_dim, error_msg, error_code = await encode_coro + try: + nbytes, embedding_len, embedding_dim, error_msg, error_code = await encode_coro + except asyncio.TimeoutError: + time_stats.trace_ctx.abort(abort_info={"reason": "encoder batch timed out"}) + raise if error_msg: + time_stats.trace_ctx.abort(abort_info={"reason": error_msg}) # zmq backends still forward an error EmbeddingData to P so it # doesn't block; send failures here are swallowed. try: @@ -2490,6 +2505,8 @@ async def _dp_worker_encode_and_send( enc.embedding_to_send.pop(req_id, None) raise MMError(error_msg, code=error_code or HTTPStatus.INTERNAL_SERVER_ERROR) + time_stats.set_mm_encode_end_time() + if backend == "mooncake": request.pop("mm_items", None) request.update( @@ -3275,6 +3292,13 @@ def launch_server(server_args: ServerArgs): dist_init_method = NetworkAddress( server_args.host or "127.0.0.1", port_args.nccl_port ).to_tcp() + if server_args.enable_trace: + process_tracing_init( + server_args.otlp_traces_endpoint, + "sglang", + trace_modules=server_args.trace_modules, + ) + trace_set_thread_info("Encoder") for rank in range(1, server_args.tp_size): schedule_path = f"ipc:///tmp/{ipc_path_prefix}_schedule_{rank}" send_sockets.append( @@ -3414,7 +3438,12 @@ async def get_condition(rid): async def handle_encode_request(request: dict): req_id = request["req_id"] start_time = time.monotonic() + time_stats_json = request.pop("time_stats_json", None) + time_stats = EncoderReqTimeStats() if dp_dispatcher is not None: + if time_stats_json: + request = dict(request) + request["time_stats_json"] = time_stats_json try: result = await dp_dispatcher.dispatch(request) except MMError as e: @@ -3497,12 +3526,19 @@ async def handle_encode_request(request: dict): async with encoder.encode_dispatch_lock: request.update({"enter_time": time.time()}) modality = Modality.from_str(request["modality"]) + if time_stats_json: + time_stats.decode_json(time_stats_json) + + time_stats.set_mm_encode_start_time() if encoder_scheduler is not None and modality in _BATCHABLE_MODALITIES: try: nbytes, embedding_len, embedding_dim, error_msg, error_code = ( await encoder_scheduler.submit(request) ) except asyncio.TimeoutError: + time_stats.trace_ctx.abort( + abort_info={"reason": "encoder batch timed out"} + ) return ORJSONResponse( status_code=HTTPStatus.GATEWAY_TIMEOUT, content={ @@ -3518,6 +3554,11 @@ async def handle_encode_request(request: dict): await encoder.encode_request(request, modality) ) + if error_msg: + time_stats.trace_ctx.abort(abort_info={"reason": error_msg}) + else: + time_stats.set_mm_encode_end_time() + if error_msg: if encoder.server_args.encoder_transfer_backend == "zmq_to_scheduler": if request["embedding_port"] is None: diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 357e3c467..aca197603 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -2860,7 +2860,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): "zmq_to_scheduler", "mooncake", ]: - self.mm_receiver.send_encode_request(obj) + time_stats_json = None + if self.server_args.enable_trace: + state = self.rid_to_state.get(obj.rid) + if state is not None: + time_stats_json = state.time_stats.encode_json() + + self.mm_receiver.send_encode_request( + obj, time_stats_json=time_stats_json + ) else: obj.need_wait_for_mm_inputs = False diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index f672045ad..1003291d0 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -145,6 +145,12 @@ class RequestStage: metrics_is_observed=True, ) + # EPD disaggregation Encode process + MM_ENCODE = RequestStageConfig( + "mm_encode", + level=1, + ) + # disaggregation prefill PREFILL_PREPARE = RequestStageConfig( "prefill_prepare", @@ -306,14 +312,34 @@ class ReqTimeStatsBase: def __getstate__(self) -> object: # The object is propagated to other processes via serialization and deserialization methods, # requiring the metric collector to be reconfigured. + trace_ctx_state = ( + self.trace_ctx.__getstate__() + if self.trace_ctx.tracing_enable + else {"tracing_enable": False} + ) return { - "disagg_mode": self.disagg_mode, + "disagg_mode": self.disagg_mode.value if self.disagg_mode else None, "enable_metrics": False, - "trace_ctx": self.trace_ctx, + "trace_ctx": trace_ctx_state, "diff_realtime_monotonic": global_diff_realtime_monotonic, } def __setstate__(self, state: object): + # Reconstruct disagg_mode from string value if needed + disagg_mode_val = state.get("disagg_mode") + if isinstance(disagg_mode_val, str): + state["disagg_mode"] = DisaggregationMode(disagg_mode_val) + + # Reconstruct trace_ctx from serialized dict if needed + trace_ctx_state = state.get("trace_ctx") + if isinstance(trace_ctx_state, dict): + if trace_ctx_state.get("tracing_enable"): + trace_ctx = object.__new__(TraceReqContext) + trace_ctx.__setstate__(trace_ctx_state) + state["trace_ctx"] = trace_ctx + else: + state["trace_ctx"] = TraceNullContext() + for key in state.keys(): if key.endswith("time"): state[key] = convert_time_cross_thread( @@ -323,6 +349,12 @@ class ReqTimeStatsBase: ) self.__dict__.update(state) + def encode_json(self) -> Dict[str, Any]: + return self.__getstate__() + + def decode_json(self, state: Dict[str, Any]): + self.__setstate__(state) + @dataclass class APIServerReqTimeStats(ReqTimeStatsBase): @@ -353,6 +385,13 @@ class APIServerReqTimeStats(ReqTimeStatsBase): if self.trace_ctx.tracing_enable: self.trace_ctx.trace_req_start(convert_time_to_realtime_ns(ts)) + # Start tokenize span early so that EPD encode dispatch can capture + # it as the predecessor span context when serializing trace_ctx. + self.trace_ctx.trace_slice_start( + RequestStage.TOKENIZE.stage_name, + RequestStage.TOKENIZE.level, + convert_time_to_realtime_ns(ts), + ) def set_finished_time(self, ts=None): ts = ts or time.perf_counter() @@ -374,8 +413,13 @@ class APIServerReqTimeStats(ReqTimeStatsBase): ts = ts or time.perf_counter() self.tokenize_finish_time = ts - stage = RequestStage.TOKENIZE - self.trace_slice(stage, self.created_time, ts) + # tokenize span was started in set_created_time(); end it here. + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_slice_end( + RequestStage.TOKENIZE.stage_name, + RequestStage.TOKENIZE.level, + convert_time_to_realtime_ns(ts), + ) def set_api_server_dispatch_time(self, ts=None): ts = ts or time.perf_counter() @@ -1141,6 +1185,34 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): return f"{convert_time_to_realtime(perf_counter_time):.3f}" +@dataclass +class EncoderReqTimeStats(ReqTimeStatsBase): + mm_encode_start_time: float = 0.0 + mm_encode_end_time: float = 0.0 + + def set_mm_encode_start_time(self, ts=None): + ts = ts or time.perf_counter() + self.mm_encode_start_time = ts + if self.trace_ctx.tracing_enable: + self.trace_ctx.rebuild_thread_context() + self.trace_ctx.trace_slice_start( + RequestStage.MM_ENCODE.stage_name, + RequestStage.MM_ENCODE.level, + convert_time_to_realtime_ns(ts), + ) + + def set_mm_encode_end_time(self, ts=None): + ts = ts or time.perf_counter() + self.mm_encode_end_time = ts + if self.trace_ctx.tracing_enable: + self.trace_ctx.trace_slice_end( + RequestStage.MM_ENCODE.stage_name, + RequestStage.MM_ENCODE.level, + convert_time_to_realtime_ns(ts), + thread_finish_flag=True, + ) + + def set_schedule_time_batch(batch: ScheduleBatch): # only for tracing if not get_global_tracing_enabled():