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

[Minor] fix type annotations and invalid method calls in constrained … (#20132)

Browse Source
This commit is contained in:
zzhpro
2026-03-13 07:42:46 +08:00
committed by GitHub
parent 78a467c74a
commit c21ddbc785
6 changed files with 35 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
python/sglang/srt/constrained/base_grammar_backend.py
View File

@@ -15,9 +15,8 @@
import logging
import time
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass, field
from threading import Event
from typing import Dict, List, Optional, Tuple
import torch
@@ -120,43 +119,36 @@ class BaseGrammarObject:
INVALID_GRAMMAR_OBJ = BaseGrammarObject()
@dataclass
class CacheEntry:
value: BaseGrammarObject
event: Event
class BaseGrammarBackend:
def __init__(self):
self.executor = ThreadPoolExecutor()
self.cache: Dict[Tuple[str, str], CacheEntry] = {}
self.cache: Dict[Tuple[str, str], BaseGrammarObject] = {}
def _not_supported(self, key_type: str, key_string: str) -> None:
def _not_supported(self, key_type: str, key_string: str) -> BaseGrammarObject:
logger.warning(f"Skip unsupported {key_type=}, {key_string=}")
return INVALID_GRAMMAR_OBJ
def dispatch_fallback(
self, key_type: str, key_string: str
) -> Optional[BaseGrammarObject]:
def dispatch_fallback(self, key_type: str, key_string: str) -> BaseGrammarObject:
"""
This function should not be reached in any case.
"""
raise ValueError(f"Invalid key_type: {key_type}={key_string}")
def dispatch_json(self, key_string: str) -> Optional[BaseGrammarObject]:
def dispatch_json(self, key_string: str) -> BaseGrammarObject:
return self._not_supported("json", key_string)
def dispatch_regex(self, key_string: str) -> Optional[BaseGrammarObject]:
def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
return self._not_supported("regex", key_string)
def dispatch_ebnf(self, key_string: str) -> Optional[BaseGrammarObject]:
def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
return self._not_supported("ebnf", key_string)
def dispatch_structural_tag(self, key_string: str) -> Optional[BaseGrammarObject]:
def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject:
return self._not_supported("structural_tag", key_string)
def _init_value_dispatch(
self, key: Tuple[str, str], require_reasoning: bool
) -> Optional[BaseGrammarObject]:
) -> BaseGrammarObject:
s = time.perf_counter()
key_type, key_string = key
if key_type == "json":
@@ -167,10 +159,6 @@ class BaseGrammarBackend:
grammar = self.dispatch_ebnf(key_string)
elif key_type == "structural_tag":
grammar = self.dispatch_structural_tag(key_string)
elif key_type == "structural_pattern":
grammar = self.dispatch_structural_pattern(key_string)
elif key_type == "structural_pattern_v2":
grammar = self.dispatch_structural_pattern_v2(key_string)
else:
grammar = self.dispatch_fallback(key_type, key_string)
@@ -180,7 +168,7 @@ class BaseGrammarBackend:
def get_cached_or_future_value(
self, key: Tuple[str, str], require_reasoning: bool
) -> Optional[BaseGrammarObject]:
) -> Tuple[BaseGrammarObject | Future[BaseGrammarObject], bool]:
value = self.cache.get(key)
if value:
copied_value = value.copy()

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

@@ -60,7 +60,7 @@ class GrammarManager:
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:
if isinstance(req.grammar, futures.Future) and req.grammar:
req.grammar.cancel()
req.set_finish_with_abort("Aborted by AbortReq.")
@@ -115,6 +115,7 @@ class GrammarManager:
ready_reqs = intersect(ready_reqs_all)
failed_reqs = union(failed_reqs_all)
"""
assert self.grammar_backend
ready_req_idxs: set[int] = set()
failed_req_idxs: set[int] = set()
@@ -170,6 +171,7 @@ class GrammarManager:
if req.finished() or req.grammar is None: # It is aborted by AbortReq
continue
assert isinstance(req.grammar, futures.Future) and req.grammar_key
req.grammar = req.grammar.result()
self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy())
if req.grammar is INVALID_GRAMMAR_OBJ:
@@ -181,6 +183,7 @@ class GrammarManager:
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, INVALID_GRAMMAR_OBJ)
error_msg = f"Grammar preprocessing timed out: {req.grammar_key=}"

10
python/sglang/srt/constrained/llguidance_backend.py
View File

@@ -122,7 +122,7 @@ class GuidanceBackend(BaseGrammarBackend):
self.whitespace_pattern = whitespace_pattern
self.llguidance_tokenizer = from_tokenizer(self.tokenizer, n_vocab)
def _from_serialized(self, serialized_grammar) -> Optional[GuidanceGrammar]:
def _from_serialized(self, serialized_grammar) -> BaseGrammarObject:
try:
return GuidanceGrammar(
llguidance_tokenizer=self.llguidance_tokenizer,
@@ -132,7 +132,7 @@ class GuidanceBackend(BaseGrammarBackend):
logger.error(f"Hit invalid grammar: {serialized_grammar=}, {e=}")
return INVALID_GRAMMAR_OBJ
def dispatch_json(self, key_string: str) -> Optional[GuidanceGrammar]:
def dispatch_json(self, key_string: str) -> BaseGrammarObject:
try:
serialized_grammar = LLMatcher.grammar_from_json_schema(
key_string,
@@ -146,11 +146,11 @@ class GuidanceBackend(BaseGrammarBackend):
return INVALID_GRAMMAR_OBJ
return self._from_serialized(serialized_grammar)
def dispatch_regex(self, key_string: str) -> Optional[GuidanceGrammar]:
def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
serialized_grammar = grammar_from("regex", key_string)
return self._from_serialized(serialized_grammar)
def dispatch_ebnf(self, key_string: str) -> Optional[GuidanceGrammar]:
def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
try:
serialized_grammar = grammar_from("ebnf", key_string)
return self._from_serialized(serialized_grammar)
@@ -158,7 +158,7 @@ class GuidanceBackend(BaseGrammarBackend):
logger.error(f"Hit invalid ebnf: {key_string=}, {e=}")
return INVALID_GRAMMAR_OBJ
def dispatch_structural_tag(self, key_string: str) -> Optional[GuidanceGrammar]:
def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject:
try:
structural_tag = json.loads(key_string)
assert is_legacy_structural_tag(structural_tag)

2
python/sglang/srt/constrained/outlines_backend.py
View File

@@ -142,7 +142,7 @@ class OutlinesGrammarBackend(BaseGrammarBackend):
)
self.whitespace_pattern = whitespace_pattern
def _compile_regex(self, regex: str) -> Optional[OutlinesGrammar]:
def _compile_regex(self, regex: str) -> BaseGrammarObject:
try:
if hasattr(RegexGuide, "from_regex"):
# outlines >= 0.1.1

18
python/sglang/srt/constrained/xgrammar_backend.py
View File

@@ -113,8 +113,6 @@ class XGrammarGrammar(BaseGrammarObject):
apply_token_bitmask_inplace_cuda(logits, vocab_mask)
else:
apply_token_bitmask_inplace_triton(logits, vocab_mask)
elif logits.device.type == "cpu" and self.apply_vocab_mask_cpu:
self.apply_vocab_mask_cpu(logits, vocab_mask)
else:
raise RuntimeError(f"Unsupported device: {logits.device.type}")
@@ -124,15 +122,17 @@ class XGrammarGrammar(BaseGrammarObject):
max_rollback_tokens=MAX_ROLLBACK_TOKENS,
override_stop_tokens=self.override_stop_tokens,
)
if grammar_stats := self.grammar_stats:
grammar_stats = dataclasses.replace(
grammar_stats, is_cache_hit=True, tree_traversal_time=[]
)
return XGrammarGrammar(
matcher,
self.vocab_size,
self.ctx,
self.override_stop_tokens,
self.key_string,
dataclasses.replace(
self.grammar_stats, is_cache_hit=True, tree_traversal_time=[]
),
grammar_stats,
)
def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
@@ -254,7 +254,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
grammar_stats,
)
def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]:
def dispatch_json(self, key_string: str) -> BaseGrammarObject:
try:
if key_string == "$$ANY$$":
# Note: This builtin JSON grammar includes *all* valid JSON (including, for example, arrays at the root)
@@ -269,7 +269,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
return INVALID_GRAMMAR_OBJ
return self._from_context(ctx, key_string, GrammarStats(dispatch_type="json"))
def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]:
def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject:
try:
ctx = self.grammar_compiler.compile_grammar(key_string)
except RuntimeError as e:
@@ -277,7 +277,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
return INVALID_GRAMMAR_OBJ
return self._from_context(ctx, key_string, GrammarStats(dispatch_type="ebnf"))
def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]:
def dispatch_regex(self, key_string: str) -> BaseGrammarObject:
try:
ctx = self.grammar_compiler.compile_regex(key_string)
except RuntimeError as e:
@@ -285,7 +285,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
return INVALID_GRAMMAR_OBJ
return self._from_context(ctx, key_string, GrammarStats(dispatch_type="regex"))
def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]:
def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject:
try:
# TODO(dark): it's REALLY stupid to construct object from string and decode it again
structural_tag = json.loads(key_string)

7
python/sglang/srt/managers/schedule_batch.py
View File

@@ -39,6 +39,7 @@ import copy
import dataclasses
import logging
import re
from concurrent.futures import Future
from enum import Enum, auto
from functools import lru_cache
from http import HTTPStatus
@@ -725,8 +726,10 @@ class Req(ReqDllmMixin):
self.embedding = None
# Constrained decoding
self.grammar_key: Optional[str] = None
self.grammar: Optional[BaseGrammarObject] = None
self.grammar_key: Optional[Tuple[str, str]] = None
self.grammar: Optional[Union[BaseGrammarObject, Future[BaseGrammarObject]]] = (
None
)
self.grammar_wait_ct = 0
# The number of cached tokens that were already cached in the KV cache