199 lines
7.8 KiB
Python
199 lines
7.8 KiB
Python
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 (
|
|
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
|
|
|
|
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
|
|
|
|
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.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 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 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.
|
|
|
|
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)
|
|
"""
|
|
assert self.grammar_backend
|
|
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
|
|
|
|
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:
|
|
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)
|
|
|
|
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=}"
|
|
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
|