llm-engine-kitty 0.1.0.dev0__py3-none-any.whl

llm_engine/kitty/server.py

@@ -0,0 +1,408 @@
+# llm_engine/kitty/server.py
+
+"""
+KittyServer：在子进程中 driving 事件循环 + 接受 socket 连接的服务端。
+
+- 一个进程内只有一个 _KittyEngine 实例被多个连接共享
+- 每个连接维护自己的 _ConnectionContext，包含其携带的 overrides 以及它提交的 task_id 集合
+- 帧级别请求→响应模型，submit 立即返回 snapshot，wait 在服务端阻塞 done_event
+- 连接断开时：持有的未完成任务被取消，未持久化的已完成结果被回收
+"""
+
+import asyncio
+import signal
+import socket
+import sys
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+import kitty_logger
+
+from ..schemas import InferenceRequest, TaskStatus
+from ..utils import gen_unique_id
+from .config import KittyEngineConfig, KittyEngineOverrides
+from .protocol import MsgType, aread_frame, encode_frame, parse_endpoint
+from .engine import _KittyEngine
+
+logger = kitty_logger.getLogger(__name__)
+
+
+@dataclass
+class _ConnectionContext:
+    connection_id: str
+    overrides: Optional[KittyEngineOverrides] = None
+    task_ids: set[str] = field(default_factory=set)
+
+
+class KittyServer:
+    def __init__(self, config: KittyEngineConfig) -> None:
+        self.config = config
+        self.engine = _KittyEngine(config)
+        self._server: Optional[asyncio.AbstractServer] = None
+        self._stop_event: Optional[asyncio.Event] = None
+        self._endpoint_kind: Optional[str] = None
+        self._endpoint_addr: Optional[tuple] = None
+        # 任务 → 所属连接 的反向映射。维护时机:
+        #   - CMD_SUBMIT 成功返回前登记（必须早于 engine.submit 首次触发
+        #     callback，否则初始 SUBMITTED 事件会找不到连接而被丢弃）。
+        #   - CMD_POP / CMD_WAIT 真正消费掉结果时 pop。
+        #   - CMD_CLEAR_DONE 按 engine.clear_done() 返回的 id 列表批量 pop。
+        #   - 连接关闭时由 _handle_client 的 finally 清理该连接剩余的全部映射。
+        self._task_to_connection: dict[str, str] = {}
+        # 每连接的状态计数器（累计口径，不随 POP/WAIT 消费而减）。
+        # key 用 TaskStatus 枚举本体，回包时再转成 name 字符串给 client 侧拿。
+        self._connection_statistics: dict[str, dict[TaskStatus, int]] = {}
+        self.engine.set_transition_callback(self._on_task_transition)
+
+    # ------------------------------------------------------------------
+    # 状态统计（callback）
+    # ------------------------------------------------------------------
+
+    def _on_task_transition(
+        self,
+        task_id: str,
+        old: Optional[TaskStatus],
+        new: TaskStatus,
+    ) -> None:
+        """worker → server 的状态回调。维护 per-connection 计数器。
+
+        语义:
+            - old is None：首次落地（worker.submit 触发），只对 new 做 +1。
+            - 正常转移：old -1、new +1，保持"当前处于各状态的任务数"口径。
+              终态由于不会再转出，对应计数只增不减。
+
+        健壮性:
+            - 找不到 connection_id 或 statistics 条目时安静返回。正常路径不会
+              走到（CMD_SUBMIT 保证映射早于 worker.submit 调用），唯一命中
+              场景是"连接已关闭、finally 已经把映射/statistics 清掉后，
+              cleanup_tasks 里取消任务产生的最后一轮 transition"——这是预期
+              的 no-op 路径，因此用 debug 级而非 warning，避免连接断开时
+              刷屏（每个未完成任务都会触发一次）。
+            - old 的 -1 做 max(0, ...) 下界保护：正常路径下初始 +1 保证非负，
+              但 callback 是"绝不能抛异常也绝不能把状态搞歪"的合约，兜底更稳。
+            - 本方法被 `_WorkerTask.transition` 同步调用，必须尽快返回。
+        """
+        connection_id: str | None = self._task_to_connection.get(task_id)
+        if connection_id is None:
+            logger.debug("KittyServer: can not find conn_id for task %s (likely cleanup path)", task_id)
+            return
+        statistics = self._connection_statistics.get(connection_id)
+        if statistics is None:
+            logger.debug("KittyServer: can not find statistics for connection %s (likely cleanup path)", connection_id)
+            return
+        if old is not None:
+            statistics[old] = max(0, statistics.get(old, 0) - 1)
+        statistics[new] = statistics.get(new, 0) + 1
+
+    # ------------------------------------------------------------------
+    # 入口
+    # ------------------------------------------------------------------
+
+    async def run(self, mode: str = "standalone", auto_port: bool = False) -> None:
+        """启动 server；阻塞直到收到 stop 信号。
+
+        mode: "standalone"（命令行启动）或 "embedded"（子进程启动，stdin EOF 时自动退出）。
+        auto_port: 仅 TCP 模式有效。True 时端口被占用会自动寻找可用端口；False 时直接 raise。
+        """
+        self._stop_event = asyncio.Event()
+        await self.engine.setup()
+
+        kind, addr = parse_endpoint(self.config.listen)
+        self._endpoint_kind = kind
+        self._endpoint_addr = addr
+
+        if kind == "unix":
+            import os
+
+            path = addr[0]
+            try:
+                os.unlink(path)
+            except FileNotFoundError:
+                pass
+            self._server = await asyncio.start_unix_server(self._handle_client, path=path)
+            os.chmod(path, 0o600)
+            logger.info("KittyServer listening on unix://%s", path)
+        else:
+            host, port = addr
+            if auto_port:
+                while True:
+                    try:
+                        self._server = await asyncio.start_server(self._handle_client, host=host, port=port)
+                        break
+                    except OSError:
+                        # 端口被占用，让 OS 分配一个随机可用端口
+                        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                            s.bind((host, 0))
+                            port = s.getsockname()[1]
+                # 更新 config.listen 以便外部可查询实际端口
+                self.config = self.config.model_copy(update={"listen": f"tcp://{host}:{port}"})
+            else:
+                self._server = await asyncio.start_server(self._handle_client, host=host, port=port)
+            logger.info("KittyServer listening on tcp://%s:%d", host, port)
+
+        # 注册信号（仅主线程；子进程 spawn 启动时主线程即这里）
+        loop = asyncio.get_running_loop()
+        for sig in (signal.SIGINT, signal.SIGTERM):
+            try:
+                loop.add_signal_handler(sig, self._stop_event.set)
+            except NotImplementedError:
+                # Windows 等平台不支持
+                pass
+
+        async with self._server:
+            stop_task = asyncio.create_task(self._stop_event.wait())
+            serve_task = asyncio.create_task(self._server.serve_forever())
+            tasks = {stop_task, serve_task}
+            if mode == "embedded":
+                tasks.add(asyncio.create_task(self._watch_stdin()))
+            done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+            for t in pending:
+                t.cancel()
+            # 清理异常
+            for t in done:
+                exc = t.exception()
+                if exc is not None and not isinstance(exc, asyncio.CancelledError):
+                    logger.error("server task 异常: %r", exc)
+
+        await self.engine.teardown()
+
+        # 清理 unix sock 文件
+        if kind == "unix":
+            import os
+
+            try:
+                os.unlink(addr[0])
+            except FileNotFoundError:
+                pass
+        logger.info("KittyServer 已停止")
+
+    def stop(self) -> None:
+        """由信号/SHUTDOWN 消息触发。"""
+        if self._stop_event is not None:
+            self._stop_event.set()
+
+    async def _watch_stdin(self) -> None:
+        """embedded 模式：监听 stdin EOF，父进程退出时自动停止 server。"""
+        loop = asyncio.get_running_loop()
+        reader = asyncio.StreamReader()
+        protocol = asyncio.StreamReaderProtocol(reader)
+        # 使用 sys.stdin.buffer（二进制模式）：避免文本层缓冲导致 asyncio transport
+        # 看不到底层 fd 中已到达的字节，且与 __main__.py 中 readline 消费第一行后
+        # 底层 fd 位置保持一致。
+        await loop.connect_read_pipe(lambda: protocol, sys.stdin.buffer)
+        try:
+            while True:
+                data = await reader.read(4096)
+                if not data:  # EOF
+                    break
+        except Exception:
+            pass
+        logger.info("KittyServer: stdin EOF，embedded 模式自动退出")
+        self.stop()
+
+    # ------------------------------------------------------------------
+    # 连接处理
+    # ------------------------------------------------------------------
+
+    async def _handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        ctx = _ConnectionContext(connection_id=gen_unique_id(prefix="conn"))
+        # 初始化该连接的状态计数器。所有 TaskStatus 都显式置 0，方便 client 侧
+        # 直接按 enum 取值渲染进度条。
+        self._connection_statistics[ctx.connection_id] = {s: 0 for s in TaskStatus}
+        peer = writer.get_extra_info("peername") or writer.get_extra_info("sockname")
+        logger.info("连接建立 connection_id=%s peer=%s", ctx.connection_id, peer)
+
+        try:
+            while True:
+                try:
+                    msg_type, payload = await aread_frame(reader)
+                except (asyncio.IncompleteReadError, ConnectionError):
+                    break
+                except Exception as e:
+                    logger.warning("connection_id=%s 读帧异常: %r", ctx.connection_id, e)
+                    break
+
+                try:
+                    should_close = await self._dispatch(ctx, msg_type, payload, writer)
+                except Exception as e:
+                    logger.exception("connection_id=%s 处理 msg_type=%d 异常", ctx.connection_id, msg_type)
+                    await self._send(writer, MsgType.MSG_ERROR, {"error": f"{type(e).__name__}: {e}"})
+                    should_close = False
+
+                if should_close:
+                    break
+        finally:
+            # 先断开 task→connection 映射 + 移除 statistics，callback 对遗留任务
+            # 变为 no-op，避免 cleanup_tasks 期间的 transition 事件仍去维护
+            # 一份马上要丢弃的计数。
+            leftover = list(ctx.task_ids)
+            for task_id in leftover:
+                self._task_to_connection.pop(task_id, None)
+            self._connection_statistics.pop(ctx.connection_id, None)
+            # 清理该连接持有的任务：取消 + 等待进入终态 + 非 persist 移除
+            if leftover:
+                try:
+                    await self.engine.cleanup_tasks(leftover)
+                except Exception:
+                    logger.exception("connection_id=%s cleanup_tasks 异常", ctx.connection_id)
+            try:
+                writer.close()
+                await writer.wait_closed()
+            except Exception:
+                pass
+            logger.info("连接关闭 connection_id=%s", ctx.connection_id)
+
+    async def _send(self, writer: asyncio.StreamWriter, msg_type: int, payload: dict[str, Any]) -> None:
+        writer.write(encode_frame(msg_type, payload))
+        await writer.drain()
+
+    # ------------------------------------------------------------------
+    # 消息分发
+    # ------------------------------------------------------------------
+
+    async def _dispatch(
+        self,
+        ctx: _ConnectionContext,
+        msg_type: int,
+        payload: dict[str, Any],
+        writer: asyncio.StreamWriter,
+    ) -> bool:
+        """处理一帧消息，返回是否应关闭连接。"""
+
+        if msg_type == MsgType.CMD_HELLO:
+            ov_dict = payload.get("overrides")
+            if ov_dict:
+                ctx.overrides = KittyEngineOverrides.model_validate(ov_dict)
+            await self._send(writer, MsgType.MSG_WELCOME, {"conn_id": ctx.connection_id})  # 协议字段名，勿改
+            return False
+
+        if msg_type == MsgType.CMD_PING:
+            await self._send(writer, MsgType.MSG_PONG, {})
+            return False
+
+        if msg_type == MsgType.CMD_SUBMIT:
+            task_id = payload.get("task_id") or gen_unique_id(prefix="task")
+            request = InferenceRequest.model_validate(payload["request"])
+            persist = bool(payload.get("persist", False))
+            timeout = payload.get("timeout")
+            # 必须先登记 task→connection 映射，再调 engine.submit：
+            # engine.submit 内部会同步触发首次 transition callback，届时回调
+            # 要能找到本任务所属连接。顺序反了首次 SUBMITTED 事件会被丢。
+            self._task_to_connection[task_id] = ctx.connection_id
+            try:
+                snap = self.engine.submit(
+                    task_id=task_id,
+                    request=request,
+                    overrides=ctx.overrides,
+                    persist=persist,
+                    timeout=timeout,
+                )
+            except Exception:
+                # 登记后若 submit 自己抛了，回滚避免 _task_to_connection 泄漏
+                self._task_to_connection.pop(task_id, None)
+                raise
+            ctx.task_ids.add(task_id)
+            await self._send(writer, MsgType.MSG_SNAPSHOT, {"snapshot": snap.model_dump()})
+            return False
+
+        if msg_type == MsgType.CMD_CANCEL:
+            task_id = payload["task_id"]
+            force = bool(payload.get("force", False))
+            cancelled = self.engine.cancel(task_id, force=force)
+            await self._send(writer, MsgType.MSG_PONG, {"cancelled": cancelled})
+            return False
+
+        if msg_type == MsgType.CMD_POLL:
+            task_id = payload["task_id"]
+            snap = self.engine.snapshot(task_id)
+            await self._send(
+                writer,
+                MsgType.MSG_SNAPSHOT,
+                {"snapshot": snap.model_dump() if snap is not None else None},
+            )
+            return False
+
+        if msg_type == MsgType.CMD_POP:
+            task_id = payload["task_id"]
+            result = self.engine.pop_result(task_id)
+            if result is not None:
+                ctx.task_ids.discard(task_id)
+                self._task_to_connection.pop(task_id, None)
+            await self._send(
+                writer,
+                MsgType.MSG_RESULT,
+                {
+                    "found": result is not None,
+                    "result": result.model_dump() if result is not None else None,
+                    "timed_out": False,
+                },
+            )
+            return False
+
+        if msg_type == MsgType.CMD_WAIT:
+            task_id = payload["task_id"]
+            timeout = payload.get("timeout")
+            result, timed_out = await self.engine.wait_task(task_id, timeout)
+            if result is not None:
+                ctx.task_ids.discard(task_id)
+                self._task_to_connection.pop(task_id, None)
+            await self._send(
+                writer,
+                MsgType.MSG_RESULT,
+                {
+                    "found": result is not None,
+                    "result": result.model_dump() if result is not None else None,
+                    "timed_out": timed_out,
+                },
+            )
+            return False
+
+        if msg_type == MsgType.CMD_CLEAR_DONE:
+            cleared_ids = self.engine.clear_done()
+            # engine 一次性 GC 了所有终态任务；同步剔除两张映射 + 本连接 task_ids
+            for task_id in cleared_ids:
+                self._task_to_connection.pop(task_id, None)
+            if cleared_ids:
+                ctx.task_ids.difference_update(cleared_ids)
+            # 把 cleared_ids 一并回传给 client，client 端按 ids 同步本地句柄。
+            # 不能依赖 client 本地 snapshot.status 判断（任务可能未 poll 过状态，
+            # 本地 status 还是 RUNNING/PENDING，但 server 这边已经摘除条目了）。
+            await self._send(
+                writer,
+                MsgType.MSG_PONG,
+                {"count": len(cleared_ids), "cleared_ids": cleared_ids},
+            )
+            return False
+
+        if msg_type == MsgType.CMD_STATS:
+            # 返回本连接累计的状态分布。这是"累计口径"：SUCCESS/FAILED/TIMEOUT/CANCELLED/REJECTED
+            # 不会因 POP/WAIT 消费而减，方便 client 侧做 tqdm 进度条。
+            statistics = self._connection_statistics.get(ctx.connection_id, {})
+            counts = {s.name: int(statistics.get(s, 0)) for s in TaskStatus}
+            finished = sum(counts[s] for s in ("SUCCESS", "FAILED", "TIMEOUT", "CANCELLED", "REJECTED"))
+            total = sum(counts.values())
+            await self._send(
+                writer,
+                MsgType.MSG_STATS,
+                {"total": total, "finished": finished, "counts": counts},
+            )
+            return False
+
+        if msg_type == MsgType.CMD_BYE:
+            return True
+
+        if msg_type == MsgType.CMD_MOCK_REQUEST:
+            request = InferenceRequest.model_validate(payload["request"])
+            mr = self.engine.mock_request(request, overrides=ctx.overrides)
+            await self._send(writer, MsgType.MSG_PREPARED_REQUEST, dict(mr))
+            return False
+
+        if msg_type == MsgType.CMD_SHUTDOWN:
+            await self._send(writer, MsgType.MSG_PONG, {})
+            self.stop()
+            return True
+
+        await self._send(writer, MsgType.MSG_ERROR, {"error": f"未知 msg_type={msg_type}"})
+        return False
+
+    # (wait_task 逻辑已下沉到 _KittyEngine.wait_task)

llm_engine/model_config.py

@@ -0,0 +1,112 @@
+# llm_engine/model_config.py
+
+import json
+import threading
+
+from pathlib import Path
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from typing import Any, Dict, List, Literal, Optional
+
+import kitty_logger
+
+from .schemas import InferenceParameters
+
+logger = kitty_logger.getLogger(__name__)
+
+
+class ModelPricing(BaseModel):
+
+    input: float
+    output: float
+
+
+class ModelConfig(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+    name: str
+    model_id: str = Field(alias="model_id")
+    api_urls: List[str] = Field(alias="api_urls")
+    api_key: Optional[str] = Field(None, alias="api_key")
+    platform: Optional[str] = None
+
+    pricing: Optional[ModelPricing] = None
+    default_inference_parameters: Optional[InferenceParameters] = Field(default=None, alias="inference_parameters")
+    extra_headers: Dict[str, str] = Field(default_factory=dict)
+    extra_payload: Dict[str, Any] = Field(default_factory=dict)
+
+    # --- 私有属性：处理运行时逻辑 ---
+    _url_active_counts: Dict[str, int] = PrivateAttr()
+    _url_lock: threading.Lock = PrivateAttr(default_factory=threading.Lock)
+
+    def model_post_init(self, __context: Any) -> None:
+        """Pydantic 初始化后的钩子，初始化各 URL 的活跃请求计数器。"""
+        if not self.api_urls:
+            raise ValueError(f"模型 '{self.name}' 的 api_urls 不能为空")
+        self._url_active_counts = {url: 0 for url in self.api_urls}
+
+    def get_url(self) -> str:
+        """选取当前活跃请求数最少的 URL，并将其计数 +1。"""
+        with self._url_lock:
+            url = min(self._url_active_counts, key=lambda u: self._url_active_counts[u])
+            self._url_active_counts[url] += 1
+            return url
+
+    def release_url(self, url: str) -> None:
+        """请求完成后，将对应 URL 的活跃计数 -1。"""
+        with self._url_lock:
+            if url in self._url_active_counts:
+                self._url_active_counts[url] = max(0, self._url_active_counts[url] - 1)
+
+    def calculate_cost(self, input_tokens: int, output_tokens: int) -> float:
+        if not self.pricing:
+            return 0.0
+        return input_tokens * self.pricing.input + output_tokens * self.pricing.output
+
+
+class ModelConfigRegistry:
+
+    __slots__ = ("model_dict",)
+
+    def __init__(self):
+
+        self.model_dict: Dict[str, ModelConfig] = {}
+
+    def load_from_json(self, config_path: str, mode: Literal["reload", "merge"] = "reload") -> bool:
+
+        path = Path(config_path)
+        logger.info(f"正在加载模型配置: '{path}'")
+
+        if self.model_dict and mode == "reload":
+            self.model_dict.clear()
+
+        if not path.is_file():
+            raise FileNotFoundError(f"模型配置文件不存在或不是文件: '{path}'")
+
+        with open(path, mode="r", encoding="utf-8") as f:
+            try:
+                raw_data: Dict[str, Dict[str, Any]] = json.load(f)
+            except json.JSONDecodeError as e:
+                logger.error(f"JSON 格式错误: {e}")
+                raise
+
+        for model_name, info in raw_data.items():
+
+            if mode == "merge" and model_name in self.model_dict:
+                logger.debug(f"模型 '{model_name}' 已存在，将被新配置覆盖。")
+
+            model_config: ModelConfig = ModelConfig(name=model_name, **info)
+            self.model_dict[model_name] = model_config
+
+        logger.info(f"共加载 {len(self.model_dict)} 个模型配置。")
+
+        return True
+
+    def get(self, name: str) -> ModelConfig:
+
+        if name not in self.model_dict:
+
+            available_models: str = ", ".join(self.model_dict.keys())
+
+            raise KeyError(f"模型 '{name}' 未注册。可用模型: [{available_models}]")
+
+        return self.model_dict[name]