diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 7ed41ba8c..9cc0a6272 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -158,9 +158,8 @@ MultimodalDataInputFormat = Union[ @dataclass class GenerateReqInput: - # 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. + # Request ID(s). If omitted, generated during normalization. For batch + # requests, a string is expanded to per-item IDs using it as a prefix. 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. @@ -328,17 +327,12 @@ 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, prefix: Optional[str] = None): + def regenerate_rid(self): """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 = [new_rid() for _ in range(len(self.rid))] + self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))] else: - self.rid = new_rid() + self.rid = uuid.uuid4().hex return self.rid def _validate_rid_uniqueness(self): @@ -506,7 +500,7 @@ class GenerateReqInput: # Expand input based on type self._expand_inputs(num) - self._normalize_rid() + self._normalize_rid(num) self._normalize_lora_paths(num) self._normalize_image_data(num) self._normalize_video_data(num) @@ -616,16 +610,16 @@ class GenerateReqInput: else: # Already a list self.sampling_params = self.sampling_params * self.parallel_sample_num - def _normalize_rid(self): - """Normalize one logical request ID per original batch item.""" + def _normalize_rid(self, num): + """Normalize request IDs for batch processing.""" if self.rid is None: - self.rid = [uuid.uuid4().hex for _ in range(self.batch_size)] + self.rid = [uuid.uuid4().hex for _ in range(num)] elif isinstance(self.rid, str): - if self.batch_size == 1: - self.rid = [self.rid] - else: - self.rid = [f"{self.rid}_{i}" for i in range(self.batch_size)] + new_rids = [f"{self.rid}_{i}" for i in range(num)] + self.rid = new_rids 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." @@ -777,9 +771,8 @@ 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[logical_index], + rid=self.rid[i], 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, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 36c639f84..1983fdd1a 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -196,10 +196,6 @@ _INCREMENTAL_STREAMING_META_INFO_KEYS = ( ) -class RequestAbortedError(ValueError): - status_code = 499 - - @dataclasses.dataclass class ReqState: """Store the state a request.""" @@ -211,9 +207,6 @@ 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 @@ -573,10 +566,6 @@ 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() @@ -793,7 +782,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): f"routed_dp_rank={obj.routed_dp_rank} out of range [0, {dp_size})" ) - request_lifecycles = self._init_req_state(obj, request) + self._init_req_state(obj, request) try: if self.server_args.language_only: self._handle_epd_disaggregation_encode_request(obj) @@ -803,16 +792,13 @@ 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) @@ -830,7 +816,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, request_lifecycles) + self._discard_pending_req_states(obj) raise def _detect_input_format( @@ -1574,9 +1560,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): tokenized_obj.wrap_pickle_fields() self._dispatch_to_scheduler(tokenized_obj) dispatched = True - 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() finally: @@ -1609,10 +1592,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): self._dispatch_to_scheduler(batch_req) dispatched = True - 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") @@ -1687,7 +1666,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: - self._remove_req_state(state.obj.rid) + del self.rid_to_state[state.obj.rid] # Mark ongoing LoRA request as finished. if self.enable_lora and state.obj.lora_path: @@ -1821,7 +1800,6 @@ 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 @@ -1844,7 +1822,6 @@ 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) @@ -1865,12 +1842,9 @@ 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 @@ -1879,20 +1853,17 @@ 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(prefix=logical_rid) + tokenized_obj.rid = tmp_obj.regenerate_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_child_req_state(logical_rid, tmp_obj) + self._init_req_state(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 @@ -1901,8 +1872,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(prefix=logical_rid) - self._init_child_req_state(logical_rid, tmp_obj) + tokenized_obj.rid = tmp_obj.regenerate_rid() + self._init_req_state(tmp_obj) state = self.rid_to_state[tmp_obj.rid] tokenized_obj.time_stats = state.time_stats if tmp_obj.return_prompt_token_ids: @@ -1911,10 +1882,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): generators.append(self._wait_one_response(tmp_obj, request)) rids.append(tmp_obj.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) + self.rid_to_state[objs[i].rid].time_stats.set_finished_time() + del self.rid_to_state[objs[i].rid] # Wait for all requests is_stream = hasattr(obj, "stream") and obj.stream @@ -1974,39 +1943,14 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): if not abort_all and not rid: logger.warning("Ignore abort_request with empty rid and abort_all=False") return - 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: + if ( + not abort_all + and self.server_args.tokenizer_worker_num == 1 + and rid not in self.rid_to_state + ): return - else: - target_rids = (rid,) - - for target_rid in target_rids: - self._dispatch_to_scheduler(AbortReq(rid=target_rid, abort_all=abort_all)) + req = AbortReq(rid=rid, abort_all=abort_all) + self._dispatch_to_scheduler(req) if self.enable_metrics: # TODO: also use custom_labels from the request self.metrics_collector.observe_one_aborted_request( @@ -2493,7 +2437,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): ) ) - self._remove_req_state(rid) + del self.rid_to_state[rid] # Mark ongoing LoRA request as finished. if self.enable_lora and state.obj.lora_path: @@ -3229,7 +3173,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): "output_ids": output_ids, "meta_info": meta_info, } - self._remove_req_state(recv_obj.rid) + del self.rid_to_state[recv_obj.rid] state.out_list.append(out) state.event.set() @@ -3385,80 +3329,11 @@ 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 @@ -3489,90 +3364,29 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): for i in range(len(obj.rid)) ] - 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: + if rid in self.rid_to_state: + raise ValueError(f"Duplicate request ID detected: {rid}") time_stats = APIServerReqTimeStats(disagg_mode=self.disaggregation_mode) - state = ReqState( - [], - False, - asyncio.Event(), - sub_obj, - time_stats, - lifecycle_id=lifecycle_id if lifecycle_id is not None else object(), - ) + state = ReqState([], False, asyncio.Event(), sub_obj, time_stats) 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, - lifecycle_ids: Optional[Dict[str, object]] = None, - ): - """Drop all logical and child state owned by *obj*. + def _discard_pending_req_states(self, obj): + """Drop rid_to_state entries created by _init_req_state for *obj*. - 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. + 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. """ - 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) + 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) def _should_dispatch_to_encoder( self, obj: Union[GenerateReqInput, EmbeddingReqInput] diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index 3b46ed8e5..46427c5f8 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -305,8 +305,7 @@ 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.""" + def test_parallel_sampling_preserves_reasoning_controls(self): single = GenerateReqInput( text="Hello", rid="single", @@ -316,8 +315,11 @@ class TestGenerateReqInputNormalization(CustomTestCase): ) single.normalize_batch_and_arguments() - self.assertEqual(single.rid, ["single"]) - self.assertEqual([single[i].rid for i in range(3)], ["single"] * 3) + self.assertEqual(single.rid, ["single_0", "single_1", "single_2"]) + self.assertEqual( + [single[i].rid for i in range(3)], + ["single_0", "single_1", "single_2"], + ) self.assertTrue(all(single[i].require_reasoning for i in range(3))) self.assertEqual( [single[i].max_thinking_tokens for i in range(3)], @@ -331,10 +333,10 @@ class TestGenerateReqInputNormalization(CustomTestCase): ) batch.normalize_batch_and_arguments() - self.assertEqual(batch.rid, ["batch_0", "batch_1"]) + self.assertEqual(batch.rid, ["batch_0", "batch_1", "batch_2", "batch_3"]) self.assertEqual( [batch[i].rid for i in range(4)], - ["batch_0", "batch_1", "batch_0", "batch_1"], + ["batch_0", "batch_1", "batch_2", "batch_3"], ) def test_audio_data_handling(self): @@ -680,15 +682,6 @@ 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 ff93d8ec4..15f89dbe7 100644 --- a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py +++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py @@ -30,7 +30,6 @@ from sglang.srt.managers.io_struct import ( # noqa: E402 ) from sglang.srt.managers.tokenizer_manager import ( # noqa: E402 ReqState, - RequestAbortedError, TokenizerManager, ) from sglang.srt.observability.req_time_stats import ( # noqa: E402 @@ -121,8 +120,6 @@ 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 @@ -132,11 +129,10 @@ 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", *, dispatched: bool = False) -> ReqState: +def _make_req_state(rid: str = "test_rid") -> ReqState: """Create a minimal ReqState for testing.""" obj = Mock(spec=GenerateReqInput) obj.rid = rid @@ -150,7 +146,6 @@ def _make_req_state(rid: str = "test_rid", *, dispatched: bool = False) -> ReqSt event=asyncio.Event(), obj=obj, time_stats=APIServerReqTimeStats(), - dispatched=dispatched, ) @@ -352,19 +347,6 @@ 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.""" @@ -467,20 +449,17 @@ def _make_generate_obj(rid, is_single): class TestDiscardPendingReqStates(CustomTestCase): """Direct tests for _discard_pending_req_states.""" - def test_discard_single_aborts_scheduler_before_cleanup(self): + def test_discard_single(self): tm = _make_tokenizer_manager() rid = "d_single" - tm.rid_to_state[rid] = _make_req_state(rid, dispatched=True) + tm.rid_to_state[rid] = _make_req_state(rid) 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_unsent_batch_without_scheduler_abort(self): + def test_discard_batch_removes_all(self): tm = _make_tokenizer_manager() rids = ["d0", "d1", "d2"] for r in rids: @@ -491,7 +470,6 @@ 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.""" @@ -503,62 +481,6 @@ 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): @@ -614,40 +536,6 @@ class TestParallelStreamTaskCleanup(CustomTestCase): 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. @@ -674,7 +562,6 @@ 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) @@ -699,49 +586,6 @@ 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() diff --git a/test/registered/unit/multimodal/test_gpu_feature_transport.py b/test/registered/unit/multimodal/test_gpu_feature_transport.py index 6e802c71d..5fb029ec9 100644 --- a/test/registered/unit/multimodal/test_gpu_feature_transport.py +++ b/test/registered/unit/multimodal/test_gpu_feature_transport.py @@ -278,8 +278,6 @@ class TestCudaVmmFeatureTransport(unittest.TestCase): transport.prepare_for_dispatch.return_value = [] manager.cuda_vmm_feature_transport = transport manager._dispatch_to_scheduler = MagicMock() - state = SimpleNamespace(dispatched=False) - manager.rid_to_state = {"test-request": state} tokenized_obj = SimpleNamespace( rid="test-request", mm_inputs=None, @@ -293,7 +291,6 @@ class TestCudaVmmFeatureTransport(unittest.TestCase): manager._dispatch_to_scheduler.assert_called_once_with(tokenized_obj) transport.prepare_for_dispatch.assert_called_once_with((None,)) transport.cancel_for_dispatch.assert_not_called() - self.assertTrue(state.dispatched) def test_failed_dispatch_cancels_published_items(self): from sglang.srt.managers import tokenizer_manager @@ -308,8 +305,6 @@ class TestCudaVmmFeatureTransport(unittest.TestCase): manager._dispatch_to_scheduler = MagicMock( side_effect=RuntimeError("send failed") ) - state = SimpleNamespace(dispatched=False) - manager.rid_to_state = {"test-request": state} items = [MultimodalDataItem(modality=Modality.IMAGE, feature=torch.arange(2))] tokenized_obj = SimpleNamespace( rid="test-request", @@ -330,7 +325,6 @@ class TestCudaVmmFeatureTransport(unittest.TestCase): (tokenized_obj.mm_inputs,) ) transport.cancel_for_dispatch.assert_called_once_with(items) - self.assertFalse(state.dispatched) def test_post_dispatch_failure_does_not_cancel_published_items(self): from sglang.srt.managers import tokenizer_manager @@ -343,8 +337,6 @@ class TestCudaVmmFeatureTransport(unittest.TestCase): manager = object.__new__(tokenizer_manager.TokenizerManager) transport = MagicMock() manager._dispatch_to_scheduler = MagicMock() - state = SimpleNamespace(dispatched=False) - manager.rid_to_state = {"test-request": state} time_stats = MagicMock() time_stats.set_api_server_dispatch_finish_time.side_effect = RuntimeError( "bookkeeping failed" @@ -367,7 +359,6 @@ class TestCudaVmmFeatureTransport(unittest.TestCase): manager._dispatch_to_scheduler.assert_called_once_with(tokenized_obj) transport.cancel_for_dispatch.assert_not_called() - self.assertTrue(state.dispatched) def test_prepare_batch_cancels_prior_groups_on_failure(self): from sglang.srt.utils.cuda_vmm_transport_utils import (