weijuesen
120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
import shutil
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
|
|
from fastapi import FastAPI, File, HTTPException, Request, UploadFile
|
|
|
|
from . import config
|
|
from .detector import MockDetector, ONNXDetector
|
|
from .stream_tasks import StreamTaskWorker, is_allowed_stream_ref
|
|
|
|
app = FastAPI(title="Silk AI Service", version="0.1.0")
|
|
START_TIME = time.time()
|
|
_lock = threading.Lock()
|
|
_requests = 0
|
|
_latency_total = 0.0
|
|
stream_worker = StreamTaskWorker(config.STREAM_TASK_MAX_WORKERS)
|
|
|
|
|
|
@app.middleware("http")
|
|
async def count_requests(request: Request, call_next):
|
|
global _requests, _latency_total
|
|
start = time.time()
|
|
response = await call_next(request)
|
|
latency = (time.time() - start) * 1000
|
|
with _lock:
|
|
_requests += 1
|
|
_latency_total += latency
|
|
return response
|
|
|
|
if config.MODEL_MODE == "onnx":
|
|
detector = ONNXDetector(config.MODEL_PATH, tuple(config.MODEL_LABELS))
|
|
else:
|
|
detector = MockDetector(class_name=config.MOCK_CLASS, confidence=config.MOCK_CONFIDENCE)
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok", "model": config.MODEL_MODE}
|
|
|
|
|
|
@app.post("/detect")
|
|
async def detect(file: UploadFile = File(...)):
|
|
data = await file.read()
|
|
if not data:
|
|
raise HTTPException(status_code=400, detail="图片为空")
|
|
try:
|
|
detections = detector.detect(data)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
return {
|
|
"model": config.MODEL_MODE,
|
|
"detections": [
|
|
{
|
|
"bbox": d["bbox"],
|
|
"class": d["class_name"],
|
|
"confidence": d["confidence"],
|
|
}
|
|
for d in detections
|
|
],
|
|
}
|
|
|
|
|
|
@app.post("/stream-detect")
|
|
def stream_detect():
|
|
raise HTTPException(status_code=410, detail="请使用 POST /internal/stream-tasks")
|
|
|
|
|
|
@app.post("/internal/stream-tasks", status_code=202)
|
|
def create_stream_task(request: Request, body: dict):
|
|
if request.headers.get("x-internal-key") != config.INTERNAL_API_KEY:
|
|
raise HTTPException(status_code=401, detail="invalid internal key")
|
|
|
|
task_id = (body or {}).get("taskId", "")
|
|
stream_ref = (body or {}).get("streamRef")
|
|
if not task_id or not isinstance(stream_ref, dict):
|
|
raise HTTPException(status_code=400, detail="缺少 taskId 或 streamRef")
|
|
if not is_allowed_stream_ref(stream_ref):
|
|
raise HTTPException(status_code=400, detail="streamRef 不在允许范围")
|
|
|
|
stream_worker.submit(task_id, stream_ref)
|
|
return {"taskId": task_id, "status": "queued"}
|
|
|
|
|
|
@app.get("/internal/stream-tasks/{task_id}")
|
|
def get_stream_task(request: Request, task_id: str):
|
|
if request.headers.get("x-internal-key") != config.INTERNAL_API_KEY:
|
|
raise HTTPException(status_code=401, detail="invalid internal key")
|
|
task = stream_worker.get(task_id)
|
|
if not task:
|
|
raise HTTPException(status_code=404, detail="task not found")
|
|
return task
|
|
|
|
|
|
@app.get("/metrics")
|
|
def metrics():
|
|
"""AI 服务监控（#26 骨架）：请求量/延迟/GPU 信息（nvidia-smi 可选）。"""
|
|
with _lock:
|
|
reqs = _requests
|
|
avg = _latency_total / _requests if _requests else 0.0
|
|
gpu = None
|
|
if shutil.which("nvidia-smi"):
|
|
try:
|
|
out = subprocess.run(
|
|
["nvidia-smi", "--query-gpu=name,temperature.gpu,memory.used,memory.total", "--format=csv,noheader,nounits"],
|
|
capture_output=True, text=True, timeout=10,
|
|
)
|
|
if out.returncode == 0 and out.stdout.strip():
|
|
parts = [p.strip() for p in out.stdout.strip().split(",")]
|
|
gpu = {"name": parts[0], "tempC": parts[1], "memUsedMB": parts[2], "memTotalMB": parts[3]}
|
|
except Exception:
|
|
gpu = None
|
|
return {
|
|
"model": config.MODEL_MODE,
|
|
"uptimeSeconds": int(time.time() - START_TIME),
|
|
"requests": reqs,
|
|
"avgLatencyMs": round(avg, 2),
|
|
"gpu": gpu,
|
|
}
|