taiji-AI-PAD/services/mcp-server/fix_agent_quotas.py

#!/usr/bin/env python3
"""修复Agent资源配置和配额记录

修复两个问题：
1. 旧的自定义Agent缺少K8s格式的资源配置
2. 租户配额使用量与实际Agent资源不匹配
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + '/../services/mcp-server')

from sqlalchemy import select, func
from database import AsyncSessionLocal
from models import Agent, TenantCustomAgentQuota
import asyncio

def parse_cpu(cpu_str):
    """解析 CPU 字符串为浮点数（核心数）"""
    if not cpu_str:
        return 0.0
    cpu_str = str(cpu_str)
    if cpu_str.endswith('m'):
        return float(cpu_str[:-1]) / 1000
    return float(cpu_str)

def parse_memory(mem_str):
    """解析内存字符串为浮点数（GB）"""
    if not mem_str:
        return 0.0
    mem_str = str(mem_str)
    if mem_str.endswith('Gi'):
        return float(mem_str[:-2])
    elif mem_str.endswith('G'):
        return float(mem_str[:-1])
    elif mem_str.endswith('Mi'):
        return float(mem_str[:-2]) / 1024
    return float(mem_str)

async def fix_agent_resources():
    """修复Agent资源配置"""
    async with AsyncSessionLocal() as db:
        # 查询所有自定义Agent
        result = await db.execute(
            select(Agent).where(Agent.type == 'custom')
        )
        agents = result.scalars().all()
        
        print(f"找到 {len(agents)} 个自定义Agent")
        print("=" * 60)
        
        fixed_count = 0
        for agent in agents:
            # 检查是否需要修复（cpu_limit 不匹配 cpu）
            current_cpu_limit = parse_cpu(agent.cpu_limit)
            expected_cpu = float(agent.cpu or 0)
            
            current_memory_limit = parse_memory(agent.memory_limit)
            expected_memory = float(agent.memory or 0)
            
            # 判断是否需要更新
            needs_update = False
            if abs(current_cpu_limit - expected_cpu) > 0.01:
                needs_update = True
            if abs(current_memory_limit - expected_memory) > 0.01:
                needs_update = True
            
            if not needs_update:
                continue
            
            print(f"\n🔧 修复Agent: {agent.name} ({agent.id})")
            print(f"   所有者: {agent.owner_id}")
            print(f"   当前 CPU: {agent.cpu} 核 -> limit: {agent.cpu_limit} ({current_cpu_limit} 核)")
            print(f"   当前 Memory: {agent.memory} GB -> limit: {agent.memory_limit} ({current_memory_limit} GB)")
            
            # 更新 K8s 格式资源配置
            cpu_limit_k8s = f"{int(expected_cpu * 1000)}m"
            cpu_request_k8s = f"{max(100, int(expected_cpu * 100))}m"
            memory_limit_k8s = f"{expected_memory}Gi"
            memory_request_k8s = f"{max(0.128, expected_memory * 0.25):.3f}Gi"
            
            agent.cpu_limit = cpu_limit_k8s
            agent.cpu_request = cpu_request_k8s
            agent.memory_limit = memory_limit_k8s
            agent.memory_request = memory_request_k8s
            
            print(f"   新 CPU limit: {cpu_limit_k8s}, request: {cpu_request_k8s}")
            print(f"   新 Memory limit: {memory_limit_k8s}, request: {memory_request_k8s}")
            
            fixed_count += 1
        
        if fixed_count > 0:
            await db.commit()
            print("\n" + "=" * 60)
            print(f"✅ Agent资源配置修复完成，共修复 {fixed_count} 个Agent")
        else:
            print("\n" + "=" * 60)
            print("✅ 所有Agent资源配置正常，无需修复")
        
        print("=" * 60)

async def recalculate_quotas():
    """重新计算租户配额使用量"""
    async with AsyncSessionLocal() as db:
        # 查询所有租户配额记录
        result = await db.execute(select(TenantCustomAgentQuota))
        quotas = result.scalars().all()
        
        print(f"\n找到 {len(quotas)} 个租户配额记录")
        print("=" * 60)
        
        for quota in quotas:
            tenant_id = quota.tenant_id
            
            # 查询该租户的所有Agent
            agent_result = await db.execute(
                select(Agent).where(
                    Agent.owner_id == tenant_id,
                    Agent.type == 'custom'
                )
            )
            agents = agent_result.scalars().all()
            
            # 计算实际使用量
            total_cpu = 0.0
            total_memory = 0.0
            for agent in agents:
                total_cpu += float(agent.cpu or 0)
                total_memory += float(agent.memory or 0)
            
            # 获取当前记录的使用量
            current_cpu = float(quota.cpu_used or 0)
            current_memory = float(quota.memory_used or 0)
            
            # 检查是否需要更新
            if abs(current_cpu - total_cpu) < 0.01 and abs(current_memory - total_memory) < 0.01:
                continue
            
            print(f"\n🔧 修复租户配额: {tenant_id}")
            print(f"   Agent数量: {len(agents)}")
            print(f"   当前记录: CPU={current_cpu} 核, Memory={current_memory} GB")
            print(f"   实际使用: CPU={total_cpu} 核, Memory={total_memory} GB")
            
            # 更新配额使用量
            quota.cpu_used = total_cpu
            quota.memory_used = total_memory
            quota.agent_count = len(agents)
            
            print(f"   ✅ 已更新配额使用量")
        
        await db.commit()
        print("\n" + "=" * 60)
        print("✅ 配额使用量重新计算完成")
        print("=" * 60)

async def main():
    """主流程"""
    print("=" * 60)
    print("修复Agent资源配置和配额记录")
    print("=" * 60)
    
    # 1. 修复Agent资源配置
    await fix_agent_resources()
    
    # 2. 重新计算配额使用量
    await recalculate_quotas()
    
    print("\n✅ 所有修复完成！")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception as e:
        print(f"\n❌ 修复失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)