fix(pd): pd limit

docs/CN/source/tutorial/api_server_args.rst

@@ -87,6 +87,26 @@ PD 分离模式参数
     推理进度健康检查：当仍有在途请求，且整个 PD Master 连续 ``HEALTH_TIMEOUT`` 秒
     没有任何请求成功返回 token 时，接口将返回 HTTP 503。
 
+.. option:: --enable_pd_node_self_request_limit
+
+    在 PD 分离模式的 Prefill 和 Decode 节点上启用本地请求准入控制。启用后，节点根据动态 QPS 和
+    当前运行请求总数决定新请求是否可以进入本地 ``shm_req`` 申请流程。超过上限时，节点会主动拒绝新请求，并由 PD Master
+    向客户端返回 HTTP 429；已经准入的请求会持续等待可用的 ``shm_req`` 对象，不再进行等待超时判断。
+    服务启动后，QPS 记录器累计完成的请求数未达到 ``running_max_req_size`` 时，最大允许进入请求数直接使用
+    ``running_max_req_size``，使冷启动阶段可以快速接收请求并积累足够的完成样本。累计完成请求数达到
+    ``running_max_req_size`` 且首个 16 请求 QPS 窗口生成后，最大允许进入请求数改为
+    ``int(QPS * 平均整���时间秒数) + 6``。如果 ``running_max_req_size`` 小于 16，节点会继续使用基础容量，
+    直到 QPS 完成初始化，避免使用尚未初始化的零值 QPS 过早切换到仅 6 个探测请求。
+    稳定运行阶段额外放行 6 个请求作为探测余量，避免低流量或长时间空闲导致 QPS 降低后，
+    系统被限制在过低的并发并且难以重新爬升。
+    这里的平均整包时间表示希望一个请求从进入
+    节点到整包完成所处的时间范围，用于将完成 QPS 换算为合理的在途请求数量。Prefill 阶段通常只负责输入
+    处理和首 token，默认值为 20 秒，避免短阶段积压过多请求；Decode 阶段需要持续生成 token，默认值为
+    60 秒，以覆盖更长的整包处理时间。可以通过环境变量
+    ``LIGHTLLM_PD_REQUEST_LIMIT_MAX_ALLOWED_REQUEST_COUNT_SECONDS`` 显式设置统一值；设置后 Prefill 和 Decode
+    节点都会采用该值。
+    该参数默认关闭，在 ``normal`` 和 ``pd_master`` 模式下不生效。
+
 .. option:: --config_server_host
 
     配置服务器模式下的主机地址
@@ -148,6 +168,8 @@ PD 分离模式参数
 .. option:: --running_max_req_size
 
     同时进行前向推理的最大请求数量，默认为 ``1000``
+    在 PD 分离模式的 Decode 节点上，该限制仅在各节点本地生效；
+    PD Master 不会汇总各 Decode 节点的值作为全局请求准入上限。
 
 .. option:: --max_req_total_len
 

docs/EN/source/tutorial/api_server_args.rst

@@ -89,6 +89,28 @@ 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 admission control on Prefill and Decode nodes in PD disaggregation mode. When enabled,
+    each node uses the dynamically measured QPS and its total running request count to decide whether a new request
+    may enter local ``shm_req`` allocation. Once the limit is exceeded, the node rejects new requests through PD Master with HTTP 429. Admitted requests
+    keep waiting for an available ``shm_req`` object without an allocation timeout. While the QPS recorder has completed
+    fewer than ``running_max_req_size`` requests since startup, that base capacity is returned
+    directly so the cold-start phase can quickly collect enough completion samples. Once the cumulative completed request
+    count reaches ``running_max_req_size`` and the first 16-request QPS window is available, the maximum admitted request
+    count becomes ``int(QPS * average whole-request seconds) + 6``. If
+    ``running_max_req_size`` is below 16, the base
+    capacity remains in effect until QPS initialization, preventing an uninitialized zero QPS from reducing the limit
+    prematurely. During steady state, six extra requests are admitted as probing headroom, preventing low traffic or a
+    long idle period from trapping the service at very low concurrency.
+    This duration represents the desired
+    average time range from node admission until the whole request finishes and converts completion QPS into a
+    reasonable in-flight request count. Prefill defaults to 20 seconds because it mainly handles input processing and
+    the first token; Decode defaults to 60 seconds because incremental token generation usually keeps the request active
+    longer. Set ``LIGHTLLM_PD_REQUEST_LIMIT_MAX_ALLOWED_REQUEST_COUNT_SECONDS`` to override both node defaults with one
+    explicit value.
+    The option is disabled by default and has no effect in ``normal`` or ``pd_master`` mode.
+
 .. option:: --config_server_host
 
     Host address in configuration server mode
@@ -150,6 +172,8 @@ Memory and Batch Processing Parameters
 .. option:: --running_max_req_size
 
     Maximum number of requests for simultaneous forward inference, default is ``1000``
+    On Decode nodes in PD disaggregation mode, this limit applies locally to each node;
+    PD Master does not aggregate the Decode-node values into a global admission limit.
 
 .. option:: --max_req_total_len
 

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=(
+            "Allow Prefill and Decode nodes in PD mode to limit the total running request count "
+            "according to the dynamically measured QPS. Default: disabled."
+        ),
     )
     parser.add_argument(
         "--pd_trans_mode",

lightllm/server/core/objs/sampling_params.py

@@ -294,6 +294,8 @@ 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 为分段续跑请求设置，避免已经完成首段的请求因本地 shm_req 限流失败。
+        ("bypass_pd_node_request_limit", ctypes.c_bool),
         ("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 +339,8 @@ 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.bypass_pd_node_request_limit = False
         self.suggested_dp_index = kwargs.get("suggested_dp_index", -1)
 
         self.skip_special_tokens = kwargs.get("skip_special_tokens", SKIP_SPECIAL_TOKENS)
@@ -503,6 +507,7 @@ 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,
+            "bypass_pd_node_request_limit": self.bypass_pd_node_request_limit,
             "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

@@ -34,11 +34,12 @@
 from lightllm.server.metrics.manager import MetricClient
 from .rl_controller import HttpRlController
 from .manager_ext import HttpRlManagerHelper
+from .qps_recorder import QPSRecorder
 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.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,9 +118,15 @@ def __init__(
 
         self.pd_mode: NodeRole = NodeRole(self.args.run_mode)
         assert self.pd_mode in [NodeRole.NORMAL, NodeRole.P, NodeRole.D]
+        # 该开关只对 PD 分离模式的 Prefill/Decode 服务主节点生效。
+        # 多机 TP 的从节点不直接与 PD Master 通信，不能独立拒绝请求。
+        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.id_gen = ReqIDGenerator()
         self.first_time_costs = MovingAverage()
         self.per_token_costs = MovingAverage()
+        self.qps_recorder = QPSRecorder(self.args)
         # 有的模型的vocab size 读取tokenizer和config.json中不一致
         self.vocab_size = max(get_vocab_size(args.model_dir), self.tokenizer.vocab_size)
 
@@ -422,22 +429,21 @@ async def generate(
                 #
                 # 这样会缩小 Prefill 节点自身健康检查的覆盖范围：prompt encode、资源上报及 Decode
                 # 资源等待阶段不再计入本地推理健康状态。资源分配异常应由 PD master 侧的运行请求计数、
-                # Decode 节点健康检查和等待资源的超时逻辑负责监控，不能依赖 Prefill 推理计数判断。
+                # Decode 节点健康检查和本地准入控制负责监控，不能依赖 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 分段续跑请求可以绕过本地进入并发限制，避免
+            # 已经成功完成首段的用户请求因为下一段暂时无法进入等待流程而被 429 中断。
+            if self.pd_node_request_limit_enabled and sampling_params.bypass_pd_node_request_limit:
+                logger.info(
+                    f"PD {self.args.run_mode} node request {group_request_id} bypasses the local request "
+                    f"concurrency limit and will wait for {sampling_params.n} shm_req object(s)"
+                )
+            alloced_req_indexes = await self._alloc_shm_req_indexes(
+                sampling_params.n,
+                bypass_pd_node_request_limit=sampling_params.bypass_pd_node_request_limit,
+            )
             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,46 @@ def _count_multimodal_tokens(self, multimodal_params: MultimodalParams) -> Tuple
 
         return image_tokens, audio_tokens
 
+    async def _alloc_shm_req_indexes(
+        self,
+        req_num: int,
+        bypass_pd_node_request_limit: bool = False,
+    ) -> List[int]:
+        """为一个请求申请全部 shm_req 索引，申请失败时回滚已分配的索引。"""
+        alloced_req_indexes = []
+
+        try:
+            if self.pd_node_request_limit_enabled and not bypass_pd_node_request_limit:
+                current_request_count = self.run_reqs_count_mark.get_value()
+                # QPS 记录器累计完成的请求数未达到 running_max_req_size 时返回基础容量，
+                # 使服务冷启动后可以快速积累足够样本；达到后再根据完成 QPS 和节点平均
+                # 整包处理时间估算准入上限，并额外保留 6 个请求的探测余量，避免系统在
+                # 低 QPS 状态下恢复过慢。Prefill 默认按 20 秒估算，Decode 默认按
+                # 60 秒估算，环境变量可以覆盖对应节点的默认时间。
+                max_allowed_request_count = self.qps_recorder.get_max_allowed_request_count()
+                if current_request_count > max_allowed_request_count:
+                    logger.warning(
+                        f"PD {self.args.run_mode} node rejects a request before shm_req allocation: "
+                        f"running_request_count={current_request_count}, "
+                        f"max_allowed_request_count={max_allowed_request_count}"
+                    )
+                    raise ServerBusyError(f"PD {self.args.run_mode} node is busy")
+
+            while len(alloced_req_indexes) < req_num:
+                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 = 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", "")
@@ -777,6 +823,7 @@ async def _wait_to_token_package(
                         unfinished_count -= 1
 
                     if unfinished_count == 0:
+                        self.qps_recorder.mark_one_req_finish()
                         total_cost_time_ms = (time.time() - start_time) * 1000
                         mean_per_token_cost_time_ms = (total_cost_time_ms - first_token_cost_ms) / out_token_counter
                         self.per_token_costs.add(mean_per_token_cost_time_ms)

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

lightllm/server/httpserver/qps_recorder.py

@@ -0,0 +1,79 @@
+import time
+from collections import deque
+from threading import Lock
+from typing import Deque, Optional
+
+from lightllm.utils.envs_utils import (
+    get_pd_request_limit_max_allowed_request_count_seconds,
+)
+
+
+class QPSRecorder:
+    """根据最近完成的请求计算系统动态 QPS。"""
+
+    def __init__(self, args, ema_alpha: float = 0.1):
+        if not 0 < ema_alpha <= 1:
+            raise ValueError("ema_alpha must be in the range (0, 1]")
+
+        self.args = args
+        self.ema_alpha = float(ema_alpha)
+        # 保存最近 16 个请求的完成时间。16 个时间点之间包含 15 个完成间隔。
+        self._finished_timestamps: Deque[float] = deque(maxlen=16)
+        # 记录服务启动后已经完成的请求总数，用于判断冷启动阶段是否已收集足够样本。
+        self._finished_request_count = 0
+        self._qps = 0.0
+        self._initialized = False
+        self._last_qps_update_time: Optional[float] = None
+        self._lock = Lock()
+
+    def mark_one_req_finish(self) -> None:
+        """记录一个请求完成事件，并在样本充足时更新全局 QPS。"""
+        finished_time = time.monotonic()
+        with self._lock:
+            self._finished_timestamps.append(finished_time)
+            self._finished_request_count += 1
+            self._update_qps()
+
+    def get_qps(self) -> float:
+        """返回经过 EMA 平滑后的全局 QPS。"""
+        with self._lock:
+            if self._last_qps_update_time is not None:
+                current_time = time.monotonic()
+                if current_time - self._last_qps_update_time > 30:
+                    self._update_qps()
+            return self._qps
+
+    def get_max_allowed_request_count(self) -> int:
+        """根据冷启动样本数和动态 QPS 返回最大允许进入请求数。
+
+        服务启动后累计完成的请求数尚未达到 ``running_max_req_size`` 时，直接返回
+        节点的基础运行容量，使冷启动阶段能够快速接收请求并积累足够的 QPS 样本。
+        当 ``running_max_req_size`` 小于 16 时，还需要等待首个完整 QPS 窗口生成，
+        避免在 QPS 尚未初始化时过早切换到仅 6 个探测请求。满足两个条件后，才根据
+        完成 QPS 和平均整包时间估算允许进入的请求数，并额外放行 6 个请求作为
+        探测余量，避免系统在低 QPS 状态下恢复过慢。
+        """
+        with self._lock:
+            finished_request_count = self._finished_request_count
+            qps_initialized = self._initialized
+        if finished_request_count < self.args.running_max_req_size or not qps_initialized:
+            return self.args.running_max_req_size
+
+        return int(self.get_qps() * get_pd_request_limit_max_allowed_request_count_seconds(self.args.run_mode)) + 6
+
+    def _update_qps(self) -> None:
+        if len(self._finished_timestamps) < self._finished_timestamps.maxlen:
+            return
+
+        current_time = time.monotonic()
+        elapsed_time = current_time - self._finished_timestamps[0]
+        if elapsed_time <= 0:
+            return
+
+        average_qps = (len(self._finished_timestamps) - 1) / elapsed_time
+        if not self._initialized:
+            self._qps = average_qps
+            self._initialized = True
+        else:
+            self._qps = self.ema_alpha * average_qps + (1 - self.ema_alpha) * self._qps
+        self._last_qps_update_time = current_time