[BugFix][PD]Fix metadata_buffer_index leak when aborted in PD (#17483)

python/sglang/srt/disaggregation/prefill.py

@@ -62,6 +62,27 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+def release_req_to_metadata_buffer(
+    req: Req, allocator: ReqToMetadataIdxAllocator
+) -> None:
+    """
+    Release the metadata buffer index allocated for a request in prefill disaggregation mode.
+
+    This function safely releases the metadata buffer index if it was allocated.
+
+    Args:
+        req: The request object that may have a metadata_buffer_index allocated
+        allocator: The ReqToMetadataIdxAllocator instance to free the index
+    """
+    if (
+        hasattr(req, "metadata_buffer_index")
+        and req.metadata_buffer_index is not None
+        and req.metadata_buffer_index >= 0
+    ):
+        allocator.free(req.metadata_buffer_index)
+        req.metadata_buffer_index = -1
+
+
 class PrefillBootstrapQueue:
     """
     Store the requests in bootstrapping
@@ -592,8 +613,9 @@ class SchedulerDisaggregationPrefillMixin:
         for req in done_reqs:
             req: Req
             req.add_latency(RequestStage.PREFILL_TRANSFER_KV_CACHE)
-            self.req_to_metadata_buffer_idx_allocator.free(req.metadata_buffer_index)
-            req.metadata_buffer_index = -1
+            release_req_to_metadata_buffer(
+                req, self.req_to_metadata_buffer_idx_allocator
+            )
             trace_slice(
                 RequestStage.PREFILL_TRANSFER_KV_CACHE, req.rid, thread_finish_flag=True
             )

python/sglang/srt/managers/scheduler.py

@@ -47,6 +47,7 @@ from sglang.srt.disaggregation.encode_receiver import MMReceiverHTTP
 from sglang.srt.disaggregation.prefill import (
     PrefillBootstrapQueue,
     SchedulerDisaggregationPrefillMixin,
+    release_req_to_metadata_buffer,
 )
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
@@ -2726,6 +2727,11 @@ class Scheduler(
             # For disaggregation decode mode, the request in the waiting queue has KV cache allocated.
             if self.disaggregation_mode == DisaggregationMode.DECODE:
                 release_kv_cache(req, self.tree_cache)
+            # For disaggregation prefill mode, free the metadata buffer index
+            if self.disaggregation_mode == DisaggregationMode.PREFILL:
+                release_req_to_metadata_buffer(
+                    req, self.req_to_metadata_buffer_idx_allocator
+                )
 
             # For mamba radix cache
             if (