AiOpsWechat/run.py

import subprocess
import sys
import os
import time
import signal

processes = []

def start_fastapi():
    """ 启动 FastAPI 服务 """
    environment = os.environ.get('environment', 'default')
    if environment == 'default':
        process = subprocess.Popen(["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "80"])
    else:
        process = subprocess.Popen(["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "5000"])
    return process

def start_celery_worker():
    """ 启动 Celery Worker """
    if sys.platform == "win32":
        process = subprocess.Popen(["celery", "-A", "celery_app", "worker", "--loglevel=info", "-P", "solo"],
                                   stdout=None, stderr=None)
    else:
        process = subprocess.Popen(["celery", "-A", "celery_app", "worker", "--loglevel=info"],
                                   stdout=None, stderr=None)
    return process

def start_celery_beat():
    """ 启动 Celery Beat，使用 RedBeat 作为调度器 """
    process = subprocess.Popen(["celery", "-A", "celery_app", "beat", "--scheduler", "redbeat.RedBeatScheduler", "--loglevel=info"],
                               stdout=None, stderr=None)
    return process

def signal_handler(sig, frame):
    """处理退出信号，确保子进程也被终止"""
    print('正在关闭所有服务...')
    for process in processes:
        process.terminate()  # 尝试优雅终止
    # 给进程一点时间来优雅终止
    time.sleep(2)
    # 检查是否有进程仍在运行，如果有则强制终止
    for process in processes:
        if process.poll() is None:  # 进程仍在运行
            process.kill()  # 强制终止
    sys.exit(0)

if __name__ == "__main__":
    # 注册信号处理程序
    signal.signal(signal.SIGINT, signal_handler)  # Ctrl+C
    signal.signal(signal.SIGTERM, signal_handler)  # 终止信号

    # 启动 Celery Worker 和 Celery Beat
    celery_worker_process = start_celery_worker()
    celery_beat_process = start_celery_beat()

    # 将 Celery 进程添加到监控列表
    processes.append(celery_worker_process)
    processes.append(celery_beat_process)

    # 启动并监控 FastAPI 进程
    while True:
        fastapi_process = start_fastapi()
        processes.append(fastapi_process)
        print("FastAPI 服务已启动。")

        # 等待 FastAPI 进程结束
        fastapi_process.wait()

        # 如果 FastAPI 进程结束，移除它并记录日志
        processes.remove(fastapi_process)
        print("FastAPI 服务意外终止，正在尝试重启...")

        # 等待一段时间再重启，避免频繁重启
        time.sleep(5)