minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

feat(grpc): add generation request semantics (#32588)

Browse Source 
Signed-off-by: Connor Carpenter <connorc@nvidia.com>
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
This commit is contained in:
Connor Carpenter
2026-08-04 18:53:45 -07:00
committed by GitHub
co-authored by ishandhanani Alex Nails
parent 29831d58ef
commit a0b04dbe4c
10 changed files with 982 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proto/sglang/runtime/v1/sglang.proto
+24 -2
View File
@@ -63,8 +63,24 @@ message SamplingParams {
repeated int32 stop_token_ids = 11;
optional bool ignore_eos = 12;
optional int32 n = 13;
optional string json_schema = 14;
optional string regex = 15;
optional string json_schema = 14 [deprecated = true];
optional string regex = 15 [deprecated = true];
optional int64 seed = 16;
optional GuidedDecoding guided_decoding = 17;
}
message GuidedDecoding {
oneof constraint {
string json_schema = 1;
string regex = 2;
string ebnf = 3;
ChoiceConstraint choice = 4;
string structural_tag = 5;
}
}
message ChoiceConstraint {
repeated string values = 1;
}
// ---- Text-based generate (text in, text out) ----
@@ -84,6 +100,9 @@ message TextGenerateRequest {
map<string, string> trace_headers = 12;
optional string session_id = 13;
optional DisaggregatedParams disaggregated_params = 14;
optional int32 priority = 15;
optional bool require_reasoning = 16;
optional uint32 max_thinking_tokens = 17;
}
message TextGenerateResponse {
@@ -108,6 +127,9 @@ message GenerateRequest {
map<string, string> trace_headers = 11;
optional string session_id = 12;
optional DisaggregatedParams disaggregated_params = 13;
optional int32 priority = 14;
optional bool require_reasoning = 15;
optional uint32 max_thinking_tokens = 16;
}
message GenerateResponse {
python/sglang/srt/entrypoints/grpc_bridge.py
+25 -5
View File
@@ -293,14 +293,23 @@ class RuntimeHandle:
async def _run_generate(self, obj, chunk_callback, stream: bool, request):
ready_event = None
gen = None
try:
ready_event = self._install_on_ready(chunk_callback) if stream else None
ready_event = self._install_on_ready(chunk_callback)
gen = self.tokenizer_manager.generate_request(obj, request=request)
if stream:
completed_choices = set()
expected_choices = obj.batch_size * obj.parallel_sample_num
async for chunk in gen:
finished = (
choice_finished = (
chunk.get("meta_info", {}).get("finish_reason") is not None
)
if choice_finished:
choice_id = chunk.get(
"index", chunk.get("meta_info", {}).get("id")
)
completed_choices.add(choice_id)
finished = len(completed_choices) >= expected_choices
keep_going = await self._send_with_backpressure(
chunk_callback,
ready_event,
@@ -314,15 +323,26 @@ class RuntimeHandle:
self._safe_callback(chunk_callback, {}, finished=True)
else:
result = await gen.__anext__()
self._safe_callback(chunk_callback, result, finished=True)
chunks = result if isinstance(result, list) else [result]
for index, chunk in enumerate(chunks):
keep_going = await self._send_with_backpressure(
chunk_callback,
ready_event,
chunk,
finished=index == len(chunks) - 1,
timeout_abort_rid=obj.rid,
)
if not keep_going:
return
except StopAsyncIteration:
self._safe_callback(chunk_callback, {}, finished=True)
except Exception as e:
logger.error("gRPC generate error for rid=%s: %s", obj.rid, e)
self._send_native_error(chunk_callback, str(e))
finally:
if stream:
self._uninstall_on_ready(chunk_callback)
if gen is not None:
await gen.aclose()
self._uninstall_on_ready(chunk_callback)
async def _run_embed(self, obj, chunk_callback, request):
try:
python/sglang/srt/entrypoints/http_server.py
+3 -2
View File
@@ -895,7 +895,7 @@ async def generate_request(obj: GenerateReqInput, request: Request):
"error": {
"message": str(e),
"type": "invalid_request_error",
"code": 400,
"code": getattr(e, "status_code", 400),
"retryable": False,
}
}
@@ -2048,7 +2048,8 @@ async def vertex_generate(
def _create_error_response(e):
return ORJSONResponse(
{"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST
{"error": {"message": str(e)}},
status_code=getattr(e, "status_code", HTTPStatus.BAD_REQUEST),
)
python/sglang/srt/managers/io_struct.py
+26 -14
View File
@@ -158,8 +158,9 @@ MultimodalDataInputFormat = Union[
@dataclass
class GenerateReqInput:
# Request ID(s). If omitted, generated during normalization. For batch
# requests, a string is expanded to per-item IDs using it as a prefix.
# Logical request ID(s). If omitted, generated during normalization. For
# batch requests, a string is expanded to one ID per original batch item.
# Parallel-sampling child IDs are internal to TokenizerManager.
rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True)
# Stable identity shared by requests in the same session. Unlike
# session_params, this does not alter or reconstruct the prompt.
@@ -276,6 +277,9 @@ class GenerateReqInput:
background: bool = False
# Require reasoning for the request (hybrid reasoning model only)
require_reasoning: bool = False
# Per-request thinking budget. Requires strict thinking so the runtime can
# enforce the limit rather than silently treating it as metadata.
max_thinking_tokens: Optional[int] = None
# Priority for the request
priority: Optional[int] = None
@@ -319,12 +323,17 @@ class GenerateReqInput:
# Batch-level: List[List[int]] (one per request). After __getitem__: List[int].
multi_item_delimiter_indices: Optional[Union[List[List[int]], List[int]]] = None
def regenerate_rid(self):
def regenerate_rid(self, prefix: Optional[str] = None):
"""Generate a new request ID and return it."""
def new_rid() -> str:
suffix = uuid.uuid4().hex
return f"{prefix}_{suffix}" if prefix is not None else suffix
if isinstance(self.rid, list):
self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))]
self.rid = [new_rid() for _ in range(len(self.rid))]
else:
self.rid = uuid.uuid4().hex
self.rid = new_rid()
return self.rid
def _validate_rid_uniqueness(self):
@@ -480,7 +489,7 @@ class GenerateReqInput:
# Expand input based on type
self._expand_inputs(num)
self._normalize_rid(num)
self._normalize_rid()
self._normalize_lora_paths(num)
self._normalize_image_data(num)
self._normalize_video_data(num)
@@ -590,16 +599,16 @@ class GenerateReqInput:
else: # Already a list
self.sampling_params = self.sampling_params * self.parallel_sample_num
def _normalize_rid(self, num):
"""Normalize request IDs for batch processing."""
def _normalize_rid(self):
"""Normalize one logical request ID per original batch item."""
if self.rid is None:
self.rid = [uuid.uuid4().hex for _ in range(num)]
self.rid = [uuid.uuid4().hex for _ in range(self.batch_size)]
elif isinstance(self.rid, str):
new_rids = [f"{self.rid}_{i}" for i in range(num)]
self.rid = new_rids
if self.batch_size == 1:
self.rid = [self.rid]
else:
self.rid = [f"{self.rid}_{i}" for i in range(self.batch_size)]
elif isinstance(self.rid, list):
# Note: the length of rid shall be the same as the batch_size,
# as the rid would be expanded for parallel sampling in tokenizer_manager
if len(self.rid) != self.batch_size:
raise ValueError(
"The specified rids length mismatch with the batch_size for batch processing."
@@ -751,8 +760,9 @@ class GenerateReqInput:
cache = self.__dict__.setdefault("_sub_obj_cache", {})
if i in cache:
return cache[i]
logical_index = i % self.batch_size
sub = GenerateReqInput(
rid=self.rid[i],
rid=self.rid[logical_index],
session_id=self.session_id,
text=self.text[i] if self.text is not None else None,
input_ids=self.input_ids[i] if self.input_ids is not None else None,
@@ -813,6 +823,8 @@ class GenerateReqInput:
disagg_prefill_dp_rank=self.disagg_prefill_dp_rank,
conversation_id=self.conversation_id,
http_worker_ipc=self.http_worker_ipc,
require_reasoning=self.require_reasoning,
max_thinking_tokens=self.max_thinking_tokens,
priority=self.priority,
extra_key=self.extra_key[i] if self.extra_key is not None else None,
no_logs=self.no_logs,
python/sglang/srt/managers/tokenizer_manager.py
+266 -37
View File
@@ -195,6 +195,10 @@ _INCREMENTAL_STREAMING_META_INFO_KEYS = (
)
class RequestAbortedError(ValueError):
status_code = 499
@dataclasses.dataclass
class ReqState:
"""Store the state a request."""
@@ -206,6 +210,9 @@ class ReqState:
# For performance metrics
time_stats: APIServerReqTimeStats
abort_requested: bool = False
lifecycle_id: object = dataclasses.field(default_factory=object)
dispatched: bool = False
last_completion_tokens: int = 1
ttft_observed: bool = False
@@ -545,6 +552,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
def init_running_status(self):
# Request states
self.rid_to_state: Dict[str, ReqState] = {}
# Parallel sampling keeps one caller-visible logical RID per original
# prompt while the scheduler operates on separate prefix/sample RIDs.
self.logical_rid_to_child_rids: Dict[str, set[str]] = {}
self.child_rid_to_logical_rid: Dict[str, str] = {}
self.event_loop = None
self.asyncio_tasks = set()
@@ -740,6 +751,15 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# Normalize the request
obj.normalize_batch_and_arguments()
self._set_default_priority(obj)
if (
isinstance(obj, GenerateReqInput)
and obj.max_thinking_tokens is not None
and not self.server_args.enable_strict_thinking
):
raise ValueError(
"max_thinking_tokens requires the server to be launched with "
"--enable-strict-thinking"
)
if isinstance(obj, GenerateReqInput) and obj.routed_dp_rank is not None:
dp_size = self.elastic_worker_count
@@ -752,7 +772,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
f"routed_dp_rank={obj.routed_dp_rank} out of range [0, {dp_size})"
)
self._init_req_state(obj, request)
request_lifecycles = self._init_req_state(obj, request)
try:
if self.server_args.language_only:
self._handle_epd_disaggregation_encode_request(obj)
@@ -762,13 +782,16 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
async with self.is_pause_cond:
await self.is_pause_cond.wait_for(lambda: not self.is_pause)
self._raise_if_logical_request_aborted(obj)
async with self.model_update_lock.reader_lock:
await self._validate_and_resolve_lora(obj)
self._raise_if_logical_request_aborted(obj)
# Tokenize the request and send it to the scheduler
if obj.is_single:
tokenized_obj = await self._tokenize_one_request(obj)
self._raise_if_logical_rid_aborted(obj.rid)
state = self.rid_to_state[obj.rid]
if obj.return_prompt_token_ids:
state.prompt_token_ids = list(tokenized_obj.input_ids)
@@ -778,7 +801,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
else:
async for response in self._handle_batch_request(obj, request):
yield response
except Exception:
except BaseException:
# _init_req_state created a rid_to_state entry per (sub-)request up
# front. The normal remover is the scheduler-response path
# (_handle_batch_output), so a failure *before* a request reaches the
@@ -786,7 +809,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# request -- would otherwise leak those entries forever. Drop any that
# are still pending; entries already removed on the normal completion
# path are left untouched (pop is a no-op).
self._discard_pending_req_states(obj)
self._discard_pending_req_states(obj, request_lifecycles)
raise
def _detect_input_format(
@@ -1308,6 +1331,11 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
sampling_kwargs = {**self.preferred_sampling_params, **obj.sampling_params}
else:
sampling_kwargs = obj.sampling_params
if isinstance(obj, GenerateReqInput) and obj.max_thinking_tokens is not None:
sampling_kwargs = dict(sampling_kwargs)
custom_params = dict(sampling_kwargs.get("custom_params") or {})
custom_params["thinking_budget"] = obj.max_thinking_tokens
sampling_kwargs["custom_params"] = custom_params
sampling_params = self.sampling_params_class(**sampling_kwargs)
sampling_params.normalize(self.tokenizer)
sampling_params.verify(self.model_config.vocab_size)
@@ -1518,6 +1546,9 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
time_stats = tokenized_obj.time_stats
tokenized_obj.wrap_pickle_fields()
self._dispatch_to_scheduler(tokenized_obj)
state = self.rid_to_state.get(tokenized_obj.rid)
if state is not None:
state.dispatched = True
tokenized_obj.time_stats = time_stats
tokenized_obj.time_stats.set_api_server_dispatch_finish_time()
@@ -1539,6 +1570,10 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
batch_req = BatchTokenizedEmbeddingReqInput(batch=tokenized_objs)
self._dispatch_to_scheduler(batch_req)
for tokenized_obj in tokenized_objs:
state = self.rid_to_state.get(tokenized_obj.rid)
if state is not None:
state.dispatched = True
for tokenized_obj, time_stat in zip(tokenized_objs, time_stats):
tokenized_obj.time_stats = time_stat
set_time_batch(tokenized_objs, "set_api_server_dispatch_finish_time")
@@ -1610,7 +1645,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# Delete the key to prevent resending abort request to the scheduler and
# to ensure aborted request state is cleaned up.
if state.obj.rid in self.rid_to_state:
del self.rid_to_state[state.obj.rid]
self._remove_req_state(state.obj.rid)
# Mark ongoing LoRA request as finished.
if self.enable_lora and state.obj.lora_path:
@@ -1744,6 +1779,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
if getattr(obj, "parallel_sample_num", 1) == 1:
if self._should_use_batch_tokenization(batch_size, obj):
tokenized_objs = await self._batch_tokenize_and_process(batch_size, obj)
self._raise_if_logical_request_aborted(obj)
self._send_batch_request(tokenized_objs)
# Set up generators for each request in the batch
@@ -1766,6 +1802,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
for i in range(batch_size):
tmp_obj = obj[i]
tokenized_obj = await self._tokenize_one_request(tmp_obj)
self._raise_if_logical_rid_aborted(tmp_obj.rid)
state = self.rid_to_state[tmp_obj.rid]
if tmp_obj.return_prompt_token_ids:
state.prompt_token_ids = list(tokenized_obj.input_ids)
@@ -1786,9 +1823,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_objs = await asyncio.gather(
*(self._tokenize_one_request(obj) for obj in objs)
)
self._raise_if_logical_request_aborted(obj)
# Cache the common prefix for parallel sampling
for i in range(batch_size):
logical_rid = objs[i].rid
self._raise_if_logical_rid_aborted(logical_rid)
tmp_obj = copy.copy(objs[i])
tokenized_obj = copy.copy(tokenized_objs[i])
# Ensure independent mm_items so wrap_shm_features won't mutate the original
@@ -1797,17 +1837,20 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_obj.mm_inputs.mm_items = [
copy.copy(item) for item in tokenized_obj.mm_inputs.mm_items
]
tokenized_obj.rid = tmp_obj.regenerate_rid()
tokenized_obj.rid = tmp_obj.regenerate_rid(prefix=logical_rid)
tokenized_obj.sampling_params = copy.copy(tokenized_obj.sampling_params)
tokenized_obj.sampling_params.max_new_tokens = 0
tokenized_obj.stream = False
self._init_req_state(tmp_obj)
self._init_child_req_state(logical_rid, tmp_obj)
self._send_one_request(tokenized_obj)
await self._wait_one_response(tmp_obj, request).__anext__()
self._raise_if_logical_rid_aborted(logical_rid)
# Expand requests, assign new rids for them, and send them
for i in range(batch_size):
logical_rid = objs[i].rid
for _ in range(obj.parallel_sample_num):
self._raise_if_logical_rid_aborted(logical_rid)
tmp_obj = copy.copy(objs[i])
tokenized_obj = copy.copy(tokenized_objs[i])
# Ensure independent mm_items so wrap_shm_features won't mutate the original
@@ -1816,8 +1859,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
tokenized_obj.mm_inputs.mm_items = [
copy.copy(item) for item in tokenized_obj.mm_inputs.mm_items
]
tokenized_obj.rid = tmp_obj.regenerate_rid()
self._init_req_state(tmp_obj)
tokenized_obj.rid = tmp_obj.regenerate_rid(prefix=logical_rid)
self._init_child_req_state(logical_rid, tmp_obj)
state = self.rid_to_state[tmp_obj.rid]
tokenized_obj.time_stats = state.time_stats
if tmp_obj.return_prompt_token_ids:
@@ -1826,17 +1869,38 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
generators.append(self._wait_one_response(tmp_obj, request))
rids.append(tmp_obj.rid)
self.rid_to_state[objs[i].rid].time_stats.set_finished_time()
del self.rid_to_state[objs[i].rid]
parent_state = self.rid_to_state.get(logical_rid)
if parent_state is not None:
parent_state.time_stats.set_finished_time()
self._remove_req_state(logical_rid)
# Wait for all requests
is_stream = hasattr(obj, "stream") and obj.stream
if not is_stream:
outputs = await asyncio.gather(*(gen.__anext__() for gen in generators))
outputs = await self._collect_batch_responses(generators)
yield outputs
else:
rid_to_index = {rid: i for i, rid in enumerate(rids)}
task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators}
async for response in self._stream_batch_responses(generators, rids):
yield response
async def _collect_batch_responses(self, generators):
tasks = [asyncio.create_task(gen.__anext__()) for gen in generators]
try:
return await asyncio.gather(*tasks)
finally:
for task in tasks:
if not task.done():
task.cancel()
await asyncio.gather(*tasks, return_exceptions=True)
await asyncio.gather(
*(gen.aclose() for gen in generators),
return_exceptions=True,
)
async def _stream_batch_responses(self, generators, rids):
rid_to_index = {rid: i for i, rid in enumerate(rids)}
task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators}
try:
while task_map:
done, _ = await asyncio.wait(
task_map.keys(), return_when=asyncio.FIRST_COMPLETED
@@ -1852,20 +1916,55 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
task_map[new_task] = gen
except StopAsyncIteration:
pass
finally:
pending_tasks = list(task_map)
for task in pending_tasks:
task.cancel()
if pending_tasks:
await asyncio.gather(*pending_tasks, return_exceptions=True)
await asyncio.gather(
*(gen.aclose() for gen in generators),
return_exceptions=True,
)
def abort_request(self, rid: str = "", abort_all: bool = False):
# Empty rid would startswith-match every request on the scheduler.
if not abort_all and not rid:
logger.warning("Ignore abort_request with empty rid and abort_all=False")
return
if (
not abort_all
and self.server_args.tokenizer_worker_num == 1
and rid not in self.rid_to_state
):
if abort_all:
for state_rid, state in self.rid_to_state.items():
if state_rid not in self.child_rid_to_logical_rid:
state.abort_requested = True
target_rids = (rid,)
elif rid in self.child_rid_to_logical_rid:
# Preserve direct child aborts for internal callers.
target_rids = (rid,)
elif rid in self.rid_to_state:
state = self.rid_to_state[rid]
state.abort_requested = True
parallel_sample_num = getattr(state.obj, "parallel_sample_num", None)
if parallel_sample_num is None:
sampling_params = getattr(state.obj, "sampling_params", None)
parallel_sample_num = (
sampling_params.get("n", 1)
if isinstance(sampling_params, dict)
else 1
)
if parallel_sample_num > 1:
# Snapshot because scheduler abort echoes remove child ownership.
target_rids = tuple(sorted(self.logical_rid_to_child_rids.get(rid, ())))
else:
target_rids = (rid,)
elif child_rids := self.logical_rid_to_child_rids.get(rid):
target_rids = tuple(sorted(child_rids))
elif self.server_args.tokenizer_worker_num == 1:
return
req = AbortReq(rid=rid, abort_all=abort_all)
self._dispatch_to_scheduler(req)
else:
target_rids = (rid,)
for target_rid in target_rids:
self._dispatch_to_scheduler(AbortReq(rid=target_rid, abort_all=abort_all))
if self.enable_metrics:
# TODO: also use custom_labels from the request
self.metrics_collector.observe_one_aborted_request(
@@ -2352,7 +2451,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
)
)
del self.rid_to_state[rid]
self._remove_req_state(rid)
# Mark ongoing LoRA request as finished.
if self.enable_lora and state.obj.lora_path:
@@ -3088,7 +3187,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
"output_ids": output_ids,
"meta_info": meta_info,
}
del self.rid_to_state[recv_obj.rid]
self._remove_req_state(recv_obj.rid)
state.out_list.append(out)
state.event.set()
@@ -3244,11 +3343,80 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
obj.lora_id[i] if isinstance(obj.lora_id, list) else obj.lora_id
)
@staticmethod
def _logical_rids(obj) -> List[str]:
if not hasattr(obj, "is_single") or obj.is_single:
return [obj.rid]
return list(obj.rid)
def _register_child_rid(self, logical_rid: str, child_rid: str) -> None:
if child_rid == logical_rid:
raise ValueError(
"Parallel-sampling child RID must differ from its logical RID"
)
owner = self.child_rid_to_logical_rid.get(child_rid)
if owner is not None and owner != logical_rid:
raise ValueError(
f"Request ID {child_rid} is already owned by logical request {owner}"
)
self.child_rid_to_logical_rid[child_rid] = logical_rid
self.logical_rid_to_child_rids.setdefault(logical_rid, set()).add(child_rid)
def _init_child_req_state(
self,
logical_rid: str,
obj: Union[GenerateReqInput, EmbeddingReqInput],
request: Optional[fastapi.Request] = None,
) -> None:
self._raise_if_logical_rid_aborted(logical_rid)
logical_state = self.rid_to_state[logical_rid]
self._init_req_state(
obj,
request,
lifecycle_id=logical_state.lifecycle_id,
)
try:
self._register_child_rid(logical_rid, obj.rid)
except BaseException:
self._remove_req_state(obj.rid)
raise
def _remove_req_state(
self,
rid: str,
lifecycle_id: Optional[object] = None,
) -> Optional[ReqState]:
"""Remove a request state and its parallel-sampling ownership."""
state = self.rid_to_state.get(rid)
if state is None or (
lifecycle_id is not None and state.lifecycle_id is not lifecycle_id
):
return None
self.rid_to_state.pop(rid)
logical_rid = self.child_rid_to_logical_rid.pop(rid, None)
if logical_rid is not None:
children = self.logical_rid_to_child_rids.get(logical_rid)
if children is not None:
children.discard(rid)
if not children:
self.logical_rid_to_child_rids.pop(logical_rid, None)
return state
def _raise_if_logical_rid_aborted(self, logical_rid: str) -> None:
state = self.rid_to_state.get(logical_rid)
if state is None or state.abort_requested:
raise RequestAbortedError(f"Request {logical_rid} was aborted")
def _raise_if_logical_request_aborted(self, obj) -> None:
for logical_rid in self._logical_rids(obj):
self._raise_if_logical_rid_aborted(logical_rid)
def _init_req_state(
self,
obj: Union[GenerateReqInput, EmbeddingReqInput],
request: Optional[fastapi.Request] = None,
):
lifecycle_id: Optional[object] = None,
) -> Dict[str, object]:
created_time = obj.received_time
external_trace_header = None
@@ -3279,29 +3447,90 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
for i in range(len(obj.rid))
]
for rid, sub_obj, bootstrap_room in items:
if rid in self.rid_to_state:
rids = [rid for rid, _, _ in items]
seen_rids = set()
for rid in rids:
if rid in seen_rids:
raise ValueError(f"Duplicate request ID detected: {rid}")
seen_rids.add(rid)
if (
rid in self.rid_to_state
or rid in self.logical_rid_to_child_rids
or rid in self.child_rid_to_logical_rid
):
raise ValueError(f"Duplicate request ID detected: {rid}")
# Mutate only after every RID passes duplicate validation so a rejected
# batch cannot leave a partial rid_to_state insertion behind.
lifecycle_ids = {}
for rid, sub_obj, bootstrap_room in items:
time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode)
state = ReqState([], False, asyncio.Event(), sub_obj, time_stats)
state = ReqState(
[],
False,
asyncio.Event(),
sub_obj,
time_stats,
lifecycle_id=lifecycle_id if lifecycle_id is not None else object(),
)
self.rid_to_state[rid] = state
lifecycle_ids[rid] = state.lifecycle_id
if self.enable_trace:
time_stats.init_trace_ctx(rid, bootstrap_room, external_trace_header)
time_stats.set_created_time(created_time)
return lifecycle_ids
def _discard_pending_req_states(self, obj):
"""Drop rid_to_state entries created by _init_req_state for *obj*.
def _discard_pending_req_states(
self,
obj,
lifecycle_ids: Optional[Dict[str, object]] = None,
):
"""Drop all logical and child state owned by *obj*.
Safe to call after a partial/failed dispatch: only entries still present
are removed, and the scheduler-response path looks up state with
``.get(...)`` so a later output for a discarded rid is ignored, not fatal.
Safe to call after a partial/failed dispatch: only requests known to have
reached the scheduler are aborted, all owned state is removed, and a later
output for a discarded RID is ignored by the scheduler-response path.
"""
if not hasattr(obj, "is_single") or obj.is_single:
rids = [obj.rid]
else:
rids = obj.rid
for rid in rids:
self.rid_to_state.pop(rid, None)
if lifecycle_ids is None:
lifecycle_ids = {
logical_rid: state.lifecycle_id
for logical_rid in self._logical_rids(obj)
if (state := self.rid_to_state.get(logical_rid)) is not None
}
for logical_rid in self._logical_rids(obj):
lifecycle_id = lifecycle_ids.get(logical_rid)
if lifecycle_id is None:
continue
child_rids = tuple(
child_rid
for child_rid in self.logical_rid_to_child_rids.get(logical_rid, ())
if (
(state := self.rid_to_state.get(child_rid)) is not None
and state.lifecycle_id is lifecycle_id
)
)
logical_state = self.rid_to_state.get(logical_rid)
owns_logical_state = (
logical_state is not None and logical_state.lifecycle_id is lifecycle_id
)
target_rids = tuple(
rid for rid in child_rids if self.rid_to_state[rid].dispatched
)
if not child_rids and owns_logical_state and logical_state.dispatched:
target_rids = (logical_rid,)
for target_rid in target_rids:
try:
self._dispatch_to_scheduler(
AbortReq(rid=target_rid, abort_all=False)
)
except Exception:
logger.exception(
"Failed to abort request rid=%s",
target_rid,
)
for child_rid in child_rids:
self._remove_req_state(child_rid, lifecycle_id)
self._remove_req_state(logical_rid, lifecycle_id)
def _should_dispatch_to_encoder(
self, obj: Union[GenerateReqInput, EmbeddingReqInput]
rust/sglang-grpc/src/server.rs
+2 -2
View File
@@ -229,7 +229,7 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl {
.rid
.clone()
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let req_dict = build_text_generate_dict(&rid, &req);
let req_dict = build_text_generate_dict(&rid, &req).map_err(Status::invalid_argument)?;
let mut receiver = self
.bridge
@@ -298,7 +298,7 @@ impl proto::sglang_service_server::SglangService for SglangServiceImpl {
.rid
.clone()
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let req_dict = build_generate_dict(&rid, &req);
let req_dict = build_generate_dict(&rid, &req).map_err(Status::invalid_argument)?;
let mut receiver = self
.bridge
rust/sglang-grpc/src/utils/request_utils.rs
+235 -17
View File
@@ -2,8 +2,25 @@ use std::collections::HashMap;
use crate::proto;
fn regex_escape_literal(value: &str) -> String {
let mut escaped = String::with_capacity(value.len());
for character in value.chars() {
if matches!(
character,
'.' | '+' | '*' | '?' | '^' | '$' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '\\'
) {
escaped.push('\\');
}
escaped.push(character);
}
escaped
}
/// Convert proto SamplingParams to a serde_json map (used as Python dict via PyO3).
fn sampling_params_to_map(params: &Option<proto::SamplingParams>) -> serde_json::Value {
#[allow(deprecated)]
fn sampling_params_to_map(
params: &Option<proto::SamplingParams>,
) -> Result<serde_json::Value, String> {
match params {
Some(p) => {
let mut map = serde_json::Map::new();
@@ -46,15 +63,90 @@ fn sampling_params_to_map(params: &Option<proto::SamplingParams>) -> serde_json:
if let Some(v) = p.n {
map.insert("n".into(), serde_json::json!(v));
}
if let Some(ref v) = p.json_schema {
map.insert("json_schema".into(), serde_json::json!(v));
if let Some(v) = p.seed {
map.insert("sampling_seed".into(), serde_json::json!(v));
}
if let Some(ref v) = p.regex {
map.insert("regex".into(), serde_json::json!(v));
if p.guided_decoding.is_some() && (p.json_schema.is_some() || p.regex.is_some()) {
return Err(
"legacy json_schema/regex cannot be combined with guided_decoding".into(),
);
}
serde_json::Value::Object(map)
if let Some(guided) = p.guided_decoding.as_ref() {
use proto::guided_decoding::Constraint;
match guided.constraint.as_ref() {
Some(Constraint::JsonSchema(value)) if !value.is_empty() => {
map.insert("json_schema".into(), serde_json::json!(value));
}
Some(Constraint::Regex(value)) if !value.is_empty() => {
map.insert("regex".into(), serde_json::json!(value));
}
Some(Constraint::Ebnf(value)) if !value.is_empty() => {
map.insert("ebnf".into(), serde_json::json!(value));
}
Some(Constraint::Choice(choice))
if !choice.values.is_empty()
&& choice.values.iter().all(|value| !value.is_empty()) =>
{
let alternatives = choice
.values
.iter()
.map(|value| regex_escape_literal(value))
.collect::<Vec<_>>()
.join("|");
map.insert(
"regex".into(),
serde_json::json!(format!("(?:{alternatives})")),
);
}
Some(Constraint::StructuralTag(value)) if !value.is_empty() => {
map.insert("structural_tag".into(), serde_json::json!(value));
}
Some(Constraint::Choice(_)) => {
return Err("guided choice must contain only non-empty values".into());
}
Some(_) => return Err("guided decoding constraint must not be empty".into()),
None => return Err("guided decoding constraint must be specified".into()),
}
} else {
if let Some(value) = p.json_schema.as_ref() {
if value.is_empty() {
return Err("legacy json_schema must not be empty".into());
}
map.insert("json_schema".into(), serde_json::json!(value));
}
if let Some(value) = p.regex.as_ref() {
if value.is_empty() {
return Err("legacy regex must not be empty".into());
}
map.insert("regex".into(), serde_json::json!(value));
}
}
Ok(serde_json::Value::Object(map))
}
None => serde_json::Value::Object(serde_json::Map::new()),
None => Ok(serde_json::Value::Object(serde_json::Map::new())),
}
}
fn insert_generation_controls(
d: &mut HashMap<String, serde_json::Value>,
priority: Option<i32>,
require_reasoning: Option<bool>,
max_thinking_tokens: Option<u32>,
) {
if let Some(priority) = priority {
d.insert("priority".into(), serde_json::json!(priority));
}
if let Some(require_reasoning) = require_reasoning {
d.insert(
"require_reasoning".into(),
serde_json::json!(require_reasoning),
);
}
if let Some(max_thinking_tokens) = max_thinking_tokens {
d.insert(
"max_thinking_tokens".into(),
serde_json::json!(max_thinking_tokens),
);
}
}
@@ -111,13 +203,13 @@ pub(crate) fn extract_model_path(json_info: &str) -> String {
pub(crate) fn build_text_generate_dict(
rid: &str,
req: &proto::TextGenerateRequest,
) -> HashMap<String, serde_json::Value> {
) -> Result<HashMap<String, serde_json::Value>, String> {
let mut d = HashMap::new();
d.insert("rid".into(), serde_json::json!(rid));
d.insert("text".into(), serde_json::json!(req.text));
d.insert(
"sampling_params".into(),
sampling_params_to_map(&req.sampling_params),
sampling_params_to_map(&req.sampling_params)?,
);
d.insert(
"stream".into(),
@@ -151,25 +243,31 @@ pub(crate) fn build_text_generate_dict(
if let Some(ref session_id) = req.session_id {
d.insert("session_id".into(), serde_json::json!(session_id));
}
insert_generation_controls(
&mut d,
req.priority,
req.require_reasoning,
req.max_thinking_tokens,
);
insert_disaggregated_params(&mut d, &req.disaggregated_params);
if let Some(trace) = trace_headers_to_json(&req.trace_headers) {
d.insert("external_trace_header".into(), trace);
}
d.insert("received_time".into(), serde_json::json!(now_timestamp()));
d
Ok(d)
}
/// Build a request dict for GenerateReqInput from proto GenerateRequest (tokenized).
pub(crate) fn build_generate_dict(
rid: &str,
req: &proto::GenerateRequest,
) -> HashMap<String, serde_json::Value> {
) -> Result<HashMap<String, serde_json::Value>, String> {
let mut d = HashMap::new();
d.insert("rid".into(), serde_json::json!(rid));
d.insert("input_ids".into(), serde_json::json!(req.input_ids));
d.insert(
"sampling_params".into(),
sampling_params_to_map(&req.sampling_params),
sampling_params_to_map(&req.sampling_params)?,
);
d.insert(
"stream".into(),
@@ -199,12 +297,18 @@ pub(crate) fn build_generate_dict(
if let Some(ref session_id) = req.session_id {
d.insert("session_id".into(), serde_json::json!(session_id));
}
insert_generation_controls(
&mut d,
req.priority,
req.require_reasoning,
req.max_thinking_tokens,
);
insert_disaggregated_params(&mut d, &req.disaggregated_params);
if let Some(trace) = trace_headers_to_json(&req.trace_headers) {
d.insert("external_trace_header".into(), trace);
}
d.insert("received_time".into(), serde_json::json!(now_timestamp()));
d
Ok(d)
}
/// Build a request dict for EmbeddingReqInput from proto TextEmbedRequest.
@@ -267,6 +371,7 @@ pub(crate) fn build_classify_dict(
}
#[cfg(test)]
#[allow(deprecated)]
mod tests {
use super::*;
@@ -283,11 +388,15 @@ mod tests {
};
assert_eq!(
build_text_generate_dict("request-1", &text_req).get("session_id"),
build_text_generate_dict("request-1", &text_req)
.unwrap()
.get("session_id"),
Some(&serde_json::json!("session-1"))
);
assert_eq!(
build_generate_dict("request-2", &token_req).get("session_id"),
build_generate_dict("request-2", &token_req)
.unwrap()
.get("session_id"),
Some(&serde_json::json!("session-1"))
);
}
@@ -312,6 +421,7 @@ mod tests {
build_text_generate_dict("request-1", &text_req),
build_generate_dict("request-2", &token_req),
] {
let request = request.unwrap();
assert_eq!(
request.get("bootstrap_host"),
Some(&serde_json::json!("10.0.0.1"))
@@ -330,8 +440,9 @@ mod tests {
#[test]
fn generate_dicts_omit_disaggregated_params_when_absent() {
let text_request =
build_text_generate_dict("request-1", &proto::TextGenerateRequest::default());
let token_request = build_generate_dict("request-2", &proto::GenerateRequest::default());
build_text_generate_dict("request-1", &proto::TextGenerateRequest::default()).unwrap();
let token_request =
build_generate_dict("request-2", &proto::GenerateRequest::default()).unwrap();
for request in [text_request, token_request] {
assert!(!request.contains_key("bootstrap_host"));
@@ -339,4 +450,111 @@ mod tests {
assert!(!request.contains_key("bootstrap_room"));
}
}
#[test]
fn generate_dicts_preserve_optional_generation_controls() {
let sampling_params = proto::SamplingParams {
seed: Some(42),
..Default::default()
};
let text_request = proto::TextGenerateRequest {
sampling_params: Some(sampling_params.clone()),
priority: Some(3),
require_reasoning: Some(false),
max_thinking_tokens: Some(128),
..Default::default()
};
let token_request = proto::GenerateRequest {
sampling_params: Some(proto::SamplingParams {
seed: Some(42),
..Default::default()
}),
priority: Some(3),
require_reasoning: Some(false),
max_thinking_tokens: Some(128),
..Default::default()
};
for mapped in [
build_text_generate_dict("text-request", &text_request).unwrap(),
build_generate_dict("token-request", &token_request).unwrap(),
] {
assert_eq!(mapped["priority"], serde_json::json!(3));
assert_eq!(mapped["require_reasoning"], serde_json::json!(false));
assert_eq!(mapped["max_thinking_tokens"], serde_json::json!(128));
assert_eq!(
mapped["sampling_params"]["sampling_seed"],
serde_json::json!(42)
);
}
for mapped in [
build_text_generate_dict("text-request", &Default::default()).unwrap(),
build_generate_dict("token-request", &Default::default()).unwrap(),
] {
assert!(!mapped.contains_key("priority"));
assert!(!mapped.contains_key("require_reasoning"));
assert!(!mapped.contains_key("max_thinking_tokens"));
}
}
#[test]
fn guided_choice_maps_to_escaped_regex() {
let request = proto::GenerateRequest {
sampling_params: Some(proto::SamplingParams {
guided_decoding: Some(proto::GuidedDecoding {
constraint: Some(proto::guided_decoding::Constraint::Choice(
proto::ChoiceConstraint {
values: vec!["a+b".into(), "x.y".into()],
},
)),
}),
..Default::default()
}),
..Default::default()
};
let mapped = build_generate_dict("request", &request).unwrap();
assert_eq!(
mapped["sampling_params"]["regex"],
serde_json::json!("(?:a\\+b|x\\.y)")
);
}
#[test]
fn invalid_guidance_combinations_are_rejected() {
let conflicting = proto::GenerateRequest {
sampling_params: Some(proto::SamplingParams {
regex: Some("[a-z]+".into()),
guided_decoding: Some(proto::GuidedDecoding {
constraint: Some(proto::guided_decoding::Constraint::Regex("[0-9]+".into())),
}),
..Default::default()
}),
..Default::default()
};
let empty_choice = proto::GenerateRequest {
sampling_params: Some(proto::SamplingParams {
guided_decoding: Some(proto::GuidedDecoding {
constraint: Some(proto::guided_decoding::Constraint::Choice(
proto::ChoiceConstraint { values: vec![] },
)),
}),
..Default::default()
}),
..Default::default()
};
let empty_legacy_regex = proto::GenerateRequest {
sampling_params: Some(proto::SamplingParams {
regex: Some(String::new()),
..Default::default()
}),
..Default::default()
};
for request in [conflicting, empty_choice, empty_legacy_regex] {
assert!(build_generate_dict("request", &request).is_err());
}
}
}
test/registered/unit/entrypoints/test_grpc_bridge.py
+115
View File
@@ -0,0 +1,115 @@
import asyncio
import enum
import unittest
from types import SimpleNamespace
from sglang.srt.entrypoints.grpc_bridge import RuntimeHandle
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
class _ChunkStatus(enum.Enum):
Ready = 1
Pending = 2
Closed = 3
class _RecordingCallback:
def __init__(self):
self.calls = []
def __call__(self, payload, *, finished=False, error=None):
self.calls.append((payload, finished, error))
return _ChunkStatus.Ready
class _FakeTokenizerManager:
def __init__(self, responses):
self.responses = responses
def generate_request(self, obj, request=None):
async def generate():
for response in self.responses:
yield response
return generate()
def _make_runtime_handle(responses):
handle = RuntimeHandle.__new__(RuntimeHandle)
handle.tokenizer_manager = _FakeTokenizerManager(responses)
return handle
class TestNativeGrpcParallelResponses(CustomTestCase):
def test_non_streaming_returns_every_choice_before_finishing(self):
callback = _RecordingCallback()
responses = [
[
{"output_ids": [1], "meta_info": {"id": "choice-0"}},
{"output_ids": [2], "meta_info": {"id": "choice-1"}},
]
]
handle = _make_runtime_handle(responses)
obj = SimpleNamespace(rid="logical", batch_size=1, parallel_sample_num=2)
asyncio.run(
handle._run_generate(
obj,
callback,
stream=False,
request=None,
)
)
self.assertEqual([call[0]["output_ids"] for call in callback.calls], [[1], [2]])
self.assertEqual([call[1] for call in callback.calls], [False, True])
def test_streaming_first_finished_choice_is_not_batch_terminal(self):
callback = _RecordingCallback()
responses = [
{
"index": 0,
"output_ids": [1],
"meta_info": {"id": "choice-0", "finish_reason": None},
},
{
"index": 0,
"output_ids": [2],
"meta_info": {
"id": "choice-0",
"finish_reason": {"type": "stop"},
},
},
{
"index": 1,
"output_ids": [3],
"meta_info": {
"id": "choice-1",
"finish_reason": {"type": "stop"},
},
},
]
handle = _make_runtime_handle(responses)
obj = SimpleNamespace(rid="logical", batch_size=1, parallel_sample_num=2)
asyncio.run(
handle._run_generate(
obj,
callback,
stream=True,
request=None,
)
)
self.assertEqual(
[call[0]["output_ids"] for call in callback.calls],
[[1], [2], [3]],
)
self.assertEqual([call[1] for call in callback.calls], [False, False, True])
if __name__ == "__main__":
unittest.main(verbosity=2)
test/registered/unit/managers/test_io_struct.py
+41
View File
@@ -305,6 +305,38 @@ class TestGenerateReqInputNormalization(CustomTestCase):
# Modalities should be set for all 3 examples
self.assertEqual(req.modalities, ["image", "image", "image"])
def test_parallel_sampling_keeps_one_logical_rid_per_prompt(self):
"""Test logical RID and reasoning control preservation across parallel samples."""
single = GenerateReqInput(
text="Hello",
rid="single",
sampling_params={"n": 3},
require_reasoning=True,
max_thinking_tokens=128,
)
single.normalize_batch_and_arguments()
self.assertEqual(single.rid, ["single"])
self.assertEqual([single[i].rid for i in range(3)], ["single"] * 3)
self.assertTrue(all(single[i].require_reasoning for i in range(3)))
self.assertEqual(
[single[i].max_thinking_tokens for i in range(3)],
[128] * 3,
)
batch = GenerateReqInput(
text=["Hello", "World"],
rid="batch",
sampling_params={"n": 2},
)
batch.normalize_batch_and_arguments()
self.assertEqual(batch.rid, ["batch_0", "batch_1"])
self.assertEqual(
[batch[i].rid for i in range(4)],
["batch_0", "batch_1", "batch_0", "batch_1"],
)
def test_audio_data_handling(self):
"""Test handling of audio_data."""
req = copy.deepcopy(self.base_req)
@@ -648,6 +680,15 @@ class TestGenerateReqInputNormalization(CustomTestCase):
self.assertNotEqual(original_rid, new_rid)
self.assertEqual(req.rid, new_rid)
def test_regenerate_rid_with_parent_prefix(self):
"""Test RID regeneration with a logical parent prefix."""
req = GenerateReqInput(text="Hello", rid="logical")
req.normalize_batch_and_arguments()
new_rid = req.regenerate_rid(prefix="logical")
self.assertTrue(new_rid.startswith("logical_"))
def test_error_cases(self):
"""Test various error cases."""
# Test when neither text, input_ids, nor input_embeds is provided
test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
+245 -7
View File
@@ -23,9 +23,19 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.io_struct import AbortReq, BatchStrOutput, GenerateReqInput
from sglang.srt.managers.tokenizer_manager import ReqState, TokenizerManager
from sglang.srt.observability.req_time_stats import APIServerReqTimeStats
from sglang.srt.managers.io_struct import ( # noqa: E402
AbortReq,
BatchStrOutput,
GenerateReqInput,
)
from sglang.srt.managers.tokenizer_manager import ( # noqa: E402
ReqState,
RequestAbortedError,
TokenizerManager,
)
from sglang.srt.observability.req_time_stats import ( # noqa: E402
APIServerReqTimeStats,
)
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
@@ -111,6 +121,8 @@ def _make_tokenizer_manager() -> TokenizerManager:
tm.server_args.dp_size = 1
tm.disaggregation_mode = "none"
tm.rid_to_state = {}
tm.logical_rid_to_child_rids = {}
tm.child_rid_to_logical_rid = {}
tm.enable_metrics = False
tm.enable_trace = False
tm.enable_lora = False
@@ -120,10 +132,11 @@ def _make_tokenizer_manager() -> TokenizerManager:
tm.dump_requests_folder = ""
tm.crash_dump_folder = ""
tm.send_to_scheduler = MagicMock()
tm._dispatch_to_scheduler = Mock()
return tm
def _make_req_state(rid: str = "test_rid") -> ReqState:
def _make_req_state(rid: str = "test_rid", *, dispatched: bool = False) -> ReqState:
"""Create a minimal ReqState for testing."""
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
@@ -137,6 +150,7 @@ def _make_req_state(rid: str = "test_rid") -> ReqState:
event=asyncio.Event(),
obj=obj,
time_stats=APIServerReqTimeStats(),
dispatched=dispatched,
)
@@ -338,6 +352,19 @@ class TestInitReqStateDuplicateDetection(CustomTestCase):
tm._init_req_state(obj)
self.assertIn(rid, tm.rid_to_state)
def test_batch_duplicate_preflight_does_not_insert_partial_state(self):
tm = _make_tokenizer_manager()
existing_rid = "existing"
existing_state = _make_req_state(existing_rid)
tm.rid_to_state[existing_rid] = existing_state
obj = _make_generate_obj(["new", existing_rid], is_single=False)
with self.assertRaisesRegex(ValueError, "Duplicate request ID"):
tm._init_req_state(obj)
self.assertNotIn("new", tm.rid_to_state)
self.assertIs(tm.rid_to_state[existing_rid], existing_state)
class TestResubmitAfterCompletion(CustomTestCase):
"""End-to-end test: complete a request, then resubmit with the same rid."""
@@ -409,6 +436,7 @@ def _make_tm_for_generate() -> TokenizerManager:
tm = _make_tokenizer_manager()
tm.server_args.language_only = False
tm.server_args.tokenizer_worker_num = 1
tm.server_args.enable_strict_thinking = False
tm.auto_create_handle_loop = Mock()
tm._set_default_priority = Mock()
tm.request_logger = Mock()
@@ -429,6 +457,7 @@ def _make_generate_obj(rid, is_single):
obj.received_time = 0.0
obj.external_trace_header = None
obj.bootstrap_room = None
obj.max_thinking_tokens = None
obj.normalize_batch_and_arguments = Mock()
if not is_single:
obj.__getitem__.side_effect = lambda i: Mock()
@@ -438,17 +467,20 @@ def _make_generate_obj(rid, is_single):
class TestDiscardPendingReqStates(CustomTestCase):
"""Direct tests for _discard_pending_req_states."""
def test_discard_single(self):
def test_discard_single_aborts_scheduler_before_cleanup(self):
tm = _make_tokenizer_manager()
rid = "d_single"
tm.rid_to_state[rid] = _make_req_state(rid)
tm.rid_to_state[rid] = _make_req_state(rid, dispatched=True)
obj = Mock(spec=GenerateReqInput)
obj.is_single = True
obj.rid = rid
tm._discard_pending_req_states(obj)
self.assertNotIn(rid, tm.rid_to_state)
abort_req = tm._dispatch_to_scheduler.call_args.args[0]
self.assertEqual(abort_req.rid, rid)
self.assertFalse(abort_req.abort_all)
def test_discard_batch_removes_all(self):
def test_discard_unsent_batch_without_scheduler_abort(self):
tm = _make_tokenizer_manager()
rids = ["d0", "d1", "d2"]
for r in rids:
@@ -459,6 +491,7 @@ class TestDiscardPendingReqStates(CustomTestCase):
tm._discard_pending_req_states(obj)
for r in rids:
self.assertNotIn(r, tm.rid_to_state)
tm._dispatch_to_scheduler.assert_not_called()
def test_discard_ignores_already_removed(self):
"""Popping a rid that is no longer present must not raise."""
@@ -470,6 +503,150 @@ class TestDiscardPendingReqStates(CustomTestCase):
tm._discard_pending_req_states(obj) # must not raise
self.assertNotIn("p1", tm.rid_to_state)
def test_parallel_cleanup_aborts_children_and_allows_parent_reuse(self):
tm = _make_tokenizer_manager()
parent = _make_generate_obj("parent", is_single=True)
lifecycle_ids = tm._init_req_state(parent)
child_rids = {"prefix", "choice_0", "choice_1"}
for child_rid in child_rids:
child = _make_generate_obj(child_rid, is_single=True)
tm._init_child_req_state("parent", child)
tm.rid_to_state[child_rid].dispatched = True
tm._remove_req_state("parent")
tm._discard_pending_req_states(parent, lifecycle_ids)
aborted_rids = {
call.args[0].rid for call in tm._dispatch_to_scheduler.call_args_list
}
self.assertEqual(aborted_rids, child_rids)
self.assertFalse(tm.rid_to_state)
self.assertFalse(tm.logical_rid_to_child_rids)
self.assertFalse(tm.child_rid_to_logical_rid)
tm._init_req_state(_make_generate_obj("parent", is_single=True))
self.assertIn("parent", tm.rid_to_state)
def test_stale_cleanup_does_not_remove_reused_rid(self):
tm = _make_tokenizer_manager()
old_obj = _make_generate_obj("reused", is_single=True)
old_lifecycle_ids = tm._init_req_state(old_obj)
tm._remove_req_state("reused")
replacement = _make_generate_obj("reused", is_single=True)
tm._init_req_state(replacement)
replacement_state = tm.rid_to_state["reused"]
tm._discard_pending_req_states(old_obj, old_lifecycle_ids)
self.assertIs(tm.rid_to_state["reused"], replacement_state)
tm._dispatch_to_scheduler.assert_not_called()
class TestParallelAbortRouting(CustomTestCase):
def test_parent_abort_fans_out_to_children(self):
tm = _make_tokenizer_manager()
tm.server_args.tokenizer_worker_num = 1
tm._register_child_rid("parent", "choice_0")
tm._register_child_rid("parent", "choice_1")
tm.abort_request("parent")
requests = [call.args[0] for call in tm._dispatch_to_scheduler.call_args_list]
self.assertEqual(
{request.rid for request in requests}, {"choice_0", "choice_1"}
)
self.assertTrue(all(not request.abort_all for request in requests))
class TestParallelStreamTaskCleanup(CustomTestCase):
def test_failing_choice_cancels_and_closes_sibling_waiters(self):
tm = _make_tokenizer_manager()
async def drive():
sibling_closed = asyncio.Event()
async def failing_choice():
await asyncio.sleep(0)
raise RuntimeError("choice failed")
yield # pragma: no cover
async def blocked_choice():
try:
await asyncio.Event().wait()
yield # pragma: no cover
finally:
sibling_closed.set()
stream = tm._stream_batch_responses(
[failing_choice(), blocked_choice()],
["choice-0", "choice-1"],
)
with self.assertRaisesRegex(RuntimeError, "choice failed"):
await stream.__anext__()
self.assertTrue(sibling_closed.is_set())
asyncio.run(drive())
def test_failing_non_stream_choice_cancels_and_closes_sibling_waiters(self):
tm = _make_tokenizer_manager()
async def drive():
sibling_closed = asyncio.Event()
async def failing_choice():
await asyncio.sleep(0)
raise RuntimeError("choice failed")
yield # pragma: no cover
async def blocked_choice():
try:
await asyncio.Event().wait()
yield # pragma: no cover
finally:
sibling_closed.set()
with self.assertRaisesRegex(RuntimeError, "choice failed"):
await tm._collect_batch_responses([failing_choice(), blocked_choice()])
self.assertTrue(sibling_closed.is_set())
asyncio.run(drive())
class TestParallelRidReuse(CustomTestCase):
def test_completed_n2_request_can_repeat_the_same_logical_rid(self):
tm = _make_tokenizer_manager()
async def complete_child(rid):
await tm._handle_batch_output(_make_batch_str_output(rid))
for _ in range(2):
logical = GenerateReqInput(
text="hello",
rid="repeat-n2",
sampling_params={"n": 2},
)
logical.normalize_batch_and_arguments()
tm._init_req_state(logical)
prefix = GenerateReqInput(text="hello", rid="prefix")
prefix.normalize_batch_and_arguments()
tm._init_child_req_state("repeat-n2", prefix)
asyncio.run(complete_child("prefix"))
for child_rid in ("choice-0", "choice-1"):
child = GenerateReqInput(text="hello", rid=child_rid)
child.normalize_batch_and_arguments()
tm._init_child_req_state("repeat-n2", child)
tm._remove_req_state("repeat-n2")
asyncio.run(complete_child("choice-0"))
asyncio.run(complete_child("choice-1"))
self.assertFalse(tm.rid_to_state)
self.assertFalse(tm.logical_rid_to_child_rids)
self.assertFalse(tm.child_rid_to_logical_rid)
class TestGenerateRequestCleanupOnDispatchFailure(CustomTestCase):
"""generate_request must not leak rid_to_state when dispatch fails.
@@ -497,6 +674,7 @@ class TestGenerateRequestCleanupOnDispatchFailure(CustomTestCase):
# Got past _init_req_state (which created the entry) ...
tm._tokenize_one_request.assert_awaited_once()
tm._send_one_request.assert_not_called()
tm._dispatch_to_scheduler.assert_not_called()
# ... and the entry was cleaned up rather than leaked.
self.assertNotIn(rid, tm.rid_to_state)
@@ -521,6 +699,66 @@ class TestGenerateRequestCleanupOnDispatchFailure(CustomTestCase):
# All sub-request entries created by _init_req_state are cleaned up.
for r in rids:
self.assertNotIn(r, tm.rid_to_state)
tm._dispatch_to_scheduler.assert_not_called()
def test_interrupted_parallel_tokenization_prevents_child_dispatch(self):
for remove_state in (False, True):
with self.subTest(remove_state=remove_state):
tm = _make_tm_for_generate()
tm._send_one_request = Mock()
obj = GenerateReqInput(
text="hello",
rid="interrupted-during-tokenization",
sampling_params={"n": 2},
)
async def drive():
tokenization_started = asyncio.Event()
allow_tokenization = asyncio.Event()
async def blocked_tokenization(_obj):
tokenization_started.set()
await allow_tokenization.wait()
return MagicMock()
tm._tokenize_one_request = blocked_tokenization
response = tm.generate_request(obj)
task = asyncio.create_task(response.__anext__())
await tokenization_started.wait()
if remove_state:
tm._remove_req_state("interrupted-during-tokenization")
else:
tm.abort_request("interrupted-during-tokenization")
allow_tokenization.set()
with self.assertRaisesRegex(
RequestAbortedError, "interrupted-during-tokenization"
):
await task
asyncio.run(drive())
tm._send_one_request.assert_not_called()
tm._dispatch_to_scheduler.assert_not_called()
self.assertFalse(tm.rid_to_state)
self.assertFalse(tm.logical_rid_to_child_rids)
self.assertFalse(tm.child_rid_to_logical_rid)
def test_thinking_budget_rejects_runtime_without_strict_thinking(self):
tm = _make_tm_for_generate()
obj = GenerateReqInput(
text="hello",
rid="thinking-budget",
sampling_params={},
max_thinking_tokens=32,
)
async def drive():
await tm.generate_request(obj).__anext__()
with self.assertRaisesRegex(ValueError, "--enable-strict-thinking"):
asyncio.run(drive())
self.assertFalse(tm.rid_to_state)
if __name__ == "__main__":