agent_management/scripts/aggregate_agents_resources.py

#!/usr/bin/env python3
"""
汇总测试平台中由 agent-manager 管理的 agents 数量及资源使用（请求/限制/实际使用）

用法示例:
  python3 scripts/aggregate_agents_resources.py --namespace ai-agents

注意:
  - 需要在运行环境中能访问 Kubernetes 集群（in-cluster 或提供 KUBECONFIG）
  - 若想使用实际资源使用（usage），需要集群安装 metrics-server
"""
import argparse
import sys
import os

# ensure repo root is on sys.path so we can import k8s_manager when running from scripts/
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from k8s_manager import K8sManager
from typing import Optional


def parse_cpu(s: Optional[str]) -> float:
    if not s:
        return 0.0
    s = str(s).strip()
    try:
        if s.endswith('m'):
            return float(s[:-1]) / 1000.0
        return float(s)
    except ValueError:
        return 0.0


def parse_memory(s: Optional[str]) -> float:
    """解析内存字符串并返回Mi单位的浮点数值"""
    if not s:
        return 0.0
    s = str(s).strip()
    units = {
        'Ki': 1024.0,
        'Mi': 1024.0 ** 2,
        'Gi': 1024.0 ** 3,
        'Ti': 1024.0 ** 4,
        'K': 1000.0,
        'M': 1000.0 ** 2,
        'G': 1000.0 ** 3,
    }
    # 直接 numeric
    try:
        return float(s) / (1024.0 ** 2)
    except Exception:
        pass

    for u, factor in units.items():
        if s.endswith(u):
            try:
                num = float(s[:-len(u)])
                # 返回 Mi 为单位
                return (num * factor) / (1024.0 ** 2)
            except Exception:
                return 0.0
    # 未知单位，尝试移除非数字字符
    num = ''.join(ch for ch in s if (ch.isdigit() or ch == '.' ))
    try:
        return float(num) / (1024.0 ** 2)
    except Exception:
        return 0.0


def human_mem(mib: float) -> str:
    if mib >= 1024:
        return f"{mib/1024:.2f} GiB"
    return f"{mib:.1f} MiB"


def human_cpu(cores: float) -> str:
    if cores < 1:
        return f"{int(cores*1000)} m"
    return f"{cores:.3f} cores"


def aggregate(namespace: str, kubeconfig: Optional[str], detailed: bool = False):
    mgr = K8sManager(namespace=namespace, kubeconfig_path=kubeconfig)

    pods = mgr.list_pods()
    total = len(pods)

    sum_req_cpu = 0.0
    sum_lim_cpu = 0.0
    sum_usage_cpu = 0.0
    have_usage_cpu = False

    sum_req_mem = 0.0
    sum_lim_mem = 0.0
    sum_usage_mem = 0.0
    have_usage_mem = False

    details = []
    for p in pods:
        name = p.get('name')
        status = mgr.get_pod_status(name)

        # requests/limits
        resources = status.get('resources', {})
        requests = resources.get('requests', {}) or {}
        limits = resources.get('limits', {}) or {}

        r_cpu = parse_cpu(requests.get('cpu'))
        l_cpu = parse_cpu(limits.get('cpu'))
        sum_req_cpu += r_cpu
        sum_lim_cpu += l_cpu

        r_mem = parse_memory(requests.get('memory'))
        l_mem = parse_memory(limits.get('memory'))
        sum_req_mem += r_mem
        sum_lim_mem += l_mem

        usage = resources.get('usage') or {}
        u_cpu = parse_cpu(usage.get('cpu'))
        u_mem = parse_memory(usage.get('memory'))
        if u_cpu:
            have_usage_cpu = True
            sum_usage_cpu += u_cpu
        if u_mem:
            have_usage_mem = True
            sum_usage_mem += u_mem

        details.append({
            'name': name,
            'status': status.get('status'),
            'template': status.get('template'),
            'req_cpu': r_cpu,
            'lim_cpu': l_cpu,
            'use_cpu': u_cpu,
            'req_mem_mi': r_mem,
            'lim_mem_mi': l_mem,
            'use_mem_mi': u_mem,
        })

    # 输出
    print(f"Agents 总数: {total}")
    print("")
    print("CPU 总计:")
    print(f"  请求 (requests): {human_cpu(sum_req_cpu)}")
    print(f"  限制 (limits): {human_cpu(sum_lim_cpu)}")
    if have_usage_cpu:
        print(f"  实际使用 (usage): {human_cpu(sum_usage_cpu)}")
    else:
        print("  实际使用 (usage): 未获取（需安装 metrics-server 或 无法访问 metrics API）")

    print("")
    print("内存 总计:")
    print(f"  请求 (requests): {human_mem(sum_req_mem)}")
    print(f"  限制 (limits): {human_mem(sum_lim_mem)}")
    if have_usage_mem:
        print(f"  实际使用 (usage): {human_mem(sum_usage_mem)}")
    else:
        print("  实际使用 (usage): 未获取（需安装 metrics-server 或 无法访问 metrics API）")

    if detailed:
        print('\n每个 Pod 详情:')
        for d in details:
            print(f"- {d['name']}: status={d['status']}, template={d['template']}, req={human_cpu(d['req_cpu'])}/{human_mem(d['req_mem_mi'])}, lim={human_cpu(d['lim_cpu'])}/{human_mem(d['lim_mem_mi'])}, use={human_cpu(d['use_cpu'])}/{human_mem(d['use_mem_mi'])}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--namespace', '-n', default='ai-agents', help='Kubernetes namespace')
    parser.add_argument('--kubeconfig', '-k', default=None, help='可选 kubeconfig 文件路径')
    parser.add_argument('--detailed', '-d', action='store_true', help='输出每个 pod 的详细资源信息')
    args = parser.parse_args()

    aggregate(args.namespace, args.kubeconfig, args.detailed)


if __name__ == '__main__':
    main()