[EPD] Add notification mechanism to fix server hang and add timeout env var (#18229)

python/sglang/srt/disaggregation/encode_receiver.py

@@ -3,9 +3,11 @@ import logging
 import pickle
 import random
 import threading
+import time
 import uuid
 from abc import ABC, abstractmethod
 from enum import IntEnum
+from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional
 
 import aiohttp
@@ -18,6 +20,7 @@ from sglang.srt.distributed.parallel_state import (
     GroupCoordinator,
     get_mooncake_transfer_engine,
 )
+from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import TokenizedGenerateReqInput
 from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
 from sglang.srt.managers.schedule_batch import Req
@@ -106,6 +109,7 @@ class WaitingImageRequestStatus(IntEnum):
     FAIL = -1
     PENDING = 0
     SUCCESS = 1
+    TIMEOUT = -2
 
 
 # For zmq_to_scheduler
@@ -138,6 +142,7 @@ class WaitingImageRequest:
         self.status = WaitingImageRequestStatus.PENDING
         self.error_msg = None
         self.error_code = None
+        self.start_time = time.time()
 
     def send_encode_request(self):
         async def _send_single_request(session, url, payload):
@@ -300,6 +305,7 @@ class MMReceiverHTTP(MMReceiverBase):
             self.hostname = get_local_ip_auto()
             self.waiting_list: List[WaitingImageRequest] = []
             self.scheduler = scheduler
+            self.wait_timeout = envs.SGLANG_ENCODER_RECV_TIMEOUT.get()
             if hf_config is not None:
                 transport_mode = _determine_tensor_transport_mode(server_args)
                 import_processors("sglang.srt.multimodal.processors")
@@ -394,9 +400,12 @@ class MMReceiverHTTP(MMReceiverBase):
         if len(self.waiting_list) == 0:
             return new_recv_reqs, []
 
+        current_time = time.time()
         local_status = []
         for waiting_req in self.waiting_list:
             waiting_req._try_recv_mm_data()
+            if current_time - waiting_req.start_time > self.wait_timeout:
+                waiting_req.status = WaitingImageRequestStatus.TIMEOUT
             local_status.append(waiting_req.status)
 
         local_status = torch.tensor(local_status, device="cpu", dtype=torch.int32)
@@ -424,6 +433,17 @@ class MMReceiverHTTP(MMReceiverBase):
                         waiting_req.error_code,
                     )
                 )
+            elif status_value == WaitingImageRequestStatus.TIMEOUT:
+                logger.error(
+                    f"Timed out waiting for image embeddings for request {waiting_req.rid}"
+                )
+                abort_reqs.append(
+                    (
+                        self.create_req(waiting_req.recv_req),
+                        f"Timeout waiting for image embedding after {self.wait_timeout}s",
+                        HTTPStatus.REQUEST_TIMEOUT,
+                    )
+                )
             else:  # status_value == WaitingImageRequestStatus.PENDING
                 new_waiting.append(waiting_req)
 

python/sglang/srt/disaggregation/encode_server.py

@@ -29,6 +29,7 @@ from sglang.srt.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import initialize_dp_attention
 from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType
 from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
@@ -55,6 +56,8 @@ rid_lock = asyncio.Lock()
 rid_to_receive_endpoint: Dict[str, List[str]] = dict()
 rid_to_receive_count: Dict[str, int] = dict()
 rid_to_err_msg: Dict[str, str] = dict()
+cond_dict_lock = asyncio.Lock()
+rid_to_cond: Dict[str, asyncio.Condition] = {}
 
 use_image_processor_gpu = (
     int(os.getenv("SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU", "0")) == 1
@@ -197,6 +200,7 @@ class MMEncoder:
         self.io_executor = concurrent.futures.ThreadPoolExecutor(
             max_workers=int(os.environ.get("SGLANG_ENCODER_MM_LOAD_WORKERS", 4))
         )
+        self.send_timeout = envs.SGLANG_ENCODER_SEND_TIMEOUT.get()
 
         if schedule_path is not None:
             self.schedule_socket = get_zmq_socket(
@@ -475,7 +479,8 @@ class MMEncoder:
         sent_urls: Set[str] = set()
         all_tasks: List[Tuple[asyncio.Task, str]] = []
         start_time = asyncio.get_running_loop().time()
-        timeout = 60.0
+        timeout = self.send_timeout
+        cond = await get_condition(req_id)
 
         try:
             while True:
@@ -504,14 +509,18 @@ class MMEncoder:
                         f"All {expected_count} endpoints initiated for {req_id}. Breaking loop."
                     )
                     break
-
+                remaining = timeout - (asyncio.get_running_loop().time() - start_time)
-                if asyncio.get_running_loop().time() - start_time > timeout:
+                if remaining <= 0:
                     logger.error(
-                        f"Timeout waiting for all endpoints for {req_id}. Initiated {len(sent_urls)}/{expected_count}"
+                        f"[{req_id}] Timeout! Sent {len(sent_urls)}/{expected_count}"
                     )
                     break
 
-                await asyncio.sleep(0.001)
+                async with cond:
+                    try:
+                        await asyncio.wait_for(cond.wait(), timeout=remaining)
+                    except asyncio.TimeoutError:
+                        continue
 
             if all_tasks:
                 logger.info(
@@ -535,6 +544,8 @@ class MMEncoder:
             async with rid_lock:
                 rid_to_receive_endpoint.pop(req_id, None)
                 rid_to_receive_count.pop(req_id, None)
+            async with cond_dict_lock:
+                rid_to_cond.pop(req_id, None)
             self.embedding_to_send.pop(req_id, None)
 
     async def get_embedding_port(self, prefill_url):
@@ -674,6 +685,13 @@ def launch_server(server_args: ServerArgs):
     uvicorn.run(app, host=server_args.host, port=server_args.port)
 
 
+async def get_condition(rid):
+    async with cond_dict_lock:
+        if rid not in rid_to_cond:
+            rid_to_cond[rid] = asyncio.Condition()
+        return rid_to_cond[rid]
+
+
 @app.post("/encode")
 async def handle_encode_request(request: dict):
     req_id = request["req_id"]
@@ -789,6 +807,9 @@ async def handle_scheduler_receive_url_request(request: dict):
             rid_to_receive_count[rid] = request["receive_count"]
         assert rid_to_receive_count[rid] == request["receive_count"]
         rid_to_receive_endpoint[rid].add(request["receive_url"])
+    cond = await get_condition(rid)
+    async with cond:
+        cond.notify_all()
 
 
 @app.get("/health")

python/sglang/srt/environ.py

@@ -466,6 +466,10 @@ class Envs:
     SGLANG_USE_AITER_FP8_PER_TOKEN = EnvBool(False)
     # fmt: on
 
+    # EPD
+    SGLANG_ENCODER_RECV_TIMEOUT = EnvFloat(180.0)
+    SGLANG_ENCODER_SEND_TIMEOUT = EnvFloat(180.0)
+
 
 envs = Envs()
 EnvField._allow_set_name = False

python/sglang/srt/managers/scheduler_metrics_mixin.py

@@ -224,6 +224,8 @@ class SchedulerMetricsMixin:
         else:
             msg += f"input throughput (token/s): {self.last_input_throughput:.2f}, "
 
+        if self.server_args.language_only:
+            msg += f"waiting-image-req: {len(self.mm_receiver.waiting_list)}, "
         graph_backend = defaultdict(
             lambda: "cuda graph",
             {
@@ -398,6 +400,9 @@ class SchedulerMetricsMixin:
             msg += f"#transfer-req: {len(self.disagg_decode_transfer_queue.queue)}, "
             msg += f"#retracted-req: {len(self.disagg_decode_prealloc_queue.retracted_queue)}, "
 
+        if self.server_args.language_only:
+            msg += f"waiting-image-req: {len(self.mm_receiver.waiting_list)}, "
+
         graph_backend = defaultdict(
             lambda: "cuda graph",
             {