354 lines
13 KiB
Python
354 lines
13 KiB
Python
from fastapi import FastAPI
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from contextlib import asynccontextmanager
|
||
from sqlalchemy import text
|
||
import redis.asyncio as redis
|
||
import subprocess
|
||
import sys
|
||
import os
|
||
|
||
from app.config import settings
|
||
from app.database import engine
|
||
from app.routes import router
|
||
from app.logger import LoggerManager, get_main_logger
|
||
|
||
|
||
# 初始化日志系统
|
||
LoggerManager.setup_logging()
|
||
logger = get_main_logger()
|
||
|
||
# Celery Worker 进程管理
|
||
celery_worker_process = None
|
||
|
||
|
||
def start_celery_worker():
|
||
"""启动 Celery Worker"""
|
||
try:
|
||
logger.info("🌿 启动Celery Worker...")
|
||
|
||
# 使用subprocess启动独立的celery worker进程
|
||
subprocess.run([
|
||
sys.executable, "-m", "celery",
|
||
"-A", "app.celery_app", # 指定celery应用模块
|
||
"worker",
|
||
'--loglevel=info',
|
||
'--concurrency=4',
|
||
'--prefetch-multiplier=1',
|
||
'--max-tasks-per-child=1000',
|
||
'--time-limit=300', # 5分钟任务超时
|
||
'--soft-time-limit=240', # 4分钟软超时
|
||
], check=True)
|
||
except Exception as e:
|
||
logger.error(f"❌ Celery Worker 启动失败: {e}")
|
||
|
||
|
||
def start_celery_beat():
|
||
"""启动 Celery Beat 调度器"""
|
||
try:
|
||
logger.info("📅 启动Celery Beat调度器...")
|
||
|
||
# 使用跨平台的调度文件路径
|
||
import tempfile
|
||
import os
|
||
schedule_file = os.path.join(tempfile.gettempdir(), 'celerybeat-schedule')
|
||
|
||
# 使用subprocess启动独立的celery beat进程
|
||
subprocess.run([
|
||
sys.executable, "-m", "celery",
|
||
"-A", "app.celery_app", # 指定celery应用模块
|
||
"beat",
|
||
'--loglevel=info',
|
||
f'--schedule={schedule_file}',
|
||
], check=True)
|
||
except Exception as e:
|
||
logger.error(f"❌ Celery Beat 启动失败: {e}")
|
||
|
||
|
||
def start_flower():
|
||
"""启动 Flower 监控服务"""
|
||
try:
|
||
logger.info("📊 启动Flower监控服务...")
|
||
|
||
# 构建Flower启动命令 - 独立进程启动
|
||
flower_cmd = [
|
||
sys.executable, "-m", "celery",
|
||
"-A", "app.celery_app", # 指定celery应用模块
|
||
f"--broker={settings.celery_broker_url}",
|
||
"flower",
|
||
f"--port={settings.flower_port}"
|
||
]
|
||
|
||
# 添加基础认证（如果配置了）
|
||
if settings.flower_basic_auth:
|
||
flower_cmd.append(f"--basic_auth={settings.flower_basic_auth}")
|
||
|
||
# 添加URL前缀（如果配置了）
|
||
if settings.flower_url_prefix:
|
||
flower_cmd.append(f"--url_prefix={settings.flower_url_prefix}")
|
||
|
||
subprocess.run(flower_cmd, check=True)
|
||
except Exception as e:
|
||
logger.error(f"❌ Flower 监控服务启动失败: {e}")
|
||
|
||
|
||
def stop_celery_service(service_name):
|
||
"""停止 Celery 服务"""
|
||
try:
|
||
logger.info(f"🛑 正在停止 {service_name} 服务...")
|
||
|
||
# 查找并停止相关进程
|
||
if service_name == "worker":
|
||
# 停止 Celery Worker
|
||
subprocess.run([
|
||
sys.executable, "-m", "celery",
|
||
"-A", "app.celery_app",
|
||
"control",
|
||
"shutdown"
|
||
], check=False)
|
||
logger.info("✅ Celery Worker 停止命令已发送")
|
||
|
||
elif service_name == "beat":
|
||
# 查找并停止 beat 进程
|
||
try:
|
||
if os.name == 'nt': # Windows
|
||
# 使用 PowerShell 查找并停止相关进程
|
||
subprocess.run([
|
||
"powershell", "-Command",
|
||
"Get-Process python | Where-Object {$_.ProcessName -eq 'python' -and $_.MainWindowTitle -like '*beat*'} | Stop-Process -Force"
|
||
], check=False)
|
||
# 备用方案：查找包含 celery beat 的进程
|
||
subprocess.run([
|
||
"powershell", "-Command",
|
||
"Get-WmiObject Win32_Process | Where-Object {$_.Name -eq 'python.exe' -and $_.CommandLine -like '*beat*'} | ForEach-Object {Stop-Process -Id $_.ProcessId -Force}"
|
||
], check=False)
|
||
else: # Linux/Mac
|
||
subprocess.run([
|
||
"pkill", "-f", "celery.*beat"
|
||
], check=False)
|
||
|
||
logger.info(f"✅ Celery Beat 停止命令已发送")
|
||
except Exception as e:
|
||
logger.warning(f"⚠️ 停止 {service_name} 时出现警告: {e}")
|
||
|
||
elif service_name == "flower":
|
||
# 查找并停止 flower 进程
|
||
try:
|
||
if os.name == 'nt': # Windows
|
||
# 使用 PowerShell 查找并停止相关进程
|
||
subprocess.run([
|
||
"powershell", "-Command",
|
||
"Get-Process python | Where-Object {$_.ProcessName -eq 'python' -and $_.MainWindowTitle -like '*flower*'} | Stop-Process -Force"
|
||
], check=False)
|
||
# 备用方案：查找包含 celery flower 的进程
|
||
subprocess.run([
|
||
"powershell", "-Command",
|
||
"Get-WmiObject Win32_Process | Where-Object {$_.Name -eq 'python.exe' -and $_.CommandLine -like '*flower*'} | ForEach-Object {Stop-Process -Id $_.ProcessId -Force}"
|
||
], check=False)
|
||
else: # Linux/Mac
|
||
subprocess.run([
|
||
"pkill", "-f", "celery.*flower"
|
||
], check=False)
|
||
|
||
logger.info(f"✅ Flower 监控服务停止命令已发送")
|
||
except Exception as e:
|
||
logger.warning(f"⚠️ 停止 {service_name} 时出现警告: {e}")
|
||
|
||
elif service_name == "all":
|
||
# 停止所有 Celery 相关服务
|
||
logger.info("🛑 正在停止所有 Celery 服务...")
|
||
stop_celery_service("worker")
|
||
stop_celery_service("beat")
|
||
stop_celery_service("flower")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 停止 {service_name} 服务失败: {e}")
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
# 启动时初始化
|
||
logger.info("🚀 应用启动中...")
|
||
|
||
# Redis 连接对象
|
||
redis_client = None
|
||
|
||
try:
|
||
# 验证数据库连接
|
||
logger.info("📊 验证数据库连接...")
|
||
try:
|
||
async with engine.begin() as conn:
|
||
# 执行简单的查询来验证连接
|
||
result = await conn.execute(text("SELECT 1"))
|
||
connection_test = result.scalar()
|
||
if connection_test == 1:
|
||
logger.info("✅ 数据库连接验证成功")
|
||
else:
|
||
raise Exception("数据库连接测试失败")
|
||
except Exception as db_error:
|
||
logger.error(f"❌ 数据库连接验证失败: {db_error}")
|
||
raise
|
||
|
||
# 验证 Redis 连接
|
||
logger.info("🔴 验证 Redis 连接...")
|
||
try:
|
||
redis_client = redis.from_url(settings.celery_broker_url)
|
||
# 执行 ping 命令验证连接
|
||
await redis_client.ping()
|
||
logger.info("✅ Redis 连接验证成功")
|
||
|
||
# 存储到应用状态中供其他组件使用
|
||
app.state.redis_client = redis_client
|
||
|
||
except Exception as redis_error:
|
||
logger.error(f"❌ Redis 连接验证失败: {redis_error}")
|
||
raise
|
||
|
||
logger.info("✅ 服务初始化完成，FastAPI应用启动中...")
|
||
|
||
logger.info(f"🎉 {settings.app_name} 启动完成!")
|
||
yield
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ 应用启动失败: {e}")
|
||
raise
|
||
|
||
finally:
|
||
# 关闭时清理
|
||
logger.info("🛑 应用关闭中...")
|
||
|
||
# 关闭 Redis 连接
|
||
if redis_client:
|
||
try:
|
||
await redis_client.close()
|
||
logger.info("🔴 Redis 连接已关闭")
|
||
except Exception as e:
|
||
logger.warning(f"⚠️ 关闭 Redis 连接时出现警告: {e}")
|
||
|
||
logger.info("👋 应用已关闭")
|
||
|
||
|
||
app = FastAPI(
|
||
title=settings.app_name,
|
||
version="1.0.0",
|
||
lifespan=lifespan,
|
||
docs_url=(
|
||
"/docs"
|
||
if not settings.disable_docs and settings.environment != "production"
|
||
else None
|
||
),
|
||
redoc_url=(
|
||
"/redoc"
|
||
if not settings.disable_docs and settings.environment != "production"
|
||
else None
|
||
),
|
||
openapi_url="/openapi.json" if not settings.disable_docs else None,
|
||
)
|
||
|
||
# 添加CORS中间件
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
# 注册路由
|
||
app.include_router(router)
|
||
|
||
|
||
@app.get("/")
|
||
async def root():
|
||
# 检查环境，生产环境下禁用根接口
|
||
if settings.environment == "production":
|
||
logger.warning("🚫 生产环境下禁止访问根接口")
|
||
from fastapi import HTTPException
|
||
|
||
raise HTTPException(status_code=404, detail="Not Found")
|
||
|
||
logger.info("📝 根接口被访问")
|
||
return {"message": f"Welcome to {settings.app_name}"}
|
||
|
||
|
||
@app.get("/health")
|
||
async def health_check():
|
||
# 检查环境，生产环境下禁用健康检查接口
|
||
if settings.environment == "production":
|
||
logger.warning("🚫 生产环境下禁止访问健康检查接口")
|
||
from fastapi import HTTPException
|
||
|
||
raise HTTPException(status_code=404, detail="Not Found")
|
||
|
||
logger.debug("💓 健康检查接口被访问")
|
||
|
||
return {"status": "healthy"}
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import uvicorn
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="AI Talk Callback API")
|
||
parser.add_argument("--mode", choices=["api", "worker", "beat", "flower"],
|
||
help="启动模式: api(仅API), worker(仅Celery Worker), beat(仅Celery Beat), flower(仅Flower监控)")
|
||
parser.add_argument("--stop", choices=["worker", "beat", "flower", "all"],
|
||
help="停止服务: worker(停止Worker), beat(停止Beat), flower(停止Flower), all(停止所有Celery服务)")
|
||
args = parser.parse_args()
|
||
|
||
# 如果没有传递任何参数，输出完整的提示信息
|
||
if not args.mode and not args.stop:
|
||
print("🚀 AI Talk Callback API 管理指南")
|
||
print("=" * 50)
|
||
print("\n📋 启动模式:")
|
||
print(" api - 启动 FastAPI Web 应用服务 (端口: 8000)")
|
||
print(" worker - 启动 Celery Worker 任务处理器")
|
||
print(" beat - 启动 Celery Beat 定时任务调度器")
|
||
print(" flower - 启动 Flower 监控服务")
|
||
print("\n🛑 停止服务:")
|
||
print(" worker - 停止 Celery Worker 服务")
|
||
print(" beat - 停止 Celery Beat 调度器")
|
||
print(" flower - 停止 Flower 监控服务")
|
||
print(" all - 停止所有 Celery 相关服务")
|
||
print("\n🔧 使用示例:")
|
||
print(" python main.py --mode=api # 启动 Web API 服务")
|
||
print(" python main.py --mode=worker # 启动任务处理器")
|
||
print(" python main.py --mode=beat # 启动定时任务调度器")
|
||
print(f" python main.py --mode=flower # 启动监控服务 (访问: {settings.flower_url})")
|
||
print("\n python main.py --stop=worker # 停止 Worker 服务")
|
||
print(" python main.py --stop=beat # 停止 Beat 调度器")
|
||
print(" python main.py --stop=flower # 停止 Flower 监控")
|
||
print(" python main.py --stop=all # 停止所有 Celery 服务")
|
||
print("\n🌐 服务地址:")
|
||
print(" API 服务: http://localhost:8000")
|
||
print(" API 文档: http://localhost:8000/docs")
|
||
print(f" 任务监控界面: {settings.flower_url}")
|
||
print("\n💡 提示:")
|
||
print(" - 请确保 Redis 和 PostgreSQL 服务已启动")
|
||
print(" - 生产环境请根据需要调整配置文件")
|
||
print(" - 建议在多个终端中分别启动不同服务")
|
||
sys.exit(0)
|
||
|
||
# 处理停止服务请求
|
||
if args.stop:
|
||
logger.info(f"🛑 正在停止 {args.stop} 服务...")
|
||
stop_celery_service(args.stop)
|
||
sys.exit(0)
|
||
|
||
# 处理启动服务请求
|
||
if args.mode == "api":
|
||
# 仅启动 FastAPI 应用
|
||
logger.info("🚀 启动FastAPI应用...")
|
||
uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=settings.debug)
|
||
elif args.mode == "worker":
|
||
# 仅启动 Celery Worker
|
||
logger.info("🌿 启动Celery Worker...")
|
||
start_celery_worker()
|
||
elif args.mode == "beat":
|
||
# 仅启动 Celery Beat
|
||
logger.info("📅 启动Celery Beat调度器...")
|
||
start_celery_beat()
|
||
elif args.mode == "flower":
|
||
# 仅启动 Flower 监控服务
|
||
logger.info("📊 启动Flower监控服务...")
|
||
start_flower()
|