diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 07b2f4e0b..19c537f7d 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -821,6 +821,9 @@ class DecodePreallocQueue: self.scheduler.stream_output( [decode_req.req], decode_req.req.return_logprob ) + if decode_req.kv_receiver is not None: + decode_req.kv_receiver.clear() + decode_req.kv_receiver = None failed_reqs.append(decode_req) continue @@ -1275,6 +1278,9 @@ class DecodeTransferQueue: ) # release pre-allocated kv cache, but don't insert into the tree since it's failed release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) + if decode_req.kv_receiver is not None: + decode_req.kv_receiver.clear() + decode_req.kv_receiver = None decode_reqs_to_free_metadata.append(decode_req) if self.scheduler.enable_metrics: self.scheduler.metrics_collector.increment_transfer_failed_reqs() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7f3f53ef3..c0d1c67db 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -3318,6 +3318,7 @@ class Scheduler( if self.disaggregation_mode == DisaggregationMode.DECODE: idle &= len(self.disagg_decode_prealloc_queue.queue) == 0 + idle &= len(self.disagg_decode_prealloc_queue.retracted_queue) == 0 idle &= len(self.disagg_decode_transfer_queue.queue) == 0 if not for_health_check: diff --git a/test/registered/unit/disaggregation/test_decode_queue_compaction.py b/test/registered/unit/disaggregation/test_decode_queue_compaction.py index 7a8bbb30d..28e217843 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_compaction.py +++ b/test/registered/unit/disaggregation/test_decode_queue_compaction.py @@ -24,6 +24,7 @@ from sglang.srt.disaggregation.utils import ( from sglang.srt.managers.scheduler_output_processor_mixin import ( SchedulerOutputProcessorMixin, ) +from sglang.srt.managers.scheduler import Scheduler from sglang.srt.managers.schedule_batch import FINISH_ABORT from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -214,9 +215,10 @@ class TestDecodeQueueCompaction(CustomTestCase): kv_receiver=cast(Any, FakeReceiver()), metadata_buffer_index=11, ) + failed_receiver = FakeReceiver(should_fail=True) failed = DecodeRequest( req=cast(Any, FakeReq("failed", 3)), - kv_receiver=cast(Any, FakeReceiver(should_fail=True)), + kv_receiver=cast(Any, failed_receiver), metadata_buffer_index=12, ) skipped = DecodeRequest( @@ -256,6 +258,8 @@ class TestDecodeQueueCompaction(CustomTestCase): self.assertEqual(released, [("failed", False)]) self.assertEqual(allocator.freed, [11, 12]) self.assertEqual(queue.queue, [keep, skipped]) + self.assertEqual(failed_receiver.clear_calls, 1) + self.assertIsNone(failed.kv_receiver) def test_decode_transfer_queue_keeps_metadata_waiters(self): allocator = FakeAllocator() @@ -556,9 +560,10 @@ class TestDecodeQueueCompaction(CustomTestCase): ) failed_req = FakeReq("failed", 2) failed_req.finished_reason = FINISH_ABORT("boom") + failed_receiver = FakeReceiver() failed = DecodeRequest( req=cast(Any, failed_req), - kv_receiver=cast(Any, FakeReceiver()), + kv_receiver=cast(Any, failed_receiver), waiting_for_input=True, ) tail = DecodeRequest( @@ -576,6 +581,8 @@ class TestDecodeQueueCompaction(CustomTestCase): self.assertEqual([req.req.rid for req in failed_reqs], ["failed"]) self.assertEqual(streamed, ["failed"]) self.assertEqual(queue.queue, [blocked, tail]) + self.assertEqual(failed_receiver.clear_calls, 1) + self.assertIsNone(failed.kv_receiver) def test_pop_preallocated_compacts_queue_in_one_pass(self): streamed = [] @@ -623,9 +630,10 @@ class TestDecodeQueueCompaction(CustomTestCase): ) failed_req = FakeReq("failed", 2) failed_req.finished_reason = FINISH_ABORT("boom") + failed_receiver = FakeReceiver() failed = DecodeRequest( req=cast(Any, failed_req), - kv_receiver=cast(Any, FakeReceiver()), + kv_receiver=cast(Any, failed_receiver), waiting_for_input=True, ) waiting = DecodeRequest( @@ -666,6 +674,29 @@ class TestDecodeQueueCompaction(CustomTestCase): self.assertEqual(streamed, ["failed"]) self.assertEqual(success.metadata_buffer_index, 7) self.assertEqual(queue.queue, [skipped, waiting, blocked, tail]) + self.assertEqual(failed_receiver.clear_calls, 1) + self.assertIsNone(failed.kv_receiver) + + def test_retracted_decode_requests_keep_scheduler_non_idle(self): + scheduler = Scheduler.__new__(Scheduler) + scheduler.running_batch = SimpleNamespace(is_empty=lambda: True) + scheduler.chunked_req = None + scheduler.dllm_manager = SimpleNamespace(any_staging_reqs=lambda: False) + scheduler.last_batch = None + scheduler.cur_batch = None + scheduler.enable_overlap = False + scheduler.pp_size = 1 + scheduler.running_mbs = [] + scheduler.waiting_queue = [] + scheduler.disaggregation_mode = DisaggregationMode.DECODE + scheduler.disagg_decode_prealloc_queue = SimpleNamespace( + queue=[], retracted_queue=[object()] + ) + scheduler.disagg_decode_transfer_queue = SimpleNamespace(queue=[]) + scheduler.grammar_manager = SimpleNamespace(grammar_queue=[]) + scheduler.enable_hierarchical_cache = False + + self.assertFalse(scheduler.is_fully_idle()) def test_get_new_prebuilt_batch_slices_waiting_queue_prefix(self): allocator = FakeAllocator()