fix(pd): pd limit

docs/CN/source/tutorial/api_server_args.rst

@@ -87,6 +87,19 @@ PD 分离模式参数
     推理进度健康检查：当仍有在途请求，且整个 PD Master 连续 ``HEALTH_TIMEOUT`` 秒
     没有任何请求成功返回 token 时，接口将返回 HTTP 503。
 
+.. option:: --enable_pd_node_self_request_limit
+
+    在 Prefill/Decode 节点上启用本地请求限流。PD Master 当前不执行请求准入限流。
+    HTTP server 申请本地 ``shm_req`` 对象的超时时间由
+    ``LIGHTLLM_PD_NODE_SHM_REQ_ALLOC_TIMEOUT_SECONDS`` 控制（默认 20 秒）；请求进入 Router 后等待
+    进入推理系统的超时时间由 ``LIGHTLLM_PD_NODE_ROUTER_WAIT_TIMEOUT_SECONDS`` 控制（默认 20 秒）。
+    超时会导致 ``Server is busy``；其中已进入 Router 但仍未进入推理系统的请求会主动标记为 aborted，
+    由 PD Master 转换为 HTTP 429。未开启限流时请求会持续等待资源；PD 高优先级请求
+    （分段续跑请求或预计输入 cache 命中率高于 0.8 的请求）由 PD Master 通过
+    ``pd_high_priority_request_time_out_seconds`` 下发一个统一的超时时间下限。P/D 节点分别取
+    该值与本地 ``shm_req``、Router 超时的较大值；该字段为 0 时不延长本地超时。PD Master 下发值由
+    ``LIGHTLLM_PD_HIGH_PRIORITY_REQUEST_TIMEOUT_SECONDS`` 控制，默认 60 秒。该参数默认关闭。
+
 .. option:: --config_server_host
 
     配置服务器模式下的主机地址

docs/EN/source/tutorial/api_server_args.rst

@@ -89,6 +89,20 @@ PD disaggregation Mode Parameters
     the endpoints return HTTP 503 if no request on the PD Master successfully returns a token for
     ``HEALTH_TIMEOUT`` consecutive seconds.
 
+.. option:: --enable_pd_node_self_request_limit
+
+    Enable local request limiting on Prefill/Decode nodes. PD Master does not currently perform request admission
+    limiting. The local ``shm_req`` allocation timeout is controlled by
+    ``LIGHTLLM_PD_NODE_SHM_REQ_ALLOC_TIMEOUT_SECONDS`` (20 seconds by default), while the timeout from Router entry
+    to inference entry is controlled by ``LIGHTLLM_PD_NODE_ROUTER_WAIT_TIMEOUT_SECONDS`` (20 seconds by default).
+    A timeout reports ``Server is busy``; a request that has entered the Router but not inference is proactively
+    marked aborted, and PD Master converts this to HTTP 429. Requests continue waiting when local admission is
+    disabled. For PD high-priority requests (segmented continuation requests or requests whose estimated input
+    cache hit rate is above 0.8), PD Master supplies a shared timeout floor through
+    ``pd_high_priority_request_time_out_seconds``. Each P/D node uses the greater of this value and its local
+    ``shm_req`` or Router timeout; zero does not extend the local timeout. The value supplied by PD Master is controlled by
+    ``LIGHTLLM_PD_HIGH_PRIORITY_REQUEST_TIMEOUT_SECONDS`` and defaults to 60 seconds. Disabled by default.
+
 .. option:: --config_server_host
 
     Host address in configuration server mode

lightllm/server/api_cli.py

@@ -69,9 +69,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
         ),
     )
     parser.add_argument(
-        "--disable_pd_master_decode_capacity_limit",
+        "--enable_pd_node_self_request_limit",
         action="store_true",
-        help="Disable PD master admission control based on the total capacity of registered decode nodes.",
+        help=(
+            "Enable local request limiting on Prefill/Decode nodes by enforcing shm_req allocation and Router "
+            "scheduling wait timeouts. PD Master admission limiting is not currently enabled. Default: disabled."
+        ),
     )
     parser.add_argument(
         "--pd_trans_mode",

lightllm/server/core/objs/req.py

@@ -65,6 +65,9 @@ def is_finished_length(self):
     def is_finished_error(self):
         return self.status == self.FINISHED_ERROR
 
+    def is_error_finished(self):
+        return self.status in (self.FINISHED_ABORTED, self.FINISHED_ERROR)
+
     def get_finish_reason(self):
         if self.status == self.FINISHED_STOP:
             return "stop"
@@ -83,6 +86,10 @@ class Req(ctypes.Structure):
         ("index_in_shm_mem", ctypes.c_int),
         ("ref_count", ctypes.c_int),  # 个人不要操作这个计数  # 个人不要操作这个引用计数
         ("recv_time", ctypes.c_double),  # 用于记录请求到达服务的时间，主要用于调试
+        # Router 收到请求和请求被调度为新 batch 的单调时钟时间戳，用于 HTTP server
+        # 判断请求是否在 Router 等待过久。时间戳写入共享内存，供不同进程读取。
+        ("router_arrival_time", ctypes.c_double),
+        ("infer_start_time", ctypes.c_double),
         ("request_id", ctypes.c_int64),  # 引用计数
         ("group_req_id", ctypes.c_int64),
         ("input_len", ctypes.c_int),
@@ -157,6 +164,8 @@ def init(
         self.index_in_shm_mem: int = self.index_in_shm_mem
         self.ref_count: int = self.ref_count
         self.recv_time: float = time.time()
+        self.router_arrival_time = 0.0
+        self.infer_start_time = 0.0
 
         self.request_id = request_id
         self.group_req_id = convert_sub_id_to_group_id(request_id)

lightllm/server/core/objs/sampling_params.py

@@ -294,6 +294,12 @@ class SamplingParams(ctypes.Structure):
         ("stop_sequences", StopSequenceGroups),
         ("exponential_decay_length_penalty", ExponentialDecayLengthPenalty),
         ("group_request_id", ctypes.c_int64),  # p d mode used params
+        # 由 PD Master 为分段续跑或预计 cache 命中率较高的请求设置，表示请求需
+        # 以高优先级插入 Router 调度队列。
+        ("pd_high_priority_request", ctypes.c_bool),
+        # PD 高优先级请求在开启本地限流的 P/D 节点上的等待时间下限。节点分别取
+        # 该值与本地超时的较大值，用于 shm_req 申请和 Router 等待进入推理系统。
+        ("pd_high_priority_request_time_out_seconds", ctypes.c_int),
         ("suggested_dp_index", ctypes.c_int),  # suggest dp index, deepseekv2 dp mode, use to suggest used dp_index
         # in pd split mode, use to keep the id of pd master
         ("pd_master_node_id", NodeUUId),
@@ -337,6 +343,9 @@ def init(self, tokenizer, **kwargs):
         self.min_new_tokens = kwargs.get("min_new_tokens", 1)
         self.input_penalty = kwargs.get("input_penalty", DEFAULT_INPUT_PENALTY)
         self.group_request_id = kwargs.get("group_request_id", -1)
+        # 这两个字段是 PD Master 的内部调度信息，不能由外部请求参数开启或修改。
+        self.pd_high_priority_request = False
+        self.pd_high_priority_request_time_out_seconds = 0
         self.suggested_dp_index = kwargs.get("suggested_dp_index", -1)
 
         self.skip_special_tokens = kwargs.get("skip_special_tokens", SKIP_SPECIAL_TOKENS)
@@ -503,6 +512,8 @@ def to_dict(self):
             "allowed_token_ids": self.allowed_token_ids.to_list(),
             "invalid_token_ids": self.invalid_token_ids.to_list(),
             "group_request_id": self.group_request_id,
+            "pd_high_priority_request": self.pd_high_priority_request,
+            "pd_high_priority_request_time_out_seconds": self.pd_high_priority_request_time_out_seconds,
             "skip_special_tokens": self.skip_special_tokens,
             "add_special_tokens": self.add_special_tokens,
             "add_spaces_between_special_tokens": self.add_spaces_between_special_tokens,

lightllm/server/core/objs/start_args_type.py

@@ -22,7 +22,7 @@ class StartArgs:
     pd_master_ip: str = field(default="0.0.0.0")
     pd_master_port: int = field(default=1212)
     pd_master_mode: str = field(default="elastic")
-    disable_pd_master_decode_capacity_limit: bool = field(default=False)
+    enable_pd_node_self_request_limit: bool = field(default=False)
     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/manager.py

@@ -36,9 +36,13 @@
 from .manager_ext import HttpRlManagerHelper
 from lightllm.utils.statics_utils import MovingAverage
 from lightllm.utils.config_utils import get_vocab_size
-from lightllm.utils.envs_utils import get_unique_server_name
+from lightllm.utils.envs_utils import (
+    get_pd_node_router_wait_timeout_seconds,
+    get_pd_node_shm_req_alloc_timeout_seconds,
+    get_unique_server_name,
+)
 from lightllm.utils.shm_port_args import get_shm_port_args
-from lightllm.utils.error_utils import ClientDisconnected, PDPrefillNodeStopGenToken
+from lightllm.utils.error_utils import ClientDisconnected, PDPrefillNodeStopGenToken, ServerBusyError
 from rpyc.utils.classic import obtain
 
 logger = init_logger(__name__)
@@ -117,6 +121,13 @@ def __init__(
 
         self.pd_mode: NodeRole = NodeRole(self.args.run_mode)
         assert self.pd_mode in [NodeRole.NORMAL, NodeRole.P, NodeRole.D]
+        # HTTP server 只负责在本地 shm_req 或 Router 等待过久时快速返回繁忙，PD Master 负责 QPS 准入限流。
+        # 该开关控制 P/D 节点是否启用这两类本地等待超时；多机 TP 从节点不独立拒绝请求。
+        self.pd_node_request_limit_enabled: bool = (
+            self.args.enable_pd_node_self_request_limit and self.pd_mode.is_P_or_D() and not self.is_multinode_tp_slave
+        )
+        self.pd_node_shm_req_alloc_timeout_seconds = get_pd_node_shm_req_alloc_timeout_seconds()
+        self.pd_node_router_wait_timeout_seconds = get_pd_node_router_wait_timeout_seconds()
         self.id_gen = ReqIDGenerator()
         self.first_time_costs = MovingAverage()
         self.per_token_costs = MovingAverage()
@@ -422,22 +433,17 @@ async def generate(
                 #
                 # 这样会缩小 Prefill 节点自身健康检查的覆盖范围：prompt encode、资源上报及 Decode
                 # 资源等待阶段不再计入本地推理健康状态。资源分配异常应由 PD master 侧的运行请求计数、
-                # Decode 节点健康检查和等待资源的超时逻辑负责监控，不能依赖 Prefill 推理计数判断。
+                # Decode 节点健康检查和本地 shm_req 等待超时负责监控，不能依赖 Prefill 推理计数判断。
                 await self._register_running_request()
                 running_request_registered = True
 
-            # 申请资源并存储
-            alloced_req_indexes = []
-            while len(alloced_req_indexes) < sampling_params.n:
-                alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
-                sleep_time = 0.1
-                while alloc_req_index is None:
-                    await asyncio.sleep(sleep_time)
-                    sleep_time *= 1.1
-                    sleep_time = min(1, sleep_time)
-
-                    alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
-                alloced_req_indexes.append(alloc_req_index)
+            # 申请资源并存储。PD 高优先级请求仍以更短的间隔抢占资源；开启本地限流时，
+            # 使用 PD Master 下发的较长超时时间，避免资源异常时一直等待。
+            alloced_req_indexes = await self._alloc_shm_req_indexes(
+                sampling_params.n,
+                pd_high_priority_request=sampling_params.pd_high_priority_request,
+                pd_high_priority_request_time_out_seconds=sampling_params.pd_high_priority_request_time_out_seconds,
+            )
             req_objs: List[Req] = []
             for i, req_index in enumerate(alloced_req_indexes):
                 req_obj = await self.shm_req_manager.async_get_req_obj_by_index(req_index)
@@ -543,6 +549,55 @@ def _count_multimodal_tokens(self, multimodal_params: MultimodalParams) -> Tuple
 
         return image_tokens, audio_tokens
 
+    async def _alloc_shm_req_indexes(
+        self,
+        req_num: int,
+        pd_high_priority_request: bool = False,
+        pd_high_priority_request_time_out_seconds: int = 0,
+    ) -> List[int]:
+        """为一个请求申请全部 shm_req 索引，申请失败时回滚已分配的索引。
+
+        未开启本地限流时无限等待。开启限流后，普通请求使用节点的 shm_req 申请
+        超时时间；高优先级请求取本地超时与 PD Master 下发值中的较大值。
+        """
+        alloced_req_indexes = []
+        alloc_timeout_seconds = None
+        if self.pd_node_request_limit_enabled:
+            alloc_timeout_seconds = self.pd_node_shm_req_alloc_timeout_seconds
+            if pd_high_priority_request:
+                alloc_timeout_seconds = max(
+                    alloc_timeout_seconds,
+                    pd_high_priority_request_time_out_seconds,
+                )
+        alloc_deadline = time.monotonic() + alloc_timeout_seconds if alloc_timeout_seconds is not None else None
+
+        try:
+            while len(alloced_req_indexes) < req_num:
+                alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
+                # 保持相同的退避起点，仅通过系数让高优先级请求更快地重新尝试获取 shm_req。
+                sleep_time_factor = 0.2 if pd_high_priority_request else 1
+                sleep_time = 0.1
+                while alloc_req_index is None:
+                    if alloc_deadline is not None and time.monotonic() >= alloc_deadline:
+                        logger.warning(
+                            f"{self.args.run_mode} node shm_req allocation timed out after "
+                            f"{alloc_timeout_seconds} seconds"
+                        )
+                        raise ServerBusyError(
+                            f"PD {self.args.run_mode} node is busy: unable to allocate a shm_req object "
+                            f"within {alloc_timeout_seconds} seconds"
+                        )
+                    await asyncio.sleep(sleep_time * sleep_time_factor)
+                    sleep_time = min(1, sleep_time * 1.1)
+                    alloc_req_index = await self.shm_req_manager.async_alloc_req_index()
+                alloced_req_indexes.append(alloc_req_index)
+            return alloced_req_indexes
+        except BaseException:
+            # 批量申请中途失败时，释放已申请的索引，避免 shm_req 资源泄漏。
+            for req_index in alloced_req_indexes:
+                await self.shm_req_manager.async_release_req_index(req_index)
+            raise
+
     async def _log_req_header(self, request_headers, group_request_id: int):
         x_request_id = request_headers.get("X-Request-Id", "")
         x_session_id = request_headers.get("X-Session-Id", "")
@@ -736,6 +791,16 @@ async def _wait_to_token_package(
             except asyncio.TimeoutError:
                 pass
 
+            if (
+                self.pd_node_request_limit_enabled
+                and is_first_token
+                and req_status.has_timed_out_waiting_for_inference(self.pd_node_router_wait_timeout_seconds)
+            ):
+                raise ServerBusyError(
+                    f"PD {self.args.run_mode} node is busy: request did not enter inference "
+                    f"within {self.pd_node_router_wait_timeout_seconds} seconds"
+                )
+
             if request is not None and await request.is_disconnected():
                 await self.abort(group_request_id)
                 raise ClientDisconnected(
@@ -1043,6 +1108,26 @@ def __init__(self, group_request_id, multimodal_params, req_objs: List[Req], sta
         )
         self.out_token_info_list = []
 
+    def has_timed_out_waiting_for_inference(self, timeout_seconds: float) -> bool:
+        """判断请求组是否已在 Router 中等待进入推理系统超时。"""
+        current_time = time.monotonic()
+        reqs = self.group_req_objs.shm_req_objs
+        # 组内任一请求已经进入新 batch，说明整个请求组已经开始执行，不能再按 Router 等待超时清理。
+        if any(req.infer_start_time > 0 for req in reqs):
+            return False
+        # 高优先级请求取本地 Router 超时与 PD Master 下发值中的较大值，既保证
+        # 它比普通请求拥有更充足的等待机会，也避免资源异常时永久滞留。
+        if any(req.sample_params.pd_high_priority_request for req in reqs):
+            timeout_seconds = max(
+                timeout_seconds,
+                reqs[0].sample_params.pd_high_priority_request_time_out_seconds,
+            )
+
+        for req in reqs:
+            if req.router_arrival_time > 0 and current_time - req.router_arrival_time >= timeout_seconds:
+                return True
+        return False
+
     def can_release(self):
         for req in self.group_req_objs.shm_req_objs:
             if not req.can_release():

lightllm/server/httpserver/pd_loop.py

@@ -20,7 +20,7 @@
 from ..pd_io_struct import PD_Master_Obj
 from lightllm.server.core.objs import StartArgs
 from lightllm.server.core.objs import SamplingParams
-from lightllm.utils.error_utils import PDPrefillNodeStopGenToken
+from lightllm.utils.error_utils import PDPrefillNodeStopGenToken, ServerBusyError
 from lightllm.utils.shm_port_args import get_shm_port_args
 
 logger = init_logger(__name__)
@@ -247,6 +247,13 @@ async def _pd_process_generate(
             await forwarding_queue.put((sub_req_id, request_output, metadata, finish_status))
     except PDPrefillNodeStopGenToken as e:
         logger.info(f"pd prefill node stop gen token for group_request_id {e.group_request_id}")
+    except ServerBusyError as e:
+        group_request_id = sampling_params.group_request_id
+        logger.warning(f"pd node rejected request {group_request_id}: {e.message}")
+        try:
+            await pd_upload_websocket.send(pickle.dumps((ObjType.PD_UPLOAD_SERVER_BUSY, group_request_id, e.message)))
+        except Exception:
+            logger.exception(f"report pd node request rejection failed, group_request_id: {group_request_id}")
     except asyncio.CancelledError:
         # PD master 主动 abort 或连接断开清理任务时会走取消路径，不需要反向重复上报。
         pass