diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index 99f123e66..5bcae0191 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -670,6 +670,8 @@ def register_pd(self, pd_info_json, websocket): # stale history does not bias CacheAware toward the zero-valued node. for prefill_node in self.prefill_nodes: prefill_node.dispatched_prompt_chars = 0 + prefill_node.recent_dispatched_chars = 0.0 + prefill_node.last_decay_ts = 0.0 elif pd_client.mode == "decode": self.decode_nodes = [e for e in self.decode_nodes if e.client_ip_port != pd_client.client_ip_port] self.decode_nodes.append(pd_client) diff --git a/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py b/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py index b7b2d525e..a4b8c627a 100644 --- a/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py +++ b/lightllm/server/httpserver_for_pd_master/pd_selector/cache_aware.py @@ -10,13 +10,15 @@ - 用前缀树(见 PromptCacheTree)记录「历史 prompt -> 处理它的 worker」; - 树中的 prefill_node 对应 worker.client_ip_port; - prompt 会按 sample_stride 抽稀后再插入/匹配,降低树的深度与内存; - - 用 worker.dispatched_prompt_chars(累计派发的 prompt 字符数)做粗粒度均衡。 + - 用 worker.recent_dispatched_chars(按 balance_half_life_secs 半衰期衰减的近期派发 + prompt 字符数)做粗粒度均衡,避免冷启动累计值掩盖「曾经忙、现在闲」的节点。 选点流程见 CacheAwarePolicy.select_worker。 """ from __future__ import annotations +import time from dataclasses import dataclass from typing import List, Optional @@ -37,6 +39,7 @@ class CacheAwareConfig: cache_threshold: float = 0.5 # 派发量不均衡判定:max > min * balance_rel_threshold 时强制选派发量最少的节点。 balance_rel_threshold: float = 1.2 + balance_half_life_secs: float = 60.0 # 前缀树允许的最大节点数(不含 root)。 max_node_count: int = 1_000_000 # 每次 LRU 驱逐的叶节点数量。 @@ -66,13 +69,23 @@ def __init__(self, config: Optional[CacheAwareConfig] = None) -> None: recursion_limit=self.config.recursion_limit, ) + def _decay_recent(self, workers: List[PD_Client_Obj]) -> None: + now = time.monotonic() + half_life = self.config.balance_half_life_secs + for worker in workers: + if half_life > 0 and worker.last_decay_ts > 0: + elapsed = now - worker.last_decay_ts + if elapsed > 0: + worker.recent_dispatched_chars *= 0.5 ** (elapsed / half_life) + worker.last_decay_ts = now + def _select_worker_min_dispatched( self, workers: List[PD_Client_Obj], request_text: str, ) -> PD_Client_Obj: - """派发量优先兜底:选择累计 dispatched_prompt_chars 最小的 worker,并写入前缀树。""" - min_dispatched_worker = min(workers, key=lambda worker: worker.dispatched_prompt_chars) + """派发量优先兜底:选择近期 dispatched 最小的 worker,并写入前缀树。""" + min_dispatched_worker = min(workers, key=lambda worker: worker.recent_dispatched_chars) self.prompt_cache_tree.insert(request_text, min_dispatched_worker.client_ip_port) return min_dispatched_worker @@ -89,28 +102,30 @@ def select_worker(self, workers: List[PD_Client_Obj], request_text: str) -> Opti 决策顺序: 1) workers 为空 -> 返回 None; - 2) 若 max(dispatched) > min(dispatched) * balance_rel_threshold, - 认为派发不均衡,直接选派发量最少的节点; - 3) 否则对 request_text 做前缀匹配,计算 + 2) 先按半衰期衰减各 worker 的 recent_dispatched_chars; + 3) 若 max(recent) > min(recent) * balance_rel_threshold, + 认为近期派发不均衡,直接选近期派发量最少的节点; + 4) 否则对 request_text 做前缀匹配,计算 match_rate = matched_char_count / input_char_count; - 4) match_rate > cache_threshold 且命中 prefill_node 仍在线 -> 路由到该节点并更新树; - 5) 未命中阈值或 prefill_node 不在当前 workers 中 -> 回退到派发量最少选择。 + 5) match_rate > cache_threshold 且命中 prefill_node 仍在线 -> 路由到该节点并更新树; + 6) 未命中阈值或 prefill_node 不在当前 workers 中 -> 回退到近期派发量最少选择。 """ if not workers: return None if len(request_text) <= 1: raise ValueError(f"request_text length must be > 1, got {len(request_text)}") - # ---- 1. 派发均衡门闩:差距过大时不再追求 cache 亲和 ---- - dispatched_chars = [worker.dispatched_prompt_chars for worker in workers] + self._decay_recent(workers) + dispatched_chars = [worker.recent_dispatched_chars for worker in workers] min_dispatched = min(dispatched_chars) if dispatched_chars else 0 max_dispatched = max(dispatched_chars) if dispatched_chars else 0 is_imbalanced = max_dispatched > (min_dispatched * self.config.balance_rel_threshold) logger.info( - f"CacheAwarePolicy: min_dispatched={min_dispatched}, max_dispatched={max_dispatched}, " + f"CacheAwarePolicy: min_dispatched={min_dispatched:.0f}, max_dispatched={max_dispatched:.0f}, " f"balance_rel_threshold={self.config.balance_rel_threshold:.4f}, " + f"balance_half_life_secs={self.config.balance_half_life_secs:.1f}, " f"is_imbalanced={is_imbalanced}" ) diff --git a/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py b/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py index d1315ae82..9d125f603 100644 --- a/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py +++ b/lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py @@ -89,6 +89,7 @@ def select_p_d_node( # 累计派发字符数,供后续 cache-aware 做派发均衡判断。 p_node.dispatched_prompt_chars += len(prompt) + p_node.recent_dispatched_chars += len(prompt) logger.info( f"LoadBalancedCacheAwareSelector: selected p_node={p_node.client_ip_port}, " diff --git a/lightllm/server/pd_io_struct.py b/lightllm/server/pd_io_struct.py index d0f32419d..61268530f 100644 --- a/lightllm/server/pd_io_struct.py +++ b/lightllm/server/pd_io_struct.py @@ -58,6 +58,8 @@ class PD_Client_Obj: run_status: _PD_Client_RunStatus = field(default_factory=_PD_Client_RunStatus) # cache-aware 选点用:累计派发到该节点的 prompt 字符数(只增不减,非实时负载)。 dispatched_prompt_chars: int = 0 + recent_dispatched_chars: float = 0.0 + last_decay_ts: float = 0.0 def __post_init__(self): if self.mode not in ["prefill", "decode"]: diff --git a/unit_tests/server/test_cache_aware_balance.py b/unit_tests/server/test_cache_aware_balance.py new file mode 100644 index 000000000..f623cfa07 --- /dev/null +++ b/unit_tests/server/test_cache_aware_balance.py @@ -0,0 +1,55 @@ +"""cache-aware PD prefill 选点器均衡门闩的单测:验证近期派发量(半衰期衰减)能打破 cache 亲和导致的单点饥饿。""" +import time + +import pytest + +from lightllm.server.pd_io_struct import PD_Client_Obj +from lightllm.server.httpserver_for_pd_master.pd_selector.cache_aware import ( + CacheAwareConfig, + CacheAwarePolicy, +) + + +def _worker(name: str) -> PD_Client_Obj: + return PD_Client_Obj(node_id=0, client_ip_port=name, mode="prefill", start_args={}) + + +PROMPT = "Explain paged KV cache. " * 200 # long enough to span a sample stride + + +def test_decay_math(): + cfg = CacheAwareConfig(balance_half_life_secs=60.0) + policy = CacheAwarePolicy(cfg) + w = _worker("A:1") + w.recent_dispatched_chars = 1000.0 + w.last_decay_ts = time.monotonic() - 120.0 # two half-lives idle + policy._decay_recent([w]) + assert w.recent_dispatched_chars == pytest.approx(250.0, rel=1e-3) + + +def test_starvation_redirects_to_idle_node(): + cfg = CacheAwareConfig(balance_rel_threshold=1.2, balance_half_life_secs=60.0) + policy = CacheAwarePolicy(cfg) + a, b = _worker("A:1"), _worker("B:1") + policy.prompt_cache_tree.insert(PROMPT, a.client_ip_port) + a.recent_dispatched_chars = 1000.0 + b.recent_dispatched_chars = 1000.0 + now = time.monotonic() + a.last_decay_ts = now + b.last_decay_ts = now - 120.0 + chosen = policy.select_worker([a, b], request_text=PROMPT) + assert chosen.client_ip_port == "B:1" + + +def test_balanced_keeps_cache_affinity(): + cfg = CacheAwareConfig(balance_rel_threshold=1.2, balance_half_life_secs=60.0) + policy = CacheAwarePolicy(cfg) + a, b = _worker("A:1"), _worker("B:1") + policy.prompt_cache_tree.insert(PROMPT, a.client_ip_port) + a.recent_dispatched_chars = 1000.0 + b.recent_dispatched_chars = 1000.0 + now = time.monotonic() + a.last_decay_ts = now + b.last_decay_ts = now + chosen = policy.select_worker([a, b], request_text=PROMPT) + assert chosen.client_ip_port == "A:1"