diff --git a/main.py b/main.py index 30a022d..09f8132 100644 --- a/main.py +++ b/main.py @@ -1,3 +1,6 @@ +import signal +import tempfile +import time from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from contextlib import asynccontextmanager @@ -18,79 +21,6 @@ LoggerManager.setup_logging() logger = get_main_logger() -def start_celery_worker(): - """启动 Celery Worker""" - try: - logger.info("🌿 启动Celery Worker...") - - # 构建Worker启动命令 - 独立进程启动 - worker_cmd = [ - sys.executable, "-m", "celery", - "-A", "app.celery_app", # 指定celery应用模块 - "worker", - '--loglevel=info', - '--pool=solo', - '--concurrency=1', - '--time-limit=300', # 5分钟任务超时 - '--soft-time-limit=240' # 4分钟软超时 - ] - subprocess.run(worker_cmd, check=True) - except Exception as e: - logger.error(f"❌ Celery Worker 启动失败: {e}") - return None - - -def start_celery_beat(): - """启动 Celery Beat 调度器""" - try: - logger.info("📅 启动Celery Beat调度器...") - - # 使用跨平台的调度文件路径 - import tempfile - import os - schedule_file = os.path.join(tempfile.gettempdir(), 'celerybeat-schedule') - - # 构建Beat启动命令 - 独立进程启动 - beat_cmd = [ - sys.executable, "-m", "celery", - "-A", "app.celery_app", # 指定celery应用模块 - "beat", - '--loglevel=info', - f'--schedule={schedule_file}' - ] - - subprocess.run(beat_cmd, 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}") - - @asynccontextmanager async def lifespan(app: FastAPI): # 启动时初始化 @@ -211,55 +141,134 @@ async def health_check(): return {"status": "healthy"} -if __name__ == "__main__": - import uvicorn - import argparse +# 全局变量存储进程 +processes = [] + +def signal_handler(signum, frame): + """信号处理器，用于优雅关闭所有服务""" + print(f"\n🛑 接收到信号 {signum}，正在关闭所有服务...") - 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监控)") - args = parser.parse_args() + # 逆序关闭进程（最后启动的最先关闭） + for i, process in enumerate(reversed(processes)): + if process and process.poll() is None: # 进程仍在运行 + print(f"🔄 正在关闭进程 {len(processes) - i}...") + try: + process.terminate() # 发送 SIGTERM 信号 + process.wait(timeout=10) # 等待最多10秒 + print(f"✅ 进程已关闭") + except subprocess.TimeoutExpired: + print(f"⚠️ 进程未在10秒内响应，强制关闭...") + process.kill() # 强制杀死进程 + except Exception as e: + print(f"❌ 关闭进程时出错: {e}") - # 如果没有传递任何参数，输出完整的提示信息 - if not args.mode: - 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(" 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("👋 所有服务已关闭") + sys.exit(0) + +def start_all_services(): + """启动所有服务""" + global processes + + print("\n🚀 AI Talk Callback API 一键启动所有服务") + print("=" * 60) + + # 注册信号处理器 + signal.signal(signal.SIGINT, signal_handler) # Ctrl+C + signal.signal(signal.SIGTERM, signal_handler) # 终止信号 + + try: + # 1. 启动 Celery Worker + print("🌿 启动 Celery Worker...") + worker_process = subprocess.Popen([ + sys.executable, "-m", "celery", + "-A", "app.celery_app", + "worker", + '--loglevel=info', + '--pool=solo', + '--concurrency=1', + '--time-limit=300', # 5分钟任务超时 + '--soft-time-limit=240' # 4分钟软超时 + ]) + processes.append(worker_process) + time.sleep(2) # 等待 Worker 启动 + + # 2. 启动 Celery Beat + print("\n📅 启动 Celery Beat...") + beat_process = subprocess.Popen([ + sys.executable, "-m", "celery", + "-A", "app.celery_app", + "beat", + '--loglevel=info', + f'--schedule={os.path.join(tempfile.gettempdir(), "celerybeat-schedule")}' + ]) + processes.append(beat_process) + time.sleep(2) # 等待 Beat 启动 + + # 3. 启动 Flower 监控（如果启用） + if settings.flower_enabled: + print("\n📊 启动 Flower 监控服务...") + flower_cmd = [ + sys.executable, "-m", "celery", + "-A", "app.celery_app", + 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}") + if settings.flower_url_prefix: + flower_cmd.append(f"--url_prefix={settings.flower_url_prefix}") + + flower_process = subprocess.Popen(flower_cmd) + processes.append(flower_process) + time.sleep(2) # 等待 Flower 启动 + + # 4. 启动 FastAPI 应用 + print("\n🚀 启动 FastAPI 应用...") + api_cmd = [ + sys.executable, "-m", "uvicorn", + "main:app", + "--host", "0.0.0.0", + "--port", "8000" + ] + + # 添加调试模式（如果配置了） + if settings.debug: + api_cmd.append("--reload") + + api_process = subprocess.Popen(api_cmd) + processes.append(api_process) + time.sleep(2) # 等待 API 启动 + + print("\n" + "=" * 60) + print("✅ 所有服务启动完成！") 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.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() + print(" 🚀 API 服务: http://localhost:8000") + if settings.environment != "production" and not settings.disable_docs: + print(" 📖 API 文档: http://localhost:8000/docs") + if settings.flower_enabled: + print(f" 📊 监控界面: {settings.flower_url}") + print("\n💡 使用 Ctrl+C 可以优雅关闭所有服务") + print("=" * 60) + + # 等待所有进程 + while True: + # 检查是否有进程异常退出 + for i, process in enumerate(processes): + if process and process.poll() is not None: + print(f"❌ 进程 {i+1} 异常退出，退出码: {process.returncode}") + signal_handler(signal.SIGINT, None) + return + + time.sleep(1) # 每秒检查一次 + + except KeyboardInterrupt: + signal_handler(signal.SIGINT, None) + except Exception as e: + print(f"❌ 启动服务时出错: {e}") + signal_handler(signal.SIGINT, None) + +if __name__ == "__main__": + # 直接启动所有服务 + start_all_services()