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

[BugFix] Fix rid_to_state leak for aborted queued requests (#24070)

Browse Source
This commit is contained in:
Yuhong Guo
2026-05-20 01:32:44 -07:00
committed by GitHub
parent e99f87c974
commit 24d27c2035
2 changed files with 392 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/tokenizer_manager.py
+2
View File
@@ -2488,6 +2488,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
"output_ids": output_ids,
"meta_info": meta_info,
}
del self.rid_to_state[recv_obj.rid]
state.out_list.append(out)
state.event.set()
test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py
+390
View File
@@ -0,0 +1,390 @@
"""
Unit tests for rid_to_state cleanup in TokenizerManager.
Verifies that request IDs are properly removed from rid_to_state after
completion or abort, allowing resubmission with the same rid without
triggering "Duplicate request ID detected" errors.
Covers:
- _handle_abort_req cleans up rid_to_state
- _handle_batch_output cleans up rid_to_state on finished requests
- _init_req_state rejects duplicate rids
- Resubmission succeeds after cleanup
"""
import asyncio
import dataclasses
import unittest
from unittest.mock import MagicMock, Mock
from sglang.test.ci.ci_register import register_cpu_ci
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
register_cpu_ci(est_time=15, suite="base-a-test-cpu")
_NOT_FINISHED = object() # Sentinel: request has not finished yet
# ---------------------------------------------------------------------------
# Per-request field defaults for BatchStrOutput construction.
# Categorised by value shape so that _make_batch_str_output can assign
# type-appropriate defaults without hardcoding every field name.
# When a field is renamed upstream, the old name simply won't appear in
# dataclasses.fields() and the new name will fall through to the
# pattern-matching or safe fallback — no test breakage.
# ---------------------------------------------------------------------------
_PER_REQUEST_INT_FIELDS = frozenset(
{
"prompt_tokens",
"completion_tokens",
"reasoning_tokens",
"cached_tokens",
"retraction_counts",
# Speculative-decoding int-scalar fields (current and historical names)
"spec_verify_ct",
"spec_accepted_drafts",
"spec_num_correct_drafts",
}
)
_PER_REQUEST_FLOAT_FIELDS = frozenset(
{
"output_token_entropy_val",
}
)
_PER_REQUEST_NESTED_LIST_FIELDS = frozenset(
{
"output_ids",
# Logprob fields
"input_token_logprobs_val",
"input_token_logprobs_idx",
"output_token_logprobs_val",
"output_token_logprobs_idx",
"input_top_logprobs_val",
"input_top_logprobs_idx",
"output_top_logprobs_val",
"output_top_logprobs_idx",
"input_token_ids_logprobs_val",
"input_token_ids_logprobs_idx",
"output_token_ids_logprobs_val",
"output_token_ids_logprobs_idx",
# Speculative-decoding histogram fields (current and historical names)
"spec_acceptance_histogram",
"spec_correct_drafts_histogram",
}
)
_PER_REQUEST_OPTIONAL_FIELDS = frozenset(
{
"output_hidden_states",
"routed_experts",
"indexer_topk",
"placeholder_tokens_idx",
"placeholder_tokens_val",
}
)
def _make_tokenizer_manager() -> TokenizerManager:
"""Create a TokenizerManager with mocked dependencies, bypassing __init__."""
tm = TokenizerManager.__new__(TokenizerManager)
tm.server_args = MagicMock()
tm.server_args.enable_trace = False
tm.server_args.enable_metrics = False
tm.server_args.enable_lora = False
tm.server_args.speculative_algorithm = None
tm.server_args.incremental_streaming_output = False
tm.server_args.skip_tokenizer_init = False
tm.server_args.batch_notify_size = 1
tm.server_args.weight_version = "1"
tm.server_args.crash_dump_folder = ""
tm.server_args.dp_size = 1
tm.disaggregation_mode = "none"
tm.rid_to_state = {}
tm.enable_metrics = False
tm.dump_requests_folder = ""
tm.crash_dump_folder = ""
tm.send_to_scheduler = MagicMock()
return tm
def _make_req_state(rid: str = "test_rid") -> ReqState:
"""Create a minimal ReqState for testing."""
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
obj.stream = False
obj.return_logprob = False
obj.lora_path = None
obj.log_metrics = False
return ReqState(
out_list=[],
finished=False,
event=asyncio.Event(),
obj=obj,
time_stats=APIServerReqTimeStats(),
)
def _make_abort_req(rid: str, abort_message: str = "Aborted") -> AbortReq:
"""Create an AbortReq for testing."""
return AbortReq(
rid=rid,
abort_all=False,
finished_reason={"type": "abort", "message": abort_message},
abort_message=abort_message,
)
def _make_batch_str_output(rid: str, finished_reason=None) -> BatchStrOutput:
"""Create a minimal BatchStrOutput for a single request.
Uses dataclass field introspection so that new or renamed fields in
BatchStrOutput don't break this test. Only the fields that matter for
test logic (rids, finished_reasons, output_strs) are set explicitly;
all others receive type-appropriate defaults based on naming patterns.
Fields with class-level defaults are left alone automatically.
"""
if finished_reason is _NOT_FINISHED:
fr = None
elif finished_reason is None:
fr = {"type": "length"}
else:
fr = finished_reason
kwargs = {}
for f in dataclasses.fields(BatchStrOutput):
if f.name == "rids":
kwargs[f.name] = [rid]
elif f.name == "finished_reasons":
kwargs[f.name] = [fr]
elif f.name == "output_strs":
kwargs[f.name] = ["hello"]
elif f.name in _PER_REQUEST_INT_FIELDS:
kwargs[f.name] = [0]
elif f.name in _PER_REQUEST_FLOAT_FIELDS:
kwargs[f.name] = [0.0]
elif f.name in _PER_REQUEST_NESTED_LIST_FIELDS:
kwargs[f.name] = [[]]
elif f.name in _PER_REQUEST_OPTIONAL_FIELDS:
kwargs[f.name] = [None]
# Fields with class defaults — skip, let the default be used
elif (
f.default is not dataclasses.MISSING
or f.default_factory is not dataclasses.MISSING
):
continue
# Unknown required field — provide a safe per-request default.
# Most BatchStrOutput fields are per-request lists; [[]] works for
# List[List[...]] and is unlikely to crash on [i] indexing for
# List[int] either (the inner [] just means "no data").
else:
kwargs[f.name] = [[]]
return BatchStrOutput(**kwargs)
class TestRidToStateCleanupOnAbort(CustomTestCase):
"""Test that _handle_abort_req removes rid from rid_to_state."""
def test_abort_removes_rid_from_state(self):
"""After _handle_abort_req, rid should be removed from rid_to_state."""
tm = _make_tokenizer_manager()
rid = "abort_test_rid"
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
abort_req = _make_abort_req(rid)
tm._handle_abort_req(abort_req)
self.assertNotIn(rid, tm.rid_to_state)
def test_abort_allows_resubmit_same_rid(self):
"""After abort, _init_req_state should accept the same rid again."""
tm = _make_tokenizer_manager()
rid = "resubmit_after_abort_rid"
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
abort_req = _make_abort_req(rid)
tm._handle_abort_req(abort_req)
# Resubmit with the same rid — should not raise
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
obj.is_single = True
obj.received_time = 0.0
obj.external_trace_header = None
obj.bootstrap_room = None
tm._init_req_state(obj)
self.assertIn(rid, tm.rid_to_state)
def test_abort_sets_finished_and_notifies(self):
"""_handle_abort_req should mark state as finished and set the event."""
tm = _make_tokenizer_manager()
rid = "abort_notify_rid"
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
abort_req = _make_abort_req(rid)
tm._handle_abort_req(abort_req)
self.assertTrue(state.finished)
self.assertTrue(state.event.is_set())
self.assertEqual(len(state.out_list), 1)
self.assertEqual(
state.out_list[0]["meta_info"]["finish_reason"]["type"], "abort"
)
class TestRidToStateCleanupOnBatchOutput(CustomTestCase):
"""Test that _handle_batch_output removes rid from rid_to_state on completion."""
def test_batch_output_removes_rid_on_finish(self):
"""When a request finishes in _handle_batch_output, rid should be removed."""
tm = _make_tokenizer_manager()
rid = "batch_finish_rid"
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
batch_output = _make_batch_str_output(rid)
asyncio.run(tm._handle_batch_output(batch_output))
self.assertNotIn(rid, tm.rid_to_state)
def test_batch_output_allows_resubmit_after_finish(self):
"""After a request finishes, the same rid can be resubmitted."""
tm = _make_tokenizer_manager()
rid = "batch_resubmit_rid"
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
batch_output = _make_batch_str_output(rid)
asyncio.run(tm._handle_batch_output(batch_output))
# Resubmit with the same rid — should not raise
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
obj.is_single = True
obj.received_time = 0.0
obj.external_trace_header = None
obj.bootstrap_room = None
tm._init_req_state(obj)
self.assertIn(rid, tm.rid_to_state)
def test_batch_output_keeps_rid_when_not_finished(self):
"""When a request is not yet finished, rid should remain in rid_to_state."""
tm = _make_tokenizer_manager()
rid = "batch_ongoing_rid"
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
# finished_reason=_NOT_FINISHED means the request is still ongoing
batch_output = _make_batch_str_output(rid, finished_reason=_NOT_FINISHED)
asyncio.run(tm._handle_batch_output(batch_output))
self.assertIn(rid, tm.rid_to_state)
class TestInitReqStateDuplicateDetection(CustomTestCase):
"""Test that _init_req_state raises ValueError for duplicate rids."""
def test_duplicate_rid_raises_error(self):
"""_init_req_state should raise ValueError if rid already exists."""
tm = _make_tokenizer_manager()
rid = "duplicate_rid"
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
obj.is_single = True
obj.received_time = 0.0
obj.external_trace_header = None
obj.bootstrap_room = None
with self.assertRaises(ValueError) as ctx:
tm._init_req_state(obj)
self.assertIn("Duplicate request ID", str(ctx.exception))
def test_unique_rid_succeeds(self):
"""_init_req_state should succeed with a unique rid."""
tm = _make_tokenizer_manager()
rid = "unique_rid"
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
obj.is_single = True
obj.received_time = 0.0
obj.external_trace_header = None
obj.bootstrap_room = None
tm._init_req_state(obj)
self.assertIn(rid, tm.rid_to_state)
class TestResubmitAfterCompletion(CustomTestCase):
"""End-to-end test: complete a request, then resubmit with the same rid."""
def test_complete_then_resubmit_same_rid(self):
"""A request that completes normally should allow resubmission with the same rid."""
tm = _make_tokenizer_manager()
rid = "complete_resubmit_rid"
# Phase 1: simulate a request in rid_to_state, then complete it
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
batch_output = _make_batch_str_output(rid, finished_reason={"type": "length"})
asyncio.run(tm._handle_batch_output(batch_output))
# rid should be cleaned up
self.assertNotIn(rid, tm.rid_to_state)
# Phase 2: resubmit with the same rid — should succeed
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
obj.is_single = True
obj.received_time = 0.0
obj.external_trace_header = None
obj.bootstrap_room = None
tm._init_req_state(obj)
self.assertIn(rid, tm.rid_to_state)
def test_abort_then_resubmit_same_rid(self):
"""An aborted request should allow resubmission with the same rid."""
tm = _make_tokenizer_manager()
rid = "abort_resubmit_rid"
# Phase 1: simulate a request, then abort it
state = _make_req_state(rid)
tm.rid_to_state[rid] = state
abort_req = _make_abort_req(rid)
tm._handle_abort_req(abort_req)
self.assertNotIn(rid, tm.rid_to_state)
# Phase 2: resubmit with the same rid — should succeed
obj = Mock(spec=GenerateReqInput)
obj.rid = rid
obj.is_single = True
obj.received_time = 0.0
obj.external_trace_header = None
obj.bootstrap_room = None
tm._init_req_state(obj)
self.assertIn(rid, tm.rid_to_state)
if __name__ == "__main__":
unittest.main(verbosity=2)