run_server.py

import os
import sys
import math
import time
import signal
import sdnotify  # Systemd 通知库
import threading
import uvicorn
import logging
import subprocess
import multiprocessing
from pathlib import Path
from dotenv import load_dotenv

load_dotenv(dotenv_path=Path(__file__).parent / ".env")  # 默认override=False
import torch  # CUDA_VISIBLE_DEVICES must be set before importing torch


# 允许分配器利用碎片内存扩展段
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

# 配置与日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("LTX-Launcher")


# =========================
# 配置管理
# =========================
class AppConfig:
    MODEL_PATH = os.environ.get(
        "MODEL_PATH", "/home/yufangjie/code/Z-Image/ckpts/Z-Image-Turbo"
    )
    DEFAULT_YAML = "LTX-Video/configs/ltxv-13b-0.9.8-distilled.yaml"
    HOST = os.environ.get("HOST", "0.0.0.0")
    PORT = int(os.environ.get("PORT", 8000))
    MANUAL_WORKERS = os.environ.get("MANUAL_WORKERS", None)  # 允许手动覆盖

    # 计数器文件路径，用于多进程协调 GPU 分配 (需与 ltx-video_api.py 一致)
    COUNTER_FILE = os.environ.get("COUNTER_FILE", "worker_counter.lock")
    COUNTER_FILE = Path(COUNTER_FILE).resolve().as_posix()

    # 假设服务支持的最大分辨率 (像素数) 用于估算显存
    MAX_PIXELS_LIMIT = 2048 * 2048

    # 单卡显存预留 (GB)
    SYSTEM_RESERVED_GB = 2.0


def check_env_vars():
    """检查必要的环境变量"""
    key = os.environ.get("SERVER_API_KEY")
    if not key:
        logger.warning("=" * 60)
        logger.warning("WARNING: 'SERVER_API_KEY' not found in environment variables!")

        # 方案 A: 强行报错退出 (生产环境推荐)
        logger.error("Service cannot start without a secure API Key.")
        sys.exit(1)

        # 方案 B: 自动生成一个临时 Key (测试方便)
        temp_key = secrets.token_urlsafe(32)
        os.environ["SERVER_API_KEY"] = temp_key
        logger.warning(f"Auto-generated temporary API Key: {temp_key}")
        logger.warning("Please set a persistent key in Systemd or .env for production.")
        logger.warning("=" * 60)
    else:
        # 处于安全考虑，只打印前几位
        masked_key = key[:4] + "*" * (len(key) - 4)
        logger.info(f"API Key loaded: {masked_key}")


class MultiGPUManager:
    _mps_started_by_me = False

    @staticmethod
    def init_counter_file():
        """初始化计数器文件，清零"""
        with open(AppConfig.COUNTER_FILE, "w") as f:
            f.write("0")
        # 赋予权限，防止权限问题
        os.chmod(AppConfig.COUNTER_FILE, 0o666)
        logger.info(f"Initialized worker counter at {AppConfig.COUNTER_FILE}")

    @classmethod
    def ensure_mps(cls):
        try:
            #  ps aux | grep mps
            os.system("pkill -f nvidia-cuda-mps-control || true")
            # os.environ.pop("CUDA_MPS_PIPE_DIRECTORY", None)
            # os.environ.pop("CUDA_MPS_LOG_DIRECTORY", None)
        except Exception as e:
            logger.error(f"Failed to clear existing MPS processes: {e}")
            sys.exit(1)

        """启动 MPS (对多卡环境，只需启动一次控制守护进程)"""
        if not torch.cuda.is_available():
            logger.warning("CUDA not available. Skipping MPS start.")
            return

        # === 配置 MPS 日志路径到当前用户目录下 ===
        mps_log_dir = os.path.join(os.getcwd(), "mps_logs")
        os.makedirs(mps_log_dir, exist_ok=True)
        # 设置环境变量，告诉 MPS 不要写 /var/log，而是写这里
        os.environ["CUDA_MPS_LOG_DIRECTORY"] = mps_log_dir

        # 检查 MPS 是否已在运行
        check_cmd = "ps -ef | grep nvidia-cuda-mps-control | grep -v grep"
        rc = subprocess.call(check_cmd, shell=True, stdout=subprocess.DEVNULL)

        if rc == 0:
            logger.info("NVIDIA MPS is already running.")
            cls._mps_started_by_me = False  # 不是我启动的，我不负责关闭
            return

        logger.info("NVIDIA MPS not found. Starting MPS daemon...")
        try:
            # 启动守护进程
            subprocess.run(["nvidia-cuda-mps-control", "-d"], check=True)
            logger.info("NVIDIA MPS started successfully.")
            cls._mps_started_by_me = True  # 【标记】是我启动的

            # (可选) 如果需要为每个 GPU 设置专有 MPS 管道，通常不需要，默认配置即可
            # 只要 CUDA_VISIBLE_DEVICES 包含了所有卡，MPS 就会接管
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to start MPS: {e}")
            logger.warning("Performance may be degraded in multi-process mode.")

    @staticmethod
    def _dry_run_measure(max_width=1280, max_height=720):
        """
        子进程执行函数：加载模型 -> Warmup -> 测量峰值 -> 放入队列
        """
        # print(f"dry run resolutions: {max_width}*{max_height}")

        # 延迟导入，确保在子进程中初始化
        from diffusers import ZImagePipeline

        # 1. 初始化 Context
        # torch.cuda.reset_peak_memory_stats()
        torch.cuda.empty_cache()

        # 2. 加载模型 (BF16)
        pipe = ZImagePipeline.from_pretrained(
            AppConfig.MODEL_PATH, torch_dtype=torch.bfloat16
        ).to("cuda")

        # 3. Warmup (极限分辨率测试)
        with torch.inference_mode():
            _ = pipe(
                prompt="benchmark",
                width=max_width,
                height=max_height,
                num_inference_steps=1,
            )

        # 4. 获取峰值显存 (MB)
        peak_bytes = torch.cuda.max_memory_allocated()
        peak_mb = peak_bytes / 1024 / 1024
        return peak_mb

    @classmethod
    def _calculate_workers_in_subprocess(cls, queue, width=1280, height=720):
        """计算所有卡能承载的总 Workers 数"""
        try:
            if not torch.cuda.is_available():
                logger.warning("No GPU found. Defaulting to 1 worker.")
                queue.put(1)
                return

            gpu_count = torch.cuda.device_count()
            visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
            logger.info(f"Detected {gpu_count} GPUs. {visible_devices}")

            # peak_mb = 20000  # 假设测出来峰值是 20000 MB
            try:
                t0 = time.time()
                peak_mb = cls._dry_run_measure(width, height)
                logger.info(
                    f"Dry run resolutions {width}*{height} took {time.time() - t0:.2f}s, Peak VRAM: {peak_mb:.0f} MB"
                )
            except Exception as e:
                logger.error(f"Dry run measurement failed: {e}")
                queue.put(1)
                return

            # 计算逻辑
            total_mem_mb = (
                torch.cuda.get_device_properties(0).total_memory / 1024 / 1024
            )
            reserved_mb = AppConfig.SYSTEM_RESERVED_GB * 1024
            available_mb = total_mem_mb - reserved_mb

            # 增加 10% 的 Buffer 应对显存碎片
            buffer = 0.05
            safe_peak_mb = peak_mb * (1 + buffer)

            optimal_workers = math.floor(available_mb / safe_peak_mb)
            optimal_workers = max(1, optimal_workers)

            logger.info("-" * 40)
            logger.info(f"GPU Total     : {total_mem_mb:.0f} MB")
            logger.info(f"Peak/Worker   : {peak_mb:.0f} MB (Allocated)")
            logger.info(
                f"Safe Estimate : {safe_peak_mb:.0f} MB (w/ {buffer * 100}% buffer)"
            )
            logger.info(f"Avail for App : {available_mb:.0f} MB")
            logger.info(f"Auto-Calc     : {optimal_workers} Workers")
            logger.info("-" * 40)

            total_workers = optimal_workers * gpu_count
            queue.put(total_workers)
        except Exception as e:
            logger.error(f"Error calculating workers: {e}")
            queue.put(1)  # 兜底：每卡至少跑1个

    @classmethod
    def calculate_total_workers(cls, width=1280, height=720) -> int:
        """
        启动一个干净的 Spawn 进程来计算 Worker 数，
        确保主进程不加载 CUDA Context。
        """
        if AppConfig.MANUAL_WORKERS:
            logger.info(
                f"Using manual worker count override: {AppConfig.MANUAL_WORKERS}"
            )
            return int(AppConfig.MANUAL_WORKERS)

        logger.info("Starting dry-run process to measure VRAM usage...")

        ctx = multiprocessing.get_context("spawn")
        queue = ctx.Queue()

        # 启动子进程进行计算
        p = ctx.Process(
            target=cls._calculate_workers_in_subprocess,
            args=(queue, width, height),
        )
        p.start()
        p.join()

        # 获取结果
        if not queue.empty():
            try:
                workers = queue.get()
            except Exception as e:
                logger.error(
                    f"Dry run failed: {e}. Defaulting to safe value (1 worker)."
                )
                return 1
            finally:
                # 显式关闭队列，回收信号量资源
                queue.close()
                queue.join_thread()
        else:
            logger.error("Dry run process returned no result. Defaulting to 1 worker.")
            return 1

        return workers

    @classmethod
    def stop_mps(cls):
        """关闭 MPS (用于服务退出时清理)"""
        if cls._mps_started_by_me:
            logger.info("Stopping NVIDIA MPS...")
            try:
                subprocess.run(
                    ["echo quit | nvidia-cuda-mps-control"], shell=True, check=False
                )
            except Exception:
                pass
        else:
            logger.info("Skipping MPS shutdown (Not started by this script).")


def watchdog_thread(notifier):
    interval = 180  # 默认值
    if os.environ.get("WATCHDOG_USEC"):
        timeout = int(os.environ.get("WATCHDOG_USEC")) / 1000000
        interval = max(timeout / 2, 30.0)  # interval 下限保护, 避免过于频繁 notify
        logger.info(
            f"[Systemd] Thread Watchdog active. Timeout={timeout}s, Ping every {interval:.1f}s\n"
        )
    else:
        logger.info("[Systemd] No watchdog detected")
    while True:
        # 发送心跳
        try:
            notifier.notify("WATCHDOG=1")
            logger.info("[Systemd] watchdog Heartbeat sent.")
        except Exception as e:
            logger.error(f"[Systemd] watchdog Heartbeat failed: {e}")

        # 线程睡眠
        time.sleep(interval)


def enable_systemd_watchdog():
    # 1. 发送 READY 信号 (告诉 Systemd 启动已完成)
    notifier = sdnotify.SystemdNotifier()
    notifier.notify("READY=1")
    # 2. 启动心跳线程
    # daemon=True 表示主程序退出时，这个线程也会自动退出
    wd_thread = threading.Thread(target=watchdog_thread, args=(notifier,), daemon=True)
    wd_thread.start()
    logger.info("[Systemd] Main process watchdog started.")


def main():
    # 注册信号处理 (Graceful Shutdown)
    def signal_handler(sig, frame):
        logger.info(f"Received signal {sig}. Shutting down...")
        MultiGPUManager.stop_mps()
        sys.exit(0)

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    try:
        # 1. 检查 Key
        check_env_vars()

        # 2. 初始化环境
        logger.info(f"Initializing Service. DEFAULT CONFIG: {AppConfig.DEFAULT_YAML}")
        MultiGPUManager.init_counter_file()
        # MultiGPUManager.ensure_mps()  # 使用不当可能会干扰device id mapping

        # 3. 计算总 Worker
        workers = MultiGPUManager.calculate_total_workers(2048, 2048)
        logger.info(
            f"Starting Server on {AppConfig.HOST}:{AppConfig.PORT} with Total {workers} Workers across all GPUs."
        )

        if os.environ.get("WATCHDOG_USEC"):
            enable_systemd_watchdog()

        # 4. 启动 Uvicorn (会阻塞主线程，直到服务停止)
        uvicorn.run(
            "ltx-video_api:app",
            host=AppConfig.HOST,
            port=AppConfig.PORT,
            workers=workers,
            log_level="info",  # "info" or "warning"
            timeout_keep_alive=30,
            access_log=True,
        )
    except Exception as e:
        logger.critical(f"Service startup failed: {e}", exc_info=True)
        MultiGPUManager.stop_mps()
        sys.exit(1)
    finally:
        # 正常退出时的清理
        MultiGPUManager.stop_mps()


if __name__ == "__main__":
    main()