diff --git a/proxy/proxy.py b/proxy/proxy.py index 3b50cce..04b15a2 100644 --- a/proxy/proxy.py +++ b/proxy/proxy.py @@ -135,17 +135,33 @@ async def _run_cp(): ) ) - # 1) register（失败不应阻塞业务启动：允许 proxy 单独跑） - try: - reg = await client.register( - proxy_id=PROXY_ID, - host=PROXY_ADVERTISE_HOST, - port=PROXY_ADVERTISE_PORT, - endpoints=["chat/completions", "completions"], - meta={"version": "proxy_v1"}, - max_capacity=PROXY_MAX_CAPACITY, - instance_count=PROXY_INSTANCE_COUNT, - kv_mem_per_instance_gb=PROXY_KV_MEM_PER_INSTANCE_GB, + # 1) register（失败不应阻塞业务启动：允许 proxy 单独跑） + try: + # CacheRoute 第二阶段： + # 可选注入 KDN->Proxy 静态拓扑信息，供 Scheduler 词典序策略使用。 + # 环境变量示例： + # PROXY_KDN_LINKS_JSON='{"kdn_a":{"bandwidth_tier":3,"latency_tier":1}}' + proxy_meta: Dict[str, Any] = {"version": "proxy_v1"} + raw_links = os.environ.get("PROXY_KDN_LINKS_JSON", "").strip() + if raw_links: + try: + parsed = json.loads(raw_links) + if isinstance(parsed, dict): + proxy_meta["kdn_links"] = parsed + else: + logger.warning("[Proxy] PROXY_KDN_LINKS_JSON is not dict, ignored") + except Exception as e: + logger.warning("[Proxy] parse PROXY_KDN_LINKS_JSON failed: %s", e) + + reg = await client.register( + proxy_id=PROXY_ID, + host=PROXY_ADVERTISE_HOST, + port=PROXY_ADVERTISE_PORT, + endpoints=["chat/completions", "completions"], + meta=proxy_meta, + max_capacity=PROXY_MAX_CAPACITY, + instance_count=PROXY_INSTANCE_COUNT, + kv_mem_per_instance_gb=PROXY_KV_MEM_PER_INSTANCE_GB, kv_cache_update_policy=PROXY_KV_CACHE_UPDATE_POLICY, ) # 用 scheduler 建议的心跳周期覆盖本地默认 diff --git a/scheduler/strategy/cacheroute.py b/scheduler/strategy/cacheroute.py index d51082f..92d1d3c 100644 --- a/scheduler/strategy/cacheroute.py +++ b/scheduler/strategy/cacheroute.py @@ -21,8 +21,11 @@ class CacheRouteStrategy(ProxySelectionStrategy): - 优先非过载 KDN； - 再比 kv_ready 覆盖长度； - 最后比轻负载与稳定顺序。 - 3) Proxy 选择先保持保守：优先当前 inflight 最小者，避免一次性把 - 拓扑/历史偏好也同时引入。后续版本再把 KDN 锚点拓扑和知识亲和性接进来。 + 3) Proxy 选择（第二阶段）升级为“非加权词典序”： + - 先在 chosen_kdn 锚点下选择拓扑最优组（带宽层级优先，时延层级其次）； + - 再做负载安全窗口过滤（避免单机过载）； + - 再比较知识历史偏好（倾向同一 LLM 连续处理相近知识）； + - 最后按当前负载 + 稳定轮转打破并列。 备注： - 这里不依赖 Injection_type 做模式分支；正式策略默认先满足文本可服务性， @@ -40,8 +43,17 @@ def __init__(self) -> None: self._kdn_qps_overload_th = float(os.environ.get("SCHEDULER_CACHEROUTE_KDN_QPS_OVERLOAD_TH", "0").strip() or 0) self._kdn_items_overload_th = int(os.environ.get("SCHEDULER_CACHEROUTE_KDN_ITEMS_OVERLOAD_TH", "0").strip() or 0) self._kdn_pending_overload_th = int(os.environ.get("SCHEDULER_CACHEROUTE_KDN_PENDING_OVERLOAD_TH", "0").strip() or 0) + # Proxy 侧第二阶段参数：负载安全窗口（非加权过滤） + self._proxy_inflight_delta = int(os.environ.get("SCHEDULER_CACHEROUTE_PROXY_INFLIGHT_DELTA", "2").strip() or 0) + self._proxy_gpu_delta = float(os.environ.get("SCHEDULER_CACHEROUTE_PROXY_GPU_DELTA", "0").strip() or 0.0) + # 历史偏好衰减（每次选择后应用），默认 0.9 + self._affinity_decay = float(os.environ.get("SCHEDULER_CACHEROUTE_AFFINITY_DECAY", "0.9").strip() or 0.9) + # 每个 proxy 只保留最近最有价值的 top-k kid，避免状态无限增长 + self._affinity_topk = int(os.environ.get("SCHEDULER_CACHEROUTE_AFFINITY_TOPK", "256").strip() or 256) # 记录最近一次决策快照，供 /debug/strategy 读取。 self._last_decision: Dict[str, Any] = {} + # 粗粒度知识历史偏好：proxy_id -> kid -> decayed_score + self._proxy_kid_affinity: Dict[str, Dict[str, float]] = {} @staticmethod def _kdn_addr(kdn: Dict[str, Any]) -> str: @@ -73,6 +85,28 @@ def _is_overloaded(kdn: Dict[str, Any]) -> bool: meta = kdn.get('meta') or {} return bool(meta.get('overloaded', False)) + @staticmethod + def _topology_tiers_for_proxy(proxy: Dict[str, Any], chosen_kdn: Dict[str, Any]) -> Tuple[int, int]: + """ + 从 proxy.meta.kdn_links 中读取 chosen_kdn 的拓扑分层： + - bandwidth_tier: 越大越好（默认 0） + - latency_tier: 越小越好（默认 999） + 支持 key 为 kdn_id 或 kdn://host:port。 + """ + meta = proxy.get("meta") or {} + links = meta.get("kdn_links") or {} + if not isinstance(links, dict): + return 0, 999 + + kdn_id = str(chosen_kdn.get("kdn_id") or "") + kdn_addr = f"kdn://{chosen_kdn.get('host')}:{int(chosen_kdn.get('port') or 0)}" + item = links.get(kdn_id) or links.get(kdn_addr) or {} + if not isinstance(item, dict): + return 0, 999 + bw_tier = int(item.get("bandwidth_tier", 0) or 0) + lat_tier = int(item.get("latency_tier", 999) or 999) + return bw_tier, lat_tier + def _is_overloaded_by_threshold(self, kdn: Dict[str, Any]) -> bool: """ 第一阶段过载判定（非加权）： @@ -174,16 +208,64 @@ def _select_kdn( self._last_decision["chosen_kdn_id"] = str((chosen or {}).get("kdn_id", "")) return chosen - def _select_proxy(self, proxies: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + def _affinity_score( + self, + proxy_id: str, + knowledge_list: List[str], + length_by_kid: Dict[str, int], + ) -> float: + table = self._proxy_kid_affinity.get(proxy_id) or {} + if not table or not knowledge_list: + return 0.0 + score = 0.0 + for kid in knowledge_list: + w = float(length_by_kid.get(kid, 1) or 1) + score += float(table.get(kid, 0.0)) * w + return score + + def _select_proxy( + self, + proxies: List[Dict[str, Any]], + chosen_kdn: Optional[Dict[str, Any]], + knowledge_list: List[str], + length_by_kid: Dict[str, int], + ) -> Optional[Dict[str, Any]]: if not proxies: return None - - # 第一阶段保持简单： - # - 优先 inflight 最小 - # - 再比 qps_1m / gpu_util - # - 最后使用本地 cursor 做稳定轮换，避免完全固定在一个 proxy 上 + # Step1: 以 chosen_kdn 为锚点选拓扑最优组（无 chosen_kdn 时退化为全量候选） + topology_rows: List[Tuple[Tuple[int, int], Dict[str, Any]]] = [] + for p in proxies: + bw_tier, lat_tier = self._topology_tiers_for_proxy(p, chosen_kdn) if chosen_kdn else (0, 999) + # 词典序：带宽 tier 越大越好，时延 tier 越小越好 + topology_rows.append(((-bw_tier, lat_tier), p)) + topology_rows.sort(key=lambda x: x[0]) + best_topo = topology_rows[0][0] + topo_group = [p for key, p in topology_rows if key == best_topo] + + # Step2: 负载安全窗口过滤（避免把知识亲和性放在过载 proxy 上） + min_inflight = min(self._proxy_inflight(p) for p in topo_group) + safe_group = [ + p for p in topo_group + if self._proxy_inflight(p) <= (min_inflight + max(0, self._proxy_inflight_delta)) + ] + if self._proxy_gpu_delta > 0 and safe_group: + min_gpu = min(self._proxy_gpu(p) for p in safe_group) + safe_group = [ + p for p in safe_group + if self._proxy_gpu(p) <= (min_gpu + self._proxy_gpu_delta) + ] or safe_group + + # Step3: 在安全窗口内按知识历史偏好选最优 + affinity_rows: List[Tuple[float, Dict[str, Any]]] = [] + for p in safe_group: + pid = str(p.get("proxy_id", "")) + affinity_rows.append((self._affinity_score(pid, knowledge_list, length_by_kid), p)) + max_aff = max((x[0] for x in affinity_rows), default=0.0) + aff_group = [p for sc, p in affinity_rows if sc == max_aff] + + # Step4: 按当前负载打破并列，再用 cursor 稳定轮换 ranked = sorted( - proxies, + aff_group, key=lambda p: ( self._proxy_inflight(p), self._proxy_qps(p), @@ -192,46 +274,52 @@ def _select_proxy(self, proxies: List[Dict[str, Any]]) -> Optional[Dict[str, Any ), ) best = ranked[0] - best_key = ( - self._proxy_inflight(best), - self._proxy_qps(best), - self._proxy_gpu(best), - ) - tied = [p for p in ranked if ( - self._proxy_inflight(p), - self._proxy_qps(p), - self._proxy_gpu(p), - ) == best_key] - - if len(tied) == 1: - with self._lock: - self._last_decision["proxy_candidates"] = [ - { - "proxy_id": str(p.get("proxy_id", "")), - "inflight": self._proxy_inflight(p), - "qps_1m": self._proxy_qps(p), - "gpu_util": self._proxy_gpu(p), - } - for p in ranked - ] - self._last_decision["chosen_proxy_id"] = str(best.get("proxy_id", "")) - return best + best_key = (self._proxy_inflight(best), self._proxy_qps(best), self._proxy_gpu(best)) + tied = [p for p in ranked if (self._proxy_inflight(p), self._proxy_qps(p), self._proxy_gpu(p)) == best_key] + # 更新亲和性：选中 proxy 对本次知识集合增益；其他 proxy 衰减 + chosen = best with self._lock: - idx = self._proxy_cursor % len(tied) - self._proxy_cursor += 1 - chosen = tied[idx] + if len(tied) > 1: + idx = self._proxy_cursor % len(tied) + self._proxy_cursor += 1 + chosen = tied[idx] + + for p in proxies: + pid = str(p.get("proxy_id", "")) + table = self._proxy_kid_affinity.setdefault(pid, {}) + if not table: + continue + for kid in list(table.keys()): + table[kid] = float(table[kid]) * self._affinity_decay + if table[kid] < 1e-6: + table.pop(kid, None) + + chosen_pid = str(chosen.get("proxy_id", "")) + chosen_table = self._proxy_kid_affinity.setdefault(chosen_pid, {}) + for kid in knowledge_list: + chosen_table[kid] = float(chosen_table.get(kid, 0.0)) + 1.0 + # 只保留 top-k，避免长期实验状态膨胀 + if self._affinity_topk > 0 and len(chosen_table) > self._affinity_topk: + kept = sorted(chosen_table.items(), key=lambda kv: kv[1], reverse=True)[: self._affinity_topk] + self._proxy_kid_affinity[chosen_pid] = dict(kept) + self._last_decision["proxy_candidates"] = [ { "proxy_id": str(p.get("proxy_id", "")), "inflight": self._proxy_inflight(p), "qps_1m": self._proxy_qps(p), "gpu_util": self._proxy_gpu(p), + "topology_tier": { + "bandwidth_tier": self._topology_tiers_for_proxy(p, chosen_kdn)[0] if chosen_kdn else 0, + "latency_tier": self._topology_tiers_for_proxy(p, chosen_kdn)[1] if chosen_kdn else 999, + }, + "affinity_score": self._affinity_score(str(p.get("proxy_id", "")), knowledge_list, length_by_kid), } - for p in ranked + for p in proxies ] - self._last_decision["chosen_proxy_id"] = str(chosen.get("proxy_id", "")) - return chosen + self._last_decision["chosen_proxy_id"] = chosen_pid + return chosen def get_debug_snapshot(self) -> Dict[str, Any]: with self._lock: @@ -249,7 +337,22 @@ def select( ctx = request_ctx or {} knowledge_list = [str(x).strip().lower() for x in (ctx.get('knowledge_list') or []) if str(x).strip()] knowledge_index = ctx.get('kdn_knowledge_index') or {} + # 统一使用知识索引中的 length，作为亲和性打分的长度权重来源 + length_by_kid: Dict[str, int] = {} + for _kdn_id, by_kid in (knowledge_index or {}).items(): + if not isinstance(by_kid, dict): + continue + for kid, meta in by_kid.items(): + if not isinstance(meta, dict): + continue + if kid not in length_by_kid: + length_by_kid[kid] = int(meta.get("length", 0) or 0) chosen_kdn = self._select_kdn(kdns=kdns, knowledge_list=knowledge_list, knowledge_index=knowledge_index) - chosen_proxy = self._select_proxy(proxies=proxies) + chosen_proxy = self._select_proxy( + proxies=proxies, + chosen_kdn=chosen_kdn, + knowledge_list=knowledge_list, + length_by_kid=length_by_kid, + ) return chosen_kdn, chosen_proxy diff --git a/test/demo_proxy.py b/test/demo_proxy.py index 6bfd4db..4d52436 100644 --- a/test/demo_proxy.py +++ b/test/demo_proxy.py @@ -24,16 +24,24 @@ parser = argparse.ArgumentParser(description="Run CacheRoute Proxy") parser.add_argument("--host", type=str, default=None, help="proxy listen host (default from config/env)") parser.add_argument("--port", type=int, default=None, help="proxy listen port (default from config/env)") - parser.add_argument( - "--strategy", - type=str, - default=None, - help="instance scheduling strategy (e.g., round_robin, least_inflight)", - ) - args = parser.parse_args() - - if args.strategy: - os.environ["PROXY_INSTANCE_STRATEGY"] = args.strategy + parser.add_argument( + "--strategy", + type=str, + default=None, + help="instance scheduling strategy (e.g., round_robin, least_inflight)", + ) + parser.add_argument( + "--kdn-links-json", + type=str, + default=None, + help="optional JSON string for PROXY_KDN_LINKS_JSON (CacheRoute topology tiers)", + ) + args = parser.parse_args() + + if args.strategy: + os.environ["PROXY_INSTANCE_STRATEGY"] = args.strategy + if args.kdn_links_json: + os.environ["PROXY_KDN_LINKS_JSON"] = args.kdn_links_json # 选择一个与 Scheduler 不同的端口，例如 8001 uvicorn.run(proxy, host=config.PROXY_DP_HOST, port=config.PROXY_DP_PORT, reload=False)