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

llm_engine/kitty/protocol.py

@@ -0,0 +1,213 @@
+# llm_engine/kitty/protocol.py
+
+"""
+KittyEngine 的线缆协议。
+
+帧格式:
+  [4 bytes BE uint32: payload_len (含 msg_type 字节)]
+  [1 byte: msg_type]
+  [payload_len - 1 bytes: msgpack 编码的 dict]
+
+序列化:
+  msgpack（安全、紧凑；对比 pickle 不支持任意对象反序列化，消除 RCE 风险）。
+  所有 payload 均由 pydantic .model_dump() 产出，已是纯 dict/list/scalar 结构。
+
+传输:
+  UDS (AF_UNIX + SOCK_STREAM) 或 TCP (AF_INET + SOCK_STREAM)。
+"""
+
+import socket
+import struct
+import asyncio
+from enum import IntEnum
+from typing import Any
+from urllib.parse import urlparse
+
+import msgpack
+
+MAX_FRAME_SIZE = 16 * 1024 * 1024  # 16 MB
+_HEADER_LEN = 4
+
+
+class MsgType(IntEnum):
+    # client → server
+    CMD_HELLO = 1
+    CMD_SUBMIT = 2
+    CMD_CANCEL = 3
+    CMD_POLL = 4
+    CMD_POP = 5
+    CMD_CLEAR_DONE = 6
+    CMD_PING = 7
+    CMD_BYE = 8
+    CMD_SHUTDOWN = 9  # 仅 embedded 模式客户端用，通知 server 整体退出
+    CMD_WAIT = 10  # 阻塞等待任务终态
+    CMD_STATS = 11  # 查询本连接所提交任务的状态分布（用于进度条）
+    CMD_MOCK_REQUEST = 12  # 组装请求但不发送，返回 PreparedRequest
+
+    # server → client
+    MSG_WELCOME = 101
+    MSG_RESULT = 102
+    MSG_SNAPSHOT = 103
+    MSG_PONG = 104
+    MSG_ERROR = 105
+    MSG_STATS = 106
+    MSG_PREPARED_REQUEST = 107
+
+
+def encode_frame(msg_type: int, payload: dict[str, Any]) -> bytes:
+    """将消息类型和 payload 编码为帧字节流。
+
+    Args:
+        msg_type (int): 消息类型，取 MsgType 枚举值。
+        payload (dict[str, Any]): 消息体，需可被 msgpack 序列化。
+
+    Returns:
+        bytes: 4 字节大端长度头 + 1 字节消息类型 + msgpack 编码的 payload。
+    """
+    # use_bin_type=True（msgpack>=1.0 默认）：str → msgpack str，bytes → msgpack bin
+    body = msgpack.packb(payload, use_bin_type=True)
+    if body is None:
+        raise ValueError("msgpack.packb returned None")
+    body_len = 1 + len(body)
+    if body_len > MAX_FRAME_SIZE:
+        raise ValueError(f"帧超限: {body_len} > {MAX_FRAME_SIZE}")
+    return struct.pack(">I", body_len) + bytes([int(msg_type)]) + body
+
+
+def _decode_body(body: bytes) -> dict[str, Any]:
+    # raw=False：msgpack str 解码为 Python str；strict_map_key=False：允许非字符串 key（我方不产生，但宽容处理）
+    payload = msgpack.unpackb(body, raw=False, strict_map_key=False)
+    if not isinstance(payload, dict):
+        raise ValueError(f"非法 payload 类型: {type(payload).__name__}")
+    return payload
+
+
+def recv_exact(sock: socket.socket, n: int) -> bytes:
+    """精确读 n 字节；EOF 抛 ConnectionError。"""
+    chunks: list[bytes] = []
+    remaining = n
+    while remaining > 0:
+        chunk = sock.recv(remaining)
+        if not chunk:
+            raise ConnectionError("socket EOF")
+        chunks.append(chunk)
+        remaining -= len(chunk)
+    return b"".join(chunks)
+
+
+async def aread_frame(reader: "asyncio.StreamReader") -> tuple[int, dict[str, Any]]:
+    """从 asyncio.StreamReader 读一帧。EOF/错误抛异常。"""
+    header = await reader.readexactly(_HEADER_LEN)
+    (body_len,) = struct.unpack(">I", header)
+    if body_len <= 0 or body_len > MAX_FRAME_SIZE:
+        raise ValueError(f"非法帧长: {body_len}")
+    body = await reader.readexactly(body_len)
+    msg_type = body[0]
+    return msg_type, _decode_body(body[1:])
+
+
+def recv_frame(sock: socket.socket) -> tuple[int, dict[str, Any]]:
+    """读一帧，返回 (msg_type, payload)。EOF/错误抛异常。"""
+    header = recv_exact(sock, _HEADER_LEN)
+    (body_len,) = struct.unpack(">I", header)
+    if body_len <= 0 or body_len > MAX_FRAME_SIZE:
+        raise ValueError(f"非法帧长: {body_len}")
+    body = recv_exact(sock, body_len)
+    msg_type = body[0]
+    return msg_type, _decode_body(body[1:])
+
+
+# —— Endpoint 解析 —— #
+
+
+def parse_endpoint(endpoint: str) -> tuple[str, tuple]:
+    """
+    解析 endpoint 字符串。返回 (kind, addr)。
+      "unix:///path/to.sock"       → ("unix", ("/path/to.sock",))
+      "tcp://host:port"            → ("tcp",  (host, port))
+    """
+    parsed = urlparse(endpoint)
+    if parsed.scheme == "unix":
+        # urlparse("unix:///tmp/x.sock") → path="/tmp/x.sock"
+        path = parsed.path
+        if not path:
+            raise ValueError(f"unix endpoint 缺路径: {endpoint}")
+        return "unix", (path,)
+    if parsed.scheme == "tcp":
+        host = parsed.hostname or "127.0.0.1"
+        port = parsed.port
+        if port is None:
+            raise ValueError(f"tcp endpoint 缺端口: {endpoint}")
+        return "tcp", (host, port)
+    raise ValueError(f"不支持的 endpoint scheme: {parsed.scheme}")
+
+
+def create_listen_socket(endpoint: str, backlog: int = 128) -> socket.socket:
+    """按 endpoint 创建监听 socket。UDS 自动 unlink 已有文件。"""
+    kind, addr = parse_endpoint(endpoint)
+    if kind == "unix":
+        import os
+
+        path = addr[0]
+        try:
+            os.unlink(path)
+        except FileNotFoundError:
+            pass
+        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        sock.bind(path)
+        os.chmod(path, 0o600)
+        sock.listen(backlog)
+        return sock
+    else:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind(addr)
+        sock.listen(backlog)
+        return sock
+
+
+def connect_endpoint(endpoint: str, connect_timeout: float = 30.0, retry_interval: float = 0.1) -> socket.socket:
+    """
+    连 endpoint，带退避重试（embedded 模式下 worker 启动需要时间）。
+    超时或重试耗尽抛 TimeoutError。
+    - connect_timeout <= 0：表示只尝试一次，不做任何等待/重试。
+    """
+    import time
+
+    kind, addr = parse_endpoint(endpoint)
+    attempts = 0
+    last_err: Exception | None = None
+
+    def _try_once() -> socket.socket:
+        if kind == "unix":
+            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        else:
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+        # connect 失败时必须关闭 sock，否则在重试循环里每次失败都会泄漏一个 fd
+        # （embedded 模式下 worker 启动慢时尤为明显）。
+        try:
+            if kind == "unix":
+                sock.connect(addr[0])
+            else:
+                sock.connect(addr)
+        except Exception:
+            sock.close()
+            raise
+        return sock
+
+    # 至少尝试一次（即便 connect_timeout<=0），避免边界条件下直接报 TimeoutError
+    deadline = time.time() + connect_timeout if connect_timeout > 0 else None
+    interval = retry_interval
+    while True:
+        attempts += 1
+        try:
+            return _try_once()
+        except (FileNotFoundError, ConnectionRefusedError, OSError) as e:
+            last_err = e
+            if deadline is None or time.time() >= deadline:
+                break
+            time.sleep(interval)
+            interval = min(interval * 2, 1.0)
+
+    raise TimeoutError(f"连接 {endpoint} 超时（尝试 {attempts} 次，connect_timeout={connect_timeout}s）: {last_err!r}")

llm_engine/kitty/schemas.py

@@ -0,0 +1,89 @@
+# llm_engine/kitty/schemas.py
+
+"""
+KittyEngine 专属数据结构：TaskSnapshot 与 WireResult。
+
+- TaskSnapshot：跨进程可序列化的任务快照，替代 GeneralEngine 的 TaskHandle。
+  不嵌 result，避免每次 poll 都回传大对象。
+- WireResult：线缆上传输的瘦身 Result，不含 request 字段。客户端收到后用本地
+  cache 的原始 request 拼回完整 InferenceRequestResult。
+- TransitionRecord：一次状态转移记录。由 worker 产生，随 snapshot 跨进程回传。
+"""
+
+from datetime import datetime
+from typing import Any, Optional, TypedDict
+
+from pydantic import BaseModel, Field
+
+from ..schemas import ModelOutput, TaskStatus
+
+
+class TransitionRecord(BaseModel):
+    """一次状态转移记录。
+
+    - ``status``：进入的状态
+    - ``timestamp``：进入该状态时刻的人类可读本地时间，形如 '2026-05-11 10:30:45.123456'
+    - ``duration``：该状态实际持续秒数；在进入下一状态时由调用方回填；末条为 None
+    - ``desc``：进入该状态的原因
+    """
+
+    status: TaskStatus
+    timestamp: str
+    duration: Optional[float] = None
+    desc: str = ""
+
+    @property
+    def epoch(self) -> float:
+        return datetime.fromisoformat(self.timestamp).timestamp()
+
+
+class TaskSnapshot(BaseModel):
+    """任务状态快照。跨进程传输用。
+
+    所有时间点/时长信息由 ``transitions`` 流水账承载，不在本结构上再挂
+    submit_time / start_time / end_time / duration 等派生字段，避免：
+      - pydantic ``@property`` 不会进 ``model_dump()`` 造成序列化不一致；
+      - 同一信息在"字段"和"派生"两处同时存在带来的同步负担。
+    如需取某阶段的时间点/时长，调用方自行遍历 ``transitions``。
+    """
+
+    task_id: str
+    status: TaskStatus
+    transitions: list[TransitionRecord] = Field(default_factory=list)
+    task_timeout: Optional[float] = None
+    persist: bool = False
+    error: Optional[str] = None
+    has_result: bool = False
+
+    def is_finished(self) -> bool:
+        return self.status in (
+            TaskStatus.SUCCESS,
+            TaskStatus.FAILED,
+            TaskStatus.TIMEOUT,
+            TaskStatus.CANCELLED,
+            TaskStatus.REJECTED,
+        )
+
+
+class WireResult(BaseModel):
+    """线缆上的 Result。不含 request，收端自行拼回。
+
+    ``status`` 无默认值：调用方必须显式指定终态，避免"忘记传 → 失败被记成 SUCCESS"
+    这类语义漂移。
+    """
+
+    success: bool
+    task_id: str
+    status: TaskStatus  # 终态；client 用于同步 snapshot.status（区分 SUCCESS/FAILED/TIMEOUT/REJECTED/CANCELLED）
+    model_output: Optional[ModelOutput] = None
+    error_message: Optional[str] = None
+    duration: Optional[float] = None
+
+
+class MockRequest(TypedDict):
+    """发给大模型服务商的完整请求数据，仅用于调试，不发出网络调用。"""
+
+    urls: list[str]
+    model: str
+    headers: dict[str, str]
+    payload: dict[str, Any]