105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
"""FastAPI 应用装配:OpenAI 路由 + 引擎调参 API + 极简 admin 页。"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import HTMLResponse
|
|
|
|
from mlx_streaming.server.admin import ADMIN_HTML
|
|
from mlx_streaming.server.openai import make_openai_router
|
|
from mlx_streaming.server.state import EngineManager, rss_bytes
|
|
from mlx_streaming.tui.backend import FakeBackend
|
|
|
|
|
|
def create_app(args, backend=None) -> FastAPI:
|
|
"""装配应用。
|
|
|
|
args: argparse Namespace 风格对象(至少含 model/k/max_tokens/expert_slots,
|
|
真后端还需要 expert_dir/mtp_out/qn_config/spec_slots)。
|
|
backend: 显式注入的后端(测试用);None 时按 SPARKLE_FAKE=1 选 FakeBackend,
|
|
否则 MLXBackend。
|
|
"""
|
|
if backend is None:
|
|
if os.environ.get("SPARKLE_FAKE") == "1":
|
|
backend = FakeBackend()
|
|
else:
|
|
from mlx_streaming.tui.backend import MLXBackend
|
|
|
|
backend = MLXBackend(args)
|
|
mgr = EngineManager(args, backend)
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app):
|
|
# 加载放线程且不等待:真模型以分钟计,不能阻塞 uvicorn 启动;
|
|
# 就绪前 loaded=False,chat 请求返回 503。
|
|
loop = asyncio.get_running_loop()
|
|
app.state.load_future = loop.run_in_executor(None, mgr.load_initial)
|
|
yield
|
|
|
|
app = FastAPI(title="sparkle server", lifespan=lifespan)
|
|
# 本机工具服务:放开 CORS,让浏览器能直连 /v1/models、/api/* 做探测
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
app.state.mgr = mgr
|
|
app.include_router(make_openai_router(mgr))
|
|
|
|
@app.get("/api/stats")
|
|
async def api_stats():
|
|
return {"tok_per_s": mgr.last_tok_per_s,
|
|
"rss_bytes": rss_bytes(),
|
|
"expert_slots": mgr.expert_slots,
|
|
"loaded": mgr.loaded,
|
|
# 专家池「分不到槽被迫落 0 号槽」的累计数。>0 = 有前向拿错专家权重算过,
|
|
# 输出不可信(调小 PREFILL_CHUNK 或调大 expert_slots)。健康值恒为 0。
|
|
"unplaced_experts": mgr.unplaced_experts(),
|
|
# 专家池命中/读盘指标:hit_rate 低就是慢的直接原因(每 miss 一个要读 ~3.3MB)
|
|
"pool": mgr.pool_metrics(),
|
|
"model": str(mgr.args.model)}
|
|
|
|
@app.get("/api/engine/config")
|
|
async def get_engine_config():
|
|
return {"expert_slots": mgr.expert_slots,
|
|
"k": mgr.k,
|
|
"max_tokens": mgr.max_tokens}
|
|
|
|
@app.post("/api/engine/config")
|
|
async def set_engine_config(req: Request):
|
|
"""有多少改多少:k/max_tokens 即时生效;expert_slots 变化触发后台重建。"""
|
|
body = await req.json()
|
|
if "k" in body:
|
|
mgr.k = int(body["k"])
|
|
if "max_tokens" in body:
|
|
mgr.max_tokens = int(body["max_tokens"])
|
|
reloading = False
|
|
if "expert_slots" in body:
|
|
slots = int(body["expert_slots"])
|
|
if slots != mgr.expert_slots:
|
|
mgr.expert_slots = slots
|
|
mgr.reloading = True # 即刻对 chat 可见(503),不等后台任务起跑
|
|
reloading = True
|
|
|
|
async def _bg_rebuild():
|
|
async with mgr.lock: # 等当前生成结束再换引擎
|
|
await asyncio.get_running_loop().run_in_executor(
|
|
None, mgr.rebuild)
|
|
|
|
asyncio.create_task(_bg_rebuild())
|
|
return {"ok": True, "reloading": reloading,
|
|
"expert_slots": mgr.expert_slots,
|
|
"k": mgr.k, "max_tokens": mgr.max_tokens}
|
|
|
|
@app.get("/admin", response_class=HTMLResponse)
|
|
async def admin_page():
|
|
return ADMIN_HTML
|
|
|
|
return app
|