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:
@@ -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()
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user