Fix grammar sync across TP ranks (#17100)

Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>

python/sglang/srt/constrained/grammar_manager.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-import random
+import time
 from concurrent import futures
 from typing import TYPE_CHECKING, List
 
@@ -18,8 +18,6 @@ if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.managers.scheduler import Scheduler
 
-GRAMMAR_TIMEOUT = envs.SGLANG_GRAMMAR_TIMEOUT.get()
-GRAMMAR_POLL_INTERVAL = envs.SGLANG_GRAMMAR_POLL_INTERVAL.get()
 logger = logging.getLogger(__name__)
 
 
@@ -43,6 +41,11 @@ class GrammarManager:
         self.grammar_sync_entry = scheduler.dp_tp_group.first_rank
         self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank
 
+        self.SGLANG_GRAMMAR_POLL_INTERVAL = envs.SGLANG_GRAMMAR_POLL_INTERVAL.get()
+        self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = (
+            envs.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS.get()
+        )
+
     def __len__(self):
         return len(self.grammar_queue)
 
@@ -101,89 +104,92 @@ class GrammarManager:
 
         return add_to_grammar_queue
 
-    def finalize_grammar(self, req: Req):
-        self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
-        if req.grammar is INVALID_GRAMMAR_OBJ:
-            error_msg = f"Invalid grammar request: {req.grammar_key=}"
+    def get_ready_grammar_requests(self) -> List[Req]:
+        """
+        Move requests whose grammar objects are ready from grammar_queue to waiting_queue.
+
+        Rank i returns two sets ready_reqs_i, failed_reqs_i
+        ready_reqs_all = all_gather(ready_reqs_i)
+        failed_reqs_all = all_gather(failed_reqs_i)
+
+        ready_reqs = intersect(ready_reqs_all)
+        failed_reqs = union(failed_reqs_all)
+        """
+        ready_req_idxs: set[int] = set()
+        failed_req_idxs: set[int] = set()
+
+        # Poll for ready requests
+        start_time = time.perf_counter()
+        while time.perf_counter() - start_time < self.SGLANG_GRAMMAR_POLL_INTERVAL:
+            for i, req in enumerate(self.grammar_queue):
+                if i in ready_req_idxs:
+                    continue
+
+                if req.finished() or req.grammar is None:  # It is aborted by AbortReq
+                    ready_req_idxs.add(i)
+                    continue
+
+                assert isinstance(req.grammar, futures.Future), f"{req=}"
+                if req.grammar.done():
+                    ready_req_idxs.add(i)
+
+            # Sleep a bit to avoid busy waiting
+            time.sleep(self.SGLANG_GRAMMAR_POLL_INTERVAL / 10)
+
+        # Check failed requests
+        for i, req in enumerate(self.grammar_queue):
+            if i not in ready_req_idxs:
+                self.grammar_queue[i].grammar_wait_ct += 1
+                if (
+                    self.grammar_queue[i].grammar_wait_ct
+                    >= self.SGLANG_GRAMMAR_MAX_POLL_ITERATIONS
+                ):
+                    # Timeout after max poll iterations
+                    # The actual waiting time is SGLANG_GRAMMAR_MAX_POLL_ITERATIONS * max(SGLANG_GRAMMAR_POLL_INTERVAL, GPU_forward_batch_latency)
+                    failed_req_idxs.add(i)
+
+        # Sync ready and failed requests across all ranks
+        if self.grammar_sync_size == 1:
+            synced_ready_req_idxs = ready_req_idxs
+            synced_failed_req_idxs = failed_req_idxs
+        else:
+            all_gather_output = [None] * self.grammar_sync_size
+            torch.distributed.all_gather_object(
+                all_gather_output,
+                (ready_req_idxs, failed_req_idxs),
+                group=self.grammar_sync_group,
+            )
+            synced_ready_req_idxs = set.intersection(*[x[0] for x in all_gather_output])
+            synced_failed_req_idxs = set.union(*[x[1] for x in all_gather_output])
+
+        # Return ready requests
+        return_reqs: List[Req] = []
+        for i in synced_ready_req_idxs:
+            req = self.grammar_queue[i]
+            return_reqs.append(req)
+            if req.finished() or req.grammar is None:  # It is aborted by AbortReq
+                continue
+
+            req.grammar = req.grammar.result()
+            self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
+            if req.grammar is INVALID_GRAMMAR_OBJ:
+                error_msg = f"Invalid grammar request: {req.grammar_key=}"
+                req.set_finish_with_abort(error_msg)
+
+        # Return failed requests
+        for i in synced_failed_req_idxs:
+            req = self.grammar_queue[i]
+            return_reqs.append(req)
+
+            req.grammar.cancel()
+            self.grammar_backend.set_cache(req.grammar_key, INVALID_GRAMMAR_OBJ)
+            error_msg = f"Grammar preprocessing timed out: {req.grammar_key=}"
             req.set_finish_with_abort(error_msg)
 
-    def get_ready_grammar_requests(self) -> List[Req]:
-        """Move requests whose grammar objects are ready from grammar_queue to waiting_queue."""
-        num_ready_reqs = 0
-        num_timeout_pt = 0
-        sim_timeout_prob = (
-            -1
-            if self.is_grammar_sync_entry
-            else envs.SGLANG_GRAMMAR_SIMULATE_TIMEOUT.get()
-        )  # Entry rank never simulates timeout
-        timeout_ct = GRAMMAR_TIMEOUT / GRAMMAR_POLL_INTERVAL
-
-        for req in self.grammar_queue:
-            try:
-                if req.finished():  # It is aborted by AbortReq
-                    num_ready_reqs += 1
-                    continue
-
-                if sim_timeout_prob > 0 and random.random() < sim_timeout_prob:
-                    # Simulate timeout for non-entry ranks in TP sync group for testing
-                    logger.warning(
-                        f"Simulating grammar timeout on {self.scheduler.tp_rank=}"
-                    )
-                    raise futures._base.TimeoutError()
-
-                req.grammar = req.grammar.result(timeout=GRAMMAR_POLL_INTERVAL)
-                self.finalize_grammar(req)
-                num_ready_reqs += 1
-                num_timeout_pt = num_ready_reqs
-            except futures._base.TimeoutError:
-                req.grammar_wait_ct += 1
-                # NOTE(lianmin): this timeout is the waiting time of the above line. It is
-                # not the waiting time from it enters the grammar queue.
-                if sim_timeout_prob > 0 or req.grammar_wait_ct > timeout_ct:
-                    num_timeout_pt = num_ready_reqs + 1
-                break
-
-        if self.grammar_sync_size > 1:
-            # Sync across TP ranks to make sure they have the same number of ready requests
-            tensor = torch.tensor(
-                [num_ready_reqs, -num_ready_reqs, -num_timeout_pt], dtype=torch.int32
-            )
-            torch.distributed.all_reduce(
-                tensor, op=torch.distributed.ReduceOp.MIN, group=self.grammar_sync_group
-            )
-            num_ready_reqs_min = tensor[0].item()
-            num_ready_reqs_max = -tensor[1].item()
-            num_timeout_pt_max = -tensor[2].item()
-        else:
-            num_ready_reqs_min = num_ready_reqs
-            num_ready_reqs_max = num_ready_reqs
-            num_timeout_pt_max = num_timeout_pt
-
-        if envs.SGLANG_GRAMMAR_SIMULATE_TIMEOUT.get() > 0:
-            # NOTE: in simulation timeout mode, if only some TP ranks report a timeout,
-            # we still treat those requests as ready instead of canceling them.
-            for i in range(num_ready_reqs_min, num_ready_reqs_max):
-                req = self.grammar_queue[i]
-                if req.finished():
-                    continue
-                if isinstance(req.grammar, futures.Future):
-                    req.grammar = req.grammar.result()
-                    self.finalize_grammar(req)
-
-            num_ready_reqs_min = num_ready_reqs_max
-
-        # NOTE: in non-simulation mode, cancel any request that times out on a subset of
-        # TP ranks because timeouts can diverge across ranks.
-        for i in range(num_ready_reqs_min, num_timeout_pt_max):
-            req = self.grammar_queue[i]
-            if req.finished():
-                continue
-            if isinstance(req.grammar, futures.Future):
-                req.grammar.cancel()
-            req.grammar = INVALID_GRAMMAR_OBJ
-            self.finalize_grammar(req)
-
-        ready_grammar_reqs = self.grammar_queue[:num_timeout_pt_max]
-        self.grammar_queue = self.grammar_queue[num_timeout_pt_max:]
-
-        return ready_grammar_reqs
+        # Remove finished requests from grammar_queue
+        self.grammar_queue = [
+            req
+            for i, req in enumerate(self.grammar_queue)
+            if i not in synced_ready_req_idxs and i not in synced_failed_req_idxs
+        ]
+        return return_reqs

python/sglang/srt/entrypoints/openai/serving_rerank.py

@@ -212,9 +212,6 @@ class OpenAIServingRerank(OpenAIServingBase):
         self._yes_token_id, self._no_token_id = _get_yes_no_token_ids(
             tokenizer_manager.tokenizer
         )
-        logger.info(
-            f"Reranker yes/no token IDs: yes={self._yes_token_id}, no={self._no_token_id}"
-        )
 
     # NOTE: /v1/rerank is not an official OpenAI endpoint. This module may be moved
     # to another module in the future.

python/sglang/srt/environ.py

@@ -172,8 +172,8 @@ class Envs:
     SGLANG_TEST_MAX_RETRY = EnvInt(None)
 
     # Constrained Decoding (Grammar)
-    SGLANG_GRAMMAR_SIMULATE_TIMEOUT = EnvFloat(-1)
-    SGLANG_GRAMMAR_POLL_INTERVAL = EnvFloat(0.03)
+    SGLANG_GRAMMAR_POLL_INTERVAL = EnvFloat(0.005)
+    SGLANG_GRAMMAR_MAX_POLL_ITERATIONS = EnvInt(10000)
 
     # Test & Debug
     SGLANG_DETECT_SLOW_RANK = EnvBool(False)
@@ -251,10 +251,6 @@ class Envs:
     SGLANG_USE_MESSAGE_QUEUE_BROADCASTER = EnvBool(True)
     SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS = EnvBool(False)
 
-    # Constrained Decoding
-    SGLANG_DISABLE_OUTLINES_DISK_CACHE = EnvBool(True)
-    SGLANG_GRAMMAR_TIMEOUT = EnvFloat(300)
-
     # Tool Calling
     SGLANG_FORWARD_UNKNOWN_TOOLS = EnvBool(False)