from __future__ import annotations import logging import time from concurrent import futures from typing import TYPE_CHECKING, List import torch from sglang.srt.constrained.base_grammar_backend import ( InvalidGrammarObject, create_grammar_backend, ) from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.environ import envs if TYPE_CHECKING: from sglang.srt.managers.io_struct import AbortReq from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.scheduler import Scheduler logger = logging.getLogger(__name__) class GrammarManager: def __init__(self, scheduler: Scheduler): self.scheduler = scheduler self.server_args = scheduler.server_args self.grammar_queue: List[Req] = [] if not self.server_args.skip_tokenizer_init: self.grammar_backend = create_grammar_backend( self.server_args, scheduler.tokenizer, scheduler.model_config.vocab_size, scheduler.model_config.hf_eos_token_id, think_end_id=scheduler.model_config.think_end_id, ) else: self.grammar_backend = None self._enable_strict_thinking = ( self.grammar_backend.enable_strict_thinking if self.grammar_backend is not None else False ) 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 self.pp_rank = scheduler.ps.pp_rank self.pp_size = scheduler.ps.pp_size self.pp_group = scheduler.pp_group self.grammar_pp_sync_work_list = [] 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) def clear(self): if self.grammar_backend: self.grammar_backend.reset() def has_waiting_grammars(self) -> bool: return len(self.grammar_queue) > 0 def _drain_pp_sync_work(self): for p2p_work in self.grammar_pp_sync_work_list: p2p_work.work.wait() self.grammar_pp_sync_work_list.clear() def _pp_sync_ready_failed( self, ready_req_idxs: set[int], failed_req_idxs: set[int], ) -> tuple[set[int], set[int]]: """ Synchronize ready/failed grammar request indexes across the PP pipeline. PP0 provides the data. Each later PP rank receives it from the previous rank and asynchronously forwards it to the next rank. """ if self.pp_size <= 1 or self.pp_group is None: return ready_req_idxs, failed_req_idxs self._drain_pp_sync_work() data = (ready_req_idxs, failed_req_idxs) if self.pp_rank > 0: data = self.pp_group.recv_object( src=self.pp_rank - 1, tag=P2PTag.GRAMMAR_PP_SYNC, ) if self.pp_rank + 1 < self.pp_size: self.grammar_pp_sync_work_list.extend( self.pp_group.send_object( data, dst=self.pp_rank + 1, async_send=True, tag=P2PTag.GRAMMAR_PP_SYNC, ) ) return data def abort_requests(self, recv_req: AbortReq): for req in self.grammar_queue: if recv_req.abort_all or req.rid.startswith(recv_req.rid): logger.debug(f"Abort grammar queue request. {req.rid=}") if isinstance(req.grammar, futures.Future) and req.grammar: req.grammar.cancel() req.set_finish_with_abort("Aborted by AbortReq.") def _get_request_thinking_budget(self, req: Req) -> int | None: custom_params = req.sampling_params.custom_params if not isinstance(custom_params, dict): return None thinking_budget = custom_params.get("thinking_budget") return thinking_budget if isinstance(thinking_budget, int) else None def _apply_request_reasoning_budget(self, req: Req) -> None: thinking_budget = self._get_request_thinking_budget(req) if thinking_budget is None: return if isinstance(req.grammar, ReasonerGrammarObject): req.grammar.max_think_tokens = thinking_budget def process_req_with_grammar(self, req: Req) -> bool: # Init grammar cache for this request add_to_grammar_queue = False if ( req.sampling_params.json_schema is not None or req.sampling_params.regex is not None or req.sampling_params.ebnf is not None or req.sampling_params.structural_tag is not None ): if self.grammar_backend is None: error_msg = "Grammar-based generation (json_schema, regex, ebnf, structural_tag) is not supported when the server is launched with --grammar-backend none" req.set_finish_with_abort(error_msg) else: if req.sampling_params.json_schema is not None: key = ("json", req.sampling_params.json_schema) elif req.sampling_params.regex is not None: key = ("regex", req.sampling_params.regex) elif req.sampling_params.ebnf is not None: key = ("ebnf", req.sampling_params.ebnf) elif req.sampling_params.structural_tag: key = ("structural_tag", req.sampling_params.structural_tag) value, cache_hit = self.grammar_backend.get_cached_or_future_value( key, req.require_reasoning ) req.grammar = value if not cache_hit: req.grammar_key = key add_to_grammar_queue = True else: if isinstance( value, InvalidGrammarObject ): # We hit a cached invalid grammar. error_msg = ( f"Failed to compile {key[0]} grammar: {value.error_message}" ) req.set_finish_with_abort(error_msg) else: self._apply_request_reasoning_budget(req) elif self._enable_strict_thinking: grammar_obj = self.grammar_backend.init_strict_reasoning_grammar( req.require_reasoning ) if grammar_obj is not None: req.grammar = grammar_obj self._apply_request_reasoning_budget(req) if add_to_grammar_queue: self.grammar_queue.append(req) return add_to_grammar_queue def get_ready_grammar_requests(self) -> List[Req]: """ Move requests whose grammar objects are ready from grammar_queue to waiting_queue. For PP0, DP/TP group rank i returns two sets ready_reqs_i, failed_reqs_i. ready_reqs_all = all_gather(ready_reqs_i) within PP0's DP/TP group. failed_reqs_all = all_gather(failed_reqs_i) within PP0's DP/TP group. ready_reqs = intersect(ready_reqs_all) failed_reqs = union(failed_reqs_all) PP0 then propagates the synced result to later PP ranks. Later PP ranks receive and apply the propagated ready/failed decision. """ assert self.grammar_backend ready_req_idxs: set[int] = set() failed_req_idxs: set[int] = set() if self.pp_rank == 0: # 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) if len(ready_req_idxs) == len(self.grammar_queue): break # 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: # grammar_wait_ct is only updated on PP0; later PP ranks # receive PP0's ready/failed decision through PP sync. 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 TP ranks in PP0. 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]) else: synced_ready_req_idxs = ready_req_idxs synced_failed_req_idxs = failed_req_idxs # Propagate PP0's grammar queue decision to later PP ranks. ( synced_ready_req_idxs, synced_failed_req_idxs, ) = self._pp_sync_ready_failed( synced_ready_req_idxs, synced_failed_req_idxs, ) # 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 assert isinstance(req.grammar, futures.Future) and req.grammar_key try: req.grammar = req.grammar.result() except Exception as e: logger.error( f"Grammar compilation raised an exception: {e}, " f"grammar_key={req.grammar_key}" ) req.grammar = InvalidGrammarObject(f"Grammar compilation failed: {e}") self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy()) self._apply_request_reasoning_budget(req) if isinstance(req.grammar, InvalidGrammarObject): error_msg = f"Failed to compile {req.grammar_key[0]} grammar: {req.grammar.error_message}" 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) assert isinstance(req.grammar, futures.Future) and req.grammar_key req.grammar.cancel() self.grammar_backend.set_cache( req.grammar_key, InvalidGrammarObject("Grammar preprocessing timed out") ) error_msg = f"Grammar preprocessing timed out: {req.grammar_key=}" req.set_finish_with_abort(error_msg) # 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