From 4f64e089b2bffc4b18dfe5496d290fd56ec35bd6 Mon Sep 17 00:00:00 2001 From: Heyao zhang <121943721+zhy1658858023@users.noreply.github.com> Date: Tue, 24 Mar 2026 15:32:33 +0800 Subject: [PATCH] Start cacheroute phase1: explicit context + overload thresholds --- core/request.py | 1265 +++++++++++----------- scheduler/knowledge/kdn_sync.py | 636 ++++++----- scheduler/scheduler.py | 1644 +++++++++++++++-------------- scheduler/strategy/base.py | 56 +- scheduler/strategy/cacheroute.py | 255 +++++ scheduler/strategy/factory.py | 43 +- scheduler/strategy/round_robin.py | 95 +- test/demo_scheduler.py | 112 +- 8 files changed, 2257 insertions(+), 1849 deletions(-) create mode 100644 scheduler/strategy/cacheroute.py diff --git a/core/request.py b/core/request.py index 9397d30..5062ed9 100644 --- a/core/request.py +++ b/core/request.py @@ -1,624 +1,641 @@ -# 构建任务和任务批次的信息结构 -import time -from dataclasses import dataclass, asdict -from core.tokenizer_registry import estimate_tokens -from core.config import ( - DEFAULT_EMBED_MODEL, - SCHEDULER_RETRIEVAL_TOP_K, - SCHEDULER_RETRIEVAL_MIN_SCORE, - SCHEDULER_RETRIEVAL_MIN_RATIO, -) -# from CacheRoute.util import parse_host_port -from typing import Dict, Any, List, Tuple, Optional -from store import EmbeddingModel, KnowledgeTable -from util import parse_stream_flag - -""" -定义调度器-本地代理的交互数据结构体 -> Class Request，调度器通过调度策略构建request并发给本地Proxy，随后Proxy按需求执行对应动作完成选池推理 - - Class Request: - | Request_ID:用户任务的唯一标识ID - |__Request_type：请求类型，如request，control，update等 - | - |__Class Prompt: 记录用户问题的基本信息，如具体问题，模型类型，问题长度等 - | | model 模型类型 - | | user_prompt 用户问题 - | | token_length 问题的token长度 - | | max_token 最长生成token数 - | | stream 是否启用流式传输 - | | temperature 采样温度，默认1.0 - | |__top_p nucleus sampling 截断阈值，默认1.0 - | - |__Class Service： 映射用户问题的服务需求，如是否支持PD分离，是否支持知识注入，TTFT，E2E和TPOT的SLO需求 - | | Enable_PD_Disaggregation 是否启用PD分离 - | | Enable_know_injection 是否启用知识注入 - | | Injection_type 知识注入类型（文本orKV） - | | Enable_compress 是否允许对KVCache压缩 - | | Compress_factor 压缩比例 - | | Enable_security 是否启用安全模式（预留接口） - | | Knowledge_block_num 知识块数量 - | | Knowledge_List 问题所需知识清单 - | | Knowledge_length 知识注入长度 - | | SLO_TTFT 首token生成时间 - | | SLO_E2E 端到端时延 - | | SLO_TPOT 每秒token生成 - | |__Endpoint_type 转发给vLLM的请求类型，包括 chat/completions 或 completions，对话类或补全类 - | - |__Class Task： 记录调度相关的任务信息，如KDN知识服务器的IP地址，PD池代理的IP地址，端口号 - | User_addr 用户IP地址 - | KDN_server_addr KDN服务器地址 - | default_know_addr 默认文本注入服务器地址 - | P_proxy_id 选择代理的ID号 - | P_proxy_addr P池代理IP地址 - | D_proxy_addr D池代理IP地址 - | P_proxy_port P池代理端口号 - | D_proxy_port D池代理端口号 - | prefill_instance 该请求分配的Prefill实例IP地址及端口 - | decode_instance 该请求分配的Decode实例IP地址及端口 - |__batch_order 该请求在其所分配实例中的批次顺序 -""" - -# ======================================================================================================================== -# ------------------------------------------------------Prompt------------------------------------------------------------ -# ======================================================================================================================== -@dataclass -class Prompt: - """ - 定义用户问题的基本信息 - model<任务模型，str> - user_prompt<用户问题，str> - token_length<问题的token长度> - bs<任务组的batch_size，通常情况下默认1> - max_token<任务支持生成的最大token数> - stream<是否启用流式传输> - temperature<采样温度，默认1.0> - top_p - """ - model: str - user_prompt: str - token_length: int - bs: int - - max_tokens: int - stream: bool - temperature: float - top_p: float - - - @classmethod - def extract_prompt_info(cls, model: str, user_prompt: str) -> int: - """ - 根据 user_prompt 自动计算 token_length并返回。 - 输入：model，user_prompt - 输出：token_length - """ - seq_length = estimate_tokens(user_prompt, model) - print(f"[Prompt-tokenizer]: get task prompt_length complete, prompt_length={seq_length}") - - return seq_length - - -# ======================================================================================================================== -# ------------------------------------------------------Service----------------------------------------------------------- -# ======================================================================================================================== -@dataclass -class Service: - """ - 定义问题服务的SLO基本信息，通过用户IP地址映射具体服务等级，支持映射模块出于安全、个性化等方面的扩展 - Enable_PD_Disaggregation<是否允许问题进行PD分离处理，默认为True> - Enable_know_injection<是否允许调用远端知识，默认为True> - Injection_type<知识注入类型，默认为text，由proxy后续根据策略调整> - Enable_compress<是都允许进行KVCache的压缩，默认为True> - Compress_factor - Enable_security<是否启用安全模式，默认false，为后续安全内容提供接口> - Knowledge_block_num<任务注入知识块的top数量，默认为3> - Knowledge_List<知识块列表，里面的元素是Knowledge_ID,表征一个具体的知识块> - Knowledge_length<任务注入知识库的token长度，默认为0> - SLO_TTFT<任务组的TTFT SLO需求，即问题开始推理至产生提一个token所需要的时间，默认2000ms> - SLO_E2E<任务从开始Prefill到结束Decode所需的完整时间，ms> - SLO_TPOT<任务组的TPOT SLO需求，即自回归推理阶段平均生成默认20ms> - Endpoint_type<转发给vLLM的请求类型，包括 chat/completions 或 completions，对话类或补全类> - """ - Enable_PD_Disaggregation: bool - Enable_know_injection: bool - Injection_type: str - Enable_compress: bool - Compress_factor: float - Enable_security: bool - Knowledge_block_num: int - Knowledge_List: List[str] - Knowledge_length: int - SLO_TTFT: int - SLO_E2E: int - SLO_TPOT: int - Endpoint_type: str - - - @classmethod - def mapping_slo_info(cls, user_addr: str) -> Dict[str, Any]: - """ - 通过用户的IP地址映射出具体的服务SLO和功能启用 - 输入：user_addr - 输出：服务SLO，如TTFT、TPOT、E2E，以及启用功能，是否启用PD分离，是否允许压缩，是否开启安全 - """ - if user_addr.startswith("10.0."): - return { - "Enable_PD_Disaggregation": False, - "SLO_TTFT": 5000, - "SLO_TPOT": 50, - "SLO_E2E": 55000, - } - elif user_addr.startswith("192.168."): - return { - "Enable_PD_Disaggregation": False, - "SLO_TTFT": 10000, - "SLO_TPOT": 150, - "SLO_E2E": 160000, - } - else: - return { - "Enable_PD_Disaggregation": True, - "SLO_TTFT": 15000, - "SLO_TPOT": 200, - "SLO_E2E": 215000, - } - - - @classmethod - def knowledge_retriever( - cls, - user_prompt: str, - knowledge_block_num: int, - embedder: EmbeddingModel, - knowledge_table: KnowledgeTable, - ) -> Tuple[List[str], int]: - """ - 根据用户问题，从知识库中检索出最相关的知识清单。 - - 参数： - user_prompt：str ->用户问句，用于生成查询 embedding - knowledge_block_num：int ->希望返回的知识块数量（top-k） - embedder： EmbeddingModel ->外部注入的embedding模型 - knowledge_table: KnowledgeTable ->实例，内部已经构建好 FAISS 索引 - - 输出： - knowledge_ids: 检索到的知识块 ID 列表（长度 <= knowledge_block_num） - total_knowledge_length: 这些知识块 length 字段的总和 - """ - user_prompt = (user_prompt or "").strip() - if not user_prompt: - # 空问题，直接返回空列表 - return [], 0 - - if knowledge_block_num <= 0: - return [], 0 - - print(f"[Knowledge retriever]: use embedder={DEFAULT_EMBED_MODEL}") - - # 1) 将 user_prompt 编码为 embedding - query_embedding = embedder.encode_vector([user_prompt])[0] - - # 2) 在知识库中检索 top-k 知识块（内部优先用 FAISS，失败则回退 Python 扫描） - results = knowledge_table.search_by_embedding( - query_embedding=query_embedding, - top_k=knowledge_block_num, - min_score=float(SCHEDULER_RETRIEVAL_MIN_SCORE), - min_ratio=float(SCHEDULER_RETRIEVAL_MIN_RATIO), - ) - - # 3) 提取 ID 列表和总长度 - knowledge_ids: List[str] = [] - total_len = 0 - for kid, unit, score in results: - knowledge_ids.append(kid) - # length 字段在你构建 knowledge_base.yaml 时已经写好 - total_len += int(getattr(unit, "length", 0)) - - # 简单 debug - print( - "[Service.knowledge_retriever] user_prompt_len={}, " - "top_k={}, hit={}, knowledge_list_ids={}, total_len={}".format( - len(user_prompt), - knowledge_block_num, - len(knowledge_ids), - knowledge_ids, - total_len, - ) - ) - - return knowledge_ids, total_len - - -def normalize_injection_type(value: Any) -> str: - """ - 规范化知识注入类型： - - 未传 / 空值：默认 kvcache - - 接受常见别名：kvcache / kv / kv_cache -> kvcache - - text / prompt -> text - - 非法值：回退 kvcache - """ - if value is None: - return "kvcache" - - s = str(value).strip().lower() - if not s: - return "kvcache" - - if s in ("kvcache", "kv", "kv_cache", "kv-cache"): - return "kvcache" - - if s in ("text", "prompt"): - return "text" - - print(f"[Request] Unknown Injection_type={value!r}, fallback to 'kvcache'") - return "kvcache" -# ======================================================================================================================== -# ------------------------------------------------------Task-------------------------------------------------------------- -# ======================================================================================================================== -@dataclass -class Task: - """ - 定义任务调度基本信息 - User_addr：用户的IP地址，表征用户身份 - KDN_server_addr<任务根据知识需求挑选出的最合适的KDN服务器地址> - default_know_addr<默认的知识注入服务器，采用文本的注入方式> - P_proxy_id<处理任务的代理ID，用于追踪流> - P_proxy_addr<处理任务的P池代理IP地址，默认为本地换回地址> - P_proxy_port<处理任务的P池代理端口号，默认为8001> - D_proxy_addr<处理任务的D池代理IP地址，默认为本地换回地址> - D_proxy_port<处理任务的D池代理端口号，默认为8001> - prefill_instance<该请求分配的Prefill实例IP地址及端口> - decode_instance<该请求分配的Decode实例IP地址及端口> - batch_order<该请求在其所分配实例中的批次顺序> - """ - User_addr: str - KDN_server_addr: str - default_know_addr: str - P_proxy_id: str - P_proxy_addr: str - P_proxy_port: int - D_proxy_addr: str - D_proxy_port: int - prefill_instance: str - decode_instance: str - batch_order: int - User_url_path: str - - -# ======================================================================================================================== -# ------------------------------------------------------Request----------------------------------------------------------- -# ======================================================================================================================== -@dataclass -class Request: - """ - 主结构体，用于描述任务的所有需求 - Request_ID：用户任务的唯一标识ID - Request_type：请求类型，如request，control，update等 - Prompt<用于处理用户内容的相关信息> - Service<用于处理用户服务的相关信息> - Task<用于记录调度任务时的相关信息> - """ - Request_ID: int - Request_type: str - Prompt: Prompt - Service: Service - Task: Task - - - @classmethod - def build_request( - cls, - url_path: str, - payload: Dict, - user_addr: str, - request_id: int, - embedder: Optional[EmbeddingModel] = None, - knowledge_table: Optional[KnowledgeTable] = None, - proxies: Optional[List[Dict[str, Any]]] = None, - kdns: Optional[List[Dict[str, Any]]] = None, - strategy: Optional[Any] = None, - ) -> "Request": - """ - 将原始用户请求信息payload，以及转换为完整的 Request 对象。 - url_path:用户发送的url地址，例如 http://127.0.0.1:7001/v1/chat/completions - payload兼容的格式： - 1）旧版TCP协议纯payload：{"model": "...", "user_prompt": "..."} - 2）/v1/chat/completions 风格：{ - "model": "...", - "messages": [ - {"role": "system", "content": "..."}, - {"role": "user", "content": "..."}, - ...], - "max_tokens": 1024, - "temperature": 0.7, - "top_p": 0.95, - "stream": true, - ...} - 3）/v1/completions 风格：{ - "model": "...", - "prompt": "...." 或 ["...","..."], - "max_tokens": 1024, - "temperature": 0.7, - "top_p": 0.95, - "stream": false, - ...} - user_addr：调度器从 TCP / HTTP 层获取的客户端 IP - request_id：调度器分配的请求 ID（可保持默认 0，然后调度器后续覆盖） - """ - - # print("[BuildRequest] Receive user request payload:", payload) - model = payload.get("model") - if not model: - raise ValueError("Request payload 缺少必需字段 'model'") - - #----------------- 为兼容多种格式，设置user_prompt判断提取 ----------------------- - user_prompt: Optional[str] = None - t0 = time.perf_counter() - # Type A: 旧版协议，直接基于TCP给user_prompt - if "user_prompt" in payload: - user_prompt = str(payload["user_prompt"]) - - # Type B: chat/completions 风格，使用 messages - elif "messages" in payload: - msgs = payload.get("messages") or [] - if not isinstance(msgs, list): - raise ValueError("'messages' 字段必须是列表") - - last_user: Optional[str] = None - for m in msgs: - if not isinstance(m, dict): - continue - role = m.get("role", "") - content = m.get("content", "") - if role == "user": - # 不断覆盖，最终保留“最后一条 user 消息” - last_user = str(content) - - if last_user is not None: - user_prompt = last_user - else: - # 没有 user 角色时，退化为把所有消息拼起来 - pieces: List[str] = [] - for m in msgs: - if not isinstance(m, dict): - continue - r = m.get("role", "") - c = m.get("content", "") - pieces.append(f"{r}: {c}") - if pieces: - user_prompt = "\n".join(pieces) - - # Type C: completions 风格，使用 prompt - elif "prompt" in payload: - prompt = payload.get("prompt") - if isinstance(prompt, str): - user_prompt = prompt - elif isinstance(prompt, list) and prompt: - # 这里简单取第一个元素，你后面可以根据需要改成拼接 - user_prompt = str(prompt[0]) - - if not user_prompt: - raise ValueError("无法从请求中解析出 user_prompt（缺少 user_prompt/messages/prompt）") - - # ---- 通用请求参数：max_tokens / stream / temperature / top_p ---- - # 兼容没传的情况，给默认值 - max_tokens = int(payload.get("max_tokens", 1000) or 1000) - stream = parse_stream_flag(payload.get("stream")) - rag = parse_stream_flag(payload.get("RAG")) - injection_type = normalize_injection_type(payload.get("Injection_type")) - - try: - temperature = float(payload.get("temperature", 1.0)) - except (TypeError, ValueError): - temperature = 1.0 - try: - top_p = float(payload.get("top_p", 1.0)) - except (TypeError, ValueError): - top_p = 1.0 - # 做一下简单夹紧，避免写乱 - if temperature <= 0: - temperature = 1.0 - if not (0 < top_p <= 1.0): - top_p = 1.0 - - # ---- Request_type，目前统一标记为 "request"，后续可扩展 control/update ---- - request_type = payload.get("Request_type", "request") - - print("=============================\n" - " Start to build request.\n" - "=============================") - # print("[Request] get request ID:", request_id) - # print("[Request] model:", model) - # print("[Request] user_prompt:", user_prompt) - - # ========================= - # ---- 构造 Prompt 对象 ---- - # ========================= - # 构造prompt主要是 - tokens = Prompt.extract_prompt_info(model, user_prompt) - prompt_obj = Prompt( - model=model, - user_prompt=user_prompt, - bs=1, - token_length=tokens, - max_tokens=max_tokens, - stream=stream, - temperature=temperature, - top_p=top_p, - ) - - # =========================== - # ---- 构造 Service 对象 ---- - # =========================== - slo_mapping = Service.mapping_slo_info(user_addr) - service_obj = Service( - Enable_PD_Disaggregation = slo_mapping["Enable_PD_Disaggregation"], - Enable_know_injection=rag, - SLO_TTFT = slo_mapping["SLO_TTFT"], - SLO_TPOT = slo_mapping["SLO_TPOT"], - SLO_E2E = slo_mapping["SLO_E2E"], - Knowledge_block_num = int(SCHEDULER_RETRIEVAL_TOP_K), - Knowledge_List= [], - Knowledge_length= 0, - Enable_security = False, - Compress_factor = 0.3, - Enable_compress = True, - Injection_type=injection_type, - Endpoint_type="", - ) - - # 确认会话类型 - if "messages" in payload: - service_obj.Endpoint_type = "chat/completions" - else: - service_obj.Endpoint_type = "completions" - - # 知识检索：填充 Knowledge_List / Knowledge_length - if ( - service_obj.Enable_know_injection - and embedder is not None - and knowledge_table is not None - ): - try: - knowledge_ids, total_len = Service.knowledge_retriever( - user_prompt=user_prompt, - knowledge_block_num=service_obj.Knowledge_block_num, - embedder=embedder, - knowledge_table=knowledge_table, - ) - - service_obj.Knowledge_List = knowledge_ids - service_obj.Knowledge_length = total_len - except Exception as e: - print(f"[Service] knowledge_retriever failed: {e}") - # 出错时维持默认值 - service_obj.Knowledge_List = [] - service_obj.Knowledge_length = 0 - else: - # 未开启知识注入或未初始化知识库，保持默认值 - if not service_obj.Enable_know_injection: - print("[Service] Ban knowledge task, skip retriever!") - elif knowledge_table is None: - print("[Service] Not found available knowledge table, skip retriever, please check whether KDN servers is up!") - service_obj.Knowledge_List = [] - service_obj.Knowledge_length = 0 - - - # ========================== - # ---- 8) 构造 Task 对象 ---- - # ========================== - task_obj = Task( - User_addr=user_addr, - KDN_server_addr = "", - default_know_addr = "", # 这里暂时用 KDN 作为默认知识服务器，可按需调整 - P_proxy_id= "", - P_proxy_addr = "", - P_proxy_port = 0, - D_proxy_addr = "", - D_proxy_port = 0, - prefill_instance = "0", - decode_instance = "", - batch_order = 0, - User_url_path=url_path - ) - - # 若 scheduler 提供了策略和 proxy 列表，则在 build_request 内做选择 - try: - if strategy and hasattr(strategy, "select"): - # 统一策略接口：要求策略提供 select(proxies, request_obj_like) 或 select(proxies, payload) - # 这里为了避免 build_request 依赖 scheduler 的 Request 类型，我们先把最小上下文传给策略： - chosen_kdn, chosen_proxy = strategy.select( - kdns=kdns or [], - proxies=proxies or [], - payload=payload, - url_path=url_path, - user_addr=user_addr, - ) - - if chosen_kdn: - # 推荐统一成 http base_url，后面 KDN client 直接用 - task_obj.KDN_server_addr = f"{chosen_kdn['host']}:{int(chosen_kdn['port'])}" - - if chosen_proxy: - task_obj.P_proxy_id = chosen_proxy["proxy_id"] - task_obj.P_proxy_addr = chosen_proxy["host"] - task_obj.P_proxy_port = int(chosen_proxy["port"]) - - except Exception as e: - print(f"[Scheduler-Strategy]: select failed, fallback to default. err={e}") - - - # ========================= - # ---- 拼装最终 Request ---- - # ========================= - request_obj = Request( - Request_type=request_type, - Request_ID=request_id, - Prompt=prompt_obj, - Service=service_obj, - Task=task_obj - ) - - print("=============================\n" - " Finish building.\n" - "=============================") - - return request_obj - - - def to_payload(self) -> Dict[str, Any]: - """ - 将Request对象序列化为可以通过HTTP发送的payload。 - 默认使用 dataclasses 保留完整结构： - { - "Request_ID": ..., - "Request_type": ..., - "Prompt": { ... }, - "Service": { ... }, - "Task": { ... } - } - """ - return asdict(self) - - - def __repr__(self): - return ( - f"Request(\n" - f" Request_ID={self.Request_ID},\n" - f" Request_type={self.Request_type},\n" - f"--------Prompt--------\n" - f" model={self.Prompt.model},\n" - f" user_prompt={self.Prompt.user_prompt},\n" - f" token_length={self.Prompt.token_length},\n" - f" stream={self.Prompt.stream},\n" - f" temperature={self.Prompt.temperature},\n" - f" top_p={self.Prompt.top_p},\n" - f" max_tokens={self.Prompt.max_tokens},\n" - f"--------Service--------\n" - f" Enable_PD={self.Service.Enable_PD_Disaggregation},\n" - f" Enable_know_injection={self.Service.Enable_know_injection},\n" - f" Injection_type={self.Service.Injection_type},\n" - f" Enable_compress={self.Service.Enable_compress},\n" - f" Compress_factor={self.Service.Compress_factor},\n" - f" Enable_security={self.Service.Enable_security},\n" - f" Knowledge_block_num={self.Service.Knowledge_block_num},\n" - f" Knowledge_List={self.Service.Knowledge_List},\n" - f" Knowledge_length={self.Service.Knowledge_length},\n" - f" SLO_TTFT(ms)={self.Service.SLO_TTFT},\n" - f" SLO_E2E(ms)={self.Service.SLO_E2E},\n" - f" SLO_TPOT(ms)={self.Service.SLO_TPOT},\n" - f" Endpoint_type={self.Service.Endpoint_type},\n" - f"--------Task--------\n" - f" User_addr={self.Task.User_addr},\n" - f" KDN_server_addr={self.Task.KDN_server_addr},\n" - f" default_know_addr={self.Task.default_know_addr},\n" - f" P_proxy_addr={self.Task.P_proxy_addr},\n" - f" D_proxy_addr={self.Task.D_proxy_addr},\n" - f" P_proxy_port={self.Task.P_proxy_port},\n" - f" D_proxy_port={self.Task.D_proxy_port},\n" - f" prefill_instance={self.Task.prefill_instance},\n" - f" decode_instance={self.Task.decode_instance},\n" - f" batch_order={self.Task.batch_order},\n" - f" User_url_path={self.Task.User_url_path},\n" - f")" - ) +# 构建任务和任务批次的信息结构 +import time +from dataclasses import dataclass, asdict +from core.tokenizer_registry import estimate_tokens +from core.config import ( + DEFAULT_EMBED_MODEL, + SCHEDULER_RETRIEVAL_TOP_K, + SCHEDULER_RETRIEVAL_MIN_SCORE, + SCHEDULER_RETRIEVAL_MIN_RATIO, +) +# from CacheRoute.util import parse_host_port +from typing import Dict, Any, List, Tuple, Optional +from store import EmbeddingModel, KnowledgeTable +from util import parse_stream_flag + +""" +定义调度器-本地代理的交互数据结构体 -> Class Request，调度器通过调度策略构建request并发给本地Proxy，随后Proxy按需求执行对应动作完成选池推理 + + Class Request: + | Request_ID:用户任务的唯一标识ID + |__Request_type：请求类型，如request，control，update等 + | + |__Class Prompt: 记录用户问题的基本信息，如具体问题，模型类型，问题长度等 + | | model 模型类型 + | | user_prompt 用户问题 + | | token_length 问题的token长度 + | | max_token 最长生成token数 + | | stream 是否启用流式传输 + | | temperature 采样温度，默认1.0 + | |__top_p nucleus sampling 截断阈值，默认1.0 + | + |__Class Service： 映射用户问题的服务需求，如是否支持PD分离，是否支持知识注入，TTFT，E2E和TPOT的SLO需求 + | | Enable_PD_Disaggregation 是否启用PD分离 + | | Enable_know_injection 是否启用知识注入 + | | Injection_type 知识注入类型（文本orKV） + | | Enable_compress 是否允许对KVCache压缩 + | | Compress_factor 压缩比例 + | | Enable_security 是否启用安全模式（预留接口） + | | Knowledge_block_num 知识块数量 + | | Knowledge_List 问题所需知识清单 + | | Knowledge_length 知识注入长度 + | | SLO_TTFT 首token生成时间 + | | SLO_E2E 端到端时延 + | | SLO_TPOT 每秒token生成 + | |__Endpoint_type 转发给vLLM的请求类型，包括 chat/completions 或 completions，对话类或补全类 + | + |__Class Task： 记录调度相关的任务信息，如KDN知识服务器的IP地址，PD池代理的IP地址，端口号 + | User_addr 用户IP地址 + | KDN_server_addr KDN服务器地址 + | default_know_addr 默认文本注入服务器地址 + | P_proxy_id 选择代理的ID号 + | P_proxy_addr P池代理IP地址 + | D_proxy_addr D池代理IP地址 + | P_proxy_port P池代理端口号 + | D_proxy_port D池代理端口号 + | prefill_instance 该请求分配的Prefill实例IP地址及端口 + | decode_instance 该请求分配的Decode实例IP地址及端口 + |__batch_order 该请求在其所分配实例中的批次顺序 +""" + +# ======================================================================================================================== +# ------------------------------------------------------Prompt------------------------------------------------------------ +# ======================================================================================================================== +@dataclass +class Prompt: + """ + 定义用户问题的基本信息 + model<任务模型，str> + user_prompt<用户问题，str> + token_length<问题的token长度> + bs<任务组的batch_size，通常情况下默认1> + max_token<任务支持生成的最大token数> + stream<是否启用流式传输> + temperature<采样温度，默认1.0> + top_p + """ + model: str + user_prompt: str + token_length: int + bs: int + + max_tokens: int + stream: bool + temperature: float + top_p: float + + + @classmethod + def extract_prompt_info(cls, model: str, user_prompt: str) -> int: + """ + 根据 user_prompt 自动计算 token_length并返回。 + 输入：model，user_prompt + 输出：token_length + """ + seq_length = estimate_tokens(user_prompt, model) + print(f"[Prompt-tokenizer]: get task prompt_length complete, prompt_length={seq_length}") + + return seq_length + + +# ======================================================================================================================== +# ------------------------------------------------------Service----------------------------------------------------------- +# ======================================================================================================================== +@dataclass +class Service: + """ + 定义问题服务的SLO基本信息，通过用户IP地址映射具体服务等级，支持映射模块出于安全、个性化等方面的扩展 + Enable_PD_Disaggregation<是否允许问题进行PD分离处理，默认为True> + Enable_know_injection<是否允许调用远端知识，默认为True> + Injection_type<知识注入类型，默认为text，由proxy后续根据策略调整> + Enable_compress<是都允许进行KVCache的压缩，默认为True> + Compress_factor + Enable_security<是否启用安全模式，默认false，为后续安全内容提供接口> + Knowledge_block_num<任务注入知识块的top数量，默认为3> + Knowledge_List<知识块列表，里面的元素是Knowledge_ID,表征一个具体的知识块> + Knowledge_length<任务注入知识库的token长度，默认为0> + SLO_TTFT<任务组的TTFT SLO需求，即问题开始推理至产生提一个token所需要的时间，默认2000ms> + SLO_E2E<任务从开始Prefill到结束Decode所需的完整时间，ms> + SLO_TPOT<任务组的TPOT SLO需求，即自回归推理阶段平均生成默认20ms> + Endpoint_type<转发给vLLM的请求类型，包括 chat/completions 或 completions，对话类或补全类> + """ + Enable_PD_Disaggregation: bool + Enable_know_injection: bool + Injection_type: str + Enable_compress: bool + Compress_factor: float + Enable_security: bool + Knowledge_block_num: int + Knowledge_List: List[str] + Knowledge_length: int + SLO_TTFT: int + SLO_E2E: int + SLO_TPOT: int + Endpoint_type: str + + + @classmethod + def mapping_slo_info(cls, user_addr: str) -> Dict[str, Any]: + """ + 通过用户的IP地址映射出具体的服务SLO和功能启用 + 输入：user_addr + 输出：服务SLO，如TTFT、TPOT、E2E，以及启用功能，是否启用PD分离，是否允许压缩，是否开启安全 + """ + if user_addr.startswith("10.0."): + return { + "Enable_PD_Disaggregation": False, + "SLO_TTFT": 5000, + "SLO_TPOT": 50, + "SLO_E2E": 55000, + } + elif user_addr.startswith("192.168."): + return { + "Enable_PD_Disaggregation": False, + "SLO_TTFT": 10000, + "SLO_TPOT": 150, + "SLO_E2E": 160000, + } + else: + return { + "Enable_PD_Disaggregation": True, + "SLO_TTFT": 15000, + "SLO_TPOT": 200, + "SLO_E2E": 215000, + } + + + @classmethod + def knowledge_retriever( + cls, + user_prompt: str, + knowledge_block_num: int, + embedder: EmbeddingModel, + knowledge_table: KnowledgeTable, + ) -> Tuple[List[str], int]: + """ + 根据用户问题，从知识库中检索出最相关的知识清单。 + + 参数： + user_prompt：str ->用户问句，用于生成查询 embedding + knowledge_block_num：int ->希望返回的知识块数量（top-k） + embedder： EmbeddingModel ->外部注入的embedding模型 + knowledge_table: KnowledgeTable ->实例，内部已经构建好 FAISS 索引 + + 输出： + knowledge_ids: 检索到的知识块 ID 列表（长度 <= knowledge_block_num） + total_knowledge_length: 这些知识块 length 字段的总和 + """ + user_prompt = (user_prompt or "").strip() + if not user_prompt: + # 空问题，直接返回空列表 + return [], 0 + + if knowledge_block_num <= 0: + return [], 0 + + print(f"[Knowledge retriever]: use embedder={DEFAULT_EMBED_MODEL}") + + # 1) 将 user_prompt 编码为 embedding + query_embedding = embedder.encode_vector([user_prompt])[0] + + # 2) 在知识库中检索 top-k 知识块（内部优先用 FAISS，失败则回退 Python 扫描） + results = knowledge_table.search_by_embedding( + query_embedding=query_embedding, + top_k=knowledge_block_num, + min_score=float(SCHEDULER_RETRIEVAL_MIN_SCORE), + min_ratio=float(SCHEDULER_RETRIEVAL_MIN_RATIO), + ) + + # 3) 提取 ID 列表和总长度 + knowledge_ids: List[str] = [] + total_len = 0 + for kid, unit, score in results: + knowledge_ids.append(kid) + # length 字段在你构建 knowledge_base.yaml 时已经写好 + total_len += int(getattr(unit, "length", 0)) + + # 简单 debug + print( + "[Service.knowledge_retriever] user_prompt_len={}, " + "top_k={}, hit={}, knowledge_list_ids={}, total_len={}".format( + len(user_prompt), + knowledge_block_num, + len(knowledge_ids), + knowledge_ids, + total_len, + ) + ) + + return knowledge_ids, total_len + + +def normalize_injection_type(value: Any) -> str: + """ + 规范化知识注入类型： + - 未传 / 空值：默认 kvcache + - 接受常见别名：kvcache / kv / kv_cache -> kvcache + - text / prompt -> text + - 非法值：回退 kvcache + """ + if value is None: + return "kvcache" + + s = str(value).strip().lower() + if not s: + return "kvcache" + + if s in ("kvcache", "kv", "kv_cache", "kv-cache"): + return "kvcache" + + if s in ("text", "prompt"): + return "text" + + print(f"[Request] Unknown Injection_type={value!r}, fallback to 'kvcache'") + return "kvcache" +# ======================================================================================================================== +# ------------------------------------------------------Task-------------------------------------------------------------- +# ======================================================================================================================== +@dataclass +class Task: + """ + 定义任务调度基本信息 + User_addr：用户的IP地址，表征用户身份 + KDN_server_addr<任务根据知识需求挑选出的最合适的KDN服务器地址> + default_know_addr<默认的知识注入服务器，采用文本的注入方式> + P_proxy_id<处理任务的代理ID，用于追踪流> + P_proxy_addr<处理任务的P池代理IP地址，默认为本地换回地址> + P_proxy_port<处理任务的P池代理端口号，默认为8001> + D_proxy_addr<处理任务的D池代理IP地址，默认为本地换回地址> + D_proxy_port<处理任务的D池代理端口号，默认为8001> + prefill_instance<该请求分配的Prefill实例IP地址及端口> + decode_instance<该请求分配的Decode实例IP地址及端口> + batch_order<该请求在其所分配实例中的批次顺序> + """ + User_addr: str + KDN_server_addr: str + default_know_addr: str + P_proxy_id: str + P_proxy_addr: str + P_proxy_port: int + D_proxy_addr: str + D_proxy_port: int + prefill_instance: str + decode_instance: str + batch_order: int + User_url_path: str + + +# ======================================================================================================================== +# ------------------------------------------------------Request----------------------------------------------------------- +# ======================================================================================================================== +@dataclass +class Request: + """ + 主结构体，用于描述任务的所有需求 + Request_ID：用户任务的唯一标识ID + Request_type：请求类型，如request，control，update等 + Prompt<用于处理用户内容的相关信息> + Service<用于处理用户服务的相关信息> + Task<用于记录调度任务时的相关信息> + """ + Request_ID: int + Request_type: str + Prompt: Prompt + Service: Service + Task: Task + + + @classmethod + def build_request( + cls, + url_path: str, + payload: Dict, + user_addr: str, + request_id: int, + embedder: Optional[EmbeddingModel] = None, + knowledge_table: Optional[KnowledgeTable] = None, + 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 对象。 + url_path:用户发送的url地址，例如 http://127.0.0.1:7001/v1/chat/completions + payload兼容的格式： + 1）旧版TCP协议纯payload：{"model": "...", "user_prompt": "..."} + 2）/v1/chat/completions 风格：{ + "model": "...", + "messages": [ + {"role": "system", "content": "..."}, + {"role": "user", "content": "..."}, + ...], + "max_tokens": 1024, + "temperature": 0.7, + "top_p": 0.95, + "stream": true, + ...} + 3）/v1/completions 风格：{ + "model": "...", + "prompt": "...." 或 ["...","..."], + "max_tokens": 1024, + "temperature": 0.7, + "top_p": 0.95, + "stream": false, + ...} + user_addr：调度器从 TCP / HTTP 层获取的客户端 IP + request_id：调度器分配的请求 ID（可保持默认 0，然后调度器后续覆盖） + """ + + # print("[BuildRequest] Receive user request payload:", payload) + model = payload.get("model") + if not model: + raise ValueError("Request payload 缺少必需字段 'model'") + + #----------------- 为兼容多种格式，设置user_prompt判断提取 ----------------------- + user_prompt: Optional[str] = None + t0 = time.perf_counter() + # Type A: 旧版协议，直接基于TCP给user_prompt + if "user_prompt" in payload: + user_prompt = str(payload["user_prompt"]) + + # Type B: chat/completions 风格，使用 messages + elif "messages" in payload: + msgs = payload.get("messages") or [] + if not isinstance(msgs, list): + raise ValueError("'messages' 字段必须是列表") + + last_user: Optional[str] = None + for m in msgs: + if not isinstance(m, dict): + continue + role = m.get("role", "") + content = m.get("content", "") + if role == "user": + # 不断覆盖，最终保留“最后一条 user 消息” + last_user = str(content) + + if last_user is not None: + user_prompt = last_user + else: + # 没有 user 角色时，退化为把所有消息拼起来 + pieces: List[str] = [] + for m in msgs: + if not isinstance(m, dict): + continue + r = m.get("role", "") + c = m.get("content", "") + pieces.append(f"{r}: {c}") + if pieces: + user_prompt = "\n".join(pieces) + + # Type C: completions 风格，使用 prompt + elif "prompt" in payload: + prompt = payload.get("prompt") + if isinstance(prompt, str): + user_prompt = prompt + elif isinstance(prompt, list) and prompt: + # 这里简单取第一个元素，你后面可以根据需要改成拼接 + user_prompt = str(prompt[0]) + + if not user_prompt: + raise ValueError("无法从请求中解析出 user_prompt（缺少 user_prompt/messages/prompt）") + + # ---- 通用请求参数：max_tokens / stream / temperature / top_p ---- + # 兼容没传的情况，给默认值 + max_tokens = int(payload.get("max_tokens", 1000) or 1000) + stream = parse_stream_flag(payload.get("stream")) + rag = parse_stream_flag(payload.get("RAG")) + injection_type = normalize_injection_type(payload.get("Injection_type")) + + try: + temperature = float(payload.get("temperature", 1.0)) + except (TypeError, ValueError): + temperature = 1.0 + try: + top_p = float(payload.get("top_p", 1.0)) + except (TypeError, ValueError): + top_p = 1.0 + # 做一下简单夹紧，避免写乱 + if temperature <= 0: + temperature = 1.0 + if not (0 < top_p <= 1.0): + top_p = 1.0 + + # ---- Request_type，目前统一标记为 "request"，后续可扩展 control/update ---- + request_type = payload.get("Request_type", "request") + + print("=============================\n" + " Start to build request.\n" + "=============================") + # print("[Request] get request ID:", request_id) + # print("[Request] model:", model) + # print("[Request] user_prompt:", user_prompt) + + # ========================= + # ---- 构造 Prompt 对象 ---- + # ========================= + # 构造prompt主要是 + tokens = Prompt.extract_prompt_info(model, user_prompt) + prompt_obj = Prompt( + model=model, + user_prompt=user_prompt, + bs=1, + token_length=tokens, + max_tokens=max_tokens, + stream=stream, + temperature=temperature, + top_p=top_p, + ) + + # =========================== + # ---- 构造 Service 对象 ---- + # =========================== + slo_mapping = Service.mapping_slo_info(user_addr) + service_obj = Service( + Enable_PD_Disaggregation = slo_mapping["Enable_PD_Disaggregation"], + Enable_know_injection=rag, + SLO_TTFT = slo_mapping["SLO_TTFT"], + SLO_TPOT = slo_mapping["SLO_TPOT"], + SLO_E2E = slo_mapping["SLO_E2E"], + Knowledge_block_num = int(SCHEDULER_RETRIEVAL_TOP_K), + Knowledge_List= [], + Knowledge_length= 0, + Enable_security = False, + Compress_factor = 0.3, + Enable_compress = True, + Injection_type=injection_type, + Endpoint_type="", + ) + + # 确认会话类型 + if "messages" in payload: + service_obj.Endpoint_type = "chat/completions" + else: + service_obj.Endpoint_type = "completions" + + # 知识检索：填充 Knowledge_List / Knowledge_length + if ( + service_obj.Enable_know_injection + and embedder is not None + and knowledge_table is not None + ): + try: + knowledge_ids, total_len = Service.knowledge_retriever( + user_prompt=user_prompt, + knowledge_block_num=service_obj.Knowledge_block_num, + embedder=embedder, + knowledge_table=knowledge_table, + ) + + service_obj.Knowledge_List = knowledge_ids + service_obj.Knowledge_length = total_len + except Exception as e: + print(f"[Service] knowledge_retriever failed: {e}") + # 出错时维持默认值 + service_obj.Knowledge_List = [] + service_obj.Knowledge_length = 0 + else: + # 未开启知识注入或未初始化知识库，保持默认值 + if not service_obj.Enable_know_injection: + print("[Service] Ban knowledge task, skip retriever!") + elif knowledge_table is None: + print("[Service] Not found available knowledge table, skip retriever, please check whether KDN servers is up!") + service_obj.Knowledge_List = [] + service_obj.Knowledge_length = 0 + + + # ========================== + # ---- 8) 构造 Task 对象 ---- + # ========================== + task_obj = Task( + User_addr=user_addr, + KDN_server_addr = "", + default_know_addr = "", # 这里暂时用 KDN 作为默认知识服务器，可按需调整 + P_proxy_id= "", + P_proxy_addr = "", + P_proxy_port = 0, + D_proxy_addr = "", + D_proxy_port = 0, + prefill_instance = "0", + decode_instance = "", + batch_order = 0, + User_url_path=url_path + ) + + # 若 scheduler 提供了策略和 proxy 列表，则在 build_request 内做选择 + try: + if strategy and hasattr(strategy, "select"): + # 统一策略接口：要求策略提供 select(proxies, request_obj_like) 或 select(proxies, payload) + # 这里为了避免 build_request 依赖 scheduler 的 Request 类型，我们先把最小上下文传给策略： + request_ctx = { + "request_id": request_id, + "knowledge_list": list(service_obj.Knowledge_List or []), + "knowledge_length": int(service_obj.Knowledge_length or 0), + "endpoint_type": service_obj.Endpoint_type, + # Injection_type 当前主要用于调试；正式策略不应依赖它做硬分支。 + "injection_type": service_obj.Injection_type, + "rag_enabled": bool(service_obj.Enable_know_injection), + "prompt_token_length": int(prompt_obj.token_length or 0), + "user_prompt": user_prompt, + # 下面两个上下文允许策略复用 scheduler 已维护的数据， + # 避免重复做 embedding 或在线拼状态。 + "knowledge_table": knowledge_table, + "kdn_knowledge_index": kdn_knowledge_index or {}, + } + chosen_kdn, chosen_proxy = strategy.select( + kdns=kdns or [], + proxies=proxies or [], + payload=payload, + url_path=url_path, + user_addr=user_addr, + request_ctx=request_ctx, + ) + + if chosen_kdn: + # 推荐统一成 http base_url，后面 KDN client 直接用 + task_obj.KDN_server_addr = f"{chosen_kdn['host']}:{int(chosen_kdn['port'])}" + + if chosen_proxy: + task_obj.P_proxy_id = chosen_proxy["proxy_id"] + task_obj.P_proxy_addr = chosen_proxy["host"] + task_obj.P_proxy_port = int(chosen_proxy["port"]) + + except Exception as e: + print(f"[Scheduler-Strategy]: select failed, fallback to default. err={e}") + + + # ========================= + # ---- 拼装最终 Request ---- + # ========================= + request_obj = Request( + Request_type=request_type, + Request_ID=request_id, + Prompt=prompt_obj, + Service=service_obj, + Task=task_obj + ) + + print("=============================\n" + " Finish building.\n" + "=============================") + + return request_obj + + + def to_payload(self) -> Dict[str, Any]: + """ + 将Request对象序列化为可以通过HTTP发送的payload。 + 默认使用 dataclasses 保留完整结构： + { + "Request_ID": ..., + "Request_type": ..., + "Prompt": { ... }, + "Service": { ... }, + "Task": { ... } + } + """ + return asdict(self) + + + def __repr__(self): + return ( + f"Request(\n" + f" Request_ID={self.Request_ID},\n" + f" Request_type={self.Request_type},\n" + f"--------Prompt--------\n" + f" model={self.Prompt.model},\n" + f" user_prompt={self.Prompt.user_prompt},\n" + f" token_length={self.Prompt.token_length},\n" + f" stream={self.Prompt.stream},\n" + f" temperature={self.Prompt.temperature},\n" + f" top_p={self.Prompt.top_p},\n" + f" max_tokens={self.Prompt.max_tokens},\n" + f"--------Service--------\n" + f" Enable_PD={self.Service.Enable_PD_Disaggregation},\n" + f" Enable_know_injection={self.Service.Enable_know_injection},\n" + f" Injection_type={self.Service.Injection_type},\n" + f" Enable_compress={self.Service.Enable_compress},\n" + f" Compress_factor={self.Service.Compress_factor},\n" + f" Enable_security={self.Service.Enable_security},\n" + f" Knowledge_block_num={self.Service.Knowledge_block_num},\n" + f" Knowledge_List={self.Service.Knowledge_List},\n" + f" Knowledge_length={self.Service.Knowledge_length},\n" + f" SLO_TTFT(ms)={self.Service.SLO_TTFT},\n" + f" SLO_E2E(ms)={self.Service.SLO_E2E},\n" + f" SLO_TPOT(ms)={self.Service.SLO_TPOT},\n" + f" Endpoint_type={self.Service.Endpoint_type},\n" + f"--------Task--------\n" + f" User_addr={self.Task.User_addr},\n" + f" KDN_server_addr={self.Task.KDN_server_addr},\n" + f" default_know_addr={self.Task.default_know_addr},\n" + f" P_proxy_addr={self.Task.P_proxy_addr},\n" + f" D_proxy_addr={self.Task.D_proxy_addr},\n" + f" P_proxy_port={self.Task.P_proxy_port},\n" + f" D_proxy_port={self.Task.D_proxy_port},\n" + f" prefill_instance={self.Task.prefill_instance},\n" + f" decode_instance={self.Task.decode_instance},\n" + f" batch_order={self.Task.batch_order},\n" + f" User_url_path={self.Task.User_url_path},\n" + f")" + ) diff --git a/scheduler/knowledge/kdn_sync.py b/scheduler/knowledge/kdn_sync.py index 0c3c724..0464a76 100644 --- a/scheduler/knowledge/kdn_sync.py +++ b/scheduler/knowledge/kdn_sync.py @@ -1,290 +1,346 @@ -# integration/kdn_sync.py -""" -KDN <-> Scheduler synchronization logic. - -This module can be reused by: -- scheduler Knowledge list startup -- Knowledge list periodic refresh -- future admin-triggered refresh -""" -import logging -logger = logging.getLogger(__name__) - -import asyncio -import time -import os - -from typing import List, Dict, Tuple -from core.config import SCHEDULER_KDN_REFRESH_INTERVAL_S_DEFAULT -from store.knowledge_base import KnowledgeTable, KnowledgeUnit -from .kdn_client import fetch_kdn_snapshot, fetch_kdn_items_by_kids -from scheduler.resource.control_plane import _hb_agg - -# kid -> signature -# signature = (embed_dim, length, kv_ready, kv_updated_at, kv_dumped_keys) -_kdn_meta_cache: Dict[str, Tuple] = {} - - -def _format_kdn_addr(host: str, port: int) -> str: - return f"kdn://{host}:{int(port)}" - - -async def select_kdn_base_url(app) -> tuple[str | None, list[str]]: - """ - Select one alive KDN server from scheduler's kdn_pool. - Returns: - - base_url: "http://host:port" or None - - alive_addrs: ["kdn://host:port", ...] (for KnowledgeUnit.avail_kdn_servers) - """ - pool = getattr(app.state, "kdn_pool", None) - if pool is None: - return None, [] - - alive = await pool.list(include_dead=False) - if not alive: - return None, [] - - # Round-robin index stored on app.state (best-effort) - idx = int(getattr(app.state, "_kdn_rr_idx", 0) or 0) - chosen = alive[idx % len(alive)] - app.state._kdn_rr_idx = idx + 1 - - base_url = f"http://{chosen.host}:{int(chosen.port)}" - alive_addrs = [_format_kdn_addr(x.host, x.port) for x in alive] - app.state.kdn_alive = len(alive) - app.state.kdn_alive_addrs = alive_addrs # ["kdn://h:p", ...] - app.state.kdn_last_selected = base_url # "http://h:p" - app.state.kdn_last_selected_id = getattr(chosen, "kdn_id", "") - - return base_url, alive_addrs - - -def build_table_from_kdn_items( - items: List[dict], - dim: int, - avail_kdn_servers: List[str], -) -> KnowledgeTable: - """ - Build a KnowledgeTable from KDN snapshot items. - - Design: - - kid(str) is the external primary key - - FAISS internal id is hidden inside KnowledgeTable - - kv metadata is preserved for scheduling decisions - """ - table = KnowledgeTable(dim=dim) - - loaded = 0 - skipped_no_emb = 0 - - for it in items: - kid = str(it.get("kid") or "").strip().lower() - emb = it.get("embedding") - length = int(it.get("length") or 0) - embed_dim = int(it.get("embed_dim") or dim) - - if not kid: - continue - if emb is None: - skipped_no_emb += 1 - continue - if embed_dim != dim: - raise RuntimeError(f"Dim mismatch for kid={kid}: {embed_dim} != {dim}") - - unit = KnowledgeUnit( - embedding=list(emb), - length=length, - avail_kdn_servers=list(avail_kdn_servers or []), - text_abstract=None, - - # KV metadata - kv_ready=int(it.get("kv_ready") or 0), - kv_rel_dir=it.get("kv_rel_dir"), - kv_dumped_keys=it.get("kv_dumped_keys"), - kv_updated_at=it.get("kv_updated_at"), - ) - - table.upsert_kid(kid, unit) - loaded += 1 - - table.build_faiss_index() - return table - - -async def kdn_refresh_once(app) -> dict: - """ - Refresh knowledge_table once from KDN snapshot. - Atomic-swap design: - - NEVER mutate old_table in-place - - Build / apply changes on new_table - - Rebuild FAISS on new_table - - Swap app.state.knowledge_table at the end - """ - lock = getattr(app.state, "_kdn_refresh_lock", None) - if lock is None: - raise RuntimeError("scheduler.state._kdn_refresh_lock is not initialized") - - async with lock: - kdn_base_url, alive_addrs = await select_kdn_base_url(app) - if not kdn_base_url: - app.state.kdn_last_refresh_ok = False - app.state.kdn_last_refresh_reason = "no alive kdn in pool" - app.state.kdn_last_refresh_ts = int(time.time()) - return {"ok": False, "reason": "no alive kdn in pool"} - - # A) meta snapshot - meta_items = await fetch_kdn_snapshot( - base_url=kdn_base_url, - need_fields=["kid", "embed_dim", "length", - "kv_ready", "kv_updated_at", "kv_dumped_keys"], - ) - - added, updated, removed = diff_kdn_meta(meta_items) - - old_table = getattr(app.state, "knowledge_table", None) - - # ---- Build new_table (either first build or clone from old) ---- - if old_table is None: - embedder = getattr(app.state, "embedding_engine", None) - if embedder is None: - raise RuntimeError("embedding_engine is not initialized; cannot build KnowledgeTable") - new_table = KnowledgeTable(dim=embedder.dim) - logger.debug("[Scheduler] KnowledgeTable initialized (first build), dim=%s", embedder.dim) - else: - new_table = old_table.clone_without_index() - - changed = False - - # B) upsert on new_table - need_fetch = added + updated - if need_fetch: - full_items = await fetch_kdn_items_by_kids(kdn_base_url, need_fetch) - - for it in full_items: - kid_raw = it.get("kid") or it.get("id") # compatibility - kid = str(kid_raw or "").strip().lower() - if not kid: - continue - - emb = it.get("embedding") - if emb is None: - continue - - unit = KnowledgeUnit( - embedding=list(emb), - length=int(it.get("length") or 0), - avail_kdn_servers=list(alive_addrs), - kv_ready=int(it.get("kv_ready") or 0), - kv_rel_dir=it.get("kv_rel_dir"), - kv_dumped_keys=it.get("kv_dumped_keys"), - kv_updated_at=it.get("kv_updated_at"), - ) - new_table.upsert_kid(kid, unit) - changed = True - - # C) delete on new_table - if removed: - new_table.delete_kids(removed) - changed = True - - # D) build FAISS on new_table (always build for first build; otherwise only if changed) - if old_table is None or changed: - new_table.build_faiss_index() - - # E) atomic swap - app.state.knowledge_table = new_table - - # status - app.state.last_refresh_ts = int(time.time()) - app.state.kdn_last_refresh_ok = True - app.state.kdn_last_refresh_reason = "" - app.state.kdn_last_refresh_ts = int(time.time()) - - return { - "ok": True, - "added": len(added), - "updated": len(updated), - "removed": len(removed), - "entries": len(getattr(new_table, "_units", {})), - } - - -async def kdn_auto_refresh_loop(app, stop_event: asyncio.Event): - """ - Periodically refresh scheduler knowledge_table from KDN. - - Properties: - - non-blocking - - failure-safe (keep old table) - - atomic swap - """ - interval_s = int(os.getenv("SCHEDULER_KDN_REFRESH_INTERVAL_S", str(SCHEDULER_KDN_REFRESH_INTERVAL_S_DEFAULT))) - interval_s = max(5, interval_s) - - # kdn_base_url = getattr(app.state, "kdn_base_url", None) - # if not kdn_base_url: - # return - - while not stop_event.is_set(): - try: - r = await kdn_refresh_once(app) - try: - await _hb_agg.record_kdn_refresh(r) - except Exception: - # 输出层故障不能影响业务 - logger.debug("[Scheduler] record_kdn_refresh failed", exc_info=True) - - # 不再刷屏：降为 debug - if r.get("ok"): - logger.debug("[Scheduler] KDN auto-refresh OK: entries=%s, added=%s, updated=%s, removed=%s", - r.get("entries"), r.get("added"), r.get("updated"), r.get("removed")) - else: - logger.debug("[Scheduler] KDN auto-refresh skipped: %s", r) - except Exception as e: - logger.exception(f"[Scheduler] KDN auto-refresh FAILED: {e}") - app.state.kdn_last_refresh_ok = False - app.state.kdn_last_refresh_reason = f"exception: {type(e).__name__}" - app.state.kdn_last_refresh_ts = int(time.time()) - - try: - await asyncio.wait_for(stop_event.wait(), timeout=interval_s) - except asyncio.TimeoutError: - pass - - -def diff_kdn_meta(items: list) -> tuple[list[str], list[str], list[str]]: - """ - Return (added, updated, removed) kid lists. - """ - global _kdn_meta_cache - - seen = set() - added, updated = [], [] - - for it in items: - kid = str(it.get("kid") or "").strip().lower() - if not kid: - continue - - sig = ( - it.get("embed_dim"), - it.get("length"), - it.get("kv_ready"), - it.get("kv_updated_at"), - it.get("kv_dumped_keys"), - ) - - seen.add(kid) - if kid not in _kdn_meta_cache: - added.append(kid) - _kdn_meta_cache[kid] = sig - elif _kdn_meta_cache[kid] != sig: - updated.append(kid) - _kdn_meta_cache[kid] = sig - - removed = [k for k in _kdn_meta_cache.keys() if k not in seen] - for k in removed: - _kdn_meta_cache.pop(k, None) - - return added, updated, removed - +# integration/kdn_sync.py +""" +KDN <-> Scheduler synchronization logic. + +This module can be reused by: +- scheduler Knowledge list startup +- Knowledge list periodic refresh +- future admin-triggered refresh +""" +import logging +logger = logging.getLogger(__name__) + +import asyncio +import time +import os + +from typing import List, Dict, Tuple +from core.config import SCHEDULER_KDN_REFRESH_INTERVAL_S_DEFAULT +from store.knowledge_base import KnowledgeTable, KnowledgeUnit +from .kdn_client import fetch_kdn_snapshot, fetch_kdn_items_by_kids +from scheduler.resource.control_plane import _hb_agg + +# kid -> signature +# signature = (embed_dim, length, kv_ready, kv_updated_at, kv_dumped_keys) +_kdn_meta_cache: Dict[str, Tuple] = {} + + +def _format_kdn_addr(host: str, port: int) -> str: + return f"kdn://{host}:{int(port)}" + + +async def refresh_kdn_knowledge_index(app) -> Dict[str, Dict[str, Dict[str, int]]]: + """ + 为 scheduler 构建一个轻量级 per-KDN 知识元数据索引。 + + 返回结构： + { + : { + : {"length": 123, "kv_ready": 1}, + ... + }, + "kdn://host:port": {...} # 兼容地址查询 + } + + 说明： + - 这里只拉 metadata，不拉 embedding，避免把知识表构建逻辑和 KDN 选择逻辑混在一起。 + - 先顺序实现，保持最小改动；后续若 KDN 数量增多，可再并发化。 + """ + pool = getattr(app.state, "kdn_pool", None) + if pool is None: + app.state.kdn_knowledge_index = {} + return {} + + alive = await pool.list(include_dead=False) + if not alive: + app.state.kdn_knowledge_index = {} + return {} + + index: Dict[str, Dict[str, Dict[str, int]]] = {} + for info in alive: + base_url = f"http://{info.host}:{int(info.port)}" + items = await fetch_kdn_snapshot( + base_url=base_url, + need_fields=["kid", "length", "kv_ready"], + ) + + meta_by_kid: Dict[str, Dict[str, int]] = {} + for it in items: + kid = str(it.get("kid") or "").strip().lower() + if not kid: + continue + meta_by_kid[kid] = { + "length": int(it.get("length") or 0), + "kv_ready": int(it.get("kv_ready") or 0), + } + + index[str(info.kdn_id)] = meta_by_kid + index[_format_kdn_addr(info.host, info.port)] = meta_by_kid + + app.state.kdn_knowledge_index = index + return index + + +async def select_kdn_base_url(app) -> tuple[str | None, list[str]]: + """ + Select one alive KDN server from scheduler's kdn_pool. + Returns: + - base_url: "http://host:port" or None + - alive_addrs: ["kdn://host:port", ...] (for KnowledgeUnit.avail_kdn_servers) + """ + pool = getattr(app.state, "kdn_pool", None) + if pool is None: + return None, [] + + alive = await pool.list(include_dead=False) + if not alive: + return None, [] + + # Round-robin index stored on app.state (best-effort) + idx = int(getattr(app.state, "_kdn_rr_idx", 0) or 0) + chosen = alive[idx % len(alive)] + app.state._kdn_rr_idx = idx + 1 + + base_url = f"http://{chosen.host}:{int(chosen.port)}" + alive_addrs = [_format_kdn_addr(x.host, x.port) for x in alive] + app.state.kdn_alive = len(alive) + app.state.kdn_alive_addrs = alive_addrs # ["kdn://h:p", ...] + app.state.kdn_last_selected = base_url # "http://h:p" + app.state.kdn_last_selected_id = getattr(chosen, "kdn_id", "") + + return base_url, alive_addrs + + +def build_table_from_kdn_items( + items: List[dict], + dim: int, + avail_kdn_servers: List[str], +) -> KnowledgeTable: + """ + Build a KnowledgeTable from KDN snapshot items. + + Design: + - kid(str) is the external primary key + - FAISS internal id is hidden inside KnowledgeTable + - kv metadata is preserved for scheduling decisions + """ + table = KnowledgeTable(dim=dim) + + loaded = 0 + skipped_no_emb = 0 + + for it in items: + kid = str(it.get("kid") or "").strip().lower() + emb = it.get("embedding") + length = int(it.get("length") or 0) + embed_dim = int(it.get("embed_dim") or dim) + + if not kid: + continue + if emb is None: + skipped_no_emb += 1 + continue + if embed_dim != dim: + raise RuntimeError(f"Dim mismatch for kid={kid}: {embed_dim} != {dim}") + + unit = KnowledgeUnit( + embedding=list(emb), + length=length, + avail_kdn_servers=list(avail_kdn_servers or []), + text_abstract=None, + + # KV metadata + kv_ready=int(it.get("kv_ready") or 0), + kv_rel_dir=it.get("kv_rel_dir"), + kv_dumped_keys=it.get("kv_dumped_keys"), + kv_updated_at=it.get("kv_updated_at"), + ) + + table.upsert_kid(kid, unit) + loaded += 1 + + table.build_faiss_index() + return table + + +async def kdn_refresh_once(app) -> dict: + """ + Refresh knowledge_table once from KDN snapshot. + Atomic-swap design: + - NEVER mutate old_table in-place + - Build / apply changes on new_table + - Rebuild FAISS on new_table + - Swap app.state.knowledge_table at the end + """ + lock = getattr(app.state, "_kdn_refresh_lock", None) + if lock is None: + raise RuntimeError("scheduler.state._kdn_refresh_lock is not initialized") + + async with lock: + kdn_knowledge_index = await refresh_kdn_knowledge_index(app) + + kdn_base_url, alive_addrs = await select_kdn_base_url(app) + if not kdn_base_url: + app.state.kdn_last_refresh_ok = False + app.state.kdn_last_refresh_reason = "no alive kdn in pool" + app.state.kdn_last_refresh_ts = int(time.time()) + return {"ok": False, "reason": "no alive kdn in pool"} + + # A) meta snapshot + meta_items = await fetch_kdn_snapshot( + base_url=kdn_base_url, + need_fields=["kid", "embed_dim", "length", + "kv_ready", "kv_updated_at", "kv_dumped_keys"], + ) + + added, updated, removed = diff_kdn_meta(meta_items) + + old_table = getattr(app.state, "knowledge_table", None) + + # ---- Build new_table (either first build or clone from old) ---- + if old_table is None: + embedder = getattr(app.state, "embedding_engine", None) + if embedder is None: + raise RuntimeError("embedding_engine is not initialized; cannot build KnowledgeTable") + new_table = KnowledgeTable(dim=embedder.dim) + logger.debug("[Scheduler] KnowledgeTable initialized (first build), dim=%s", embedder.dim) + else: + new_table = old_table.clone_without_index() + + changed = False + + # B) upsert on new_table + need_fetch = added + updated + if need_fetch: + full_items = await fetch_kdn_items_by_kids(kdn_base_url, need_fetch) + + for it in full_items: + kid_raw = it.get("kid") or it.get("id") # compatibility + kid = str(kid_raw or "").strip().lower() + if not kid: + continue + + emb = it.get("embedding") + if emb is None: + continue + + unit = KnowledgeUnit( + embedding=list(emb), + length=int(it.get("length") or 0), + # 当前 KnowledgeTable 继续保留“哪些 KDN 可用”的全局视角； + # 精确到每个 KDN 的覆盖信息由 kdn_knowledge_index 维护，供 cacheroute 使用。 + avail_kdn_servers=list(alive_addrs), + kv_ready=int(it.get("kv_ready") or 0), + kv_rel_dir=it.get("kv_rel_dir"), + kv_dumped_keys=it.get("kv_dumped_keys"), + kv_updated_at=it.get("kv_updated_at"), + ) + new_table.upsert_kid(kid, unit) + changed = True + + # C) delete on new_table + if removed: + new_table.delete_kids(removed) + changed = True + + # D) build FAISS on new_table (always build for first build; otherwise only if changed) + if old_table is None or changed: + new_table.build_faiss_index() + + # E) atomic swap + app.state.knowledge_table = new_table + + # status + app.state.last_refresh_ts = int(time.time()) + app.state.kdn_last_refresh_ok = True + app.state.kdn_last_refresh_reason = "" + app.state.kdn_last_refresh_ts = int(time.time()) + + return { + "ok": True, + "added": len(added), + "updated": len(updated), + "removed": len(removed), + "entries": len(getattr(new_table, "_units", {})), + } + + +async def kdn_auto_refresh_loop(app, stop_event: asyncio.Event): + """ + Periodically refresh scheduler knowledge_table from KDN. + + Properties: + - non-blocking + - failure-safe (keep old table) + - atomic swap + """ + interval_s = int(os.getenv("SCHEDULER_KDN_REFRESH_INTERVAL_S", str(SCHEDULER_KDN_REFRESH_INTERVAL_S_DEFAULT))) + interval_s = max(5, interval_s) + + # kdn_base_url = getattr(app.state, "kdn_base_url", None) + # if not kdn_base_url: + # return + + while not stop_event.is_set(): + try: + r = await kdn_refresh_once(app) + try: + await _hb_agg.record_kdn_refresh(r) + except Exception: + # 输出层故障不能影响业务 + logger.debug("[Scheduler] record_kdn_refresh failed", exc_info=True) + + # 不再刷屏：降为 debug + if r.get("ok"): + logger.debug("[Scheduler] KDN auto-refresh OK: entries=%s, added=%s, updated=%s, removed=%s", + r.get("entries"), r.get("added"), r.get("updated"), r.get("removed")) + else: + logger.debug("[Scheduler] KDN auto-refresh skipped: %s", r) + except Exception as e: + logger.exception(f"[Scheduler] KDN auto-refresh FAILED: {e}") + app.state.kdn_last_refresh_ok = False + app.state.kdn_last_refresh_reason = f"exception: {type(e).__name__}" + app.state.kdn_last_refresh_ts = int(time.time()) + + try: + await asyncio.wait_for(stop_event.wait(), timeout=interval_s) + except asyncio.TimeoutError: + pass + + +def diff_kdn_meta(items: list) -> tuple[list[str], list[str], list[str]]: + """ + Return (added, updated, removed) kid lists. + """ + global _kdn_meta_cache + + seen = set() + added, updated = [], [] + + for it in items: + kid = str(it.get("kid") or "").strip().lower() + if not kid: + continue + + sig = ( + it.get("embed_dim"), + it.get("length"), + it.get("kv_ready"), + it.get("kv_updated_at"), + it.get("kv_dumped_keys"), + ) + + seen.add(kid) + if kid not in _kdn_meta_cache: + added.append(kid) + _kdn_meta_cache[kid] = sig + elif _kdn_meta_cache[kid] != sig: + updated.append(kid) + _kdn_meta_cache[kid] = sig + + removed = [k for k in _kdn_meta_cache.keys() if k not in seen] + for k in removed: + _kdn_meta_cache.pop(k, None) + + return added, updated, removed + diff --git a/scheduler/scheduler.py b/scheduler/scheduler.py index c8ca19b..0b053e8 100644 --- a/scheduler/scheduler.py +++ b/scheduler/scheduler.py @@ -1,789 +1,855 @@ -""" -基于 FastAPI 的简单 Scheduler HTTP 入口。 -相较于v0 scheduler.py，采用http集成库进行发包，不再采用protocol/interface - -调度器核心： - - 基于 FastAPI 提供 HTTP 接口 - - 启动时挂起控制平面，构建proxy池和kdn池 - - 控制平面接收来自proxy和kdn服务器的注册，并维护资源池 - - 接收 /v1/chat/completions 和 /v1/completions 的 POST 请求 - - 解析 URL / payload / 客户端 IP - - 分配 1~65535 循环的 request_id - - 调用 Request.build_request(...) 构建内部 Request 对象，构建Request时调用策略选择最优kdn和proxy - -调度器Workflow： - 调度层： - - Scheduler 收到 HTTP 请求（OpenAI 风格）； - - 构建request（包含分配ID，决定各项服务以及下一级调度） - - 决定具体要打哪个下游 URL。 - 转发层： - - 拿着“下游 URL + 要发的数据 + headers”， - - 用 forward_request 把请求送到 Proxy， - - 一边从下游拉流，一边原样推回给用户（StreamingResponse）。 - -后续可以在这里接 Proxy / 流式传输等逻辑。 -""" - -from __future__ import annotations -import sys -import asyncio -import os, logging -from logging.handlers import RotatingFileHandler -from datetime import datetime -from pathlib import Path - -import uvicorn -# import aiohttp -# import httpx -from typing import Any, Dict, List -from contextlib import asynccontextmanager - -from fastapi import FastAPI, Request as FastAPIRequest -from fastapi.responses import StreamingResponse, JSONResponse -from pydantic import BaseModel - -from core import TokenizerRegistry -from core import Request as SchedulerRequest -from core import forward_request -from core import config - -from model import EmbeddingEngine -# from .knowledge.kdn_client import fetch_kdn_snapshot - -from .knowledge.kdn_sync import ( - # build_table_from_kdn_items, - kdn_auto_refresh_loop, - kdn_refresh_once -) -from .resource.control_plane import control_plane, set_pool, get_pool, set_kdn_pool, get_kdn_pool, set_on_kdn_register -from .resource.proxy_pool import ProxyPool -from .resource.kdn_pool import KDNPool -from .strategy import create_strategy - -# from store import ( -# # DummyEmbeddingModel, -# init_knowledge_table, -# # KnowledgeTable, -# # KnowledgeUnit -# ) - -from util import timing - -# 如果外部没给，就用一个保底默认值（方便本地直接跑） -from core.config import DEFAULT_MODEL, DEFAULT_EMBED_MODEL, SCHEDULER_CP_PORT - - -# ======================= 请求ID分配器等基本函数======================= - -class RequestIdAllocator: - """ - 简单的 16bit 请求 ID 分配器： - - 有效范围：1 ~ max_id（默认 65535） - - 超过后从 1 重新开始 - """ - def __init__(self, max_id: int = 65535) -> None: - self._max_id = max_id - self._current = 0 - - def next_id(self) -> int: - self._current += 1 - if self._current > self._max_id: - self._current = 1 - return self._current - - -id_alloc = RequestIdAllocator() -logger = logging.getLogger("scheduler") -logger.setLevel(logging.INFO) - - -class PeekPayload(BaseModel): - kids: List[str] - need_fields: List[str] = ["length", "avail_kdn_servers", "text_abstract"] - - -def init_logging() -> str: - """ - 目标： - 1) SCHEDULER_LOG_FILE 允许配置为“目录路径”或“文件路径” - - 目录：/logs/scheduler -> /logs/scheduler/scheduler-YYYYMMDD-HHMMSS.log - - 文件：/logs/scheduler/scheduler.log -> /logs/scheduler/scheduler-YYYYMMDD-HHMMSS.log - 2) 修复 handler 重复创建（先复用已有 RotatingFileHandler，找不到才创建） - 3) 隔离 uvicorn 对 root logger 的影响：业务日志写在独立 logger 上 propagate=False - 4) 控制台只输出 ERROR（错误立刻可见） - """ - # ---------- 1) 解析用户配置：目录 or 文件 ---------- - cfg_path = os.environ.get("SCHEDULER_LOG_FILE", getattr(config, "SCHEDULER_LOG_FILE", "scheduler.log")) - p = Path(str(cfg_path)).expanduser() - - ts = datetime.now().strftime("%Y%m%d-%H%M%S") - - if p.suffix.lower() == ".log": - # 用户给的是文件路径 - base_dir = p.parent - stem = p.stem or "scheduler" - else: - # 用户给的是目录路径 - base_dir = p - stem = "scheduler" - - base_dir.mkdir(parents=True, exist_ok=True) - log_path = str(base_dir / f"{stem}-{ts}.log") - - # ---------- 2) 业务 logger（不走 root） ---------- - biz = logging.getLogger("scheduler") - cp = logging.getLogger("scheduler.control_plane") - hb = logging.getLogger("scheduler.hbreport") - - for lg in (biz, cp, hb): - lg.setLevel(logging.INFO) - lg.propagate = False - - # ---------- 3) 先复用已有 file handler（修复必修问题1） ---------- - def _find_same_file_handler(lg: logging.Logger, target: str) -> RotatingFileHandler | None: - for h in lg.handlers: - if isinstance(h, RotatingFileHandler) and getattr(h, "baseFilename", None) == target: - return h - return None - - fh = ( - _find_same_file_handler(biz, log_path) - or _find_same_file_handler(cp, log_path) - or _find_same_file_handler(hb, log_path) - ) - - if fh is None: - file_fmt = logging.Formatter( - "%(asctime)s.%(msecs)03d %(levelname)s %(name)s:%(lineno)d | %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - ) - fh = RotatingFileHandler( - log_path, - maxBytes=50 * 1024 * 1024, - backupCount=5, - encoding="utf-8", - ) - fh.setLevel(logging.INFO) - fh.setFormatter(file_fmt) - - # 挂载到三个 logger（避免重复 add） - for lg in (biz, cp, hb): - if fh not in lg.handlers: - lg.addHandler(fh) - - # ---------- 4) root 控制台：ERROR only ---------- - root = logging.getLogger() - root.setLevel(logging.INFO) - - stream_handlers = [h for h in root.handlers if isinstance(h, logging.StreamHandler)] - if not stream_handlers: - sh = logging.StreamHandler(sys.stdout) - sh.setLevel(logging.ERROR) - sh.setFormatter(logging.Formatter("%(levelname)s | %(name)s | %(message)s")) - root.addHandler(sh) - else: - for h in stream_handlers: - h.setLevel(logging.ERROR) - - # ---------- 5) 降噪（可选） ---------- - logging.getLogger("httpx").setLevel(logging.WARNING) - logging.getLogger("httpcore").setLevel(logging.WARNING) - logging.getLogger("uvicorn").setLevel(logging.WARNING) - logging.getLogger("uvicorn.error").setLevel(logging.WARNING) - logging.getLogger("uvicorn.access").setLevel(logging.WARNING) - - # ---------- 6) 自检：写一条，立刻 flush ---------- - biz.info("[Scheduler] logging initialized, log_path=%s", log_path) - for h in biz.handlers: - try: - h.flush() - except Exception: - pass - - return log_path - - -# ======================= Scheduler初始化 ======================= -@asynccontextmanager -async def lifespan(app: FastAPI): - """ - FastAPI 生命周期管理： - - startup: 预热分词器 - - shutdown: 目前没资源要清理，预留接口 - """ - log_path = init_logging() - print(f"[Scheduler] started. log_file={log_path}") - # ------------------------------------------ - # 尝试预热 tokenizer 分词器 - # ------------------------------------------ - model_path = os.getenv("SCHEDULER_MODEL_PATH", DEFAULT_MODEL) - try: - # print("[Scheduler] startup: 开始预热分词器") - TokenizerRegistry.warmup_tokenizers(model_path) - logger.info(f"[Scheduler] Warmup tokenizers, model_path={model_path!r}") - except Exception as e: - logger.info(f"[Scheduler] 分词器预热失败:{e}") - - # ------------------------------------------ - # 尝试预热Embedding模型 - # ------------------------------------------ - embedding_model_name = os.getenv("SCHEDULER_EMBEDDING_MODEL", DEFAULT_EMBED_MODEL) - if os.path.isabs(embedding_model_name) and not os.path.isdir(embedding_model_name): - raise RuntimeError( - f"SCHEDULER_EMBEDDING_MODEL is a local path but not found: {embedding_model_name}" - ) - - try: - embedder = EmbeddingEngine(model_name=embedding_model_name) - app.state.embedding_engine = embedder # type: ignore - - try: - _ = embedder.encode_vector(["__warm_up__"])[0] - logger.info( - f"[Scheduler] Warmup embedding model: {embedding_model_name!r}, " - f"dim={embedder.dim}, device={embedder.device}" - ) - - except Exception as e: - logger.warning(f"[Scheduler] Warmup embedding model failed: {e}") - - except Exception as e: - logger.exception(f"[Scheduler] 预热 Embedding 模型失败: {e}") - app.state.embedding_engine = None # type: ignore - - # ------------------------------------------------------------- - # 尝试启动控制平面，启用proxy池和KDN池，并从KDN池拉取知识清单来初始化知识库 - # ------------------------------------------------------------- - #先构建Proxy池实例 - ttl = int(os.environ.get("SCHEDULER_PROXY_TTL_S", config.CONTROL_PLANE_TTL_S)) - app.state.proxy_pool = ProxyPool(ttl_s=ttl) # type: ignore - set_pool(app.state.proxy_pool) # type: ignore 让 7002 复用这份池 - - kdn_ttl = int(os.environ.get("SCHEDULER_KDN_TTL_S", config.CONTROL_PLANE_TTL_S)) - app.state.kdn_pool = KDNPool(ttl_s=kdn_ttl) # type: ignore - set_kdn_pool(app.state.kdn_pool) # type: ignore - - app.state.knowledge_table = None # type: ignore - app.state.last_refresh_ts = 0 # type: ignore - - # KDN refresh loop infra (always on) - app.state._kdn_refresh_lock = asyncio.Lock() # type: ignore - app.state._kdn_stop_event = asyncio.Event() # type: ignore - app.state._kdn_refresh_task = asyncio.create_task( # type: ignore - kdn_auto_refresh_loop(app, app.state._kdn_stop_event) # type: ignore - ) - - logger.info("[Scheduler] KDN refresh loop started (pool-based).") - - async def _trigger_refresh_once() -> None: - # 如果 refresh 正在跑，就直接跳过（复用你的锁） - lock = getattr(app.state, "_kdn_refresh_lock", None) # type: ignore - if lock is None: - return - if lock.locked(): - return - try: - await kdn_refresh_once(app) - except Exception: - logger.exception("[Scheduler] kdn_refresh_once failed when triggered by kdn_register") - - set_on_kdn_register(_trigger_refresh_once) - logger.info("[Scheduler] on_kdn_register callback installed") - - async def _run_control_plane(app): - host = os.environ.get("SCHEDULER_CP_HOST", "0.0.0.0") - port = int(os.environ.get("SCHEDULER_CP_PORT", SCHEDULER_CP_PORT)) - - uv_cfg = uvicorn.Config( - control_plane, - host=host, - port=port, - log_level=os.environ.get("SCHEDULER_CP_LOG_LEVEL", "info"), - loop="asyncio", - lifespan="on", - access_log=False, - ) - server = uvicorn.Server(uv_cfg) - app.state._cp_server = server - await server.serve() - - app.state._cp_task = asyncio.create_task(_run_control_plane(app)) # type: ignore - - # ------------------------------------------ - # 调用具体scheduler策略 - # ------------------------------------------ - strategy_name = os.environ.get("SCHEDULER_STRATEGY", "round_robin") - app.state.proxy_strategy = create_strategy(strategy_name) # type: ignore - logger.info("[Scheduler] scheduler strategy=%s", strategy_name) - - # 这里 yield 之后是正常服务期 - logger.info("[Scheduler] startup: 初始化完成，监听服务中") - try: - yield - - finally: - # ---------- shutdown ---------- - cp_server = getattr(app.state, "_cp_server", None) # type: ignore - cp_task = getattr(app.state, "_cp_task", None) # type: ignore - ev = getattr(app.state, "_kdn_stop_event", None) # type: ignore - task = getattr(app.state, "_kdn_refresh_task", None) # type: ignore - - if ev is not None: - ev.set() - if task is not None: - task.cancel() - - if cp_server is not None: - cp_server.should_exit = True - if cp_task is not None: - try: - await cp_task - except Exception: - pass - - # shutdown 阶段，如果后面有资源要清理，可以写在这里 - logger.info("[Scheduler] shutdown: 结束服务") - - -scheduler = FastAPI( - title="CacheRoute Scheduler", - version="0.1.1", - lifespan=lifespan, -) - - - - -# ======================= 公共内部处理函数 ======================= -@timing -def _handle_client( - app: FastAPI, - url_path: str, - payload: Dict[str, Any], - client_ip: str, - proxies: List[Dict[str, Any]] | None = None, - kdns: List[Dict[str, Any]] | None = None, - strategy: Any | None = None, -) -> SchedulerRequest: - """ - 根据 HTTP 请求信息 - 构造内部 Request 对象，分配request_id - - url_path: 例如 "/v1/chat/completions" - - payload: 解析后的 JSON 字典 - - client_ip: 客户端 IP 字符串 - """ - # 分配 request_id（1~65535 循环） - request_id = id_alloc.next_id() - - # 打印调试信息 - if os.environ.get("SCHEDULER_VERBOSE_REQUEST_LOG", config.SCHEDULER_VERBOSE_REQUEST_LOG) == 1: - print("+" * 80) - print(f"[Scheduler] 收到 HTTP 请求: path={url_path}, client_ip={client_ip},\n" - f"分配 request_id={request_id},\n" - f"payload={payload}") - print("+" * 80) - - # 直接复用你改好的 build_request - req_obj = SchedulerRequest.build_request( - url_path=url_path, - payload=payload, - user_addr=client_ip, - request_id=request_id, - embedder=getattr(app.state, "embedding_engine", None), # type: ignore - knowledge_table=getattr(app.state, "knowledge_table", None), # type: ignore - proxies=proxies, - kdns=kdns, - strategy=strategy, - ) - 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" - f"[Scheduler] Endpoint_type={getattr(req_obj.Service, 'Endpoint_type', None)},\n" - f"[Scheduler] Knowledge_List={req_obj.Service.Knowledge_List},\n" - f"[Scheduler] Selected KDN={req_obj.Task.KDN_server_addr}, ProxyID={req_obj.Task.P_proxy_id}[{req_obj.Task.P_proxy_addr}:{req_obj.Task.P_proxy_port}]" - ) - print(f"[Scheduler] Injection_type={getattr(req_obj.Service, 'Injection_type', None)}") - return req_obj - - -# ======================= 调度器方法路由 ======================= -@scheduler.post("/v1/chat/completions") -async def create_chat_completions(request: FastAPIRequest): - """ - 对应： - curl http://HOST:PORT/v1/chat/completions -H "Content-Type: application/json" -d '{...}' - payload 结构参考 OpenAI chat 接口： - { - "model": "...", - "messages": [ ... ], - "max_tokens": ..., - "temperature": ..., - "top_p": ..., - "stream": true/false, - ... - } - """ - # 解析用户原始请求体 - try: - payload = await request.json() - except Exception as e: - return JSONResponse( - status_code=400, - content={"error": "invalid_json", "detail": str(e)}, - ) - - # 取客户端 IP 和路径（用于构建内部 Request） - client_ip = request.client.host if request.client else "unknown" - url_path = request.url.path - - # 把用户请求头转成 dict[str, str] - raw_headers = {k.lower(): v for k, v in request.headers.items()} - - # 构建内部 Request（包含 Prompt/Service/Task 等） - pool = get_pool() - proxy_infos = await pool.list(include_dead=False) - proxies = [{"proxy_id": p.proxy_id, "host": p.host, "port": p.port} for p in proxy_infos] - - kdn_pool = get_kdn_pool() - kdn_infos = await kdn_pool.list(include_dead=False) - kdns = [{"kdn_id": k.kdn_id, "host": k.host, "port": k.port} for k in kdn_infos] - - strategy = getattr(request.app.state, "proxy_strategy", None) - req_obj = _handle_client(request.app, url_path, payload, client_ip, proxies=proxies, kdns=kdns, strategy=strategy,) - - 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") - - # 根据调度结果选择下游 URL，更新对应proxy_id的inflight - host = req_obj.Task.P_proxy_addr - port = req_obj.Task.P_proxy_port - endpoint = getattr(req_obj.Service, "Endpoint_type", "chat/completions") - downstream_url = f"http://{host}:{port}/v1/{endpoint}" - proxy_id = req_obj.Task.P_proxy_id - if not proxy_id: - raise RuntimeError("Routing failed: missing P_proxy_id for inflight tracking") - ok = await pool.inflight_delta(proxy_id, +1) - if not ok: - raise RuntimeError(f"Proxy not found in pool: {proxy_id}") - - # 转化payload - data_for_downstream = req_obj.to_payload() - - # 处理需要透传给下游的 headers（可选择性过滤） - extra_headers = {} - - # 把用户的 Authorization 原样传给下游（如果希望由 Proxy 做鉴权） - if "authorization" in raw_headers: - extra_headers["authorization"] = raw_headers["authorization"] - - # 带一个 Scheduler分配的Request ID，便于链路追踪 - extra_headers["scheduler-request-id"] = str(req_obj.Request_ID) - - # 定义一个 async 生成器，从下游流式读取，再回给上游 - async def iter_upstream(): - try: - # forward_request 本身是一个 async generator - async for chunk in forward_request( - url=downstream_url, - data=data_for_downstream, - use_chunked=True, # chat/completions 一般是流式 - extra_headers=extra_headers, # 透传必要头 - ): - # 这里 chunk 已经是 bytes，直接 yield 出去即可 - yield chunk - finally: - await pool.inflight_delta(proxy_id, -1) - - # 用 StreamingResponse 包一层，实现用户侧的流式响应 - return StreamingResponse(iter_upstream(), media_type="application/json") - - -@scheduler.post("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/v1/completions") -async def create_completions(request: FastAPIRequest): - """ - 对应： - curl http://HOST:PORT/v1/completions -H "Content-Type: application/json" -d '{...}' - payload 结构类似 OpenAI completions 接口： - { - "model": "...", - "prompt": "...." 或 ["..."], - "max_tokens": ..., - "temperature": ..., - "top_p": ..., - "stream": true/false, - ... - } - """ - try: - payload = await request.json() - except Exception as e: - return JSONResponse( - status_code=400, - content={"error": "invalid_json", "detail": str(e)}, - ) - - # 取客户端 IP 和路径（用于构建内部 Request） - client_ip = request.client.host if request.client else "unknown" - url_path = request.url.path - - # 把用户请求头转成 dict[str, str] - raw_headers = {k.lower(): v for k, v in request.headers.items()} - - # 构建内部 Request（包含 Prompt/Service/Task 等） - pool = get_pool() - proxy_infos = await pool.list(include_dead=False) - proxies = [{"proxy_id": p.proxy_id, "host": p.host, "port": p.port} for p in proxy_infos] - - kdn_pool = get_kdn_pool() - kdn_infos = await kdn_pool.list(include_dead=False) - kdns = [{"kdn_id": k.kdn_id, "host": k.host, "port": k.port} for k in kdn_infos] - - strategy = getattr(request.app.state, "proxy_strategy", None) - req_obj = _handle_client(request.app, url_path, payload, client_ip, proxies=proxies, kdns=kdns, strategy=strategy,) - - 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") - - # 根据调度结果选择下游 URL，更新对应proxy_id的inflight - host = req_obj.Task.P_proxy_addr - port = req_obj.Task.P_proxy_port - endpoint = getattr(req_obj.Service, "Endpoint_type", "completions") - downstream_url = f"http://{host}:{port}/v1/{endpoint}" - proxy_id = req_obj.Task.P_proxy_id - if not proxy_id: - raise RuntimeError("Routing failed: missing P_proxy_id for inflight tracking") - ok = await pool.inflight_delta(proxy_id, +1) - if not ok: - raise RuntimeError(f"Proxy not found in pool: {proxy_id}") - - # 转化payload - data_for_downstream = req_obj.to_payload() - - # 处理需要透传给下游的 headers（可选择性过滤） - extra_headers = {} - - # 把用户的 Authorization 原样传给下游（如果希望由 Proxy 做鉴权） - if "authorization" in raw_headers: - extra_headers["authorization"] = raw_headers["authorization"] - - # 带一个 Scheduler分配的Request ID，便于链路追踪 - extra_headers["scheduler-request-id"] = str(req_obj.Request_ID) - - # 定义一个 async 生成器，从下游流式读取，再回给上游 - async def iter_upstream(): - try: - # forward_request 本身是一个 async generator - async for content in forward_request( - url=downstream_url, - data=data_for_downstream, - use_chunked=False, # chat/completions 一般是流式 - extra_headers=extra_headers, # 透传必要头 - ): - # 这里 chunk 已经是 bytes，直接 yield 出去即可 - yield content - finally: - await pool.inflight_delta(proxy_id, -1) - - # 用 StreamingResponse 包一层，实现用户侧的流式响应 - return StreamingResponse(iter_upstream(), media_type="application/json") - - -@scheduler.get("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/debug/status") -async def debug_status() -> Dict[str, Any]: - """ - Debug endpoint for Scheduler CLI. - - Returns: - - knowledge_loaded: whether knowledge_table is ready - - entries: number of knowledge units - - dim: embedding dim - - faiss_total: index.ntotal if built (else None) - - kdn_base_url: where snapshot comes from (if any) - - last_refresh_ts: unix seconds (if you maintain it), else None - - unit_fields: what fields a KnowledgeUnit contains (for inspection) - - sample_kids: first few kids (for quick view) - """ - table = getattr(scheduler.state, "knowledge_table", None) # type: ignore - kdn_base_url = getattr(scheduler.state, "kdn_base_url", None) # type: ignore - last_refresh_ts = getattr(scheduler.state, "last_refresh_ts", None) # type: ignore - kdn_alive = int(getattr(scheduler.state, "kdn_alive", 0) or 0) # type: ignore - kdn_alive_addrs = list(getattr(scheduler.state, "kdn_alive_addrs", []) or []) # type: ignore - kdn_last_selected = getattr(scheduler.state, "kdn_last_selected", None) # type: ignore - kdn_last_selected_id = getattr(scheduler.state, "kdn_last_selected_id", None) # type: ignore - kdn_last_refresh_ok = bool(getattr(scheduler.state, "kdn_last_refresh_ok", False)) # type: ignore - kdn_last_refresh_reason = getattr(scheduler.state, "kdn_last_refresh_reason", "") # type: ignore - kdn_last_refresh_ts = int(getattr(scheduler.state, "kdn_last_refresh_ts", 0) or 0) # type: ignore - - pool = get_pool() - proxy_infos = await pool.list(include_dead=False) - - proxy_states = [] - for p in proxy_infos: - proxy_states.append({ - "proxy_id": p.proxy_id, - "host": p.host, - "port": p.port, - "max_capacity": p.load.max_capacity, - "instance_count": p.load.instance_count, - "kv_mem_per_instance_gb": p.load.kv_mem_per_instance_gb, - "kv_cache_pool_gb": p.load.kv_cache_pool_gb, - "kv_cache_update_policy": getattr(p, "kv_cache_update_policy", "static"), - "inflight": p.load.inflight, - "qps_1m": p.load.qps_1m, - "gpu_util": p.load.gpu_util, - }) - - if table is None: - - return { - "knowledge_loaded": False, - "entries": 0, - "dim": None, - "faiss_total": None, - "kdn_base_url": kdn_base_url, - "last_refresh_ts": last_refresh_ts, - "unit_fields": [], - "sample_kids": [], - "kdn_alive": 0, - "kdn_alive_addrs": [], - "kdn_last_selected": None, - "kdn_last_selected_id": None, - "kdn_last_refresh_ok": False, - "kdn_last_refresh_reason": "", - "kdn_last_refresh_ts": 0, - "proxies": proxy_states, - } - - units = getattr(table, "_units", {}) - kids = list(units.keys()) - kids_sorted = sorted(kids)[:10] - - # 尽量从 table 拿到 dim - dim = getattr(table, "dim", None) - - # FAISS - faiss_total = None - idx = getattr(table, "_faiss_index", None) - if idx is not None: - try: - faiss_total = int(getattr(idx, "ntotal")) - except Exception: - faiss_total = None - - # KnowledgeUnit 字段探测（帮助你看 snapshot 后到底存了什么） - unit_fields: List[str] = [] - if kids_sorted: - u0 = units[kids_sorted[0]] - # dataclass 一般有 __dict__ - if hasattr(u0, "__dict__"): - unit_fields = sorted(list(u0.__dict__.keys())) - else: - # fallback：dir 过滤 - unit_fields = [x for x in dir(u0) if not x.startswith("_")] - - return { - "knowledge_loaded": True, - "entries": len(units), - "dim": dim, - "faiss_total": faiss_total, - "kdn_base_url": kdn_base_url, - "last_refresh_ts": last_refresh_ts, - "unit_fields": unit_fields, - "sample_kids": kids_sorted, - "kdn_alive": kdn_alive, - "kdn_alive_addrs": kdn_alive_addrs, - "kdn_last_selected": kdn_last_selected, - "kdn_last_selected_id": kdn_last_selected_id, - "kdn_last_refresh_ok": kdn_last_refresh_ok, - "kdn_last_refresh_reason": kdn_last_refresh_reason, - "kdn_last_refresh_ts": kdn_last_refresh_ts, - "proxies": proxy_states, - } - - -@scheduler.post("/debug/knowledge/peek") -async def debug_peek_knowledge(payload: PeekPayload) -> Dict[str, Any]: - """ - Peek knowledge units by kids. - Default fields are safe (no full embedding). - """ - table = getattr(scheduler.state, "knowledge_table", None) # type: ignore - if table is None: - return {"items": [], "miss": payload.kids} - - units = getattr(table, "_units", {}) - items = [] - miss = [] - - for kid in payload.kids: - k = (kid or "").strip().lower() - u = units.get(k) - if u is None: - miss.append(kid) - continue - - it = {"kid": k} - allow = set(payload.need_fields or []) - - if "length" in allow: - it["length"] = getattr(u, "length", None) - if "avail_kdn_servers" in allow: - it["avail_kdn_servers"] = getattr(u, "avail_kdn_servers", None) - if "text_abstract" in allow: - it["text_abstract"] = getattr(u, "text_abstract", None) - - # 可选：你若在 KnowledgeUnit/metadata 里存了 kv_ready 等，也可以在这里输出 - if "meta" in allow and hasattr(u, "meta"): - it["meta"] = getattr(u, "meta") - if "kv_ready" in allow: - it["kv_ready"] = getattr(u, "kv_ready", 0) - if "kv_rel_dir" in allow: - it["kv_rel_dir"] = getattr(u, "kv_rel_dir", None) - if "kv_dumped_keys" in allow: - it["kv_dumped_keys"] = getattr(u, "kv_dumped_keys", None) - if "kv_updated_at" in allow: - it["kv_updated_at"] = getattr(u, "kv_updated_at", None) - - items.append(it) - - return {"items": items, "miss": miss} - - -@scheduler.post("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/admin/refresh_knowledge") -async def admin_refresh_knowledge(): - try: - r = await kdn_refresh_once(scheduler) - return JSONResponse(content=r) - except Exception as e: - logger.exception(f"[Scheduler] admin refresh failed: {e}") - return JSONResponse(status_code=500, content={"ok": False, "error": str(e)}) - - -@scheduler.get("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/debug/strategy") -async def debug_strategy() -> Dict[str, Any]: - strategy = getattr(scheduler.state, "proxy_strategy", None) # type: ignore - pool = get_pool() - proxy_infos = await pool.list(include_dead=False) - - # 只返回简要信息，避免输出过大 - sample = [{ - "proxy_id": p.proxy_id, - "host": p.host, - "port": p.port, - "is_alive": True, # list(include_dead=False) 已保证 alive - } for p in proxy_infos[:10]] - - return { - "strategy": getattr(strategy, "name", None) or type(strategy).__name__ if strategy else None, - "proxy_count": len(proxy_infos), - "proxies_sample": sample, - } - -# 预留：/knowledge/update 等路由以后再加 -# @api.post("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/knowledge/update") -# async def knowledge_update(request: FastAPIRequest): -# payload = await request.json() -# client_ip = request.client.host if request.client else "unknown" -# ... -# return JSONResponse(content={"status": "ok"}) - - +""" +基于 FastAPI 的简单 Scheduler HTTP 入口。 +相较于v0 scheduler.py，采用http集成库进行发包，不再采用protocol/interface + +调度器核心： + - 基于 FastAPI 提供 HTTP 接口 + - 启动时挂起控制平面，构建proxy池和kdn池 + - 控制平面接收来自proxy和kdn服务器的注册，并维护资源池 + - 接收 /v1/chat/completions 和 /v1/completions 的 POST 请求 + - 解析 URL / payload / 客户端 IP + - 分配 1~65535 循环的 request_id + - 调用 Request.build_request(...) 构建内部 Request 对象，构建Request时调用策略选择最优kdn和proxy + +调度器Workflow： + 调度层： + - Scheduler 收到 HTTP 请求（OpenAI 风格）； + - 构建request（包含分配ID，决定各项服务以及下一级调度） + - 决定具体要打哪个下游 URL。 + 转发层： + - 拿着“下游 URL + 要发的数据 + headers”， + - 用 forward_request 把请求送到 Proxy， + - 一边从下游拉流，一边原样推回给用户（StreamingResponse）。 + +后续可以在这里接 Proxy / 流式传输等逻辑。 +""" + +from __future__ import annotations +import sys +import asyncio +import os, logging +from logging.handlers import RotatingFileHandler +from datetime import datetime +from pathlib import Path + +import uvicorn +# import aiohttp +# import httpx +from typing import Any, Dict, List +from contextlib import asynccontextmanager + +from fastapi import FastAPI, Request as FastAPIRequest +from fastapi.responses import StreamingResponse, JSONResponse +from pydantic import BaseModel + +from core import TokenizerRegistry +from core import Request as SchedulerRequest +from core import forward_request +from core import config + +from model import EmbeddingEngine +# from .knowledge.kdn_client import fetch_kdn_snapshot + +from .knowledge.kdn_sync import ( + # build_table_from_kdn_items, + kdn_auto_refresh_loop, + kdn_refresh_once +) +from .resource.control_plane import control_plane, set_pool, get_pool, set_kdn_pool, get_kdn_pool, set_on_kdn_register +from .resource.proxy_pool import ProxyPool +from .resource.kdn_pool import KDNPool +from .strategy import create_strategy + +# from store import ( +# # DummyEmbeddingModel, +# init_knowledge_table, +# # KnowledgeTable, +# # KnowledgeUnit +# ) + +from util import timing + +# 如果外部没给，就用一个保底默认值（方便本地直接跑） +from core.config import DEFAULT_MODEL, DEFAULT_EMBED_MODEL, SCHEDULER_CP_PORT + + +# ======================= 请求ID分配器等基本函数======================= + +class RequestIdAllocator: + """ + 简单的 16bit 请求 ID 分配器： + - 有效范围：1 ~ max_id（默认 65535） + - 超过后从 1 重新开始 + """ + def __init__(self, max_id: int = 65535) -> None: + self._max_id = max_id + self._current = 0 + + def next_id(self) -> int: + self._current += 1 + if self._current > self._max_id: + self._current = 1 + return self._current + + +id_alloc = RequestIdAllocator() +logger = logging.getLogger("scheduler") +logger.setLevel(logging.INFO) + + +class PeekPayload(BaseModel): + kids: List[str] + need_fields: List[str] = ["length", "avail_kdn_servers", "text_abstract"] + + +def init_logging() -> str: + """ + 目标： + 1) SCHEDULER_LOG_FILE 允许配置为“目录路径”或“文件路径” + - 目录：/logs/scheduler -> /logs/scheduler/scheduler-YYYYMMDD-HHMMSS.log + - 文件：/logs/scheduler/scheduler.log -> /logs/scheduler/scheduler-YYYYMMDD-HHMMSS.log + 2) 修复 handler 重复创建（先复用已有 RotatingFileHandler，找不到才创建） + 3) 隔离 uvicorn 对 root logger 的影响：业务日志写在独立 logger 上 propagate=False + 4) 控制台只输出 ERROR（错误立刻可见） + """ + # ---------- 1) 解析用户配置：目录 or 文件 ---------- + cfg_path = os.environ.get("SCHEDULER_LOG_FILE", getattr(config, "SCHEDULER_LOG_FILE", "scheduler.log")) + p = Path(str(cfg_path)).expanduser() + + ts = datetime.now().strftime("%Y%m%d-%H%M%S") + + if p.suffix.lower() == ".log": + # 用户给的是文件路径 + base_dir = p.parent + stem = p.stem or "scheduler" + else: + # 用户给的是目录路径 + base_dir = p + stem = "scheduler" + + base_dir.mkdir(parents=True, exist_ok=True) + log_path = str(base_dir / f"{stem}-{ts}.log") + + # ---------- 2) 业务 logger（不走 root） ---------- + biz = logging.getLogger("scheduler") + cp = logging.getLogger("scheduler.control_plane") + hb = logging.getLogger("scheduler.hbreport") + + for lg in (biz, cp, hb): + lg.setLevel(logging.INFO) + lg.propagate = False + + # ---------- 3) 先复用已有 file handler（修复必修问题1） ---------- + def _find_same_file_handler(lg: logging.Logger, target: str) -> RotatingFileHandler | None: + for h in lg.handlers: + if isinstance(h, RotatingFileHandler) and getattr(h, "baseFilename", None) == target: + return h + return None + + fh = ( + _find_same_file_handler(biz, log_path) + or _find_same_file_handler(cp, log_path) + or _find_same_file_handler(hb, log_path) + ) + + if fh is None: + file_fmt = logging.Formatter( + "%(asctime)s.%(msecs)03d %(levelname)s %(name)s:%(lineno)d | %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + fh = RotatingFileHandler( + log_path, + maxBytes=50 * 1024 * 1024, + backupCount=5, + encoding="utf-8", + ) + fh.setLevel(logging.INFO) + fh.setFormatter(file_fmt) + + # 挂载到三个 logger（避免重复 add） + for lg in (biz, cp, hb): + if fh not in lg.handlers: + lg.addHandler(fh) + + # ---------- 4) root 控制台：ERROR only ---------- + root = logging.getLogger() + root.setLevel(logging.INFO) + + stream_handlers = [h for h in root.handlers if isinstance(h, logging.StreamHandler)] + if not stream_handlers: + sh = logging.StreamHandler(sys.stdout) + sh.setLevel(logging.ERROR) + sh.setFormatter(logging.Formatter("%(levelname)s | %(name)s | %(message)s")) + root.addHandler(sh) + else: + for h in stream_handlers: + h.setLevel(logging.ERROR) + + # ---------- 5) 降噪（可选） ---------- + logging.getLogger("httpx").setLevel(logging.WARNING) + logging.getLogger("httpcore").setLevel(logging.WARNING) + logging.getLogger("uvicorn").setLevel(logging.WARNING) + logging.getLogger("uvicorn.error").setLevel(logging.WARNING) + logging.getLogger("uvicorn.access").setLevel(logging.WARNING) + + # ---------- 6) 自检：写一条，立刻 flush ---------- + biz.info("[Scheduler] logging initialized, log_path=%s", log_path) + for h in biz.handlers: + try: + h.flush() + except Exception: + pass + + return log_path + + +# ======================= Scheduler初始化 ======================= +@asynccontextmanager +async def lifespan(app: FastAPI): + """ + FastAPI 生命周期管理： + - startup: 预热分词器 + - shutdown: 目前没资源要清理，预留接口 + """ + log_path = init_logging() + print(f"[Scheduler] started. log_file={log_path}") + # ------------------------------------------ + # 尝试预热 tokenizer 分词器 + # ------------------------------------------ + model_path = os.getenv("SCHEDULER_MODEL_PATH", DEFAULT_MODEL) + try: + # print("[Scheduler] startup: 开始预热分词器") + TokenizerRegistry.warmup_tokenizers(model_path) + logger.info(f"[Scheduler] Warmup tokenizers, model_path={model_path!r}") + except Exception as e: + logger.info(f"[Scheduler] 分词器预热失败:{e}") + + # ------------------------------------------ + # 尝试预热Embedding模型 + # ------------------------------------------ + embedding_model_name = os.getenv("SCHEDULER_EMBEDDING_MODEL", DEFAULT_EMBED_MODEL) + if os.path.isabs(embedding_model_name) and not os.path.isdir(embedding_model_name): + raise RuntimeError( + f"SCHEDULER_EMBEDDING_MODEL is a local path but not found: {embedding_model_name}" + ) + + try: + embedder = EmbeddingEngine(model_name=embedding_model_name) + app.state.embedding_engine = embedder # type: ignore + + try: + _ = embedder.encode_vector(["__warm_up__"])[0] + logger.info( + f"[Scheduler] Warmup embedding model: {embedding_model_name!r}, " + f"dim={embedder.dim}, device={embedder.device}" + ) + + except Exception as e: + logger.warning(f"[Scheduler] Warmup embedding model failed: {e}") + + except Exception as e: + logger.exception(f"[Scheduler] 预热 Embedding 模型失败: {e}") + app.state.embedding_engine = None # type: ignore + + # ------------------------------------------------------------- + # 尝试启动控制平面，启用proxy池和KDN池，并从KDN池拉取知识清单来初始化知识库 + # ------------------------------------------------------------- + #先构建Proxy池实例 + ttl = int(os.environ.get("SCHEDULER_PROXY_TTL_S", config.CONTROL_PLANE_TTL_S)) + app.state.proxy_pool = ProxyPool(ttl_s=ttl) # type: ignore + set_pool(app.state.proxy_pool) # type: ignore 让 7002 复用这份池 + + kdn_ttl = int(os.environ.get("SCHEDULER_KDN_TTL_S", config.CONTROL_PLANE_TTL_S)) + app.state.kdn_pool = KDNPool(ttl_s=kdn_ttl) # type: ignore + set_kdn_pool(app.state.kdn_pool) # type: ignore + + app.state.knowledge_table = None # type: ignore + app.state.last_refresh_ts = 0 # type: ignore + + # KDN refresh loop infra (always on) + app.state._kdn_refresh_lock = asyncio.Lock() # type: ignore + app.state._kdn_stop_event = asyncio.Event() # type: ignore + app.state._kdn_refresh_task = asyncio.create_task( # type: ignore + kdn_auto_refresh_loop(app, app.state._kdn_stop_event) # type: ignore + ) + + logger.info("[Scheduler] KDN refresh loop started (pool-based).") + + async def _trigger_refresh_once() -> None: + # 如果 refresh 正在跑，就直接跳过（复用你的锁） + lock = getattr(app.state, "_kdn_refresh_lock", None) # type: ignore + if lock is None: + return + if lock.locked(): + return + try: + await kdn_refresh_once(app) + except Exception: + logger.exception("[Scheduler] kdn_refresh_once failed when triggered by kdn_register") + + set_on_kdn_register(_trigger_refresh_once) + logger.info("[Scheduler] on_kdn_register callback installed") + + async def _run_control_plane(app): + host = os.environ.get("SCHEDULER_CP_HOST", "0.0.0.0") + port = int(os.environ.get("SCHEDULER_CP_PORT", SCHEDULER_CP_PORT)) + + uv_cfg = uvicorn.Config( + control_plane, + host=host, + port=port, + log_level=os.environ.get("SCHEDULER_CP_LOG_LEVEL", "info"), + loop="asyncio", + lifespan="on", + access_log=False, + ) + server = uvicorn.Server(uv_cfg) + app.state._cp_server = server + await server.serve() + + app.state._cp_task = asyncio.create_task(_run_control_plane(app)) # type: ignore + + # ------------------------------------------ + # 调用具体scheduler策略 + # ------------------------------------------ + strategy_name = os.environ.get("SCHEDULER_STRATEGY", "round_robin") + app.state.proxy_strategy = create_strategy(strategy_name) # type: ignore + logger.info("[Scheduler] scheduler strategy=%s", strategy_name) + + # 这里 yield 之后是正常服务期 + logger.info("[Scheduler] startup: 初始化完成，监听服务中") + try: + yield + + finally: + # ---------- shutdown ---------- + cp_server = getattr(app.state, "_cp_server", None) # type: ignore + cp_task = getattr(app.state, "_cp_task", None) # type: ignore + ev = getattr(app.state, "_kdn_stop_event", None) # type: ignore + task = getattr(app.state, "_kdn_refresh_task", None) # type: ignore + + if ev is not None: + ev.set() + if task is not None: + task.cancel() + + if cp_server is not None: + cp_server.should_exit = True + if cp_task is not None: + try: + await cp_task + except Exception: + pass + + # shutdown 阶段，如果后面有资源要清理，可以写在这里 + logger.info("[Scheduler] shutdown: 结束服务") + + +scheduler = FastAPI( + title="CacheRoute Scheduler", + version="0.1.1", + lifespan=lifespan, +) + + + + +# ======================= 公共内部处理函数 ======================= +@timing +def _handle_client( + app: FastAPI, + url_path: str, + payload: Dict[str, Any], + client_ip: str, + 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 请求信息 + 构造内部 Request 对象，分配request_id + - url_path: 例如 "/v1/chat/completions" + - payload: 解析后的 JSON 字典 + - client_ip: 客户端 IP 字符串 + """ + # 分配 request_id（1~65535 循环） + request_id = id_alloc.next_id() + + # 打印调试信息 + if os.environ.get("SCHEDULER_VERBOSE_REQUEST_LOG", config.SCHEDULER_VERBOSE_REQUEST_LOG) == 1: + print("+" * 80) + print(f"[Scheduler] 收到 HTTP 请求: path={url_path}, client_ip={client_ip},\n" + f"分配 request_id={request_id},\n" + f"payload={payload}") + print("+" * 80) + + # 直接复用你改好的 build_request + req_obj = SchedulerRequest.build_request( + url_path=url_path, + payload=payload, + user_addr=client_ip, + request_id=request_id, + embedder=getattr(app.state, "embedding_engine", None), # type: ignore + knowledge_table=getattr(app.state, "knowledge_table", None), # type: ignore + 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" + f"[Scheduler] Endpoint_type={getattr(req_obj.Service, 'Endpoint_type', None)},\n" + f"[Scheduler] Knowledge_List={req_obj.Service.Knowledge_List},\n" + f"[Scheduler] Selected KDN={req_obj.Task.KDN_server_addr}, ProxyID={req_obj.Task.P_proxy_id}[{req_obj.Task.P_proxy_addr}:{req_obj.Task.P_proxy_port}]" + ) + print(f"[Scheduler] Injection_type={getattr(req_obj.Service, 'Injection_type', None)}") + return req_obj + + +# ======================= 调度器方法路由 ======================= +@scheduler.post("/v1/chat/completions") +async def create_chat_completions(request: FastAPIRequest): + """ + 对应： + curl http://HOST:PORT/v1/chat/completions -H "Content-Type: application/json" -d '{...}' + payload 结构参考 OpenAI chat 接口： + { + "model": "...", + "messages": [ ... ], + "max_tokens": ..., + "temperature": ..., + "top_p": ..., + "stream": true/false, + ... + } + """ + # 解析用户原始请求体 + try: + payload = await request.json() + except Exception as e: + return JSONResponse( + status_code=400, + content={"error": "invalid_json", "detail": str(e)}, + ) + + # 取客户端 IP 和路径（用于构建内部 Request） + client_ip = request.client.host if request.client else "unknown" + url_path = request.url.path + + # 把用户请求头转成 dict[str, str] + raw_headers = {k.lower(): v for k, v in request.headers.items()} + + # 构建内部 Request（包含 Prompt/Service/Task 等） + pool = get_pool() + proxy_infos = await pool.list(include_dead=False) + proxies = [ + { + "proxy_id": p.proxy_id, + "host": p.host, + "port": p.port, + "inflight": p.load.inflight, + "qps_1m": p.load.qps_1m, + "gpu_util": p.load.gpu_util, + "meta": dict(p.meta or {}), + } + for p in proxy_infos + ] + + kdn_pool = get_kdn_pool() + kdn_infos = await kdn_pool.list(include_dead=False) + kdns = [ + { + "kdn_id": k.kdn_id, + "host": k.host, + "port": k.port, + "items": k.load.items, + "qps_1m": k.load.qps_1m, + "meta": dict(k.meta or {}), + } + for k in kdn_infos + ] + + strategy = getattr(request.app.state, "proxy_strategy", None) + 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") + + # 根据调度结果选择下游 URL，更新对应proxy_id的inflight + host = req_obj.Task.P_proxy_addr + port = req_obj.Task.P_proxy_port + endpoint = getattr(req_obj.Service, "Endpoint_type", "chat/completions") + downstream_url = f"http://{host}:{port}/v1/{endpoint}" + proxy_id = req_obj.Task.P_proxy_id + if not proxy_id: + raise RuntimeError("Routing failed: missing P_proxy_id for inflight tracking") + ok = await pool.inflight_delta(proxy_id, +1) + if not ok: + raise RuntimeError(f"Proxy not found in pool: {proxy_id}") + + # 转化payload + data_for_downstream = req_obj.to_payload() + + # 处理需要透传给下游的 headers（可选择性过滤） + extra_headers = {} + + # 把用户的 Authorization 原样传给下游（如果希望由 Proxy 做鉴权） + if "authorization" in raw_headers: + extra_headers["authorization"] = raw_headers["authorization"] + + # 带一个 Scheduler分配的Request ID，便于链路追踪 + extra_headers["scheduler-request-id"] = str(req_obj.Request_ID) + + # 定义一个 async 生成器，从下游流式读取，再回给上游 + async def iter_upstream(): + try: + # forward_request 本身是一个 async generator + async for chunk in forward_request( + url=downstream_url, + data=data_for_downstream, + use_chunked=True, # chat/completions 一般是流式 + extra_headers=extra_headers, # 透传必要头 + ): + # 这里 chunk 已经是 bytes，直接 yield 出去即可 + yield chunk + finally: + await pool.inflight_delta(proxy_id, -1) + + # 用 StreamingResponse 包一层，实现用户侧的流式响应 + return StreamingResponse(iter_upstream(), media_type="application/json") + + +@scheduler.post("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/v1/completions") +async def create_completions(request: FastAPIRequest): + """ + 对应： + curl http://HOST:PORT/v1/completions -H "Content-Type: application/json" -d '{...}' + payload 结构类似 OpenAI completions 接口： + { + "model": "...", + "prompt": "...." 或 ["..."], + "max_tokens": ..., + "temperature": ..., + "top_p": ..., + "stream": true/false, + ... + } + """ + try: + payload = await request.json() + except Exception as e: + return JSONResponse( + status_code=400, + content={"error": "invalid_json", "detail": str(e)}, + ) + + # 取客户端 IP 和路径（用于构建内部 Request） + client_ip = request.client.host if request.client else "unknown" + url_path = request.url.path + + # 把用户请求头转成 dict[str, str] + raw_headers = {k.lower(): v for k, v in request.headers.items()} + + # 构建内部 Request（包含 Prompt/Service/Task 等） + pool = get_pool() + proxy_infos = await pool.list(include_dead=False) + proxies = [ + { + "proxy_id": p.proxy_id, + "host": p.host, + "port": p.port, + "inflight": p.load.inflight, + "qps_1m": p.load.qps_1m, + "gpu_util": p.load.gpu_util, + "meta": dict(p.meta or {}), + } + for p in proxy_infos + ] + + kdn_pool = get_kdn_pool() + kdn_infos = await kdn_pool.list(include_dead=False) + kdns = [ + { + "kdn_id": k.kdn_id, + "host": k.host, + "port": k.port, + "items": k.load.items, + "qps_1m": k.load.qps_1m, + "meta": dict(k.meta or {}), + } + for k in kdn_infos + ] + + strategy = getattr(request.app.state, "proxy_strategy", None) + 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") + + # 根据调度结果选择下游 URL，更新对应proxy_id的inflight + host = req_obj.Task.P_proxy_addr + port = req_obj.Task.P_proxy_port + endpoint = getattr(req_obj.Service, "Endpoint_type", "completions") + downstream_url = f"http://{host}:{port}/v1/{endpoint}" + proxy_id = req_obj.Task.P_proxy_id + if not proxy_id: + raise RuntimeError("Routing failed: missing P_proxy_id for inflight tracking") + ok = await pool.inflight_delta(proxy_id, +1) + if not ok: + raise RuntimeError(f"Proxy not found in pool: {proxy_id}") + + # 转化payload + data_for_downstream = req_obj.to_payload() + + # 处理需要透传给下游的 headers（可选择性过滤） + extra_headers = {} + + # 把用户的 Authorization 原样传给下游（如果希望由 Proxy 做鉴权） + if "authorization" in raw_headers: + extra_headers["authorization"] = raw_headers["authorization"] + + # 带一个 Scheduler分配的Request ID，便于链路追踪 + extra_headers["scheduler-request-id"] = str(req_obj.Request_ID) + + # 定义一个 async 生成器，从下游流式读取，再回给上游 + async def iter_upstream(): + try: + # forward_request 本身是一个 async generator + async for content in forward_request( + url=downstream_url, + data=data_for_downstream, + use_chunked=False, # chat/completions 一般是流式 + extra_headers=extra_headers, # 透传必要头 + ): + # 这里 chunk 已经是 bytes，直接 yield 出去即可 + yield content + finally: + await pool.inflight_delta(proxy_id, -1) + + # 用 StreamingResponse 包一层，实现用户侧的流式响应 + return StreamingResponse(iter_upstream(), media_type="application/json") + + +@scheduler.get("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/debug/status") +async def debug_status() -> Dict[str, Any]: + """ + Debug endpoint for Scheduler CLI. + + Returns: + - knowledge_loaded: whether knowledge_table is ready + - entries: number of knowledge units + - dim: embedding dim + - faiss_total: index.ntotal if built (else None) + - kdn_base_url: where snapshot comes from (if any) + - last_refresh_ts: unix seconds (if you maintain it), else None + - unit_fields: what fields a KnowledgeUnit contains (for inspection) + - sample_kids: first few kids (for quick view) + """ + table = getattr(scheduler.state, "knowledge_table", None) # type: ignore + kdn_base_url = getattr(scheduler.state, "kdn_base_url", None) # type: ignore + last_refresh_ts = getattr(scheduler.state, "last_refresh_ts", None) # type: ignore + kdn_alive = int(getattr(scheduler.state, "kdn_alive", 0) or 0) # type: ignore + kdn_alive_addrs = list(getattr(scheduler.state, "kdn_alive_addrs", []) or []) # type: ignore + kdn_last_selected = getattr(scheduler.state, "kdn_last_selected", None) # type: ignore + kdn_last_selected_id = getattr(scheduler.state, "kdn_last_selected_id", None) # type: ignore + kdn_last_refresh_ok = bool(getattr(scheduler.state, "kdn_last_refresh_ok", False)) # type: ignore + kdn_last_refresh_reason = getattr(scheduler.state, "kdn_last_refresh_reason", "") # type: ignore + kdn_last_refresh_ts = int(getattr(scheduler.state, "kdn_last_refresh_ts", 0) or 0) # type: ignore + + pool = get_pool() + proxy_infos = await pool.list(include_dead=False) + + proxy_states = [] + for p in proxy_infos: + proxy_states.append({ + "proxy_id": p.proxy_id, + "host": p.host, + "port": p.port, + "max_capacity": p.load.max_capacity, + "instance_count": p.load.instance_count, + "kv_mem_per_instance_gb": p.load.kv_mem_per_instance_gb, + "kv_cache_pool_gb": p.load.kv_cache_pool_gb, + "kv_cache_update_policy": getattr(p, "kv_cache_update_policy", "static"), + "inflight": p.load.inflight, + "qps_1m": p.load.qps_1m, + "gpu_util": p.load.gpu_util, + }) + + if table is None: + + return { + "knowledge_loaded": False, + "entries": 0, + "dim": None, + "faiss_total": None, + "kdn_base_url": kdn_base_url, + "last_refresh_ts": last_refresh_ts, + "unit_fields": [], + "sample_kids": [], + "kdn_alive": 0, + "kdn_alive_addrs": [], + "kdn_last_selected": None, + "kdn_last_selected_id": None, + "kdn_last_refresh_ok": False, + "kdn_last_refresh_reason": "", + "kdn_last_refresh_ts": 0, + "proxies": proxy_states, + } + + units = getattr(table, "_units", {}) + kids = list(units.keys()) + kids_sorted = sorted(kids)[:10] + + # 尽量从 table 拿到 dim + dim = getattr(table, "dim", None) + + # FAISS + faiss_total = None + idx = getattr(table, "_faiss_index", None) + if idx is not None: + try: + faiss_total = int(getattr(idx, "ntotal")) + except Exception: + faiss_total = None + + # KnowledgeUnit 字段探测（帮助你看 snapshot 后到底存了什么） + unit_fields: List[str] = [] + if kids_sorted: + u0 = units[kids_sorted[0]] + # dataclass 一般有 __dict__ + if hasattr(u0, "__dict__"): + unit_fields = sorted(list(u0.__dict__.keys())) + else: + # fallback：dir 过滤 + unit_fields = [x for x in dir(u0) if not x.startswith("_")] + + return { + "knowledge_loaded": True, + "entries": len(units), + "dim": dim, + "faiss_total": faiss_total, + "kdn_base_url": kdn_base_url, + "last_refresh_ts": last_refresh_ts, + "unit_fields": unit_fields, + "sample_kids": kids_sorted, + "kdn_alive": kdn_alive, + "kdn_alive_addrs": kdn_alive_addrs, + "kdn_last_selected": kdn_last_selected, + "kdn_last_selected_id": kdn_last_selected_id, + "kdn_last_refresh_ok": kdn_last_refresh_ok, + "kdn_last_refresh_reason": kdn_last_refresh_reason, + "kdn_last_refresh_ts": kdn_last_refresh_ts, + "proxies": proxy_states, + } + + +@scheduler.post("/debug/knowledge/peek") +async def debug_peek_knowledge(payload: PeekPayload) -> Dict[str, Any]: + """ + Peek knowledge units by kids. + Default fields are safe (no full embedding). + """ + table = getattr(scheduler.state, "knowledge_table", None) # type: ignore + if table is None: + return {"items": [], "miss": payload.kids} + + units = getattr(table, "_units", {}) + items = [] + miss = [] + + for kid in payload.kids: + k = (kid or "").strip().lower() + u = units.get(k) + if u is None: + miss.append(kid) + continue + + it = {"kid": k} + allow = set(payload.need_fields or []) + + if "length" in allow: + it["length"] = getattr(u, "length", None) + if "avail_kdn_servers" in allow: + it["avail_kdn_servers"] = getattr(u, "avail_kdn_servers", None) + if "text_abstract" in allow: + it["text_abstract"] = getattr(u, "text_abstract", None) + + # 可选：你若在 KnowledgeUnit/metadata 里存了 kv_ready 等，也可以在这里输出 + if "meta" in allow and hasattr(u, "meta"): + it["meta"] = getattr(u, "meta") + if "kv_ready" in allow: + it["kv_ready"] = getattr(u, "kv_ready", 0) + if "kv_rel_dir" in allow: + it["kv_rel_dir"] = getattr(u, "kv_rel_dir", None) + if "kv_dumped_keys" in allow: + it["kv_dumped_keys"] = getattr(u, "kv_dumped_keys", None) + if "kv_updated_at" in allow: + it["kv_updated_at"] = getattr(u, "kv_updated_at", None) + + items.append(it) + + return {"items": items, "miss": miss} + + +@scheduler.post("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/admin/refresh_knowledge") +async def admin_refresh_knowledge(): + try: + r = await kdn_refresh_once(scheduler) + return JSONResponse(content=r) + except Exception as e: + logger.exception(f"[Scheduler] admin refresh failed: {e}") + return JSONResponse(status_code=500, content={"ok": False, "error": str(e)}) + + +@scheduler.get("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/debug/strategy") +async def debug_strategy() -> Dict[str, Any]: + strategy = getattr(scheduler.state, "proxy_strategy", None) # type: ignore + pool = get_pool() + proxy_infos = await pool.list(include_dead=False) + + # 只返回简要信息，避免输出过大 + sample = [{ + "proxy_id": p.proxy_id, + "host": p.host, + "port": p.port, + "is_alive": True, # list(include_dead=False) 已保证 alive + } for p in proxy_infos[:10]] + + 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") +# async def knowledge_update(request: FastAPIRequest): +# payload = await request.json() +# client_ip = request.client.host if request.client else "unknown" +# ... +# return JSONResponse(content={"status": "ok"}) diff --git a/scheduler/strategy/base.py b/scheduler/strategy/base.py index 0227558..9ddedc1 100644 --- a/scheduler/strategy/base.py +++ b/scheduler/strategy/base.py @@ -1,26 +1,30 @@ -# scheduler/strategy/base.py -from __future__ import annotations - -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple - - -class ProxySelectionStrategy(ABC): - """ - 统一路由策略接口：给一组候选 KDNs + Proxies + 当前 request， - 返回 (chosen_kdn, chosen_proxy)。 - 后续扩展：可以在这里加更多上下文（如历史统计、请求类别、KV信息等） - """ - - name: str = "base" - - @abstractmethod - def select( - self, - kdns: List[Dict[str, Any]], - proxies: List[Dict[str, Any]], - payload: Dict[str, Any], - url_path: str, - user_addr: str, - ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: - raise NotImplementedError +# scheduler/strategy/base.py +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple + + +class ProxySelectionStrategy(ABC): + """ + 统一路由策略接口：给一组候选 KDNs + Proxies + 当前 request， + 返回 (chosen_kdn, chosen_proxy)。 + + request_ctx 采用轻量上下文形式，用来给策略传递已经在 scheduler + 内部计算完成的信息（例如 Knowledge_List、知识长度、每个 KDN 的知识元数据索引等）， + 避免策略层重复做 embedding 检索。 + """ + + name: str = "base" + + @abstractmethod + def select( + self, + kdns: List[Dict[str, Any]], + proxies: List[Dict[str, Any]], + payload: Dict[str, Any], + url_path: str, + user_addr: str, + request_ctx: Optional[Dict[str, Any]] = None, + ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + raise NotImplementedError diff --git a/scheduler/strategy/cacheroute.py b/scheduler/strategy/cacheroute.py new file mode 100644 index 0000000..d51082f --- /dev/null +++ b/scheduler/strategy/cacheroute.py @@ -0,0 +1,255 @@ +# scheduler/strategy/cacheroute.py +from __future__ import annotations + +import os +import threading +from typing import Any, Dict, List, Optional, Tuple + +from .base import ProxySelectionStrategy + + +class CacheRouteStrategy(ProxySelectionStrategy): + """ + CacheRoute 调度策略的第一阶段实现。 + + 当前版本只做最小可用能力： + 1) 基于 request 的 Knowledge_List，在每个 KDN 的知识元数据索引里判断 + - 是否能完整提供请求所需知识文本； + - 该 KDN 上 kv_ready 覆盖的总长度。 + 2) KDN 选择采用“词典序规则”，避免引入加权打分： + - 优先完整文本覆盖； + - 优先非过载 KDN； + - 再比 kv_ready 覆盖长度； + - 最后比轻负载与稳定顺序。 + 3) Proxy 选择先保持保守：优先当前 inflight 最小者，避免一次性把 + 拓扑/历史偏好也同时引入。后续版本再把 KDN 锚点拓扑和知识亲和性接进来。 + + 备注： + - 这里不依赖 Injection_type 做模式分支；正式策略默认先满足文本可服务性， + 再在可服务候选中偏向 KVCache 覆盖更高的 KDN。 + - request_ctx / kdn_knowledge_index 由 scheduler 内部准备。 + """ + + name: str = "cacheroute" + + 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: + return f"kdn://{kdn['host']}:{int(kdn['port'])}" + + @staticmethod + def _proxy_inflight(proxy: Dict[str, Any]) -> int: + return int(proxy.get('inflight', 0) or 0) + + @staticmethod + def _proxy_qps(proxy: Dict[str, Any]) -> float: + return float(proxy.get('qps_1m', 0.0) or 0.0) + + @staticmethod + def _proxy_gpu(proxy: Dict[str, Any]) -> float: + return float(proxy.get('gpu_util', 0.0) or 0.0) + + @staticmethod + def _kdn_items(kdn: Dict[str, Any]) -> int: + return int(kdn.get('items', 0) or 0) + + @staticmethod + def _kdn_qps(kdn: Dict[str, Any]) -> float: + return float(kdn.get('qps_1m', 0.0) or 0.0) + + @staticmethod + def _is_overloaded(kdn: Dict[str, Any]) -> bool: + """向后兼容：保留 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]], + knowledge_list: List[str], + knowledge_index: Dict[str, Dict[str, Dict[str, Any]]], + ) -> Optional[Dict[str, Any]]: + if not kdns: + return None + + # 没有知识需求时，退化为最轻负载 KDN，避免复杂逻辑干扰原有流程。 + if not knowledge_list: + ranked = sorted( + kdns, + key=lambda k: ( + self._is_overloaded(k), + self._kdn_qps(k), + self._kdn_items(k), + str(k.get('kdn_id', '')), + ), + ) + 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) + + # 支持两种索引键：kdn_id 优先，其次 kdn://host:port，方便后续平滑过渡。 + meta_by_kid = knowledge_index.get(kdn_id) or knowledge_index.get(kdn_addr) or {} + + text_full = True + kv_cover_len = 0 + missing_count = 0 + for kid in knowledge_list: + item = meta_by_kid.get(kid) + if item is None: + text_full = False + missing_count += 1 + continue + if int(item.get('kv_ready', 0) or 0) == 1: + kv_cover_len += int(item.get('length', 0) or 0) + + # 词典序： + # 1. 完整文本覆盖优先 + # 2. 非过载优先 + # 3. KVCache 覆盖长度越大越优 + # 4. 缺失知识数越少越优（作为非完整覆盖场景的回退） + # 5. QPS / items 越小越优 + # 6. kdn_id 稳定打破并列 + key = ( + 0 if text_full else 1, + 0 if not self._is_overloaded_by_threshold(kdn) else 1, + -kv_cover_len, + missing_count, + self._kdn_qps(kdn), + self._kdn_items(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]) + 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: + return None + + # 第一阶段保持简单： + # - 优先 inflight 最小 + # - 再比 qps_1m / gpu_util + # - 最后使用本地 cursor 做稳定轮换，避免完全固定在一个 proxy 上 + ranked = sorted( + proxies, + key=lambda p: ( + self._proxy_inflight(p), + self._proxy_qps(p), + self._proxy_gpu(p), + str(p.get('proxy_id', '')), + ), + ) + 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 + + with self._lock: + idx = self._proxy_cursor % len(tied) + self._proxy_cursor += 1 + 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, + kdns: List[Dict[str, Any]], + proxies: List[Dict[str, Any]], + payload: Dict[str, Any], + url_path: str, + user_addr: str, + request_ctx: Optional[Dict[str, Any]] = None, + ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + 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 {} + + chosen_kdn = self._select_kdn(kdns=kdns, knowledge_list=knowledge_list, knowledge_index=knowledge_index) + chosen_proxy = self._select_proxy(proxies=proxies) + return chosen_kdn, chosen_proxy diff --git a/scheduler/strategy/factory.py b/scheduler/strategy/factory.py index 6ea01c5..5d4ebe5 100644 --- a/scheduler/strategy/factory.py +++ b/scheduler/strategy/factory.py @@ -1,20 +1,23 @@ -# scheduler/strategy/factory.py -from __future__ import annotations - -from typing import Dict - -from .base import ProxySelectionStrategy -from .round_robin import RoundRobinStrategy - - -_STRATEGIES: Dict[str, type[ProxySelectionStrategy]] = { - "round_robin": RoundRobinStrategy, -} - -def create_strategy(name: str) -> ProxySelectionStrategy: - key = (name or "").strip().lower() - if not key: - key = "round_robin" - if key not in _STRATEGIES: - raise ValueError(f"Unknown strategy: {name!r}. Available={list(_STRATEGIES.keys())}") - return _STRATEGIES[key]() +# scheduler/strategy/factory.py +from __future__ import annotations + +from typing import Dict + +from .base import ProxySelectionStrategy +from .round_robin import RoundRobinStrategy +from .cacheroute import CacheRouteStrategy + + +_STRATEGIES: Dict[str, type[ProxySelectionStrategy]] = { + "round_robin": RoundRobinStrategy, + "cacheroute": CacheRouteStrategy, +} + + +def create_strategy(name: str) -> ProxySelectionStrategy: + key = (name or "").strip().lower() + if not key: + key = "round_robin" + if key not in _STRATEGIES: + raise ValueError(f"Unknown strategy: {name!r}. Available={list(_STRATEGIES.keys())}") + return _STRATEGIES[key]() diff --git a/scheduler/strategy/round_robin.py b/scheduler/strategy/round_robin.py index 694cffc..40b5ae4 100644 --- a/scheduler/strategy/round_robin.py +++ b/scheduler/strategy/round_robin.py @@ -1,47 +1,48 @@ -# scheduler/strategy/round_robin.py -from __future__ import annotations - -import threading - -from typing import Any, Dict, List, Optional, Tuple -from .base import ProxySelectionStrategy - - - -class RoundRobinStrategy(ProxySelectionStrategy): - """ - 最简单的 round-robin： - - 在“当前存活 proxy 列表”上做循环取模 - - 用 threading.Lock 保护 index，避免并发请求打乱顺序 - """ - - name: str = "round_robin" - - def __init__(self): - self._lock = threading.Lock() - self._kdn_cursor = 0 - self._proxy_cursor = 0 - - def select( - self, - kdns: List[Dict[str, Any]], - proxies: List[Dict[str, Any]], - payload: Dict[str, Any], - url_path: str, - user_addr: str, - ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: - chosen_kdn = None - chosen_proxy = None - - with self._lock: - if kdns: - ki = self._kdn_cursor % len(kdns) - self._kdn_cursor += 1 - chosen_kdn = kdns[ki] - - if proxies: - pi = self._proxy_cursor % len(proxies) - self._proxy_cursor += 1 - chosen_proxy = proxies[pi] - - return chosen_kdn, chosen_proxy +# scheduler/strategy/round_robin.py +from __future__ import annotations + +import threading + +from typing import Any, Dict, List, Optional, Tuple +from .base import ProxySelectionStrategy + + +class RoundRobinStrategy(ProxySelectionStrategy): + """ + 最简单的 round-robin： + - 在“当前存活 proxy 列表”上做循环取模 + - 用 threading.Lock 保护 index，避免并发请求打乱顺序 + - request_ctx 对该策略无用，但保留统一接口，方便 scheduler 无差别调用 + """ + + name: str = "round_robin" + + def __init__(self): + self._lock = threading.Lock() + self._kdn_cursor = 0 + self._proxy_cursor = 0 + + def select( + self, + kdns: List[Dict[str, Any]], + proxies: List[Dict[str, Any]], + payload: Dict[str, Any], + url_path: str, + user_addr: str, + request_ctx: Optional[Dict[str, Any]] = None, + ) -> Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + chosen_kdn = None + chosen_proxy = None + + with self._lock: + if kdns: + ki = self._kdn_cursor % len(kdns) + self._kdn_cursor += 1 + chosen_kdn = kdns[ki] + + if proxies: + pi = self._proxy_cursor % len(proxies) + self._proxy_cursor += 1 + chosen_proxy = proxies[pi] + + return chosen_kdn, chosen_proxy diff --git a/test/demo_scheduler.py b/test/demo_scheduler.py index baa7c7c..3fe8b4f 100644 --- a/test/demo_scheduler.py +++ b/test/demo_scheduler.py @@ -1,53 +1,59 @@ -# demo_scheduler.py -""" -Scheduler_v1启动demo： - - 引用 scheduler.py 里的 api - - 使用 uvicorn 启动 HTTP 服务 -""" -import argparse -import os, logging -import uvicorn - -from scheduler import scheduler -from core import config - -# 在 demo 里配置要预热的模型路径以及知识库路径 -MODEL_PATH = config.DEFAULT_MODEL -KNOWLEDGE_YAML_PATH = config.KNOWLEDGE_YAML_PATH -EMBEDDING_MODEL = config.EMBEDDING_MODEL -KDN_BASE_URL = config.KDN_BASE_URL -dp_port = config.SCHEDULER_DP_PORT -dp_host = config.SCHEDULER_DP_HOST - -def main(): - # logging配置 - logging.basicConfig( - level=logging.INFO, # 根 logger 级别设为 INFO - format=" [%(levelname)s] %(name)s: %(message)s", - ) - - parser = argparse.ArgumentParser() - parser.add_argument("--strategy", default="round_robin", help="proxy scheduling strategy") - args = parser.parse_args() - - # 把模型路径暴露给 scheduler（scheduler.py 里通过 os.getenv 读取） - os.environ["SCHEDULER_MODEL_PATH"] = MODEL_PATH - os.environ["SCHEDULER_TOKENIZER_MAP"]='{"llama3-70b":"/workspace/llm-stack/models/LLM-Research/Meta-Llama-3-70B-Instruct"}' - os.environ["SCHEDULER_KNOWLEDGE_YAML"] = str(KNOWLEDGE_YAML_PATH) - os.environ["SCHEDULER_KDN_BASE_URL"] = str(KDN_BASE_URL).rstrip("/") - os.environ["SCHEDULER_EMBEDDING_MODEL"] = EMBEDDING_MODEL - os.environ["HF_HUB_OFFLINE"] = "1" - os.environ["TRANSFORMERS_OFFLINE"] = "1" - os.environ["SCHEDULER_STRATEGY"] = args.strategy - - - # 配置 uvicorn.Server - config = uvicorn.Config(scheduler, host=dp_host, port=dp_port, reload=False) - server = uvicorn.Server(config) - server.run() - print("[DEMO] Scheduler stopped, demo exit.") - - -if __name__ == "__main__": - main() - +# demo_scheduler.py +""" +Scheduler_v1启动demo： + - 引用 scheduler.py 里的 api + - 使用 uvicorn 启动 HTTP 服务 +""" +import argparse +import os, logging +import uvicorn + +from scheduler import scheduler +from core import config + +# 在 demo 里配置要预热的模型路径以及知识库路径 +MODEL_PATH = config.DEFAULT_MODEL +KNOWLEDGE_YAML_PATH = config.KNOWLEDGE_YAML_PATH +EMBEDDING_MODEL = config.EMBEDDING_MODEL +KDN_BASE_URL = config.KDN_BASE_URL +dp_port = config.SCHEDULER_DP_PORT +dp_host = config.SCHEDULER_DP_HOST + + +def main(): + # logging配置 + logging.basicConfig( + level=logging.INFO, # 根 logger 级别设为 INFO + format=" [%(levelname)s] %(name)s: %(message)s", + ) + + parser = argparse.ArgumentParser() + parser.add_argument("--strategy", default="round_robin", help="proxy scheduling strategy") + parser.add_argument( + "--cacheroute", + action="store_true", + help="shortcut for --strategy cacheroute", + ) + args = parser.parse_args() + + strategy_name = "cacheroute" if args.cacheroute else args.strategy + + # 把模型路径暴露给 scheduler（scheduler.py 里通过 os.getenv 读取） + os.environ["SCHEDULER_MODEL_PATH"] = MODEL_PATH + os.environ["SCHEDULER_TOKENIZER_MAP"]='{"llama3-70b":"/workspace/llm-stack/models/LLM-Research/Meta-Llama-3-70B-Instruct"}' + os.environ["SCHEDULER_KNOWLEDGE_YAML"] = str(KNOWLEDGE_YAML_PATH) + os.environ["SCHEDULER_KDN_BASE_URL"] = str(KDN_BASE_URL).rstrip("/") + os.environ["SCHEDULER_EMBEDDING_MODEL"] = EMBEDDING_MODEL + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + os.environ["SCHEDULER_STRATEGY"] = strategy_name + + # 配置 uvicorn.Server + config = uvicorn.Config(scheduler, host=dp_host, port=dp_port, reload=False) + server = uvicorn.Server(config) + server.run() + print("[DEMO] Scheduler stopped, demo exit.") + + +if __name__ == "__main__": + main()