model-watchdog/main.py

"""model-watchdog 入口：一个 loop，按 interval_minutes 反复执行两个检测模块。

    python main.py                # 常驻循环
    python main.py --once         # 只跑一轮
    python main.py --once --dry-run   # 只检测并打印将要发送的告警，不真发

每轮执行：
    模块 A  check_offline_notice  收件箱下线通知检测 -> 转发原邮件 + 钉钉（已发送标题去重）
    模块 B  check_model_health    模型可用性检测 -> 聚合告警邮件 + 钉钉
"""
import argparse
import logging
import sys
import time
from datetime import datetime, timedelta

from config import ENDPOINTS, enabled_probe_targets, load_config
from mailer import MailChecker
from notifier import dual_alert, dual_forward
from probes import probe_all

log = logging.getLogger("watchdog")


# ========== 模块 A：模型下线通知检测 ==========

def check_offline_notice(cfg, dry_run=False):
    """检测收件箱下线通知，未转发过的立即转发 + 钉钉推送。返回本轮转发条数。"""
    checker = MailChecker(cfg["mail"], cfg["offline_notice"])
    try:
        notices = checker.find_offline_notices()
    except Exception as e:  # noqa: BLE001 - IMAP 异常不能中断整轮
        log.error("[模块A] 邮箱检测失败: %s", e)
        return 0

    if not notices:
        log.info("[模块A] 无需要转发的模型下线通知")
        return 0

    for n in notices:
        dual_forward(cfg, n, dry_run=dry_run)
        log.info("[模块A] 已处理下线通知: %s", n["message_id"])

    return len(notices)


# ========== 模块 B：模型可用性检测 ==========

def check_model_health(cfg, dry_run=False):
    """全量探测配置的模型，把失败的聚合成一条双重告警。返回失败模型数。"""
    targets = enabled_probe_targets(cfg)
    if not targets:
        log.info("[模块B] 未配置任何待探测模型")
        return 0

    results = probe_all(targets, cfg["api_key"])
    failed = [r for r in results if not r["ok"]]
    log.info("[模块B] 探测 %d 个模型，失败 %d 个", len(results), len(failed))
    if not failed:
        return 0

    segments = []
    for r in failed:
        code = r["code"] or "未知错误"
        message = r["message"] or "无错误描述"
        raw = r["raw"] or "（无响应体）"
        segments.append(
            f"{r['model']}（{r['label']}）报错 {code} {message}，完整响应 {raw}"
        )

    body = "\n".join([
        "【模型可用性告警】",
        "",
        f"检测时间: {datetime.now():%Y-%m-%d %H:%M:%S}",
        f"探测总数: {len(results)}　失败: {len(failed)}",
        "",
        "；\n".join(segments),
        "",
        "—— model-watchdog 自动告警",
    ])
    subject = f"【模型可用性告警】{len(failed)} 个模型不可用：" + "、".join(
        r["model"] for r in failed
    )
    dual_alert(cfg, subject, body, dry_run=dry_run)
    return len(failed)


# ========== 单轮 + 主循环 ==========

def run_once(cfg, dry_run=False, skip_mail=False, skip_probe=False):
    started = datetime.now()
    log.info("========== 开始第 %s 轮检测 ==========", started.strftime("%Y-%m-%d %H:%M:%S"))
    notices = 0 if skip_mail else check_offline_notice(cfg, dry_run=dry_run)
    failures = 0 if skip_probe else check_model_health(cfg, dry_run=dry_run)
    log.info(
        "========== 本轮结束，耗时 %.1fs，转发下线通知 %d 条，不可用模型 %d 个 ==========",
        (datetime.now() - started).total_seconds(), notices, failures,
    )


def run_loop(cfg, **kwargs):
    interval = cfg["interval_minutes"] * 60
    while True:
        try:
            run_once(cfg, **kwargs)
        except Exception:  # noqa: BLE001 - 单轮任何异常都不能让守护进程退出
            log.exception("本轮检测出现未捕获异常，跳过本轮")
        next_at = datetime.now() + timedelta(seconds=interval)
        log.info("休眠 %d 分钟，下次检测 %s", cfg["interval_minutes"], next_at.strftime("%Y-%m-%d %H:%M:%S"))
        time.sleep(interval)


def main():
    parser = argparse.ArgumentParser(description="阿里云百炼模型下线 & 可用性监控")
    parser.add_argument("-c", "--config", default="config.yaml", help="配置文件路径")
    parser.add_argument("--once", action="store_true", help="只执行一轮后退出")
    parser.add_argument("--dry-run", action="store_true", help="只检测并打印告警内容，不真正发送")
    parser.add_argument("--skip-mail", action="store_true", help="跳过模块 A（只测模型）")
    parser.add_argument("--skip-probe", action="store_true", help="跳过模块 B（只测邮箱）")
    parser.add_argument("--test-dingtalk", action="store_true",
                        help="只给钉钉发一条测试消息后退出（用于单独联调机器人）")
    args = parser.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    cfg = load_config(args.config)

    if args.test_dingtalk:
        from notifier import send_dingtalk, _to_markdown
        title = "【model-watchdog】钉钉连通性测试"
        body = "\n".join([
            "这是一条测试消息，说明钉钉通道配置正确。",
            "",
            f"发送时间: {datetime.now():%Y-%m-%d %H:%M:%S}",
            f"监控邮箱: {cfg['mail']['username']}",
        ])
        try:
            send_dingtalk(cfg["dingtalk"], title, _to_markdown(title, body))
            log.info("钉钉测试消息已发送，请到群里确认")
            return 0
        except Exception as e:  # noqa: BLE001
            log.error("钉钉测试失败: %s", e)
            return 1

    summary = "、".join(
        f"{ENDPOINTS[name]['label']}×{len(g['models'])}"
        for name, g in cfg["models"].items()
        if g["enabled"] and g["models"]
    )
    log.info("配置加载完成: 每 %d 分钟一轮 | 待测模型 %s | 邮箱 %s",
             cfg["interval_minutes"], summary or "无", cfg["mail"]["username"])
    log.info("下线通知匹配条件: 发件人域名=%s 关键词=%s 回溯=%d 天",
             "、".join(cfg["offline_notice"]["sender_domains"]) or "(不限，有被外部邮件触发的风险)",
             cfg["offline_notice"]["keyword"], cfg["offline_notice"]["lookback_days"])
    if args.dry_run:
        log.info("DRY-RUN 模式：不会真正发送任何告警")

    opts = {"dry_run": args.dry_run, "skip_mail": args.skip_mail, "skip_probe": args.skip_probe}
    if args.once:
        run_once(cfg, **opts)
    else:
        run_loop(cfg, **opts)


if __name__ == "__main__":
    sys.exit(main() or 0)