Release decode transfer state on abort paths

Decode prealloc and transfer queues own the receiver lifetime once a request leaves the queue. The abort and transfer-failure paths were removing requests after streaming/releasing KV state without clearing the receiver, leaving backend-specific request tracking behind. The scheduler idle check also ignored decode retracted requests, so idle housekeeping could run while decode handoff state still existed.

Constraint: This is a manual port of upstream 18989f3d48 onto a locally diverged decode queue implementation.

Rejected: Direct cherry-pick | current decode queue compaction and streamer APIs differ from upstream.

Confidence: high

Scope-risk: narrow

Directive: Any path that removes a DecodeRequest from prealloc or transfer ownership must clear its kv_receiver unless ownership is explicitly transferred.

Tested: remote cjy-glm5-new pytest test/registered/unit/disaggregation/test_decode_queue_compaction.py -q: 16 passed

Tested: remote cjy-glm5-new pytest test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/observability/test_scheduler_metrics_load.py -q: 17 passed

Tested: git diff --check on modified files

Not-tested: full disaggregated decode E2E under live traffic
This commit is contained in:
laoyao0822
2026-06-27 01:57:31 +08:00
parent ac47eb61c9
commit f3cf95c0f1
3 changed files with 41 additions and 3 deletions

View File

@@ -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()

View File

@@ -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: