forked from xiaohei/taiji-AI-PAD
400 lines
15 KiB
Python
400 lines
15 KiB
Python
"""
|
||
周期性计费任务模块
|
||
|
||
功能：
|
||
1. 定时更新运行中 Agent 的计费日志（agent_billing_records 表）
|
||
2. 增量扣款：只扣除上次计费后新增的费用，避免重复扣款
|
||
3. 余额不足时自动停止 Agent
|
||
|
||
为什么需要周期性统计：
|
||
- Agent 可能长时间运行，需要定期更新计费记录以便用户查看实时消费
|
||
- 避免 Agent 结束回调失败导致计费记录不准确
|
||
- 及时发现余额不足的用户并停止其 Agent，防止欠费
|
||
|
||
EU计算（Agent运行时长）：
|
||
- 公式：EU = ceil(duration_seconds / 10)
|
||
- 即：1 EU = 10秒运行时间，不足10秒按1 EU计算
|
||
|
||
成本计算：
|
||
- 平台 Agent: 按模板固定价格计费（$/小时）
|
||
- 自定义 Agent: 按资源使用量（CPU/内存）计费
|
||
|
||
执行周期：每小时执行一次（BILLING_INTERVAL_SECONDS = 3600）
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import math
|
||
from datetime import datetime, timedelta
|
||
from decimal import Decimal
|
||
from typing import Optional
|
||
|
||
from sqlalchemy import select, and_
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from database import AsyncSessionLocal
|
||
from models import AgentBillingRecord, User, Balance, TenantCustomAgentQuota, Agent
|
||
from app.billing import (
|
||
calculate_platform_agent_cost,
|
||
calculate_agent_cost_by_resources,
|
||
deduct_balance,
|
||
get_available_balance,
|
||
_parse_cpu_to_cores,
|
||
_parse_memory_to_gb,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 计费周期配置
|
||
BILLING_INTERVAL_SECONDS = 3600 # 每小时执行一次
|
||
QUOTA_CHECK_INTERVAL_HOURS = 24 # 配额一致性检查周期：每24小时一次
|
||
VM_EU_PER_HOUR = Decimal("0.5") # VM计算 0.5 EU/hour
|
||
|
||
# 上次配额检查时间
|
||
_last_quota_check: Optional[datetime] = None
|
||
|
||
|
||
def _calculate_eu(duration_seconds: int) -> int:
|
||
"""计算EU：1 EU = 10秒，不足10秒按1 EU计算"""
|
||
return math.ceil(duration_seconds / 10)
|
||
|
||
|
||
async def stop_user_agents(user_id: str, db: AsyncSession) -> list:
|
||
"""
|
||
停止用户所有运行中的Agent（余额不足时调用）
|
||
|
||
Args:
|
||
user_id: 用户ID
|
||
db: 数据库会话
|
||
|
||
Returns:
|
||
停止的Agent列表
|
||
"""
|
||
from app.agent_manager_client import get_agent_manager_client, AgentManagerError
|
||
|
||
stopped_agents = []
|
||
now = datetime.utcnow()
|
||
|
||
# 查询该用户所有运行中的Agent
|
||
result = await db.execute(
|
||
select(AgentBillingRecord).where(
|
||
and_(
|
||
AgentBillingRecord.user_id == user_id,
|
||
AgentBillingRecord.end_time == None
|
||
)
|
||
)
|
||
)
|
||
user_agents = result.scalars().all()
|
||
|
||
client = get_agent_manager_client()
|
||
|
||
for record in user_agents:
|
||
try:
|
||
# 调用 Agent Manager 停止 Pod
|
||
await client.delete_agent(record.agent_name)
|
||
|
||
# 更新计费记录为已结束
|
||
record.end_time = now
|
||
if record.start_time:
|
||
duration = (now - record.start_time).total_seconds()
|
||
record.duration_seconds = int(duration)
|
||
record.eu_consumed = _calculate_eu(int(duration))
|
||
|
||
stopped_agents.append(record.agent_name)
|
||
logger.info(f"余额不足，已停止Agent: {record.agent_name}, 用户: {user_id}")
|
||
|
||
except AgentManagerError as e:
|
||
logger.error(f"停止Agent失败: {record.agent_name}, 错误: {e.message}")
|
||
except Exception as e:
|
||
logger.error(f"停止Agent异常: {record.agent_name}, 错误: {str(e)}")
|
||
|
||
return stopped_agents
|
||
|
||
|
||
async def check_quota_consistency(db: AsyncSession) -> dict:
|
||
"""
|
||
检查配额一致性（假用量检测）
|
||
|
||
对比配额表中的使用量与实际运行的Agent，识别假用量问题。
|
||
如果发现不一致，记录警告日志，运维人员可以运行 fix_fake_quota.py 修复。
|
||
|
||
Returns:
|
||
检查结果统计
|
||
"""
|
||
logger.info("🔍 开始配额一致性检查（假用量检测）")
|
||
|
||
try:
|
||
# 查询所有配额记录
|
||
result = await db.execute(select(TenantCustomAgentQuota))
|
||
quotas = result.scalars().all()
|
||
|
||
inconsistencies = []
|
||
total_fake_cpu = 0.0
|
||
total_fake_memory = 0.0
|
||
|
||
for quota in quotas:
|
||
tenant_id = str(quota.tenant_id)
|
||
|
||
# 查询配额显示的使用量
|
||
quota_cpu = float(quota.cpu_used or 0)
|
||
quota_memory = float(quota.memory_used or 0)
|
||
quota_count = quota.agent_count or 0
|
||
|
||
# 查询实际运行的Agent
|
||
agents_result = await db.execute(
|
||
select(Agent)
|
||
.where(Agent.owner_id == tenant_id)
|
||
.where(Agent.type == 'custom')
|
||
.where(Agent.status == 'active')
|
||
)
|
||
agents = agents_result.scalars().all()
|
||
|
||
# 计算实际使用量
|
||
real_cpu = sum(float(a.cpu or 0) for a in agents)
|
||
real_memory = sum(float(a.memory or 0) for a in agents)
|
||
real_count = len(agents)
|
||
|
||
# 计算差异（允许小误差 0.01）
|
||
cpu_diff = quota_cpu - real_cpu
|
||
memory_diff = quota_memory - real_memory
|
||
count_diff = quota_count - real_count
|
||
|
||
if abs(cpu_diff) > 0.01 or abs(memory_diff) > 0.01 or count_diff != 0:
|
||
inconsistencies.append({
|
||
'tenant_id': tenant_id,
|
||
'quota': {'cpu': quota_cpu, 'memory': quota_memory, 'count': quota_count},
|
||
'real': {'cpu': real_cpu, 'memory': real_memory, 'count': real_count},
|
||
'diff': {'cpu': cpu_diff, 'memory': memory_diff, 'count': count_diff}
|
||
})
|
||
|
||
total_fake_cpu += cpu_diff
|
||
total_fake_memory += memory_diff
|
||
|
||
logger.warning(
|
||
f"⚠️ 配额不一致: 租户={tenant_id}, "
|
||
f"配额显示=[CPU:{quota_cpu:.2f}核, 内存:{quota_memory:.2f}GB, Agent:{quota_count}个], "
|
||
f"实际运行=[CPU:{real_cpu:.2f}核, 内存:{real_memory:.2f}GB, Agent:{real_count}个], "
|
||
f"假用量=[CPU:{cpu_diff:.2f}核, 内存:{memory_diff:.2f}GB, Agent:{count_diff}个]"
|
||
)
|
||
|
||
if inconsistencies:
|
||
logger.error(
|
||
f"❌ 发现 {len(inconsistencies)} 个租户存在配额不一致（假用量）！\n"
|
||
f" 假用量汇总: CPU={total_fake_cpu:.2f}核, 内存={total_fake_memory:.2f}GB\n"
|
||
f" 🔧 请运行修复脚本: python fix_fake_quota.py"
|
||
)
|
||
else:
|
||
logger.info("✅ 配额一致性检查通过，无假用量问题")
|
||
|
||
return {
|
||
'total_tenants': len(quotas),
|
||
'inconsistent_tenants': len(inconsistencies),
|
||
'total_fake_cpu': total_fake_cpu,
|
||
'total_fake_memory': total_fake_memory,
|
||
'details': inconsistencies
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"配额一致性检查失败: {e}")
|
||
return {'error': str(e)}
|
||
|
||
|
||
async def update_running_agent_billing(db: AsyncSession) -> dict:
|
||
"""
|
||
更新所有运行中 Agent 的计费记录
|
||
|
||
功能：
|
||
1. 更新EU消耗和成本
|
||
2. 增量扣款
|
||
3. 余额不足时自动停止Agent
|
||
|
||
Returns:
|
||
统计信息
|
||
"""
|
||
stats = {
|
||
"processed": 0,
|
||
"platform_agents": 0,
|
||
"custom_agents": 0,
|
||
"total_eu_consumed": Decimal("0"),
|
||
"total_cost": Decimal("0"),
|
||
"failed": 0,
|
||
"errors": [],
|
||
"stopped_agents": [], # 因余额不足而停止的Agent
|
||
}
|
||
|
||
try:
|
||
# 查询所有运行中的 Agent（end_time 为空）
|
||
result = await db.execute(
|
||
select(AgentBillingRecord).where(
|
||
AgentBillingRecord.end_time == None
|
||
)
|
||
)
|
||
running_agents = result.scalars().all()
|
||
|
||
now = datetime.utcnow()
|
||
|
||
for record in running_agents:
|
||
try:
|
||
# 跳过没有开始时间的记录
|
||
if record.start_time is None:
|
||
logger.warning(f"Agent {record.agent_name} 缺少 start_time，跳过计费")
|
||
stats["failed"] += 1
|
||
stats["errors"].append(f"Agent {record.agent_name}: start_time 为空")
|
||
continue
|
||
|
||
# 计算从开始到现在的运行时长
|
||
duration = (now - record.start_time).total_seconds()
|
||
duration_seconds = int(duration)
|
||
|
||
# 计算 EU 消耗
|
||
eu_consumed = _calculate_eu(duration_seconds)
|
||
|
||
# 计算成本
|
||
if record.is_platform_agent:
|
||
cost = calculate_platform_agent_cost(record.agent_type, duration_seconds)
|
||
stats["platform_agents"] += 1
|
||
else:
|
||
cpu_cores = _parse_cpu_to_cores(record.cpu_used) if record.cpu_used else 0.1
|
||
memory_gb = _parse_memory_to_gb(record.memory_used) if record.memory_used else 0.125
|
||
cost = calculate_agent_cost_by_resources(cpu_cores, memory_gb, duration_seconds)
|
||
stats["custom_agents"] += 1
|
||
|
||
# 计算本次周期需要扣除的增量
|
||
previous_cost = Decimal(str(record.cost or 0))
|
||
cost_increment = cost - previous_cost
|
||
|
||
# 更新记录
|
||
record.duration_seconds = duration_seconds
|
||
record.eu_consumed = eu_consumed
|
||
record.cost = float(cost)
|
||
|
||
# 如果有增量，进行扣款
|
||
if cost_increment > 0:
|
||
success, message = await deduct_balance(
|
||
str(record.user_id), cost_increment, db,
|
||
f"Agent周期计费: {record.agent_name}"
|
||
)
|
||
if not success:
|
||
logger.warning(
|
||
f"周期计费扣款失败: {message}, "
|
||
f"用户: {record.user_id}, Agent: {record.agent_name}"
|
||
)
|
||
|
||
# 检查用户余额，如果透支则停止该用户的所有Agent
|
||
balance, credit_limit, available = await get_available_balance(
|
||
str(record.user_id), db
|
||
)
|
||
if available < 0:
|
||
logger.warning(
|
||
f"用户 {record.user_id} 余额不足 (可用: {available})，将停止所有Agent"
|
||
)
|
||
stopped = await stop_user_agents(str(record.user_id), db)
|
||
stats["stopped_agents"].extend(stopped)
|
||
|
||
stats["processed"] += 1
|
||
stats["total_eu_consumed"] += Decimal(str(eu_consumed))
|
||
stats["total_cost"] += cost
|
||
|
||
except Exception as e:
|
||
stats["failed"] += 1
|
||
stats["errors"].append(f"Agent {record.agent_name}: {str(e)}")
|
||
logger.error(f"处理 Agent {record.agent_name} 计费失败: {e}")
|
||
|
||
await db.commit()
|
||
|
||
except Exception as e:
|
||
logger.error(f"周期性计费任务失败: {e}")
|
||
stats["errors"].append(f"全局错误: {str(e)}")
|
||
await db.rollback()
|
||
|
||
return stats
|
||
|
||
|
||
async def periodic_billing_task():
|
||
"""
|
||
周期性计费任务（后台运行）
|
||
|
||
每小时执行一次，更新所有运行中 Agent 的计费记录
|
||
每24小时执行一次配额一致性检查（假用量检测）
|
||
"""
|
||
global _last_quota_check
|
||
|
||
logger.info("周期性计费任务启动")
|
||
|
||
while True:
|
||
try:
|
||
async with AsyncSessionLocal() as db:
|
||
# 1. 执行计费更新（每小时）
|
||
stats = await update_running_agent_billing(db)
|
||
|
||
logger.info(
|
||
f"周期性计费完成: "
|
||
f"处理 {stats['processed']} 个Agent, "
|
||
f"平台Agent {stats['platform_agents']}, "
|
||
f"自定义Agent {stats['custom_agents']}, "
|
||
f"总EU消耗 {stats['total_eu_consumed']}, "
|
||
f"总成本 {stats['total_cost']}, "
|
||
f"失败 {stats['failed']}, "
|
||
f"因余额不足停止 {len(stats.get('stopped_agents', []))} 个Agent"
|
||
)
|
||
|
||
if stats["errors"]:
|
||
logger.warning(f"计费错误: {stats['errors']}")
|
||
|
||
# 2. 配额一致性检查（每24小时一次）
|
||
now = datetime.utcnow()
|
||
should_check_quota = (
|
||
_last_quota_check is None or
|
||
(now - _last_quota_check).total_seconds() >= QUOTA_CHECK_INTERVAL_HOURS * 3600
|
||
)
|
||
|
||
if should_check_quota:
|
||
logger.info(f"⏰ 触发配额一致性检查（距上次检查: {(now - _last_quota_check).total_seconds() / 3600:.1f}小时）" if _last_quota_check else "⏰ 首次执行配额一致性检查")
|
||
quota_stats = await check_quota_consistency(db)
|
||
_last_quota_check = now
|
||
|
||
# 如果发现假用量，记录到统计中
|
||
if quota_stats.get('inconsistent_tenants', 0) > 0:
|
||
stats['quota_check'] = quota_stats
|
||
|
||
except Exception as e:
|
||
logger.error(f"周期性计费任务异常: {e}")
|
||
|
||
# 等待下一个计费周期
|
||
await asyncio.sleep(BILLING_INTERVAL_SECONDS)
|
||
|
||
|
||
async def run_billing_now() -> dict:
|
||
"""
|
||
立即执行一次计费更新（手动触发）
|
||
|
||
Returns:
|
||
统计信息
|
||
"""
|
||
async with AsyncSessionLocal() as db:
|
||
return await update_running_agent_billing(db)
|
||
|
||
|
||
# 用于存储后台任务引用
|
||
_billing_task: Optional[asyncio.Task] = None
|
||
|
||
|
||
def start_periodic_billing():
|
||
"""启动周期性计费任务"""
|
||
global _billing_task
|
||
if _billing_task is None or _billing_task.done():
|
||
_billing_task = asyncio.create_task(periodic_billing_task())
|
||
logger.info("周期性计费任务已启动")
|
||
else:
|
||
logger.warning("周期性计费任务已在运行中")
|
||
|
||
|
||
def stop_periodic_billing():
|
||
"""停止周期性计费任务"""
|
||
global _billing_task
|
||
if _billing_task and not _billing_task.done():
|
||
_billing_task.cancel()
|
||
logger.info("周期性计费任务已停止")
|
||
_billing_task = None
|
||
|