TravelContentCreator/run_api.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
TravelContentCreator API 服务启动脚本（生产环境优化版）

- 基于 uvicorn 官方 API 启动 FastAPI 应用。
- 默认按「2 * CPU + 1」规则自动计算 worker 数，以充分利用多核性能。
- 可选择使用 uvloop 作为事件循环，提升 I/O 性能（如未安装 uvloop 会自动降级）。
- 支持通过环境变量或命令行参数覆盖默认配置，方便容器化部署。
"""

from __future__ import annotations

import argparse
import logging
import multiprocessing
import os
import sys
from typing import Any

import uvicorn


def build_arg_parser() -> argparse.ArgumentParser:
    """构造命令行解析器"""
    cpu_cores = multiprocessing.cpu_count()
    default_workers = cpu_cores * 2 + 1

    parser = argparse.ArgumentParser(description="TravelContentCreator API 服务（生产环境）")
    parser.add_argument("--host", default=os.getenv("HOST", "0.0.0.0"), help="监听主机地址")
    parser.add_argument(
        "--port",
        type=int,
        default=int(os.getenv("PORT", 2714)),
        help="监听端口",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=int(os.getenv("WORKERS", default_workers)),
        help="worker 进程数，默认 2*CPU + 1",
    )
    parser.add_argument(
        "--log-level",
        default=os.getenv("LOG_LEVEL", "info"),
        choices=[
            "critical",
            "error",
            "warning",
            "info",
            "debug",
            "trace",
        ],
        help="日志级别",
    )
    parser.add_argument(
        "--proxy-headers",
        action="store_true",
        help="信任代理头(例如部署在 Nginx / Traefik 之后时开启)",
    )
    parser.add_argument(
        "--loop",
        default=os.getenv("LOOP", "uvloop"),
        choices=["uvloop", "asyncio"],
        help="事件循环实现(uvloop 未安装时自动降级到 asyncio)",
    )
    parser.add_argument(
        "--timeout-keep-alive",
        type=int,
        default=int(os.getenv("TIMEOUT_KEEP_ALIVE", 5)),
        help="Keep-Alive 超时时间 (秒)",
    )
    return parser


def setup_logging(log_level: str = "info") -> None:
    """初始化日志配置"""
    logging.basicConfig(
        level=getattr(logging, log_level.upper(), logging.INFO),
        format="%(asctime)s | %(levelname)s | %(process)d | %(name)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        stream=sys.stdout,
    )


def maybe_install_uvloop(enable: bool) -> None:
    """若选择 uvloop 并且可用，则安装之"""
    if not enable:
        return
    try:
        import uvloop  # type: ignore

        uvloop.install()
        logging.getLogger(__name__).info("uvloop 已启用")
    except ModuleNotFoundError:
        logging.getLogger(__name__).warning("未检测到 uvloop，降级使用默认 asyncio 事件循环")


def run() -> None:  # noqa: C901  # (函数较长但在此处可接受)
    parser = build_arg_parser()
    args = parser.parse_args()

    setup_logging(args.log_level)
    maybe_install_uvloop(args.loop == "uvloop")

    logging.getLogger(__name__).info(
        "启动 API 服务: http://%s:%d | workers=%d | loop=%s",
        args.host,
        args.port,
        args.workers,
        args.loop,
    )

    uvicorn.run(
        "api.main:app",
        host=args.host,
        port=args.port,
        workers=args.workers,
        proxy_headers=args.proxy_headers,
        log_level=args.log_level,
        timeout_keep_alive=args.timeout_keep_alive,  # type: ignore[arg-type]
        loop="uvloop" if args.loop == "uvloop" else "asyncio",
    )


if __name__ == "__main__":
    run()