diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py new file mode 100644 index 000000000..4cfcf6e83 --- /dev/null +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import logging +from concurrent import futures +from typing import TYPE_CHECKING, List + +import torch + +from sglang.srt.constrained.base_grammar_backend import ( + INVALID_GRAMMAR_OBJ, + create_grammar_backend, +) +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 + +GRAMMAR_TIMEOUT = envs.SGLANG_GRAMMAR_TIMEOUT.get() +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, + ) + else: + self.grammar_backend = None + + 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 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 req.grammar: + req.grammar.cancel() + req.set_finish_with_abort("Aborted by AbortReq.") + + 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 value is INVALID_GRAMMAR_OBJ: # We hit a cached invalid grammar. + error_msg = f"Invalid grammar request with cache hit: {key=}" + req.set_finish_with_abort(error_msg) + + 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.""" + + num_ready_reqs = 0 + num_timeout_reqs = 0 + 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) + + num_ready_reqs += 1 + 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 + 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: + # 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 + ) + 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) + else: + num_ready_reqs_max = num_ready_reqs + num_timeout_reqs_max = num_timeout_reqs + + for i in range(num_ready_reqs, num_ready_reqs + num_timeout_reqs_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) + + 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:] + + return ready_grammar_reqs diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 51af67636..4738b032f 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -946,8 +946,10 @@ class SchedulerDisaggregationDecodeMixin: def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]: """Create a schedulebatch for fake completed prefill""" - if self.grammar_queue: - self.move_ready_grammar_requests() + if self.grammar_manager.has_waiting_grammars(): + ready_grammar_requests = self.grammar_manager.get_ready_grammar_requests() + for req in ready_grammar_requests: + self._add_request_to_queue(req) if len(self.waiting_queue) == 0: return None diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5f168595c..22e173192 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -719,6 +719,7 @@ class Req: self.embedding = None # Constrained decoding + self.grammar_key: Optional[str] = None self.grammar: Optional[BaseGrammarObject] = None self.grammar_wait_ct = 0 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a3765d261..1bf294973 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -20,7 +20,6 @@ import signal import sys import time from collections import deque -from concurrent import futures from dataclasses import dataclass from http import HTTPStatus from typing import Any, Deque, Dict, List, Optional, Tuple, Union @@ -35,10 +34,7 @@ from torch.cuda import StreamContext as CudaStreamContext from torch.distributed import barrier from sglang.srt.configs.model_config import ModelConfig -from sglang.srt.constrained.base_grammar_backend import ( - INVALID_GRAMMAR_OBJ, - create_grammar_backend, -) +from sglang.srt.constrained.grammar_manager import GrammarManager from sglang.srt.disaggregation.decode import ( DecodePreallocQueue, DecodeTransferQueue, @@ -212,7 +208,6 @@ logger = logging.getLogger(__name__) TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get() TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get() TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get() -GRAMMAR_TIMEOUT = float(os.environ.get("SGLANG_GRAMMAR_TIMEOUT", 300)) @dataclass @@ -355,9 +350,6 @@ class Scheduler( # Init chunked prefill self.init_chunked_prefill() - # Init the grammar backend for constrained generation - self.init_grammar_backend() - # Init schedule policy and new token estimation self.init_schedule_policy() @@ -379,6 +371,9 @@ class Scheduler( # Init request dispatcher self.init_request_dispatcher() + # Init the grammar backend for constrained generation + self.grammar_manager = GrammarManager(self) + self.is_initializing = False def init_model_config(self): @@ -756,18 +751,6 @@ class Scheduler( ) self.enable_dynamic_chunking = False - def init_grammar_backend(self): - self.grammar_queue: List[Req] = [] - if not self.server_args.skip_tokenizer_init: - self.grammar_backend = create_grammar_backend( - self.server_args, - self.tokenizer, - self.model_config.vocab_size, - self.model_config.hf_eos_token_id, - ) - else: - self.grammar_backend = None - def init_schedule_policy(self): # Init schedule policy and new token estimation self.policy = SchedulePolicy( @@ -1554,43 +1537,8 @@ class Scheduler( self._add_request_to_queue(req) return - # 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 value is INVALID_GRAMMAR_OBJ: # We hit a cached invalid grammar. - error_msg = f"Invalid grammar request with cache hit: {key=}" - req.set_finish_with_abort(error_msg) - - if add_to_grammar_queue: - self.grammar_queue.append(req) - else: + added_to_grammar_queue = self.grammar_manager.process_req_with_grammar(req) + if not added_to_grammar_queue: self._add_request_to_queue(req) def handle_batch_generate_request( @@ -1889,8 +1837,10 @@ class Scheduler( self, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] ) -> Optional[ScheduleBatch]: # Check if the grammar is ready in the grammar queue - if self.grammar_queue: - self.move_ready_grammar_requests() + if self.grammar_manager.has_waiting_grammars(): + ready_grammar_requests = self.grammar_manager.get_ready_grammar_requests() + for req in ready_grammar_requests: + self._add_request_to_queue(req) if self.try_preemption: # Reset batch_is_full to try preemption with a prefill adder. @@ -2369,73 +2319,6 @@ class Scheduler( self.return_health_check_ct -= 1 self.send_to_tokenizer.send_output(HealthCheckOutput()) - def move_ready_grammar_requests(self): - """Move requests whose grammar objects are ready from grammar_queue to waiting_queue.""" - - num_ready_reqs = 0 - num_timeout_reqs = 0 - 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) - - num_ready_reqs += 1 - 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 - break - - if self.server_args.enable_dp_attention: - tp_size = self.attn_tp_size - tp_group = self.attn_tp_cpu_group - else: - tp_size = self.tp_size - tp_group = self.tp_cpu_group - - if tp_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 - ) - 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) - else: - num_ready_reqs_max = num_ready_reqs - num_timeout_reqs_max = num_timeout_reqs - - for i in range(num_ready_reqs, num_ready_reqs + num_timeout_reqs_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) - - num_ready_reqs = num_ready_reqs_max + num_timeout_reqs_max - - for req in self.grammar_queue[:num_ready_reqs]: - self._add_request_to_queue(req) - self.grammar_queue = self.grammar_queue[num_ready_reqs:] - def flush_cache_wrapped(self, recv_req: FlushCacheReqInput): success = self.flush_cache() return FlushCacheReqOutput(success=success) @@ -2478,8 +2361,7 @@ class Scheduler( self.tree_cache.reset() self.req_to_token_pool.clear() self.token_to_kv_pool_allocator.clear() - if self.grammar_backend: - self.grammar_backend.reset() + self.grammar_manager.clear() self.reset_metrics() if self.draft_worker: @@ -2616,15 +2498,10 @@ class Scheduler( logger.debug(f"Abort queued request. {req.rid=}") # Delete the requests in the grammar queue - for req in self.grammar_queue: - # Abort method 2: call `set_finish_with_abort` - # The request will still run one prefill forward pass. - # In this case, we change the input_ids to be only one token to make this prefill cheap. - if recv_req.abort_all or req.rid.startswith(recv_req.rid): - logger.debug(f"Abort grammar queue request. {req.rid=}") - if req.grammar: - req.grammar.cancel() - req.set_finish_with_abort("Aborted by AbortReq.") + # Abort method 2: call `set_finish_with_abort` + # The request will still run one prefill forward pass. + # In this case, we change the input_ids to be only one token to make this prefill cheap. + self.grammar_manager.abort_requests(recv_req) # Delete requests not in the waiting queue when PD disaggregation is enabled if self.disaggregation_mode == DisaggregationMode.PREFILL: diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index 9f663bc71..c943d1886 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -227,7 +227,7 @@ class SchedulerMetricsMixin: if self.is_hybrid_ssm: self.stats.mamba_usage = mamba_usage self.stats.num_queue_reqs = len(self.waiting_queue) - self.stats.num_grammar_queue_reqs = len(self.grammar_queue) + self.stats.num_grammar_queue_reqs = len(self.grammar_manager) self.stats.cache_hit_rate = cache_hit_rate self.stats.max_total_num_tokens = self.max_total_num_tokens @@ -393,7 +393,7 @@ class SchedulerMetricsMixin: self.stats.decode_sum_seq_lens = batch.seq_lens_cpu.sum().item() self.stats.gen_throughput = self.last_gen_throughput self.stats.num_queue_reqs = len(self.waiting_queue) - self.stats.num_grammar_queue_reqs = len(self.grammar_queue) + self.stats.num_grammar_queue_reqs = len(self.grammar_manager) self.stats.cache_hit_rate = cache_hit_rate self.stats.max_total_num_tokens = self.max_total_num_tokens diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index ac11e4ff8..36b17556c 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -288,7 +288,7 @@ class SchedulerRuntimeCheckerMixin: self.stats.token_usage = round(token_usage, 2) self.stats.gen_throughput = 0 self.stats.num_queue_reqs = len(self.waiting_queue) - self.stats.num_grammar_queue_reqs = len(self.grammar_queue) + self.stats.num_grammar_queue_reqs = len(self.grammar_manager) if self.disaggregation_mode == DisaggregationMode.PREFILL: self.stats.num_prefill_prealloc_queue_reqs = len( self.disagg_prefill_bootstrap_queue.queue