zhifu-platform/scripts/verify_runtime.py

import json
import os
import subprocess
import time
from pathlib import Path

import httpx
from prometheus_client.parser import text_string_to_metric_families

ROOT = Path(__file__).resolve().parents[1]
COMPOSE = ["docker", "compose", "-f", "compose.runtime.yaml"]
ROLES = ["agent", "memory", "fulfillment", "knowledge", "channel-delivery"]


def run(*arguments, check=True):
    return subprocess.run(
        [*COMPOSE, *arguments], cwd=ROOT, check=check, capture_output=True, encoding="utf-8"
    )


def workers(expected):
    with httpx.Client(timeout=10, trust_env=False) as client:
        for _ in range(40):
            response = client.get("http://127.0.0.1:28300/metrics")
            response.raise_for_status()
            counts = {
                sample.labels["role"]: sample.value
                for family in text_string_to_metric_families(response.text)
                for sample in family.samples
                if sample.name == "zhifu_workers_alive"
            }
            if counts == dict.fromkeys(ROLES, expected):
                return counts
            time.sleep(1)
    raise RuntimeError(f"工作进程心跳不符合要求：{counts}")


def main():
    if not os.environ.get("DEEPSEEK_API_KEY"):
        raise ValueError("运行验证需要已配置的 DEEPSEEK_API_KEY")
    report = {"release_check": json.loads(run("run", "--rm", "release-check").stdout)}
    missing = run("run", "--rm", "-e", "AGENT_MODEL_NAME=", "release-check", check=False)
    if missing.returncode != 1 or "模型名称、服务地址及环境变量凭据必须齐全" not in missing.stderr:
        raise RuntimeError("发布检查未拒绝缺少模型配置的环境")
    report["missing_model_configuration_rejected"] = True
    report["workers_before"] = workers(1)
    started = time.monotonic()
    try:
        run("stop", *ROLES)
        ids = run("ps", "--all", "--quiet", *ROLES).stdout.split()
        metadata = json.loads(subprocess.check_output(["docker", "inspect", *ids]))
        if len(metadata) != len(ROLES):
            raise RuntimeError("工作进程容器数量不符合要求")
        states = {
            value["Config"]["Labels"]["com.docker.compose.service"]: {
                "exit_code": value["State"]["ExitCode"],
                "oom_killed": value["State"]["OOMKilled"],
                "running": value["State"]["Running"],
            }
            for value in metadata
        }
        if any(
            value != {"exit_code": 0, "oom_killed": False, "running": False}
            for value in states.values()
        ):
            raise RuntimeError(f"工作进程未正常停止：{states}")
        report["workers_stopped"] = states
        report["workers_after_stop"] = workers(0)
        report["stop_seconds"] = round(time.monotonic() - started, 2)
    finally:
        run("up", "-d", *ROLES)
    report["workers_after_start"] = workers(1)
    with httpx.Client(timeout=10, trust_env=False) as client:
        report["http_statuses"] = {}
        for name, url in {
            "api": "http://127.0.0.1:28000/health/ready",
            "oms": "http://127.0.0.1:28100/health",
            "inference": "http://127.0.0.1:28200/health/ready",
        }.items():
            response = client.get(url)
            response.raise_for_status()
            report["http_statuses"][name] = response.status_code
    containers = run("ps", "--quiet", "api", "oms", "operations", "inference", *ROLES)
    metadata = json.loads(
        subprocess.check_output(["docker", "inspect", *containers.stdout.split()])
    )
    report["containers"] = {
        value["Config"]["Labels"]["com.docker.compose.service"]: {
            "image": value["Image"],
            "user": value["Config"]["User"],
            "read_only": value["HostConfig"]["ReadonlyRootfs"],
            "restart_count": value["RestartCount"],
        }
        for value in metadata
    }
    if len(metadata) != 4 + len(ROLES) or any(
        value["user"] != "10001:10001" or not value["read_only"]
        for value in report["containers"].values()
    ):
        raise RuntimeError("容器数量或权限配置不符合要求")
    (ROOT / "docs/verification/container-runtime.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()