diff --git a/proto/sglang/runtime/v1/sglang.proto b/proto/sglang/runtime/v1/sglang.proto index 980506957..fb8e4f7a1 100644 --- a/proto/sglang/runtime/v1/sglang.proto +++ b/proto/sglang/runtime/v1/sglang.proto @@ -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 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 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 { diff --git a/python/sglang/srt/entrypoints/grpc_bridge.py b/python/sglang/srt/entrypoints/grpc_bridge.py index dfaf2840f..176503a50 100644 --- a/python/sglang/srt/entrypoints/grpc_bridge.py +++ b/python/sglang/srt/entrypoints/grpc_bridge.py @@ -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: diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 289d00d6b..ac38b4668 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -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), ) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 4e7a54b39..928348eb5 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -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, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 3658aa209..80351f330 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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] diff --git a/rust/sglang-grpc/src/server.rs b/rust/sglang-grpc/src/server.rs index 88ce594e9..c05736ce4 100644 --- a/rust/sglang-grpc/src/server.rs +++ b/rust/sglang-grpc/src/server.rs @@ -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 diff --git a/rust/sglang-grpc/src/utils/request_utils.rs b/rust/sglang-grpc/src/utils/request_utils.rs index 4685a8209..20a192f14 100644 --- a/rust/sglang-grpc/src/utils/request_utils.rs +++ b/rust/sglang-grpc/src/utils/request_utils.rs @@ -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) -> serde_json::Value { +#[allow(deprecated)] +fn sampling_params_to_map( + params: &Option, +) -> Result { match params { Some(p) => { let mut map = serde_json::Map::new(); @@ -46,15 +63,90 @@ fn sampling_params_to_map(params: &Option) -> 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::>() + .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, + priority: Option, + require_reasoning: Option, + max_thinking_tokens: Option, +) { + 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 { +) -> Result, 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 { +) -> Result, 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()); + } + } } diff --git a/test/registered/unit/entrypoints/test_grpc_bridge.py b/test/registered/unit/entrypoints/test_grpc_bridge.py new file mode 100644 index 000000000..dccc7b24b --- /dev/null +++ b/test/registered/unit/entrypoints/test_grpc_bridge.py @@ -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) diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index eb8042855..835c846c5 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -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 diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py index 2926a22a5..ff93d8ec4 100644 --- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py +++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py @@ -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__":