From 24d27c203547cbb9fe23ad8adab1a981c4ee71a1 Mon Sep 17 00:00:00 2001 From: Yuhong Guo Date: Wed, 20 May 2026 16:32:44 +0800 Subject: [PATCH] [BugFix] Fix rid_to_state leak for aborted queued requests (#24070) --- .../sglang/srt/managers/tokenizer_manager.py | 2 + .../test_tokenizer_manager_rid_cleanup.py | 390 ++++++++++++++++++ 2 files changed, 392 insertions(+) create mode 100644 test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 64a69018d..415236ddb 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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() diff --git a/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py new file mode 100644 index 000000000..95a8ffffd --- /dev/null +++ b/test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py @@ -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)