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

llm_engine/general_engine.py

@@ -0,0 +1,562 @@
+# llm_engine/general_engine.py
+
+"""
+GeneralEngine：自包含的任务化推理引擎。
+
+- 后台维护一个 asyncio 事件循环（独立 daemon 线程）
+- submit() 返回 TaskHandle，非阻塞
+- wait()/get() 消费结果；poll() 非消费快照
+- inference()/batch_inference() 为同步便捷接口
+- 不依赖 SimpleEngine / SimpleCoroutineEngine
+
+详细设计见 docs/general_engine_design.md
+"""
+
+import asyncio
+import httpx
+import threading
+import time
+
+import kitty_logger
+from pydantic import ValidationError
+from tqdm import tqdm
+from typing import IO, Any, TypeAlias
+
+from .model_config import ModelConfig, ModelConfigRegistry
+from .schemas import (
+    ChatCompletionChoice,
+    ChatCompletionResponse,
+    InferenceParameters,
+    InferenceRequest,
+    InferenceRequestResult,
+    Message,
+    MessageRole,
+    ModelOutput,
+    PreparedRequest,
+    TaskHandle,
+    TaskStatus,
+)
+from .utils import gen_unique_id
+
+logger = kitty_logger.getLogger(__name__)
+
+
+Key: TypeAlias = str | TaskHandle
+
+
+class GeneralEngine:
+
+    def __init__(
+        self,
+        model_registry: ModelConfigRegistry,
+        default_model: str,
+        default_api_key: str | None = None,
+        default_inference_parameters: InferenceParameters | None = None,
+        extra_headers: dict[str, str] | None = None,
+        extra_payload: dict[str, Any] | None = None,
+        stream: bool = False,
+        max_global_concurrency: int = 64,
+        default_timeout: float | None = None,
+        default_max_retries: int = 5,
+        default_base_delay: int = 2,
+    ) -> None:
+        # —— 静态配置 ——
+        self.model_registry: ModelConfigRegistry = model_registry
+        self.default_model_name: str = default_model
+        self.default_model: ModelConfig = model_registry.get(name=default_model)
+        self.default_api_key: str | None = default_api_key
+        self.default_inference_parameters: InferenceParameters | None = default_inference_parameters
+        self.extra_headers: dict[str, str] = extra_headers or {}
+        self.extra_payload: dict[str, Any] = extra_payload or {}
+        self.stream: bool = stream
+        self.max_global_concurrency: int = max_global_concurrency
+        self.default_timeout: float | None = default_timeout
+        self.default_max_retries: int = default_max_retries
+        self.default_base_delay: int = default_base_delay
+
+        # —— 运行时（由 setup 填充）——
+        self._loop: asyncio.AbstractEventLoop | None = None
+        self._loop_thread: threading.Thread | None = None
+        self._running: bool = False
+
+        self._global_sem: asyncio.Semaphore | None = None
+
+        self._handles: dict[str, TaskHandle] = {}
+        self._handles_lock: threading.Lock = threading.Lock()
+
+    # ------------------------------------------------------------------
+    # 生命周期
+    # ------------------------------------------------------------------
+
+    def setup(self) -> None:
+        """启动后台事件循环线程。幂等。"""
+        if self._running:
+            logger.warning(msg="GeneralEngine 已启动，请勿重复调用 setup()")
+            return
+
+        ready: threading.Event = threading.Event()
+
+        def _run() -> None:
+            loop: asyncio.AbstractEventLoop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            self._loop = loop
+            # 在 loop 线程中创建 asyncio 原语
+            self._global_sem = asyncio.Semaphore(value=self.max_global_concurrency)
+            ready.set()
+            try:
+                loop.run_forever()
+            finally:
+                # run_forever 结束后，把所有待清理的生成器/任务关掉
+                try:
+                    tasks: set[asyncio.Task[Any]] = asyncio.all_tasks(loop)
+                    for t in tasks:
+                        t.cancel()
+                    if tasks:
+                        loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
+                finally:
+                    loop.close()
+
+        self._loop_thread = threading.Thread(target=_run, name="GeneralEngine-Loop", daemon=True)
+        self._loop_thread.start()
+        ready.wait()
+        self._running = True
+        logger.info("GeneralEngine 已启动 (max_global_concurrency=%d)", self.max_global_concurrency)
+
+    def teardown(self, cancel_running: bool = True, timeout: float = 10.0) -> None:
+        """停止后台事件循环，回收线程。"""
+        if not self._running:
+            return
+
+        assert self._loop is not None
+        assert self._loop_thread is not None
+
+        if cancel_running:
+            with self._handles_lock:
+                handles: list[TaskHandle] = list[TaskHandle](self._handles.values())
+            for h in handles:
+                if not h.is_finished() and h._asyncio_task is not None:
+                    self._loop.call_soon_threadsafe(callback=h._asyncio_task.cancel)
+
+        self._loop.call_soon_threadsafe(callback=self._loop.stop)
+        self._loop_thread.join(timeout=timeout)
+        self._running = False
+        logger.info("GeneralEngine 已停止")
+
+    def is_alive(self) -> bool:
+        return self._running and self._loop_thread is not None and self._loop_thread.is_alive()
+
+    def __enter__(self) -> "GeneralEngine":
+        self.setup()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.teardown()
+
+    # ------------------------------------------------------------------
+    # 提交与执行
+    # ------------------------------------------------------------------
+
+    def submit(
+        self,
+        request: InferenceRequest,
+        *,
+        persist: bool = False,
+        timeout: float | None = None,
+    ) -> TaskHandle:
+        if not self._running:
+            raise RuntimeError("GeneralEngine 未启动，请先调用 setup()")
+
+        task_id: str = gen_unique_id(prefix="task")
+
+        # 超时层级：request.timeout > submit timeout > default
+        effective_timeout: float | None = request.timeout if request.timeout is not None else (timeout if timeout is not None else self.default_timeout)
+
+        handle: TaskHandle = TaskHandle(
+            task_id=task_id,
+            request=request,
+            status=TaskStatus.SUBMITTED,
+            task_timeout=effective_timeout,
+            persist=persist,
+        )
+        # threading.Event 在任意线程都能创建
+        handle._done_threading_event = threading.Event()
+
+        with self._handles_lock:
+            self._handles[task_id] = handle
+
+        assert self._loop is not None
+        loop = self._loop
+
+        # 在 loop 线程里创建 Task，并同步回填到 handle。用 threading.Event 等回执，
+        # 确保 submit() 返回后 handle._asyncio_task 一定可用（否则 cancel() 竞态）。
+        bound: threading.Event = threading.Event()
+
+        def _schedule() -> None:
+            try:
+                handle._asyncio_task = loop.create_task(self._run_task(handle))
+                handle._done_event = asyncio.Event()
+            finally:
+                bound.set()
+
+        loop.call_soon_threadsafe(_schedule)
+        bound.wait()
+        return handle
+
+    async def _run_task(self, handle: TaskHandle) -> None:
+        request = handle.request
+        # —— 验证阶段：失败 → REJECTED ——
+        try:
+            model_name: str = request.model_name or self.default_model_name
+            if not model_name:
+                raise ValueError("no model_name (request/default all empty)")
+            model: ModelConfig = self.model_registry.get(name=model_name)
+        except Exception as e:
+            handle.status = TaskStatus.REJECTED
+            handle.error = f"{type(e).__name__}: {e}"
+            handle.error_exception = e
+            handle.result = InferenceRequestResult(
+                success=False,
+                task_id=handle.task_id,
+                request=request,
+                error_message=handle.error,
+            )
+            handle.end_time = time.time()
+            if handle._done_event is not None:
+                handle._done_event.set()
+            if handle._done_threading_event is not None:
+                handle._done_threading_event.set()
+            return
+
+        # 验证通过：ACCEPTED
+        handle.status = TaskStatus.ACCEPTED
+
+        assert self._global_sem is not None
+        try:
+            # 入队等 sem
+            handle.status = TaskStatus.PENDING
+            async with self._global_sem:
+                handle.status = TaskStatus.RUNNING
+                handle.start_time = time.time()
+                try:
+                    # 用 asyncio.timeout 而非 asyncio.wait_for：超时直接抛
+                    # TimeoutError，不经内层 CancelledError 上传，因此下面的
+                    # CancelledError 分支可确定是"外部 cancel"，避免双终态写入。
+                    async with asyncio.timeout(handle.task_timeout):
+                        model_output = await self._send_request(request, model)
+                    handle.result = InferenceRequestResult(
+                        success=True,
+                        task_id=handle.task_id,
+                        request=request,
+                        model_output=model_output,
+                        duration=time.time() - handle.start_time,
+                    )
+                    handle.status = TaskStatus.SUCCESS
+                except asyncio.CancelledError:
+                    handle.status = TaskStatus.CANCELLED
+                    handle.error = "cancelled"
+                    handle.result = InferenceRequestResult(
+                        success=False,
+                        task_id=handle.task_id,
+                        request=request,
+                        error_message="cancelled",
+                    )
+                    logger.info("任务 %s 已取消", handle.task_id)
+                    raise
+                except TimeoutError:
+                    handle.status = TaskStatus.TIMEOUT
+                    handle.error = "timeout"
+                    handle.result = InferenceRequestResult(
+                        success=False,
+                        task_id=handle.task_id,
+                        request=request,
+                        error_message="timeout",
+                    )
+                except Exception as e:
+                    logger.exception("任务 %s 执行异常", handle.task_id)
+                    handle.status = TaskStatus.FAILED
+                    handle.error = f"{type(e).__name__}: {e}"
+                    handle.error_exception = e
+                    handle.result = InferenceRequestResult(
+                        success=False,
+                        task_id=handle.task_id,
+                        request=request,
+                        error_message=handle.error,
+                    )
+        finally:
+            handle.end_time = time.time()
+            if handle._done_event is not None:
+                handle._done_event.set()
+            if handle._done_threading_event is not None:
+                handle._done_threading_event.set()
+
+    # ------------------------------------------------------------------
+    # 查询与控制
+    # ------------------------------------------------------------------
+
+    def _resolve(self, key: Key) -> TaskHandle:
+        if isinstance(key, TaskHandle):
+            return key
+        if isinstance(key, str):
+            with self._handles_lock:
+                h: TaskHandle | None = self._handles.get(key)
+            if h is None:
+                raise KeyError(f"未找到 task_id='{key}' 的任务")
+            return h
+        raise TypeError(f"key 必须是 str 或 TaskHandle，实际: {type(key).__name__}")
+
+    def _maybe_consume(self, handle: TaskHandle) -> None:
+        """终态 + 非 persist 时从 _handles 删除。"""
+        if handle.is_finished() and not handle.persist:
+            with self._handles_lock:
+                self._handles.pop(handle.task_id, None)
+
+    def poll(self, key: Key) -> TaskHandle:
+        """非消费快照：返回 handle 当前状态，不删除。"""
+        return self._resolve(key)
+
+    def get(self, key: Key) -> InferenceRequestResult | None:
+        """非阻塞消费：未完成返回 None；完成则返回 result 并（若非 persist）删除。"""
+        h: TaskHandle = self._resolve(key)
+        if not h.is_finished():
+            return None
+        result: Any | None = h.result
+        self._maybe_consume(handle=h)
+        return result
+
+    def wait(self, key: Key, timeout: float | None = None) -> InferenceRequestResult:
+        """阻塞到终态，返回 result。等待超时抛 TimeoutError（任务仍在后台运行）。"""
+        h: TaskHandle = self._resolve(key)
+        if h._done_threading_event is None:
+            raise RuntimeError("handle 缺少 done event，可能未经过 submit")
+        finished: bool = h._done_threading_event.wait(timeout=timeout)
+        if not finished:
+            raise TimeoutError(f"等待任务 {h.task_id} 超时（{timeout}s），任务仍在后台运行")
+        result = h.result
+        self._maybe_consume(handle=h)
+        # 终态后 result 必然已被 _run_task 赋值（SUCCESS/FAILED/CANCELLED 均覆盖）；
+        # 若仍为 None 则说明引擎内部状态异常，直接把失败信息返回给调用方。
+        if result is None:
+            return InferenceRequestResult(
+                success=False,
+                task_id=h.task_id,
+                request=h.request,
+                error_message=h.error or "unknown error (result missing)",
+            )
+        return result
+
+    def cancel(self, key: Key, force: bool = False) -> bool:
+        h: TaskHandle = self._resolve(key)
+        if h.is_finished():
+            return False
+        if h.status == TaskStatus.RUNNING and not force:
+            return False
+        if h._asyncio_task is not None and self._loop is not None:
+            self._loop.call_soon_threadsafe(callback=h._asyncio_task.cancel)
+        # 注意：实际状态更新在 _run_task 的 except CancelledError 分支完成
+        return True
+
+    def pop(self, key: Key) -> InferenceRequestResult | None:
+        """强制从 _handles 移除（无论 persist），返回 result（若有）。"""
+        h: TaskHandle = self._resolve(key)
+        with self._handles_lock:
+            self._handles.pop(h.task_id, None)
+        return h.result
+
+    def clear_done(self) -> int:
+        """清理所有终态 handle（含 persist）。返回清理个数。"""
+        with self._handles_lock:
+            done_ids: list[str] = [tid for tid, h in self._handles.items() if h.is_finished()]
+            for tid in done_ids:
+                self._handles.pop(tid, None)
+        return len(done_ids)
+
+    # ------------------------------------------------------------------
+    # 同步便捷接口
+    # ------------------------------------------------------------------
+
+    def inference(self, request: InferenceRequest) -> InferenceRequestResult:
+        handle: TaskHandle = self.submit(request)
+        # task_timeout 在 _run_task 的 asyncio.wait_for 里生效，此处 threading 侧不再重复设置，
+        # 避免两个超时竞争导致偶发 TimeoutError
+        return self.wait(handle)
+
+    def infer(self, query: str) -> str:
+        result: InferenceRequestResult = self.inference(request=InferenceRequest(messages=[Message(role=MessageRole.USER, content=query)]))
+        if result.success and result.model_output is not None:
+            return result.model_output.content
+        return ""
+
+    def batch_inference(
+        self,
+        requests: list[InferenceRequest],
+        output_file: str | None = None,
+        silent_mode: bool = False,
+    ) -> list[InferenceRequestResult]:
+        handles: list[TaskHandle] = [self.submit(request=req) for req in requests]
+
+        pbar: tqdm | None = tqdm(total=len(requests), desc=f"Batch: {self.default_model_name}") if not silent_mode else None
+        f: IO[str] | None = open(output_file, mode="w", encoding="utf-8") if output_file else None
+        results: list[InferenceRequestResult] = []
+        try:
+            for h in handles:
+                result: InferenceRequestResult = self.wait(key=h)
+                results.append(result)
+                if f is not None:
+                    f.write(result.model_dump_json() + "\n")
+                    f.flush()
+                if pbar is not None:
+                    pbar.update(n=1)
+        finally:
+            if pbar is not None:
+                pbar.close()
+            if f is not None:
+                f.close()
+        return results
+
+    def batch_infer(self, queries: list[str], **kwargs: Any) -> list[str]:
+        reqs: list[InferenceRequest] = [InferenceRequest(messages=[Message(role=MessageRole.USER, content=q)]) for q in queries]
+        results: list[InferenceRequestResult] = self.batch_inference(requests=reqs, **kwargs)
+        return [r.model_output.content if r.success and r.model_output else "" for r in results]
+
+    # ------------------------------------------------------------------
+    # HTTP 核心（自包含）
+    # ------------------------------------------------------------------
+
+    def _get_api_key(self, request: InferenceRequest, model: ModelConfig) -> str:
+        if request.api_key is not None:
+            return request.api_key
+        if model.api_key is not None:
+            return model.api_key
+        if self.default_api_key is not None:
+            return self.default_api_key
+        logger.warning("没有找到可用的 api_key, 将传递空 api_key")
+        return ""
+
+    def _build_headers(self, request: InferenceRequest, model: ModelConfig) -> dict[str, str]:
+        headers: dict[str, str] = {
+            "Authorization": "Bearer " + self._get_api_key(request, model),
+            "Content-Type": "application/json",
+        }
+        headers.update(self.extra_headers)
+        headers.update(model.extra_headers)
+        headers.update(request.extra_headers)
+        return headers
+
+    def _build_payload(self, request: InferenceRequest, model: ModelConfig) -> dict[str, Any]:
+        payload: dict[str, Any] = {
+            "model": model.model_id,
+            "stream": request.stream if request.stream is not None else self.stream,
+            "messages": [msg.to_dict() for msg in request.messages],
+        }
+        if self.default_inference_parameters:
+            payload.update(self.default_inference_parameters.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.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:
+        if status_code == 429:
+            return 1.0
+        if status_code >= 500:
+            return float(base_delay**attempt)
+        return 5.0
+
+    async def _send_request(self, request: InferenceRequest, model: ModelConfig) -> ModelOutput:
+        """
+        发送请求，带重试。每次 attempt 重新选 URL（支持多节点故障转移）。
+        仅实现非流式；流式留待后续。
+        """
+        headers: dict[str, str] = self._build_headers(request, model)
+        payload: dict[str, Any] = self._build_payload(request, model)
+        stream = payload["stream"]
+        max_retries: int = self.default_max_retries
+        base_delay: int = self.default_base_delay
+
+        if stream:
+            raise NotImplementedError("GeneralEngine 首版暂未实现 stream 模式")
+
+        async with httpx.AsyncClient(timeout=httpx.Timeout(timeout=10.0, read=3600.0)) as client:
+            for attempt in range(max_retries):
+                is_last = attempt == max_retries - 1
+                url = model.get_url()
+                try:
+                    resp = await client.post(url, json=payload, headers=headers)
+                    if resp.status_code != 200:
+                        body = resp.text
+                        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":
+                        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:
+                    # HTTPStatusError 仅会在上面 is_last 分支抛出，不再在此捕获
+                    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:
+                    model.release_url(url)
+
+            raise RuntimeError("Unexpected end of retry loop")
+
+    # ------------------------------------------------------------------
+    # 预留：B 方案（async-native），首版不实现
+    # ------------------------------------------------------------------
+
+    async def a_submit(self, request: InferenceRequest, *, persist: bool = False, timeout: float | None = None) -> TaskHandle:
+        raise NotImplementedError("async-native 接口尚未实现")
+
+    async def a_wait(self, key: Key, timeout: float | None = None) -> InferenceRequestResult:
+        raise NotImplementedError("async-native 接口尚未实现")
+
+    async def a_inference(self, request: InferenceRequest) -> InferenceRequestResult:
+        raise NotImplementedError("async-native 接口尚未实现")
+
+    async def a_batch_inference(self, requests: list[InferenceRequest]) -> list[InferenceRequestResult]:
+        raise NotImplementedError("async-native 接口尚未实现")

llm_engine/kitty/__init__.py

@@ -0,0 +1,8 @@
+# llm_engine/kitty/__init__.py
+
+"""KittyEngine：独立进程化的 LLM 推理引擎。"""
+
+from .client import ClientTaskHandle, KittyClient
+from .server import KittyServer
+
+__all__ = ["KittyClient", "KittyServer", "ClientTaskHandle"]

llm_engine/kitty/__main__.py

@@ -0,0 +1,46 @@
+# llm_engine/kitty/__main__.py
+
+"""
+CLI 入口。
+
+用法:
+  python -m llm_engine.kitty -c /path/to/kitty_engine.yaml
+"""
+
+import argparse
+import asyncio
+import logging
+import sys
+
+import kitty_logger
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(prog="python -m llm_engine.kitty", description="KittyEngine 后端 server")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("-c", "--config", help="YAML 配置文件路径")
+    group.add_argument("--embedded", action="store_true", help="embedded 模式：从 stdin 读取 JSON config，stdin EOF 时自动退出")
+    args = parser.parse_args()
+
+    from .config import KittyEngineConfig
+
+    if args.embedded:
+        import json
+
+        config_dict = json.loads(sys.stdin.readline())
+        config = KittyEngineConfig.model_validate(config_dict)
+    else:
+        config = KittyEngineConfig.from_yaml(args.config)
+
+    level = getattr(logging, config.log_level.upper(), logging.INFO)
+    kitty_logger.setup_logging(level=level)
+
+    from .server import KittyServer
+
+    server = KittyServer(config)
+    asyncio.run(server.run(mode="embedded" if args.embedded else "standalone"))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())