"""MiniCPMO45 推理 Worker 每个 Worker 占用一张 GPU，持有一个 UnifiedProcessor 实例， 提供 Chat (HTTP) / Streaming (WebSocket) / Duplex (WebSocket) 三种推理 API。 启动方式： cd /user/sunweiyue/lib/swy-dev/minicpmo45_service CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. .venv/base/bin/python worker.py \\ --port 10031 \\ --model-path /path/to/base_model \\ --pt-path /path/to/custom.pt \\ --ref-audio-path /path/to/ref.wav """ import json import time import asyncio import argparse import logging import base64 import threading from typing import Optional, List, Dict, Any import numpy as np import uvicorn from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect from pydantic import BaseModel, Field from worker_state import WorkerState, WorkerStatus from runtime.protocol import DEFAULT_WORKER_CAPABILITIES from runtime.session import BackendRuntimeSession logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", ) logger = logging.getLogger("worker") # ============ 请求/响应模型 ============ class WorkerHealthResponse(BaseModel): """健康检查响应""" status: str worker_status: WorkerStatus gpu_id: int model_loaded: bool current_ticket_id: Optional[str] = None total_requests: int = 0 avg_inference_time_ms: float = 0.0 kv_cache_length: int = 0 # 当前 LLM KV cache token 总数 capabilities: List[str] = Field(default_factory=list) # ============ FastAPI 应用 ============ worker: Optional[Any] = None # 启动参数（通过 main() 传入） WORKER_CONFIG: Dict[str, Any] = {} class RemoteBackendWorker: """Worker host used when inference lives in backend_server.py.""" def __init__(self, *, backend_server_url: str, gpu_id: int = 0) -> None: self.backend_server_url = backend_server_url self.gpu_id = gpu_id self.processor = None self.state = WorkerState(status=WorkerStatus.IDLE) def metrics(self) -> Dict[str, Any]: return {"backend": "backend_server", "backend_server_url": self.backend_server_url} def shutdown(self) -> None: return None def _backend_server_url() -> Optional[str]: value = WORKER_CONFIG.get("backend_server_url") return str(value).rstrip("/") if value else None def _input_payload(message: Dict[str, Any]) -> Dict[str, Any]: value = message.get("input") if isinstance(value, dict): return value raise RuntimeError("input.append must carry an object `input`") def _init_payload(message: Dict[str, Any]) -> Dict[str, Any]: value = message.get("payload") if isinstance(value, dict): return value raise RuntimeError("session.init must carry an object `payload`") def _event_payload(event: Any) -> Dict[str, Any]: payload = dict(getattr(event, "payload", {}) or {}) raw_event = payload.get("event") if isinstance(raw_event, dict): return raw_event return payload @asynccontextmanager async def lifespan(app: FastAPI): """应用生命周期：启动时加载模型""" global worker config = WORKER_CONFIG backend_server_url = _backend_server_url() if backend_server_url: worker = RemoteBackendWorker( backend_server_url=backend_server_url, gpu_id=int(config.get("gpu_id", 0) or 0), ) logger.info("Worker running as backend-server runtime host: %s", backend_server_url) else: from core.processors.backend_factory import create_backend worker = create_backend(config) # 模型加载是同步操作（~15s），在线程中执行避免阻塞 await asyncio.to_thread(worker.load_model) try: yield finally: logger.info("Worker shutting down") if worker is not None: await asyncio.to_thread(worker.shutdown) app = FastAPI(title="MiniCPMO45 Worker", lifespan=lifespan) # ========== 健康检查 ========== @app.get("/health", response_model=WorkerHealthResponse) async def health(): """健康检查""" if worker is None: return WorkerHealthResponse( status="initializing", worker_status=WorkerStatus.LOADING, gpu_id=0, model_loaded=False, capabilities=[], ) avg_time = 0.0 if worker.state.total_requests > 0: avg_time = worker.state.total_inference_time_ms / worker.state.total_requests worker_metrics = worker.metrics() kv_len = int(worker_metrics.get("kv_cache_length", 0) or 0) remote_backend_url = _backend_server_url() model_loaded = bool(remote_backend_url) or worker.processor is not None or bool(getattr(worker.state, "is_idle", False)) return WorkerHealthResponse( status="healthy" if model_loaded else "error", worker_status=worker.state.status, gpu_id=worker.gpu_id, model_loaded=model_loaded, current_ticket_id=worker.state.current_ticket_id, total_requests=worker.state.total_requests, avg_inference_time_ms=avg_time, kv_cache_length=kv_len, capabilities=DEFAULT_WORKER_CAPABILITIES, ) async def _handle_remote_backend_runtime_ws( ws: WebSocket, *, mode: str, active_status: WorkerStatus, idle_status: WorkerStatus, ) -> None: """Bridge worker runtime WebSocket to backend_server.py.""" backend_url = _backend_server_url() if backend_url is None: await ws.close(code=1013, reason="Backend server URL is not configured") return if worker is None: await ws.close(code=1013, reason="Worker not ready") return if not worker.state.is_idle: await ws.close(code=1013, reason=f"Worker busy: {worker.state.status.value}") return await ws.accept() worker.state.status = active_status runtime = BackendRuntimeSession( backend_base_url=backend_url, mode=mode, ) backend_closed = False async def _send_runtime_event(event: Any) -> Dict[str, Any]: payload = _event_payload(event) await ws.send_json(payload) return payload try: first = json.loads(await ws.receive_text()) first_type = str(first.get("type") or "") pending_input: Optional[Dict[str, Any]] = None if first_type == "session.init": init_params = _init_payload(first) elif first_type == "input.append": init_params = {"mode": mode} pending_input = _input_payload(first) else: raise RuntimeError(f"first message must initialize or push input, got: {first_type}") init_params = dict(init_params) init_params.setdefault("mode", mode) await _send_runtime_event(await runtime.init(init_params)) if pending_input is not None: await runtime.push(pending_input) async def client_to_backend() -> None: nonlocal backend_closed async for raw in ws.iter_text(): msg = json.loads(raw) msg_type = str(msg.get("type") or "") if msg_type == "input.append": await runtime.push(_input_payload(msg)) continue if msg_type == "input.context": context = msg.get("context") if not isinstance(context, dict): raise RuntimeError( "input.context must carry an object `context`" ) await runtime.push({ "input_id": msg.get("input_id"), "context": context, }) continue if msg_type == "session.close": close_event = await runtime.unary("close", {"reason": str(msg.get("reason") or "client_closed")}) backend_closed = True close_payload = _event_payload(close_event) if close_payload.get("type") != "session.closed": close_payload = { "type": "session.closed", "session_id": runtime.session_id, "reason": msg.get("reason", "client_closed"), } await ws.send_json(close_payload) await ws.close(code=1000, reason="client_closed") return raise RuntimeError(f"unsupported runtime message type: {msg_type}") async def backend_to_client() -> None: nonlocal backend_closed while not backend_closed: event = await runtime.pull() payload = await _send_runtime_event(event) if payload.get("type") == "session.closed": backend_closed = True return done, pending = await asyncio.wait( [ asyncio.create_task(client_to_backend()), asyncio.create_task(backend_to_client()), ], return_when=asyncio.FIRST_COMPLETED, ) for task in pending: task.cancel() for task in done: task.result() except WebSocketDisconnect: logger.info("Remote backend runtime WebSocket disconnected") except Exception as exc: logger.error("Remote backend runtime failed: error=%s", exc, exc_info=True) try: await ws.send_json({ "type": "session.closed", "session_id": runtime.session_id, "reason": "backend_error", }) except Exception: pass finally: try: if not backend_closed: await runtime.unary("close", {"reason": "worker_disconnected"}) except Exception: logger.exception("Remote backend runtime cleanup failed") worker.state.status = idle_status worker.state.current_ticket_id = None try: await ws.close() except Exception: pass @app.websocket("/v1/worker/chat") async def worker_chat_runtime_ws(ws: WebSocket): """Worker-internal turn-based chat runtime protocol (backend-server only).""" await _handle_remote_backend_runtime_ws( ws, mode="turn_based", active_status=WorkerStatus.BUSY_CHAT, idle_status=WorkerStatus.IDLE, ) # ========== Duplex WebSocket ========== @app.websocket("/v1/worker/duplex") async def worker_duplex_runtime_ws(ws: WebSocket): """Worker-internal duplex runtime protocol (backend-server only). This endpoint is meant for gateway-worker communication and uses runtime event payloads instead of page/demo-shaped result messages. """ await _handle_remote_backend_runtime_ws( ws, mode="full_duplex", active_status=WorkerStatus.DUPLEX_ACTIVE, idle_status=WorkerStatus.IDLE, ) # ============ 缓存状态查询 ========== @app.get("/cache_info") async def cache_info(): """查询当前 Worker 的 KV Cache 状态""" if worker is None: raise HTTPException(status_code=503, detail="Worker not ready") return { "status": worker.state.status.value, "note": "KV cache state is now tracked by Gateway (cached_hash on WorkerConnection)", } @app.post("/clear_cache") async def clear_cache(): """手动清除 KV Cache（重置 Streaming 模型 session）""" if worker is None: raise HTTPException(status_code=503, detail="Worker not ready") worker.reset_half_duplex_session() return {"success": True, "message": "Cache cleared"} # ============ 入口 ============ def main(): from config import get_config cfg = get_config() parser = argparse.ArgumentParser(description="MiniCPMO45 Worker") parser.add_argument("--port", type=int, default=None, help=f"Worker port (default: from config, base={cfg.worker_base_port})") parser.add_argument("--host", type=str, default="0.0.0.0", help="Host") parser.add_argument("--model-path", type=str, default=None, help="Base model path") parser.add_argument("--pt-path", type=str, default=None, help="Custom weights path (.pt)") parser.add_argument("--ref-audio-path", type=str, default=None, help="Default ref audio path") parser.add_argument("--gpu-id", type=int, default=None, help="GPU ID (inferred from port if not set)") parser.add_argument("--worker-index", type=int, default=0, help="Worker index (0, 1, 2, ...)") parser.add_argument("--duplex-pause-timeout", type=float, default=None, help="Duplex pause timeout (s)") parser.add_argument("--backend-server-url", type=str, default=None, help="Remote backend_server.py base URL") args = parser.parse_args() port = args.port or cfg.worker_port(args.worker_index) gpu_id = args.gpu_id if args.gpu_id is not None else args.worker_index WORKER_CONFIG.update({ "model_path": args.model_path or cfg.model.model_path, "gpu_id": gpu_id, "pt_path": args.pt_path or cfg.model.pt_path, "ref_audio_path": args.ref_audio_path or cfg.ref_audio_path, "duplex_pause_timeout": args.duplex_pause_timeout or cfg.duplex_pause_timeout, "backend_server_url": args.backend_server_url, "compile": cfg.compile, "chat_vocoder": cfg.chat_vocoder, "attn_implementation": cfg.attn_implementation, }) logger.info(f"Starting Worker on port {port}, GPU {gpu_id}") # Bump WS max payload from uvicorn's 16 MiB default to 128 MiB so that # base64-encoded video attachments (commonly 30-60 MiB after inflation) # can be received without the connection being torn down with code 1009. uvicorn.run(app, host=args.host, port=port, ws_max_size=128 * 1024 * 1024) if __name__ == "__main__": main()