pyai/ai_server.py

import logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(name)s] %(levelname)s: %(message)s'
)

from contextlib import asynccontextmanager
from fastapi import FastAPI
from route import gpu, file, audio, text, image
import service.ai_task as ai_task

logger = logging.getLogger(__name__)
# 获取 uvicorn 的 logger
# logger = logging.getLogger("uvicorn.error")

@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("🚀 pyai 服务已启动...")
    await ai_task.start_worker()
    yield
    logger.info("🛑 pyai 服务已停止")
app = FastAPI(title="GPU Worker Server", lifespan=lifespan)

# 挂载子路由
app.include_router(gpu.router)
app.include_router(file.router)
app.include_router(audio.router)
app.include_router(text.router)
app.include_router(image.router)

@app.get("/")
async def root():
    return {"message": "Welcome to pyai"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7070)