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

fix grammar timeout sync across tp ranks. (#16898)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-01-14 10:26:31 +08:00
committed by GitHub
parent a4825ed588
commit e2c8a50b38
3 changed files with 207 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
python/sglang/srt/constrained/grammar_manager.py
View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
import random
from concurrent import futures
from typing import TYPE_CHECKING, List
@@ -18,6 +19,7 @@ if TYPE_CHECKING:
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__)
@@ -36,6 +38,11 @@ class GrammarManager:
else:
self.grammar_backend = None
self.grammar_sync_group = scheduler.dp_tp_cpu_group
self.grammar_sync_size = scheduler.dp_tp_group.world_size
self.grammar_sync_entry = scheduler.dp_tp_group.first_rank
self.is_grammar_sync_entry = scheduler.dp_tp_group.is_first_rank
def __len__(self):
return len(self.grammar_queue)
@@ -94,70 +101,89 @@ 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=}"
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_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
req.grammar = req.grammar.result(timeout=0.03)
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)
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 req.grammar_wait_ct > GRAMMAR_TIMEOUT / 0.03:
num_timeout_reqs = 1
if sim_timeout_prob > 0 or req.grammar_wait_ct > timeout_ct:
num_timeout_pt = num_ready_reqs + 1
break
if self.server_args.enable_dp_attention:
tp_size = self.scheduler.attn_tp_size
tp_group = self.scheduler.attn_tp_cpu_group
else:
tp_size = self.scheduler.tp_size
tp_group = self.scheduler.tp_cpu_group
if tp_size > 1:
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_timeout_reqs], dtype=torch.int32)
torch.distributed.all_reduce(
tensor, op=torch.distributed.ReduceOp.MAX, group=tp_group
tensor = torch.tensor(
[num_ready_reqs, -num_ready_reqs, -num_timeout_pt], dtype=torch.int32
)
num_ready_reqs_max, num_timeout_reqs_max = tensor.tolist()
for i in range(num_ready_reqs, num_ready_reqs_max):
req = self.grammar_queue[i]
if req.finished(): # 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)
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_reqs_max = num_timeout_reqs
num_timeout_pt_max = num_timeout_pt
for i in range(num_ready_reqs, num_ready_reqs + num_timeout_reqs_max):
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]
req.grammar.cancel()
self.grammar_backend.set_cache(req.grammar_key, INVALID_GRAMMAR_OBJ)
error_msg = f"Grammar preprocessing timed out for {req.grammar_key=}"
req.set_finish_with_abort(error_msg)
if req.finished():
continue
if isinstance(req.grammar, futures.Future):
req.grammar.cancel()
req.grammar = INVALID_GRAMMAR_OBJ
self.finalize_grammar(req)
num_ready_reqs = num_ready_reqs_max + num_timeout_reqs_max
ready_grammar_reqs = self.grammar_queue[:num_ready_reqs]
self.grammar_queue = self.grammar_queue[num_ready_reqs:]
ready_grammar_reqs = self.grammar_queue[:num_timeout_pt_max]
self.grammar_queue = self.grammar_queue[num_timeout_pt_max:]
return ready_grammar_reqs

4
python/sglang/srt/environ.py
View File

@@ -171,6 +171,10 @@ class Envs:
SGLANG_IS_IN_CI_AMD = EnvBool(False)
SGLANG_TEST_MAX_RETRY = EnvInt(None)
# Constrained Decoding (Grammar)
SGLANG_GRAMMAR_SIMULATE_TIMEOUT = EnvFloat(-1)
SGLANG_GRAMMAR_POLL_INTERVAL = EnvFloat(0.03)
# Test & Debug
SGLANG_DETECT_SLOW_RANK = EnvBool(False)
SGLANG_TEST_STUCK_DETOKENIZER = EnvFloat(0)