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

llm_engine/kitty/engine.py

@@ -0,0 +1,1077 @@
+# llm_engine/kitty/engine.py
+
+"""KittyEngine 后端进程的任务执行器层。
+
+本模块只包含两个核心对象：
+
+- `_WorkerTask`：后端进程内的任务状态容器（仅进程内使用，不跨进程）。
+- `_KittyEngine`：任务生命周期管理器，由 KittyServer 在子进程主线程的
+  asyncio loop 中驱动。
+
+分层设计:
+    - 协议/序列化：由 `llm_engine.kitty.schemas`（`TaskSnapshot` / `WireResult`
+      / `TransitionRecord`）与 `llm_engine.kitty.config` 负责；worker 不感知。
+    - 事件循环 / socket：由 KittyServer 负责。worker 只暴露任务级 API
+      （`submit` / `cancel` / `snapshot` / `pop_result` / `clear_done`
+      / `force_remove` / `wait_task` / `cleanup_tasks`），不接触 socket。
+    - HTTP / 模型路由：由 `_KittyEngine` 直接使用 `httpx.AsyncClient` +
+      `ModelConfigRegistry`。上层 client 不感知。
+
+并发模型:
+    - worker 运行在子进程的唯一 asyncio 事件循环线程里。
+    - 所有 `self._tasks` / `_WorkerTask` 字段的读写都在该线程内完成，故不需要锁。
+    - 并发度由 `asyncio.Semaphore(max_global_concurrency)` 控制；每个任务是
+      一条 `asyncio.Task`，运行 `_KittyEngine._run_task`。
+
+状态机（见 `llm_engine.schemas.TaskStatus`）::
+
+    SUBMITTED ──► ACCEPTED ──► PENDING ──► RUNNING ──► SUCCESS
+         │            │           │           │
+         │            │           │           ├──► FAILED      (异常/HTTP/解析失败)
+         │            │           │           ├──► TIMEOUT     (task_timeout 触发)
+         │            │           │           └──► CANCELLED   (运行中取消)
+         │            └───────────┴──► CANCELLED               (排队中取消)
+         └──► REJECTED（校验失败：缺 model_name、模型未注册……）
+
+- SUBMITTED / ACCEPTED：任务已落盘到 worker，但尚未入队 sem。REJECTED 从这一
+  段产生。
+- PENDING：已入队 `_sem`，等待并发名额。
+- RUNNING：持有 sem，正在发 HTTP。
+- SUCCESS / FAILED / TIMEOUT / CANCELLED / REJECTED：终态，由 `is_finished` 判定。
+
+将来若支持 "RUNNING 中异常回到 PENDING 再跑" 的重试，状态会多次经过 PENDING /
+RUNNING。`TaskSnapshot.submit_time` / `start_time` 等派生属性按"首次出现"
+语义实现，已为此预留。
+
+状态转移的唯一入口:
+    所有状态变更必须经由 `_WorkerTask.transition`。该方法会：
+
+    1. 写入新状态；
+    2. 追加一条 `TransitionRecord`；
+    3. 回填上一条记录的 `duration`（= 上一个状态的实际持续秒数）。
+
+    禁止直接 `wt.status = X`，否则时间线会丢失。
+"""
+
+import asyncio
+import time
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Callable, Optional
+
+import httpx
+import kitty_logger
+from pydantic import ValidationError
+
+from ..model_config import ModelConfig, ModelConfigRegistry
+from ..schemas import (
+    ChatCompletionResponse,
+    InferenceParameters,
+    InferenceRequest,
+    PreparedRequest,
+    ModelOutput,
+    TaskStatus,
+)
+from .config import KittyEngineConfig, KittyEngineOverrides
+from .schemas import MockRequest, TaskSnapshot, TransitionRecord, WireResult
+
+logger = kitty_logger.getLogger(__name__)
+
+
+def _now_str() -> str:
+    """返回人类可读的本地时间戳，形如 `'2026-05-11 10:30:45.123456'`。
+
+    用法约定:
+        - 仅用于 `TransitionRecord.timestamp`（字符串便于日志 / 跨进程传输 /
+          可读性）。
+        - 需要数值秒时通过 `TransitionRecord.epoch`
+          （`datetime.fromisoformat(...).timestamp()`）换算，保持
+          "字符串为真，epoch 派生"。
+
+    Returns:
+        带微秒精度的本地时间字符串。
+    """
+    return datetime.now().isoformat(sep=" ", timespec="microseconds")
+
+
+@dataclass
+class _WorkerTask:
+    """后端进程内部的任务状态容器。不跨进程传输。
+
+    为什么用 dataclass 而非 pydantic:
+        - 持有 `asyncio.Task` / `asyncio.Event` 等不可序列化对象；
+        - 仅在事件循环线程内存在，不需要 `model_validate`；
+        - 需要频繁原地变更（状态/结果/错误信息）。
+
+    时间线语义（重要）:
+        - 不保存 `submit_time / start_time / end_time` 字段。所有时间点统一
+          由 `transitions` 列表记录；每次进入新状态追加一条 `TransitionRecord`，
+          并把上一条的 `duration` 回填为上一个状态实际持续的秒数。
+        - 若未来支持 "RUNNING 失败→回到 PENDING→再次 RUNNING" 的重试路径，
+          `transitions` 会出现多个同状态条目；对应的 submit/start 等派生值
+          按"首次出现"取值。
+
+    与外部对象的协作要求:
+        - `_KittyEngine.submit` 在构造完实例之后、`create_task` 之前，必须
+          显式设置 `done_event = asyncio.Event()`，并挂上 `asyncio_task`。
+          （当前 `__post_init__` 只初始化 `transitions`，不创建 event，以避免
+          在无事件循环的线程里被误构造。）
+        - `snapshot` 对 `transitions` 做浅拷贝；`TransitionRecord` 在 worker
+          内部一旦被追加后，除回填末条 `duration` 外不应再被修改，否则外部
+          已经 pop 走的快照可能被"幽灵修改"。
+
+    Attributes:
+        task_id: 任务唯一 ID（由上层 client/engine 生成，本模块只负责查表）。
+        request: 原始 `InferenceRequest`。用于 `_run_task` 发起 HTTP。
+        overrides: 本次提交附带的 `KittyEngineOverrides`（可为 None），在模型 /
+            超时 / header / payload 解析时作为中间优先级兜底。
+        status: 当前状态。禁止直接赋值，必须通过 `transition` 更新。
+        task_timeout: 单次任务（发 HTTP + 解析）的 wall-clock 超时。由 submit
+            按优先级解析得到：`request.timeout` > submit 参数 >
+            `overrides.default_timeout` > `config.default_timeout`。
+        persist: true 则终态后仍保留在 `_tasks` 中（可多次 snapshot / 重复 pop
+            语义见 `_KittyEngine.pop_result`）；false 则一旦 pop_result 拿走
+            即删除。
+        error: 终态为 FAILED / REJECTED / CANCELLED 时的简要错误信息
+            （同样会写入 `result`）。
+        result: 终态的 `WireResult`，由 `_run_task` 在各终态分支显式构造。
+        asyncio_task: 运行 `_run_task` 的底层 `asyncio.Task`，用于 cancel。
+        done_event: 到达终态时由 `_run_task` 的 `finally` 置位。`wait_task` /
+            `cleanup_tasks` 依赖它做异步等待。
+        transitions: 状态转移流水账。`transitions[0]` 为 SUBMITTED 初始记录，
+            顺序即状态机实际走过的路径。
+    """
+
+    task_id: str
+    request: InferenceRequest
+    overrides: Optional[KittyEngineOverrides]
+    status: TaskStatus = TaskStatus.SUBMITTED
+    task_timeout: Optional[float] = None
+    persist: bool = False
+    error: Optional[str] = None
+    result: Optional[WireResult] = None
+    asyncio_task: Optional[asyncio.Task] = None
+    done_event: Optional[asyncio.Event] = None
+    transitions: list[TransitionRecord] = field(default_factory=list)
+    # 状态转移回调。由 `_KittyEngine.submit` 在构造后注入（当前只支持单订阅者）；
+    # 不持有 worker 反向引用，避免循环。签名: (task_id, old, new) -> None；
+    # `old` 为 None 表示"首次落地"（由 `_KittyEngine.submit` 显式触发，不经 transition）。
+    on_transition: Optional[Callable[[str, Optional[TaskStatus], TaskStatus], None]] = None
+
+    def __post_init__(self) -> None:
+        # 初始状态（默认 SUBMITTED）作为第一条转移记录。
+        # 注意：这里直接 append，不走 transition()，因为此时没有"上一条"需要回填 duration，
+        # 且 transition() 的语义是"从旧状态进入新状态"，不适用于首次落地。
+        if not self.transitions:
+            self.transitions.append(TransitionRecord(status=self.status, timestamp=_now_str(), desc="task created"))
+
+    # ------------------------------------------------------------------
+    # 状态转移（唯一入口）
+    # ------------------------------------------------------------------
+
+    def transition(self, new_status: TaskStatus, desc: str = "") -> None:
+        """更新 `status`、追加转移记录，并回填上一条的 `duration`。
+
+        行为:
+            1. 取当前时间 `now`（单次 `datetime.now()`）；
+            2. 把 `transitions[-1].duration` 置为 `now.timestamp() - prev.epoch`；
+            3. 将 `self.status` 置为 `new_status`；
+            4. 追加 `TransitionRecord(new_status, now.isoformat(...), desc=desc)`。
+
+        Args:
+            new_status: 要进入的新状态。
+            desc: 进入该状态的原因，便于日志 / 排障。
+
+        Note:
+            - 所有状态变更都应走这里，不要直接赋值 `self.status`，否则会丢失
+              时间线。
+            - 注意：本模块的 `transition()` 约束仅限 `_WorkerTask`；
+              `TaskHandle`（GeneralEngine）当前仍直接赋值 `status`，未接入
+              transitions 流水账。两边状态机由各自维护。
+            - 不做状态合法性校验：调用方自己确保转移顺序合理（见模块开头状态机图）。
+        """
+        # 单次取 now，避免"写入字符串 → 再 fromisoformat 解回来"的冗余解析。
+        now = datetime.now()
+        now_str = now.isoformat(sep=" ", timespec="microseconds")
+        old = self.status
+        if self.transitions:
+            self.transitions[-1].duration = now.timestamp() - self.transitions[-1].epoch
+        self.status = new_status
+        self.transitions.append(TransitionRecord(status=new_status, timestamp=now_str, desc=desc))
+        # 通知回调。callback 异常绝不能污染 worker 状态机。
+        callback = self.on_transition
+        if callback is not None:
+            try:
+                callback(self.task_id, old, new_status)
+            except Exception:
+                logger.exception("transition callback 异常 (task=%s)", self.task_id)
+
+    def is_finished(self) -> bool:
+        """判断是否已到达终态（SUCCESS / FAILED / TIMEOUT / CANCELLED / REJECTED）。
+
+        Returns:
+            到达终态返回 True，否则 False。
+        """
+        return self.status in (
+            TaskStatus.SUCCESS,
+            TaskStatus.FAILED,
+            TaskStatus.TIMEOUT,
+            TaskStatus.CANCELLED,
+            TaskStatus.REJECTED,
+        )
+
+    def snapshot(self) -> TaskSnapshot:
+        """生成可跨进程传输的 `TaskSnapshot`。
+
+        Returns:
+            当前任务状态的快照。
+
+        Note:
+            - `transitions` 做**深拷贝**（逐条 `TransitionRecord.model_copy()`）：
+              `transition()` 会原地修改 `transitions[-1].duration` 来回填上一个
+              状态的持续时长；若只做浅拷贝，则已经交给客户端的 snapshot 末条
+              记录会被后续 transition 调用"幽灵修改"。深拷贝彻底消除该隐患。
+            - 不回传 `result`：瘦身，避免每次 poll status 都搬运模型输出。
+              客户端通过 `has_result` 判断是否需要单独 pop。
+            - 首条 SUBMITTED 记录由 `__post_init__` 在构造时自动追加，此处无需
+              额外处理；REJECTED 分支里调用 `transition(REJECTED)` 时，会自然
+              回填 SUBMITTED 的 duration，等价于"校验阶段耗时"。
+        """
+        return TaskSnapshot(
+            task_id=self.task_id,
+            status=self.status,
+            # 深拷贝：TransitionRecord 可被 transition() 原地改 duration，
+            # 不深拷贝会把已发出的 snapshot 一起带跑。
+            transitions=[r.model_copy() for r in self.transitions],
+            task_timeout=self.task_timeout,
+            persist=self.persist,
+            error=self.error,
+            has_result=self.result is not None,
+        )
+
+
+class _KittyEngine:
+    """KittyEngine 后端进程的任务执行器。
+
+    职责:
+        - 维护模型注册表 + 单个 `httpx.AsyncClient` + 全局并发信号量；
+        - 维护 `_tasks: dict[task_id, _WorkerTask]`；
+        - 暴露任务级 API（submit / cancel / snapshot / pop_result / wait_task /
+          cleanup_tasks / clear_done / force_remove）。
+
+    使用约束:
+        - 必须在 asyncio 事件循环线程内调用所有方法（含同步方法：它们会创建
+          `asyncio.Task` / `asyncio.Event`，依赖当前 loop）。
+        - 生命周期：`setup()` 一次 → 正常工作 → `teardown()` 一次。重复 setup
+          / teardown 会 warning 但不抛错。
+    """
+
+    def __init__(self, config: KittyEngineConfig) -> None:
+        self.config = config
+
+        # —— 模型注册表（与 GeneralEngine 一致，reload 模式） ——
+        self.model_registry = ModelConfigRegistry()
+        self.model_registry.load_from_json(config.registry_path, mode="reload")
+
+        # —— 运行时资源，由 setup 初始化；teardown 置回 None / False ——
+        # _sem 用于 RUNNING 并发上限控制；PENDING → RUNNING 之间 await 它。
+        self._sem: Optional[asyncio.Semaphore] = None
+        # 单一 AsyncClient，连接复用；timeout 由 config 指定。
+        self._http_client: Optional[httpx.AsyncClient] = None
+        # 活跃任务表。终态任务是否保留取决于 persist。
+        self._tasks: dict[str, _WorkerTask] = {}
+        self._setup_done: bool = False
+        # 状态转移回调（单订阅者，一般由 KittyServer 注册，用于维护 per-connection
+        # 状态统计）。callback 在事件循环线程内被**同步**调用，不允许阻塞。
+        self._transition_callback: Optional[Callable[[str, Optional[TaskStatus], TaskStatus], None]] = None
+
+    def set_transition_callback(
+        self,
+        callback_function: Optional[Callable[[str, Optional[TaskStatus], TaskStatus], None]],
+    ) -> None:
+        """注册（或清除）状态转移回调。
+
+        约定:
+            - 只能有一个订阅者；多次调用以最后一次为准。
+            - 设置时机应在 `submit` 首次调用之前，否则更早的任务不会带上 callback。
+            - callback 内不能抛异常（调用端会吞掉 + 记 exception 日志）。
+        """
+        self._transition_callback = callback_function
+
+    # ------------------------------------------------------------------
+    # 生命周期
+    # ------------------------------------------------------------------
+
+    async def setup(self) -> None:
+        """初始化 sem + AsyncClient。幂等（重复调用仅 warning）。"""
+        if self._setup_done:
+            logger.warning("KittyWorker 已启动，无需重复启动。")
+            return
+        self._sem = asyncio.Semaphore(self.config.max_global_concurrency)
+        self._http_client = httpx.AsyncClient(
+            timeout=httpx.Timeout(
+                timeout=self.config.http_client_connect_timeout,
+                read=self.config.http_client_read_timeout,
+            ),
+            limits=httpx.Limits(
+                max_connections=self.config.max_global_concurrency,
+                max_keepalive_connections=self.config.max_global_concurrency,
+            ),
+        )
+        self._setup_done = True
+        logger.info(
+            "KittyWorker 启动 (max_global_concurrency=%d, models=%d)",
+            self.config.max_global_concurrency,
+            len(self.model_registry.model_dict),
+        )
+
+    async def teardown(self) -> None:
+        """取消所有未完成任务并关闭 HTTP 客户端。幂等。
+
+        流程:
+            1. 收集所有未 finished 的 `asyncio_task`；
+            2. 统一 `cancel()`；
+            3. `gather(..., return_exceptions=True)` 等它们真正退出，防止在关闭
+               client 之后还有未结束协程尝试发请求；
+            4. 关闭 `_http_client`。
+        """
+        if not self._setup_done:
+            logger.warning("KittyWorker 未启动，无需停止。")
+            return
+        # 取消未完成任务
+        pending = [t for t in self._tasks.values() if not t.is_finished() and t.asyncio_task is not None]
+        for t in pending:
+            t.asyncio_task.cancel()  # type: ignore[union-attr]
+        if pending:
+            await asyncio.gather(*[t.asyncio_task for t in pending if t.asyncio_task], return_exceptions=True)
+        if self._http_client is not None:
+            await self._http_client.aclose()
+            self._http_client = None
+        self._setup_done = False
+        logger.info("KittyWorker 已停止")
+
+    # ------------------------------------------------------------------
+    # 任务 API（均在事件循环线程内调用）
+    # ------------------------------------------------------------------
+
+    def submit(
+        self,
+        *,
+        task_id: str,
+        request: InferenceRequest,
+        overrides: Optional[KittyEngineOverrides],
+        persist: bool,
+        timeout: Optional[float],
+    ) -> TaskSnapshot:
+        """提交一个任务并立刻返回其初始 snapshot。
+
+        行为:
+            - 按 "request.timeout > submit timeout > overrides.default_timeout >
+              config.default_timeout" 的优先级解析 `task_timeout`；
+            - 构造 `_WorkerTask`（此时已经写入 SUBMITTED 记录）；
+            - 创建 `done_event` 与 `asyncio_task`（= `_run_task(wt)`）；
+            - 登记到 `_tasks`；
+            - 返回 snapshot 给调用方（此时 task 大概率还在 SUBMITTED；真正的
+              校验 / ACCEPTED / REJECTED 由 `_run_task` 在下一个 tick 写入）。
+
+        Args:
+            task_id: 上层分配的任务 ID。
+            request: 推理请求。
+            overrides: 可选的 per-submit 覆盖配置。
+            persist: 终态后是否保留在 `_tasks` 表里。
+            timeout: 单次任务超时秒数（可被 `request.timeout` 覆盖）。
+
+        Returns:
+            初始 `TaskSnapshot`。
+
+        Note:
+            - 不做同 `task_id` 的去重；上层 client/engine 负责保证唯一性。
+            - 不立刻做模型校验：校验在 `_run_task` 头部完成，REJECTED 也走统一
+              的终态路径（置 result + transition + done_event.set）。
+        """
+        # 超时层级：request.timeout > submit timeout > overrides.default_timeout > config.default_timeout
+        effective_timeout = request.timeout
+        if effective_timeout is None:
+            effective_timeout = timeout
+        if effective_timeout is None and overrides is not None:
+            effective_timeout = overrides.default_timeout
+        if effective_timeout is None:
+            effective_timeout = self.config.default_timeout
+
+        wt = _WorkerTask(
+            task_id=task_id,
+            request=request,
+            overrides=overrides,
+            task_timeout=effective_timeout,
+            persist=persist,
+        )
+        # done_event 必须在 event loop 线程内创建，因此放到此处而非 __post_init__
+        wt.done_event = asyncio.Event()
+        wt.on_transition = self._transition_callback
+        # 注意：**先**登记到 _tasks，**再**显式 fire 一次"首次落地"事件（old=None），
+        # 最后才 create_task。这样：
+        #   - 回调被触发时，`worker._tasks[task_id]` 已可见，callback 里若需要
+        #     回查 snapshot 不会 KeyError；
+        #   - create_task 调度的 _run_task 只会在当前同步段结束后才实际执行，
+        #     因此它里面后续的 transition 事件一定晚于此处首次事件。
+        # 之所以不让 `__post_init__` 走 transition()：初始记录不回填 duration、
+        # 没有"旧状态"，语义上不是一次 transition。把首次 fire 放在这里，可以保证
+        # 订阅方看到任意 task 全生命周期的第一条事件总是由 worker 主动发出。
+        self._tasks[task_id] = wt
+        if self._transition_callback is not None:
+            try:
+                self._transition_callback(wt.task_id, None, wt.status)
+            except Exception:
+                logger.exception("transition callback 异常（initial, task=%s）", wt.task_id)
+        wt.asyncio_task = asyncio.create_task(self._run_task(wt))
+        return wt.snapshot()
+
+    def cancel(self, task_id: str, force: bool = False) -> bool:
+        """取消任务。
+
+        语义:
+            - 任务不存在 / 已到终态：返回 False。
+            - RUNNING 中的任务：默认不取消（保护已经在打的 HTTP 请求），需要
+              `force=True` 才取消。
+            - 其它非终态（SUBMITTED / ACCEPTED / PENDING）：直接取消
+              `asyncio_task`；`_run_task` 的 `except CancelledError` 分支会把
+              状态推到 CANCELLED。
+
+        Args:
+            task_id: 目标任务 ID。
+            force: 是否强制取消 RUNNING 中的任务。
+
+        Returns:
+            True 表示已向 `asyncio_task` 发出 cancel。真正进入终态需等
+            `_run_task` 清理 + `done_event.set()`（可用 `wait_task` 同步等待）。
+        """
+        wt = self._tasks.get(task_id)
+        if wt is None or wt.is_finished():
+            return False
+        if wt.status == TaskStatus.RUNNING and not force:
+            return False
+        if wt.asyncio_task is not None:
+            wt.asyncio_task.cancel()
+        return True
+
+    def snapshot(self, task_id: str) -> Optional[TaskSnapshot]:
+        """返回当前任务的快照。
+
+        Args:
+            task_id: 目标任务 ID。
+
+        Returns:
+            对应任务的 `TaskSnapshot`；任务不存在返回 None。
+        """
+        wt = self._tasks.get(task_id)
+        return wt.snapshot() if wt is not None else None
+
+    def pop_result(self, task_id: str) -> Optional[WireResult]:
+        """取结果。仅在终态有效。
+
+        Args:
+            task_id: 目标任务 ID。
+
+        Returns:
+            - 任务不存在 / 未终态 → None；
+            - 终态 + `persist=False` → 返回 result 并从 `_tasks` 删除该条目；
+            - 终态 + `persist=True` → 返回 result，条目保留（可再次 pop 到同一 result）。
+        """
+        wt = self._tasks.get(task_id)
+        if wt is None or not wt.is_finished():
+            return None
+        result = wt.result
+        if not wt.persist:
+            self._tasks.pop(task_id, None)
+        return result
+
+    def clear_done(self) -> list[str]:
+        """主动清理所有终态任务（含 persist）。
+
+        典型用法：客户端做周期 GC，或上层显式要求丢弃历史。
+
+        Returns:
+            被清理的任务 ID 列表。调用方据此同步自己的索引（例如 server 的
+            `_task_to_connection` 映射）。长度即 "实际清理的任务数"。
+        """
+        done = [tid for tid, wt in self._tasks.items() if wt.is_finished()]
+        for tid in done:
+            self._tasks.pop(tid, None)
+        return done
+
+    def force_remove(self, task_id: str) -> Optional[WireResult]:
+        """无视状态强制移除条目。
+
+        与 `cancel` 的区别：cancel 只是触发取消流程、等待终态；force_remove
+        直接从 `_tasks` 摘除，不管任务有没有跑完。一般只用于"客户端已断线、
+        条目不再有意义"的兜底清理。
+
+        Args:
+            task_id: 目标任务 ID。
+
+        Returns:
+            该任务的 `WireResult`（若已有），否则 None。
+        """
+        wt = self._tasks.pop(task_id, None)
+        return wt.result if wt is not None else None
+
+    async def wait_task(self, task_id: str, timeout: Optional[float]) -> tuple[Optional[WireResult], bool]:
+        """异步等待任务终态后 pop result。
+
+        依赖：`done_event` 在 `_run_task` 的 `finally` 里被 set；所以所有终态
+        路径都能唤醒本方法。
+
+        Args:
+            task_id: 目标任务 ID。
+            timeout: 等待秒数；None 表示不限时。
+
+        Returns:
+            `(result, timed_out)`：
+
+            - 任务不存在 → `(None, False)`；
+            - 超时 → `(None, True)`（task 仍在进行，状态不变）；
+            - 正常 → `(result, False)`；`result` 的 pop 遵循 `pop_result` 的
+              persist 语义。
+        """
+        wt = self._tasks.get(task_id)
+        if wt is None:
+            return None, False
+        if not wt.is_finished():
+            if wt.done_event is None:
+                return None, False
+            try:
+                await asyncio.wait_for(wt.done_event.wait(), timeout=timeout)
+            except asyncio.TimeoutError:
+                return None, True
+        return self.pop_result(task_id), False
+
+    async def cleanup_tasks(self, task_ids, *, wait_timeout: float = 5.0) -> None:
+        """批量取消一组任务并等其真正进入终态；非 persist 的一并摘除。
+
+        用于客户端断连、engine 主动 drain 等场景。使用 `force=True` 以便 RUNNING
+        任务也会被取消。
+
+        Args:
+            task_ids: 要清理的任务 ID 序列。
+            wait_timeout: 等待全部任务进入终态的整体超时秒数。
+        """
+        tids = [t for t in task_ids if t in self._tasks]
+        if not tids:
+            return
+        events: list[asyncio.Event] = []
+        for tid in tids:
+            wt = self._tasks.get(tid)
+            if wt is None:
+                continue
+            self.cancel(tid, force=True)
+            if not wt.is_finished() and wt.done_event is not None:
+                events.append(wt.done_event)
+        if events:
+            try:
+                await asyncio.wait_for(
+                    asyncio.gather(*[e.wait() for e in events], return_exceptions=True),
+                    timeout=wait_timeout,
+                )
+            except asyncio.TimeoutError:
+                logger.warning("cleanup_tasks 等待 %d 个任务终态超时", len(events))
+        for tid in tids:
+            wt = self._tasks.get(tid)
+            if wt is not None and not wt.persist:
+                self._tasks.pop(tid, None)
+
+    # ------------------------------------------------------------------
+    # 任务执行
+    # ------------------------------------------------------------------
+
+    async def _run_task(self, wt: _WorkerTask) -> None:
+        """单个任务的完整生命周期协程。状态机在此闭合。
+
+        阶段:
+            1. 模型解析：`request.model_name` > `overrides.default_model` >
+               `config.default_model`，全空则 REJECTED。
+            2. 模型校验：在 `model_registry` 中查表，未命中则 REJECTED。
+            3. `transition(ACCEPTED)`：通过校验。
+            4. `transition(PENDING)` → `async with self._sem` →
+               `transition(RUNNING)`：排队 → 持有 sem → 开始执行。
+            5. `async with asyncio.timeout(wt.task_timeout):` 内 `await self._send_request(...)`：
+               发 HTTP + 解析；
+               - 正常 → 填 `result` + `transition(SUCCESS)`；
+               - `CancelledError` → 填 `CANCELLED` result + transition，再 raise
+                 让外层 task 正确结束；此处可确定是外部 cancel，因为
+                 `asyncio.timeout` 超时直接抛 `TimeoutError`，不经 inner cancel；
+               - `TimeoutError` → TIMEOUT；
+               - 其它异常 → FAILED，error=`type: msg`。
+            6. 外层 `except CancelledError`：覆盖 ACCEPTED / PENDING 阶段（尚未
+               进入 RUNNING）被取消的情形。若此时还不是终态，补写 CANCELLED。再
+               raise。
+            7. `finally`：无论哪条路径，只要 `done_event` 存在就 set，唤醒所有
+               `wait_task` / `cleanup_tasks`。
+
+        与其它组件的配合要点:
+            - 每个终态分支同时写 `wt.error` / `wt.result` / `transition`，
+              缺一不可：
+                - `wt.result` 要带正确的 `status` 字段，client 据此同步
+                  snapshot.status；
+                - `transition` 保证 `transitions` 的时间线闭合；
+                - `done_event.set()` 保证等待方被唤醒。
+            - 绝不在此函数外修改 `wt.status`（见 `_WorkerTask.transition` 的
+              约束）。
+
+        Args:
+            wt: 要驱动的任务实例，已由 `submit` 登记到 `_tasks`。
+        """
+        request = wt.request
+        overrides = wt.overrides
+
+        # ---- 1. 模型名解析 ----
+        model_name = request.model_name
+        if model_name is None and overrides is not None:
+            model_name = overrides.default_model
+        if model_name is None:
+            model_name = self.config.default_model
+        if model_name is None:
+            wt.error = "no model_name (request/overrides/config all empty)"
+            wt.result = WireResult(success=False, task_id=wt.task_id, status=TaskStatus.REJECTED, error_message=wt.error)
+            wt.transition(TaskStatus.REJECTED, desc="no model_name")
+            if wt.done_event is not None:
+                wt.done_event.set()
+            return
+
+        # ---- 2. 模型查表校验 ----
+        try:
+            model = self.model_registry.get(name=model_name)
+        except KeyError as e:
+            wt.error = str(e)
+            wt.result = WireResult(success=False, task_id=wt.task_id, status=TaskStatus.REJECTED, error_message=wt.error)
+            wt.transition(TaskStatus.REJECTED, desc=f"model not found: {model_name}")
+            if wt.done_event is not None:
+                wt.done_event.set()
+            return
+
+        # ---- 3. 校验通过 ----
+        # ACCEPTED 与下一行 PENDING 在正常路径下紧邻；保留两条独立转移，便于
+        # 外部观察"校验耗时"与"排队耗时"的界限。
+        wt.transition(TaskStatus.ACCEPTED, desc="validation passed")
+
+        assert self._sem is not None
+        try:
+            # ---- 4. 排队等 sem ----
+            wt.transition(TaskStatus.PENDING, desc="waiting for sem")
+            async with self._sem:
+                # ---- 5. 持有 sem，开跑 ----
+                wt.transition(TaskStatus.RUNNING, desc="sem acquired")
+                run_start = time.time()
+                try:
+                    # 用 asyncio.timeout 而非 asyncio.wait_for：前者超时时直接
+                    # 抛 TimeoutError，不会经由内层 CancelledError 上传，因此
+                    # `except asyncio.CancelledError` 分支可以确定是"外部 cancel"，
+                    # 避免 timeout 与 cancel 路径竞争写两次终态（参见 issue 历史）。
+                    async with asyncio.timeout(wt.task_timeout):
+                        model_output = await self._send_request(request, model, overrides)
+                    wt.result = WireResult(
+                        success=True,
+                        task_id=wt.task_id,
+                        status=TaskStatus.SUCCESS,
+                        model_output=model_output,
+                        duration=time.time() - run_start,
+                    )
+                    wt.transition(TaskStatus.SUCCESS, desc="completed")
+                except asyncio.CancelledError:
+                    # RUNNING 中被外部 cancel：先写终态（外层 finally 统一触发 done_event），再继续 raise。
+                    wt.error = "cancelled"
+                    wt.result = WireResult(success=False, task_id=wt.task_id, status=TaskStatus.CANCELLED, error_message="cancelled")
+                    wt.transition(TaskStatus.CANCELLED, desc="cancelled while running")
+                    logger.info("任务 %s 已取消（RUNNING）", wt.task_id)
+                    raise
+                except TimeoutError:
+                    # asyncio.timeout 超时；与 CancelledError 分支互斥，不会重复写。
+                    wt.error = "timeout"
+                    wt.result = WireResult(success=False, task_id=wt.task_id, status=TaskStatus.TIMEOUT, error_message="timeout")
+                    wt.transition(TaskStatus.TIMEOUT, desc="timeout")
+                except Exception as e:
+                    logger.exception("任务 %s 执行异常", wt.task_id)
+                    wt.error = f"{type(e).__name__}: {e}"
+                    wt.result = WireResult(success=False, task_id=wt.task_id, status=TaskStatus.FAILED, error_message=wt.error)
+                    wt.transition(TaskStatus.FAILED, desc=wt.error)
+        except asyncio.CancelledError:
+            # ACCEPTED / PENDING 阶段（等 sem 时）被取消：此处兜底，保证状态进入终态。
+            # 若 RUNNING 的分支已经写过 CANCELLED，则 is_finished() 为 True，这里不会重复写。
+            if not wt.is_finished():
+                wt.error = "cancelled"
+                wt.result = WireResult(success=False, task_id=wt.task_id, status=TaskStatus.CANCELLED, error_message="cancelled")
+                wt.transition(TaskStatus.CANCELLED, desc="cancelled while pending")
+                logger.info("任务 %s 已取消（PENDING）", wt.task_id)
+            raise
+        finally:
+            # 无论正常 / 异常 / 取消，所有终态路径都在此唤醒等待方。
+            if wt.done_event is not None:
+                wt.done_event.set()
+
+    # ------------------------------------------------------------------
+    # HTTP 核心
+    #
+    # 下面这组方法负责把 (request, model, overrides, config) 四层配置压平成
+    # 一次 HTTP 调用，并处理重试。优先级统一为：
+    #     request > model > overrides > config
+    # 仅对列表/字典采用 **merge**（后者覆盖前者）；对单值（api_key、默认 IP 等）
+    # 采用"优先级高者非空即返回"。
+    # ------------------------------------------------------------------
+
+    def _resolve_default_ip(
+        self,
+        overrides: Optional[KittyEngineOverrides],
+    ) -> Optional[InferenceParameters]:
+        """解析"默认推理参数"层，overrides 优先于 config。
+
+        Args:
+            overrides: per-submit 覆盖配置。
+
+        Returns:
+            选中的默认 `InferenceParameters`；全空则 None。
+
+        Note:
+            model 和 request 的 inference_parameters 在 `_build_payload` 里
+            另行 merge，不在此处返回。
+        """
+        if overrides is not None and overrides.default_inference_parameters is not None:
+            return overrides.default_inference_parameters
+        return self.config.default_inference_parameters
+
+    def _get_api_key(self, request: InferenceRequest, model: ModelConfig, overrides: Optional[KittyEngineOverrides]) -> str:
+        """解析 API Key，按 request > model > overrides > config 取第一个非 None。
+
+        全部为空时返回空字符串并 warning（有些自研网关允许无 key；严格环境应在
+        上游校验）。
+
+        Args:
+            request: 推理请求。
+            model: 已解析的模型配置。
+            overrides: per-submit 覆盖配置。
+
+        Returns:
+            选中的 API Key 字符串。
+        """
+        if request.api_key is not None:
+            return request.api_key
+        if model.api_key is not None:
+            return model.api_key
+        if overrides is not None and overrides.default_api_key is not None:
+            return overrides.default_api_key
+        if self.config.default_api_key is not None:
+            return self.config.default_api_key
+        logger.warning("没有找到可用的 api_key, 将传递空 api_key")
+        return ""
+
+    def _build_headers(self, request: InferenceRequest, model: ModelConfig, overrides: Optional[KittyEngineOverrides]) -> dict[str, str]:
+        """组装请求头。
+
+        顺序：基础（Authorization + Content-Type）→ config.extra_headers →
+        overrides.extra_headers → model.extra_headers → request.extra_headers。
+        后写入的 key 会覆盖前者。
+
+        Args:
+            request: 推理请求。
+            model: 已解析的模型配置。
+            overrides: per-submit 覆盖配置。
+
+        Returns:
+            合并好的 headers 字典。
+        """
+        headers: dict[str, str] = {
+            "Authorization": "Bearer " + self._get_api_key(request, model, overrides),
+            "Content-Type": "application/json",
+        }
+        headers.update(self.config.extra_headers)
+        if overrides is not None:
+            headers.update(overrides.extra_headers)
+        headers.update(model.extra_headers)
+        headers.update(request.extra_headers)
+        return headers
+
+    def _build_payload(self, request: InferenceRequest, model: ModelConfig, overrides: Optional[KittyEngineOverrides]) -> dict[str, Any]:
+        """组装 OpenAI 兼容 payload。
+
+        合并顺序（后者覆盖前者）:
+            1. 骨架：`model` / `stream` / `messages`；
+            2. default_inference_parameters（overrides / config 二选一，见
+               `_resolve_default_ip`）；
+            3. `model.default_inference_parameters`；
+            4. `request.inference_parameters`；
+            5. `config.extra_payload`；
+            6. `overrides.extra_payload`；
+            7. `model.extra_payload`；
+            8. `request.extra_payload`。
+
+        Args:
+            request: 推理请求。
+            model: 已解析的模型配置。
+            overrides: per-submit 覆盖配置。
+
+        Returns:
+            合并好的 payload 字典。
+        """
+        payload: dict[str, Any] = {
+            "model": model.model_id,
+            "stream": request.stream if request.stream is not None else False,
+            "messages": [msg.to_dict() for msg in request.messages],
+        }
+        default_ip = self._resolve_default_ip(overrides)
+        if default_ip:
+            payload.update(default_ip.to_dict())
+        if model.default_inference_parameters:
+            payload.update(model.default_inference_parameters.to_dict())
+        if request.inference_parameters:
+            payload.update(request.inference_parameters.to_dict())
+        payload.update(self.config.extra_payload)
+        if overrides is not None:
+            payload.update(overrides.extra_payload)
+        payload.update(model.extra_payload)
+        payload.update(request.extra_payload)
+        return payload
+
+    def _get_wait_time(self, status_code: int, attempt: int, base_delay: int) -> float:
+        """按 HTTP 状态码决定下一次重试前的等待秒数。
+
+        Args:
+            status_code: 上一次 HTTP 响应状态码。
+            attempt: 当前已尝试的次数（0-based）。
+            base_delay: 指数退避基数。
+
+        Returns:
+            等待秒数。
+
+            - 429（限流）：固定 1 秒，避免指数退避把 QPS 压到 0；
+            - 408（请求超时）：按指数退避；
+            - 5xx（服务端错误）：指数退避 `base_delay ** attempt`；
+            - 其它 4xx：调用方应在外层判定为不可重试，不会进到这里。
+
+        Note:
+            4xx 非 429 / 408 视为客户端错误（鉴权 / 参数 / 资源不存在等），
+            重试只是浪费时间。判断逻辑由 `_send_request` 在调用本函数之前
+            完成，本函数仅给"应当重试的状态码"算等待时长。
+        """
+        if status_code == 429:
+            return 1.0
+        if status_code == 408 or status_code >= 500:
+            return float(base_delay**attempt)
+        # 理论上不可达：调用方已用 `_is_retriable_status` 把其它 4xx 拦在外面。
+        # 留作 defensive 兜底，避免未来调用点漏判时退化成热循环。
+        return 5.0
+
+    @staticmethod
+    def _is_retriable_status(status_code: int) -> bool:
+        """HTTP 状态码是否值得重试。
+
+        - 429 / 408：限流 / 请求超时，可重试。
+        - 5xx：服务端错误，可重试。
+        - 其它 4xx：客户端错误（401/403/404/422 …），重试无意义，立刻失败。
+        """
+        return status_code in (408, 429) or status_code >= 500
+
+    def _resolve_retry_params(self, overrides: Optional[KittyEngineOverrides]) -> tuple[int, int]:
+        """解析 `(max_retries, base_delay)`，overrides 中非 None 的字段覆盖 config。
+
+        Args:
+            overrides: per-submit 覆盖配置。
+
+        Returns:
+            `(max_retries, base_delay)` 二元组。
+        """
+        max_retries = self.config.default_max_retries
+        base_delay = self.config.default_base_delay
+        if overrides is not None:
+            if overrides.default_max_retries is not None:
+                max_retries = overrides.default_max_retries
+            if overrides.default_base_delay is not None:
+                base_delay = overrides.default_base_delay
+        return max_retries, base_delay
+
+    def build_request(
+        self,
+        request: InferenceRequest,
+        overrides: Optional[KittyEngineOverrides] = None,
+    ) -> PreparedRequest:
+        """组装最终发给大模型服务商的请求，不发出任何网络调用。
+
+        用于调试：验证 URL、Headers、Payload 是否符合预期。
+
+        Args:
+            request (InferenceRequest): 推理请求对象。
+            overrides (Optional[KittyEngineOverrides]): 覆盖配置，默认为 None。
+
+        Returns:
+            PreparedRequest: 包含模型名、URLs、请求头和负载的已准备请求。
+        """
+
+        model_name = request.model_name
+        if not model_name:
+            model_name = overrides.default_model if (overrides and overrides.default_model) else self.config.default_model
+        if not model_name:
+            raise ValueError("未指定 model_name 且 config 中无 default_model")
+        model = self.model_registry.get(name=model_name)
+        if model is None:
+            raise ValueError(f"模型 '{model_name}' 不在注册表中")
+        return PreparedRequest(
+            model_name=model_name,
+            url="",
+            urls=model.api_urls,
+            headers=self._build_headers(request, model, overrides),
+            payload=self._build_payload(request, model, overrides),
+        )
+
+    def mock_request(
+        self,
+        request: InferenceRequest,
+        overrides: Optional[KittyEngineOverrides] = None,
+    ) -> MockRequest:
+        """组装最终发给大模型服务商的请求，不发出任何网络调用。
+
+        用于调试：验证 URL、Headers、Payload 是否符合预期。
+
+        Args:
+            request (InferenceRequest): 推理请求对象。
+            overrides (Optional[KittyEngineOverrides]): 覆盖配置，默认为 None。
+
+        Returns:
+            MockRequest: 包含 URLs、模型名、请求头和负载的模拟请求。
+        """
+        prepared: PreparedRequest = self.build_request(request, overrides=overrides)
+        return MockRequest(
+            urls=prepared.urls,
+            model=prepared.model_name,
+            headers=prepared.headers,
+            payload=prepared.payload,
+        )
+
+    async def _send_request(
+        self,
+        request: InferenceRequest,
+        model: ModelConfig,
+        overrides: Optional[KittyEngineOverrides],
+    ) -> ModelOutput:
+        """一次完整的"带重试的 HTTP 调用 + 响应解析"。
+
+        重试策略:
+            共 `max_retries` 次尝试（非 `max_retries + 1`）。每次尝试内：
+
+            - 网络异常 (NetworkError / TimeoutException)：
+              `wait = base_delay ** attempt` 后重试；最后一次直接 raise。
+            - HTTP 不可重试状态（4xx 非 408/429）：立即 raise
+              `httpx.HTTPStatusError`，不消耗剩余 attempts（鉴权/参数错重试
+              无意义）。
+            - HTTP 可重试状态（408 / 429 / 5xx）：按 `_get_wait_time`
+              等待后重试；最后一次 raise `httpx.HTTPStatusError`。
+            - 响应 JSON 解析失败：等 `base_delay` 后重试；最后一次 raise
+              `RuntimeError`。
+            - choices 为空（疑似风控）：等 `base_delay` 后重试；最后一次 raise
+              `RuntimeError`。
+            - finish_reason != 'stop'：warning 但不重试，视为正常返回（例如
+              length、tool_calls 等）。
+
+        与 model 的协作:
+            - `model.get_url()` 在每次 attempt 内取 URL（支持多节点轮询 /
+              负载均衡）；
+            - 必须在 `finally` 中 `model.release_url(url)` 释放该节点的占用
+              计数，否则会把节点计数泄漏，最终把模型自己的并发打爆。
+
+        外部超时:
+            本方法自身不设顶层超时；由 `_run_task` 用
+            `asyncio.wait_for(..., timeout=wt.task_timeout)` 包裹，超时会转成
+            `asyncio.TimeoutError`。
+
+        Stream:
+            首版不支持 stream；若 payload `stream=True` 直接抛
+            `NotImplementedError`。
+
+        Args:
+            request: 推理请求。
+            model: 已解析的模型配置。
+            overrides: per-submit 覆盖配置。
+
+        Returns:
+            解析完成的 `ModelOutput`。
+
+        Raises:
+            NotImplementedError: payload 中 stream=True。
+            httpx.HTTPStatusError: 最后一次 attempt 仍为非 200。
+            httpx.NetworkError: 最后一次 attempt 仍为网络异常。
+            httpx.TimeoutException: 最后一次 attempt 仍为读超时。
+            RuntimeError: 最后一次 attempt 响应解析失败 / choices 为空。
+        """
+        assert self._http_client is not None
+        client = self._http_client
+
+        prepared = self.build_request(request, overrides)
+        if prepared.payload["stream"]:
+            raise NotImplementedError("KittyEngine 暂未实现 stream 模式")
+        max_retries, base_delay = self._resolve_retry_params(overrides)
+
+        for attempt in range(max_retries):
+            is_last = attempt == max_retries - 1
+            url = model.get_url()
+            try:
+                resp = await client.post(url, json=prepared.payload, headers=prepared.headers)
+                if resp.status_code != 200:
+                    body = resp.text
+                    # 不可重试的 4xx 直接抛，不浪费剩余 attempts。
+                    if not self._is_retriable_status(resp.status_code):
+                        logger.error("HTTP %d (不可重试)，响应: %s", resp.status_code, body[:500])
+                        raise httpx.HTTPStatusError(f"HTTP {resp.status_code}", request=resp.request, response=resp)
+                    if is_last:
+                        logger.error("HTTP %d，已达最大重试次数，响应: %s", resp.status_code, body[:500])
+                        raise httpx.HTTPStatusError(f"HTTP {resp.status_code}", request=resp.request, response=resp)
+                    wait_time = self._get_wait_time(resp.status_code, attempt, base_delay)
+                    logger.warning(
+                        "HTTP %d, 第 %d 次重试，等待 %.1fs... 响应: %s",
+                        resp.status_code,
+                        attempt + 1,
+                        wait_time,
+                        body[:500],
+                    )
+                    await asyncio.sleep(wait_time)
+                    continue
+
+                try:
+                    parsed = ChatCompletionResponse.model_validate_json(resp.text)
+                except ValidationError as ve:
+                    if is_last:
+                        raise RuntimeError(f"响应 JSON 解析失败，已达最大重试次数: {ve}; body: {resp.text[:500]}") from ve
+                    logger.warning("响应 JSON 解析失败，第 %d 次重试: %r", attempt + 1, ve)
+                    await asyncio.sleep(base_delay)
+                    continue
+
+                if not parsed.choices:
+                    if is_last:
+                        raise RuntimeError(f"empty choices after {max_retries} attempts, body: {resp.text[:500]}")
+                    logger.warning("resp.choices 为空，疑似风控，第 %d 次重试...", attempt + 1)
+                    await asyncio.sleep(base_delay)
+                    continue
+                choice = parsed.choices[0]
+                if choice.finish_reason and choice.finish_reason != "stop":
+                    # 不重试：length / tool_calls / content_filter 等属业务层语义，交由上层判断。
+                    logger.warning(
+                        "finish_reason='%s' (非 stop), content_len=%d, usage=%s",
+                        choice.finish_reason,
+                        len(choice.message.content or ""),
+                        parsed.usage.model_dump() if parsed.usage else None,
+                    )
+                return ModelOutput(
+                    role=choice.message.role,
+                    content=choice.message.content,
+                    reasoning=choice.message.reasoning_content,
+                    finish_reason=choice.finish_reason,
+                    usage=parsed.usage.model_dump() if parsed.usage else None,
+                )
+            except (httpx.NetworkError, httpx.TimeoutException) as e:
+                if is_last:
+                    logger.error("达到最大重试次数，最后错误: %r", e)
+                    raise
+                wait_time = base_delay**attempt
+                logger.warning("网络异常: %r, 等待 %.1fs 后重试...", e, wait_time)
+                await asyncio.sleep(wait_time)
+            finally:
+                # 必须释放节点占用，否则 ModelConfig 的负载均衡计数会泄漏。
+                model.release_url(url)
+
+        # 理论上到不了这里：要么在循环内 return，要么最后一次 attempt raise。
+        raise RuntimeError("Unexpected end of retry loop")