Add CacheRoute strategy: per-KDN index, topology-aware proxy selection, and demo flags

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 建议的心跳周期覆盖本地默认

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

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)