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

Revert parallel request lifecycle tracking from #32588 (#34160)

Browse Source 
Signed-off-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
Lianmin Zheng
2026-08-09 03:10:02 -07:00
committed by GitHub
parent fcc5468cce
commit 71043b9dbb
5 changed files with 59 additions and 424 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/io_struct.py
+14 -21
View File
@@ -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,
python/sglang/srt/managers/tokenizer_manager.py
+33 -219
View File
@@ -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]
test/registered/unit/managers/test_io_struct.py
+8 -15
View File
@@ -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
test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
+4 -160
View File
@@ -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()
test/registered/unit/multimodal/test_gpu_feature_transport.py
-9
View File
@@ -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 (