feat(pd): add node-authoritative Decode admission

lightllm/server/api_cli.py

@@ -71,7 +71,30 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
     parser.add_argument(
         "--disable_pd_master_decode_capacity_limit",
         action="store_true",
-        help="Disable PD master admission control based on the total capacity of registered decode nodes.",
+        help=(
+            "Deprecated alias for --disable_pd_node_decode_admission. Set it consistently on PD Masters and "
+            "Decode nodes when intentionally allowing Decode nodes without authoritative admission."
+        ),
+    )
+    parser.add_argument(
+        "--disable_pd_node_decode_admission",
+        action="store_true",
+        help="Disable the bounded request-slot admission queue on PD Decode nodes.",
+    )
+    parser.add_argument(
+        "--pd_node_decode_admission_queue_size",
+        type=int,
+        default=None,
+        help=(
+            "Maximum number of Decode request slots allowed to wait on a PD Decode node. "
+            "Defaults to --running_max_req_size."
+        ),
+    )
+    parser.add_argument(
+        "--pd_node_decode_admission_timeout",
+        type=float,
+        default=5.0,
+        help="Maximum seconds a request may wait in the PD Decode node admission queue. Default: 5.",
     )
     parser.add_argument(
         "--pd_trans_mode",

lightllm/server/api_start.py

@@ -121,6 +121,11 @@ def _launch_subprocesses(args: StartArgs):
         args.mem_fraction > 0 and args.mem_fraction < 1
     ), f"Invalid mem_fraction {args.mem_fraction}, The expected value is between 0 and 1."
 
+    if args.pd_node_decode_admission_queue_size is not None and args.pd_node_decode_admission_queue_size < 0:
+        raise ValueError("pd_node_decode_admission_queue_size must be non-negative")
+    if args.pd_node_decode_admission_timeout <= 0:
+        raise ValueError("pd_node_decode_admission_timeout must be positive")
+
     if args.graph_max_len_in_batch == 0:
         args.graph_max_len_in_batch = args.max_req_total_len
 

lightllm/server/core/objs/shm_req_manager.py

@@ -82,17 +82,34 @@ def init_alloc_state_shm(self):
     # alloc_req_index 和 release_req_index 是分配资源时使用的接口。
     # 只有管理请求申请和释放的首节点才能调用这个接口。
     def alloc_req_index(self):
+        indexes = self.alloc_req_indexes(1)
+        return None if indexes is None else indexes[0]
+
+    def alloc_req_indexes(self, req_num: int):
+        """Atomically allocate all requested indexes or leave the pool unchanged."""
+        if req_num < 1:
+            raise ValueError("req_num must be positive")
+
         with self.manager_lock:
-            idx = self.linked_req_manager.alloc()
-            if idx is not None:
+            indexes = []
+            for _ in range(req_num):
+                idx = self.linked_req_manager.alloc()
+                if idx is None:
+                    for allocated_idx in reversed(indexes):
+                        self.alloc_state_shm.arr[allocated_idx] = 0
+                        self.linked_req_manager.free(allocated_idx)
+                    return None
                 assert self.alloc_state_shm.arr[idx] == 0
                 self.alloc_state_shm.arr[idx] = 1
-            return idx
-        return None
+                indexes.append(idx)
+            return indexes
 
     async def async_alloc_req_index(self):
         return self.alloc_req_index()
 
+    async def async_alloc_req_indexes(self, req_num: int):
+        return self.alloc_req_indexes(req_num)
+
     def release_req_index(self, req_index_in_mem):
         assert req_index_in_mem < self.max_req_num
         with self.manager_lock:

lightllm/server/core/objs/start_args_type.py

@@ -23,6 +23,9 @@ class StartArgs:
     pd_master_port: int = field(default=1212)
     pd_master_mode: str = field(default="elastic")
     disable_pd_master_decode_capacity_limit: bool = field(default=False)
+    disable_pd_node_decode_admission: bool = field(default=False)
+    pd_node_decode_admission_queue_size: Optional[int] = field(default=None)
+    pd_node_decode_admission_timeout: float = field(default=5.0)
     pd_trans_mode: str = field(default="nccl", metadata={"choices": ["nccl", "nixl"]})
     config_server_host: str = field(default=None)
     config_server_port: int = field(default=None)

lightllm/server/httpserver/decode_admission.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import deque
+from dataclasses import dataclass
+from typing import Callable, Deque, Optional
+
+from lightllm.utils.error_utils import ServerBusyError
+
+
+@dataclass(slots=True)
+class _Waiter:
+    slots: int
+    enqueue_time: float
+    future: asyncio.Future
+
+
+class DecodeAdmissionLease:
+    """A fixed number of request slots reserved on one Decode node."""
+
+    def __init__(self, controller: "DecodeAdmissionController", slots: int, waited_seconds: float) -> None:
+        self._controller = controller
+        self.slots = slots
+        self.waited_seconds = waited_seconds
+        self._released = False
+
+    def release(self) -> None:
+        if self._released:
+            return
+        self._released = True
+        self._controller._release(self.slots)
+
+    def split(self, slots: list[int]) -> list["DecodeAdmissionLease"]:
+        """Transfer this reservation into independently releasable child leases."""
+        if self._released:
+            raise RuntimeError("Decode admission lease has already been released")
+        if not slots or any(slot_count < 1 for slot_count in slots) or sum(slots) != self.slots:
+            raise ValueError("child lease slots must be positive and sum to the parent lease")
+
+        self._released = True
+        return [DecodeAdmissionLease(self._controller, slot_count, self.waited_seconds) for slot_count in slots]
+
+
+class DecodeAdmissionLeaseHandle:
+    """Transfers a pre-acquired lease into an asynchronously started generator."""
+
+    def __init__(self, lease: DecodeAdmissionLease) -> None:
+        self._lease: Optional[DecodeAdmissionLease] = lease
+
+    def take(self) -> DecodeAdmissionLease:
+        if self._lease is None:
+            raise RuntimeError("Decode admission lease has already been taken")
+        lease = self._lease
+        self._lease = None
+        return lease
+
+    def release(self) -> None:
+        if self._lease is None:
+            return
+        self._lease.release()
+        self._lease = None
+
+
+class DecodeAdmissionController:
+    """A bounded, cancellable FIFO for request slots owned by one Decode node."""
+
+    def __init__(
+        self,
+        capacity: int,
+        max_queued_slots: int,
+        timeout_seconds: float,
+        clock: Callable[[], float] = time.monotonic,
+    ) -> None:
+        if capacity < 1:
+            raise ValueError("decode admission capacity must be positive")
+        if max_queued_slots < 0:
+            raise ValueError("decode admission queue size must be non-negative")
+        if timeout_seconds <= 0:
+            raise ValueError("decode admission timeout must be positive")
+
+        self.capacity = capacity
+        self.max_queued_slots = max_queued_slots
+        self.timeout_seconds = timeout_seconds
+        self._clock = clock
+        self._active_slots = 0
+        self._queued_slots = 0
+        self._waiters: Deque[_Waiter] = deque()
+
+    @property
+    def active_slots(self) -> int:
+        return self._active_slots
+
+    @property
+    def queued_slots(self) -> int:
+        return self._queued_slots
+
+    @property
+    def queued_request_count(self) -> int:
+        return len(self._waiters)
+
+    async def acquire(self, slots: int) -> DecodeAdmissionLease:
+        if slots < 1:
+            raise ValueError("decode admission slots must be positive")
+        if slots > self.capacity:
+            raise ServerBusyError(f"request needs {slots} Decode slots, but the node capacity is {self.capacity}")
+
+        if not self._waiters and self._active_slots + slots <= self.capacity:
+            return self._activate(slots, waited_seconds=0.0)
+
+        if self._queued_slots + slots > self.max_queued_slots:
+            raise ServerBusyError("Decode node admission queue is full")
+
+        waiter = _Waiter(
+            slots=slots,
+            enqueue_time=self._clock(),
+            future=asyncio.get_running_loop().create_future(),
+        )
+        self._waiters.append(waiter)
+        self._queued_slots += slots
+        self._drain()
+
+        try:
+            return await asyncio.wait_for(asyncio.shield(waiter.future), timeout=self.timeout_seconds)
+        except asyncio.TimeoutError as exc:
+            lease = self._remove_waiter_or_take_lease(waiter)
+            if lease is not None:
+                lease.release()
+            raise ServerBusyError("Decode node admission queue wait timed out") from exc
+        except BaseException:
+            lease = self._remove_waiter_or_take_lease(waiter)
+            if lease is not None:
+                lease.release()
+            raise
+
+    def _activate(self, slots: int, waited_seconds: float) -> DecodeAdmissionLease:
+        self._active_slots += slots
+        return DecodeAdmissionLease(self, slots, waited_seconds)
+
+    def _release(self, slots: int) -> None:
+        self._active_slots -= slots
+        if self._active_slots < 0:
+            raise RuntimeError("Decode admission active slot count became negative")
+        self._drain()
+
+    def _drain(self) -> None:
+        while self._waiters:
+            waiter = self._waiters[0]
+            if waiter.future.done():
+                self._remove_waiter(waiter)
+                continue
+            if self._active_slots + waiter.slots > self.capacity:
+                break
+
+            self._remove_waiter(waiter)
+            waiter.future.set_result(
+                self._activate(
+                    waiter.slots,
+                    waited_seconds=max(0.0, self._clock() - waiter.enqueue_time),
+                )
+            )
+
+    def _remove_waiter(self, waiter: _Waiter) -> bool:
+        try:
+            self._waiters.remove(waiter)
+        except ValueError:
+            return False
+        self._queued_slots -= waiter.slots
+        return True
+
+    def _remove_waiter_or_take_lease(self, waiter: _Waiter) -> Optional[DecodeAdmissionLease]:
+        if self._remove_waiter(waiter):
+            waiter.future.cancel()
+            self._drain()
+            return None
+        if waiter.future.done() and not waiter.future.cancelled():
+            result = waiter.future.result()
+            if isinstance(result, DecodeAdmissionLease):
+                return result
+        return None