From 364b5f23e65f5e7041d824fed402814efd84b0ba Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Fri, 24 Jul 2026 04:37:52 -0400 Subject: [PATCH] [BugFix][EPD] Harden zmq_to_scheduler receiver failures; sync error info across TP (#31592) Co-authored-by: siyu --- .../srt/disaggregation/encode_receiver.py | 139 +++++++++++++----- .../scheduler_components/request_receiver.py | 11 +- 2 files changed, 105 insertions(+), 45 deletions(-) diff --git a/python/sglang/srt/disaggregation/encode_receiver.py b/python/sglang/srt/disaggregation/encode_receiver.py index 620d952d1..ee45b415e 100644 --- a/python/sglang/srt/disaggregation/encode_receiver.py +++ b/python/sglang/srt/disaggregation/encode_receiver.py @@ -431,7 +431,9 @@ class EmbeddingData: self.shape = list(embedding.shape) if embedding is not None else None self.cached_embedding = None self.error_msg = error_msg - self.error_code = error_code + # Coerce to plain int: this object crosses process boundaries via + # safe_pickle_loads, whose allowlist blocks http.HTTPStatus. + self.error_code = int(error_code) if error_code is not None else None # Store additional metadata (e.g., video_timestamps for qwen3_vl) for key, value in kwargs.items(): setattr(self, key, value) @@ -839,51 +841,78 @@ class WaitingImageRequest: except zmq.Again: # No data available yet, wait a bit and retry return - recv_obj: EmbeddingData = safe_pickle_loads(parts[0]) - if getattr(recv_obj, "error_msg", None) is not None: - logger.warning( - f"Received error signal from encoder for {self.rid}: {recv_obj.error_msg} {recv_obj.error_code = }" + try: + recv_obj: EmbeddingData = safe_pickle_loads(parts[0]) + if getattr(recv_obj, "error_msg", None) is not None: + logger.warning( + f"Received error signal from encoder for {self.rid}: {recv_obj.error_msg} {recv_obj.error_code = }" + ) + self.error_msg = recv_obj.error_msg + self.error_code = recv_obj.error_code + self.status = WaitingImageRequestStatus.FAIL + self.recv_socket.close() + return + + # Extract original req_id from part_req_id and drop stale payloads + # that may arrive on a reused ZMQ port after a prior request aborted. + original_req_id = extract_original_req_id(recv_obj.req_id) + if original_req_id != self.recv_req.rid: + logger.warning( + f"Dropping stale embedding data: expected rid={self.recv_req.rid}, " + f"got rid={recv_obj.req_id} (likely from ZMQ port reuse)" + ) + continue + recv_obj.req_id = original_req_id + + buffer = parts[1].buffer if hasattr(parts[1], "buffer") else parts[1] + recv_obj.embedding = ( + torch.frombuffer(buffer, dtype=recv_obj.dtype) + .reshape(recv_obj.shape) + .clone() ) - self.error_msg = recv_obj.error_msg - self.error_code = recv_obj.error_code + + if self.recv_embedding_data is None: + self.recv_embedding_data = ( + MultiModalEmbeddingData.from_embedding_data( + recv_obj, model_type=self.model_type + ) + ) + else: + self.recv_embedding_data.add(recv_obj) + except Exception as e: + # A message the scheduler cannot decode (blocked unpickle, + # bad shape/dtype, ...) must fail this request, not crash the + # scheduler event loop; FAIL still reaches the TP-wide status + # all-reduce in _process_waiting_requests. + logger.exception( + "Failed to decode embedding message for rid=%s", self.rid + ) + self.error_msg = f"Failed to decode embedding message: {e}" self.status = WaitingImageRequestStatus.FAIL + self._cleanup_gpu_buffer() self.recv_socket.close() return - # Extract original req_id from part_req_id and drop stale payloads - # that may arrive on a reused ZMQ port after a prior request aborted. - original_req_id = extract_original_req_id(recv_obj.req_id) - if original_req_id != self.recv_req.rid: - logger.warning( - f"Dropping stale embedding data: expected rid={self.recv_req.rid}, " - f"got rid={recv_obj.req_id} (likely from ZMQ port reuse)" - ) - continue - recv_obj.req_id = original_req_id - - buffer = parts[1].buffer if hasattr(parts[1], "buffer") else parts[1] - recv_obj.embedding = ( - torch.frombuffer(buffer, dtype=recv_obj.dtype) - .reshape(recv_obj.shape) - .clone() + # Assemble mm_inputs. Wrapped so an assembly failure still reaches the + # TP-wide status all-reduce in _process_waiting_requests instead of + # raising past it. + try: + recv_embedding = self.recv_embedding_data.get_embedding(is_concat=True) + mm_inputs = self.mm_processor.get_mm_data( + self.recv_req.input_text, + recv_embedding, + **self.recv_embedding_data.get_mm_extra_meta(), ) - - if self.recv_embedding_data is None: - self.recv_embedding_data = MultiModalEmbeddingData.from_embedding_data( - recv_obj, model_type=self.model_type - ) - else: - self.recv_embedding_data.add(recv_obj) - - recv_embedding = self.recv_embedding_data.get_embedding(is_concat=True) - mm_inputs = self.mm_processor.get_mm_data( - self.recv_req.input_text, - recv_embedding, - **self.recv_embedding_data.get_mm_extra_meta(), - ) - self.recv_req.mm_inputs = mm_inputs - self.recv_req.input_ids = array("q", mm_inputs.input_ids) - self.status = WaitingImageRequestStatus.SUCCESS + self.recv_req.mm_inputs = mm_inputs + self.recv_req.input_ids = array("q", mm_inputs.input_ids) + self.status = WaitingImageRequestStatus.SUCCESS + except Exception as e: + logger.exception( + "Failed to assemble multimodal inputs for rid=%s", self.rid + ) + self.status = WaitingImageRequestStatus.FAIL + self.error_msg = f"Failed to assemble multimodal inputs: {e}" + self._cleanup_gpu_buffer() self.recv_socket.close() def _cleanup_gpu_buffer(self): @@ -1177,6 +1206,10 @@ class WaitingImageRDMARequest(WaitingImageRequest): parts = self.recv_socket.recv_multipart(flags=zmq.NOBLOCK, copy=False) except zmq.Again: return + except zmq.ZMQError: + # The RDMA pipeline thread closed the socket after an encoder + # error (e.g. OOM). It already set status=FAIL; just bail. + return recv_obj: EmbeddingData = safe_pickle_loads(parts[0]) if getattr(recv_obj, "error_msg", None) is not None: @@ -1803,6 +1836,31 @@ class MMReceiverBase(ABC): ) obj.need_wait_for_mm_inputs = False + def _sync_fail_info_across_tp(self, waiting_req: WaitingImageRequest) -> None: + """Share encoder error fields across TP ranks before abort. + + The encoder sends ZMQ error signals to each TP rank's receive socket, + but they can arrive at different times. ``all_reduce`` on status makes + every rank enter FAIL together while only some ranks have populated + ``error_msg`` / ``error_code``. attn_tp_rank 0 streams the abort to the + client, so merge the best-known payload from all ranks first. + """ + if self.tp_size <= 1 or self.tp_group is None: + return + + gathered = self.tp_group.all_gather_object( + (waiting_req.error_msg, waiting_req.error_code) + ) + best_msg = waiting_req.error_msg + best_code = waiting_req.error_code + for msg, code in gathered: + if msg is not None: + best_msg = msg + if code is not None: + best_code = code + waiting_req.error_msg = best_msg + waiting_req.error_code = best_code + # For zmq_to_scheduler def _process_waiting_requests(self, recv_reqs, waiting_cls, **extra_kwargs): new_recv_reqs = [] @@ -1861,6 +1919,7 @@ class MMReceiverBase(ABC): if status_value == WaitingImageRequestStatus.SUCCESS: new_recv_reqs.append(waiting_req.recv_req) elif status_value == WaitingImageRequestStatus.FAIL: + self._sync_fail_info_across_tp(waiting_req) logger.error( f"Waiting request {waiting_req.rid} failed: {waiting_req.error_msg} {waiting_req.error_code = }" ) diff --git a/python/sglang/srt/managers/scheduler_components/request_receiver.py b/python/sglang/srt/managers/scheduler_components/request_receiver.py index fccd74c1d..11df155a3 100644 --- a/python/sglang/srt/managers/scheduler_components/request_receiver.py +++ b/python/sglang/srt/managers/scheduler_components/request_receiver.py @@ -226,11 +226,12 @@ class SchedulerRequestReceiver: ): recv_reqs, abort_reqs = self.mm_receiver.process_waiting_requests(recv_reqs) for req, error_msg, error_code in abort_reqs: - status_code = ( - HTTPStatus.BAD_REQUEST - if error_code == 400 - else HTTPStatus.INTERNAL_SERVER_ERROR - ) + if error_code is None: + status_code = HTTPStatus.INTERNAL_SERVER_ERROR + elif isinstance(error_code, HTTPStatus): + status_code = error_code + else: + status_code = HTTPStatus(int(error_code)) prepare_abort(req, error_msg, status_code=status_code) self.stream_output([req], req.return_logprob) return recv_reqs