diff --git a/core/request.py b/core/request.py index 925a018..5062ed9 100644 --- a/core/request.py +++ b/core/request.py @@ -317,6 +317,7 @@ def build_request( proxies: Optional[List[Dict[str, Any]]] = None, kdns: Optional[List[Dict[str, Any]]] = None, strategy: Optional[Any] = None, + kdn_knowledge_index: Optional[Dict[str, Dict[str, Dict[str, Any]]]] = None, ) -> "Request": """ 将原始用户请求信息payload，以及转换为完整的 Request 对象。 @@ -540,7 +541,7 @@ def build_request( # 下面两个上下文允许策略复用 scheduler 已维护的数据， # 避免重复做 embedding 或在线拼状态。 "knowledge_table": knowledge_table, - "kdn_knowledge_index": payload.get("_scheduler_kdn_knowledge_index") if isinstance(payload, dict) else None, + "kdn_knowledge_index": kdn_knowledge_index or {}, } chosen_kdn, chosen_proxy = strategy.select( kdns=kdns or [], diff --git a/scheduler/scheduler.py b/scheduler/scheduler.py index 20f6646..0b053e8 100644 --- a/scheduler/scheduler.py +++ b/scheduler/scheduler.py @@ -364,6 +364,7 @@ def _handle_client( proxies: List[Dict[str, Any]] | None = None, kdns: List[Dict[str, Any]] | None = None, strategy: Any | None = None, + kdn_knowledge_index: Dict[str, Dict[str, Dict[str, Any]]] | None = None, ) -> SchedulerRequest: """ 根据 HTTP 请求信息 @@ -394,6 +395,7 @@ def _handle_client( proxies=proxies, kdns=kdns, strategy=strategy, + kdn_knowledge_index=kdn_knowledge_index, ) if os.environ.get("SCHEDULER_VERBOSE_REQUEST_LOG", config.SCHEDULER_VERBOSE_REQUEST_LOG) == 1: print(f"[Scheduler] 构建内部 Request 成功: Request_ID={req_obj.Request_ID},\n" @@ -469,9 +471,16 @@ async def create_chat_completions(request: FastAPIRequest): ] strategy = getattr(request.app.state, "proxy_strategy", None) - payload_for_build = dict(payload) - payload_for_build["_scheduler_kdn_knowledge_index"] = getattr(request.app.state, "kdn_knowledge_index", {}) - req_obj = _handle_client(request.app, url_path, payload_for_build, client_ip, proxies=proxies, kdns=kdns, strategy=strategy,) + req_obj = _handle_client( + request.app, + url_path, + payload, + client_ip, + proxies=proxies, + kdns=kdns, + strategy=strategy, + kdn_knowledge_index=getattr(request.app.state, "kdn_knowledge_index", {}), + ) if not req_obj.Task.KDN_server_addr or not req_obj.Task.P_proxy_addr or req_obj.Task.P_proxy_port <= 0: raise RuntimeError("Routing failed: missing KDN or Proxy selection") @@ -582,9 +591,16 @@ async def create_completions(request: FastAPIRequest): ] strategy = getattr(request.app.state, "proxy_strategy", None) - payload_for_build = dict(payload) - payload_for_build["_scheduler_kdn_knowledge_index"] = getattr(request.app.state, "kdn_knowledge_index", {}) - req_obj = _handle_client(request.app, url_path, payload_for_build, client_ip, proxies=proxies, kdns=kdns, strategy=strategy,) + req_obj = _handle_client( + request.app, + url_path, + payload, + client_ip, + proxies=proxies, + kdns=kdns, + strategy=strategy, + kdn_knowledge_index=getattr(request.app.state, "kdn_knowledge_index", {}), + ) if not req_obj.Task.KDN_server_addr or not req_obj.Task.P_proxy_addr or req_obj.Task.P_proxy_port <= 0: raise RuntimeError("Routing failed: missing KDN or Proxy selection") @@ -818,11 +834,17 @@ async def debug_strategy() -> Dict[str, Any]: "is_alive": True, # list(include_dead=False) 已保证 alive } for p in proxy_infos[:10]] - return { + out = { "strategy": getattr(strategy, "name", None) or type(strategy).__name__ if strategy else None, "proxy_count": len(proxy_infos), "proxies_sample": sample, } + if strategy is not None and hasattr(strategy, "get_debug_snapshot"): + try: + out["strategy_debug"] = strategy.get_debug_snapshot() # type: ignore + except Exception: + out["strategy_debug"] = {"error": "strategy debug snapshot unavailable"} + return out # 预留：/knowledge/update 等路由以后再加 # @api.post("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/knowledge/update") @@ -831,5 +853,3 @@ async def debug_strategy() -> Dict[str, Any]: # client_ip = request.client.host if request.client else "unknown" # ... # return JSONResponse(content={"status": "ok"}) - - diff --git a/scheduler/strategy/cacheroute.py b/scheduler/strategy/cacheroute.py index fa44693..d51082f 100644 --- a/scheduler/strategy/cacheroute.py +++ b/scheduler/strategy/cacheroute.py @@ -1,6 +1,7 @@ # scheduler/strategy/cacheroute.py from __future__ import annotations +import os import threading from typing import Any, Dict, List, Optional, Tuple @@ -34,6 +35,13 @@ class CacheRouteStrategy(ProxySelectionStrategy): def __init__(self) -> None: self._lock = threading.Lock() self._proxy_cursor = 0 + # 第一阶段参数：使用阈值规则（非加权）判断 KDN 过载。 + # 值为 0 表示不启用对应阈值。 + 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) + # 记录最近一次决策快照，供 /debug/strategy 读取。 + self._last_decision: Dict[str, Any] = {} @staticmethod def _kdn_addr(kdn: Dict[str, Any]) -> str: @@ -61,14 +69,32 @@ def _kdn_qps(kdn: Dict[str, Any]) -> float: @staticmethod def _is_overloaded(kdn: Dict[str, Any]) -> bool: - """ - 第一阶段先只保留钩子，不引入新的 KDN 负载字段。 - 当前没有专门过载指标时，统一视作未过载；如果后续 heartbeat 已经开始上报， - 可以通过 meta.overloaded 直接覆盖。 - """ + """向后兼容：保留 meta.overloaded 显式覆盖。""" meta = kdn.get('meta') or {} return bool(meta.get('overloaded', False)) + def _is_overloaded_by_threshold(self, kdn: Dict[str, Any]) -> bool: + """ + 第一阶段过载判定（非加权）： + 1) meta.overloaded 显式为 True -> 过载； + 2) 若配置了阈值，按 qps/items/meta.pending_transfers 逐条判定。 + """ + if self._is_overloaded(kdn): + return True + + if self._kdn_qps_overload_th > 0 and self._kdn_qps(kdn) >= self._kdn_qps_overload_th: + return True + + if self._kdn_items_overload_th > 0 and self._kdn_items(kdn) >= self._kdn_items_overload_th: + return True + + meta = kdn.get("meta") or {} + pending = int(meta.get("pending_transfers", 0) or 0) + if self._kdn_pending_overload_th > 0 and pending >= self._kdn_pending_overload_th: + return True + + return False + def _select_kdn( self, kdns: List[Dict[str, Any]], @@ -92,6 +118,7 @@ def _select_kdn( return ranked[0] if ranked else None ranked: List[Tuple[Tuple[Any, ...], Dict[str, Any]]] = [] + decision_rows: List[Dict[str, Any]] = [] for kdn in kdns: kdn_id = str(kdn.get('kdn_id') or '') kdn_addr = self._kdn_addr(kdn) @@ -120,7 +147,7 @@ def _select_kdn( # 6. kdn_id 稳定打破并列 key = ( 0 if text_full else 1, - 0 if not self._is_overloaded(kdn) else 1, + 0 if not self._is_overloaded_by_threshold(kdn) else 1, -kv_cover_len, missing_count, self._kdn_qps(kdn), @@ -128,9 +155,24 @@ def _select_kdn( kdn_id, ) ranked.append((key, kdn)) + decision_rows.append( + { + "kdn_id": kdn_id, + "text_full": text_full, + "overloaded": self._is_overloaded_by_threshold(kdn), + "kv_cover_len": kv_cover_len, + "missing_count": missing_count, + "qps_1m": self._kdn_qps(kdn), + "items": self._kdn_items(kdn), + } + ) ranked.sort(key=lambda x: x[0]) - return ranked[0][1] if ranked else None + chosen = ranked[0][1] if ranked else None + with self._lock: + self._last_decision["kdn_candidates"] = decision_rows + 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]]: if not proxies: @@ -162,12 +204,38 @@ def _select_proxy(self, proxies: List[Dict[str, Any]]) -> Optional[Dict[str, Any ) == 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 with self._lock: idx = self._proxy_cursor % len(tied) self._proxy_cursor += 1 - return tied[idx] + chosen = tied[idx] + 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(chosen.get("proxy_id", "")) + return chosen + + def get_debug_snapshot(self) -> Dict[str, Any]: + with self._lock: + return dict(self._last_decision) def select( self,