File size: 6,898 Bytes
66916b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | """Extended admin helpers for per-user management.
What lives here
---------------
Pure collectors that the ``/admin/users`` table and the
``/admin/users/{user_id}`` drilldown page both consume, so they
stay in lockstep. No business logic — just roll-ups that read from
``BotInstance``, ``AuditLog``, ``DeploymentLog`` etc.
Returned shapes
---------------
::
collect_user_usage(db, user) -> {
"user": User,
"bots_total": int,
"bots_running": int,
"bots_by_status": {status: count, ...},
"cpu_alloc_cores": float,
"cpu_quota_cores": float,
"cpu_percent": float, # 0..100 of quota used
"ram_alloc_mb": int,
"ram_quota_mb": int,
"ram_percent": float,
"storage_used_mb": int,
"storage_quota_mb": int,
"storage_percent": float,
"bots_at_max": bool, # True if bots_total >= plan_max_bots
}
collect_user_traffic(db, user, hours=24) -> {
"hours": int,
"total_requests": int,
"error_count": int,
"avg_latency_ms": float,
}
The quota percentages are clamped to [0, 100] for the bar UI;
values above 100 indicate the user is over their plan (still useful
information — admin should see "100%+!" red bar).
"""
from __future__ import annotations
import json
from typing import Any
from sqlalchemy import select, func, and_
from sqlalchemy.ext.asyncio import AsyncSession
from models import (
AuditLog,
BotInstance,
BotStatus,
DeploymentMode,
User,
)
def _percent(used: float, quota: float) -> float:
"""Used / quota * 100, clamped to [0, 200]. ``>100`` for over-plan."""
if quota <= 0:
return 100.0 if used > 0 else 0.0
pct = round(used / quota * 100.0, 1)
return min(200.0, max(0.0, pct))
async def collect_user_usage(db: AsyncSession, user: User) -> dict[str, Any]:
"""Compute current usage vs plan for one user."""
bots = list((await db.execute(
select(BotInstance).where(BotInstance.owner_id == user.id)
)).scalars().all())
bots_by_status: dict[str, int] = {}
for b in bots:
key = b.status.value if hasattr(b.status, "value") else str(b.status)
bots_by_status[key] = bots_by_status.get(key, 0) + 1
# Live process metrics for multitentant bots; alloc for legacy ones
# uses the bot's configured cpu/ram (allocated at create-time).
cpu_alloc = 0.0
ram_alloc = 0
storage_used = 0
for b in bots:
if b.deployment_mode == DeploymentMode.MULTITENANT:
try:
import bot_runner
st = bot_runner.bot_status(b)
if st.get("running"):
cpu_alloc += float(st.get("cpu_percent") or 0.0) / 100.0 # cpu_percent -> cores (approx)
ram_alloc += int(float(st.get("rss_mb") or 0))
except Exception:
pass
else:
# Legacy / HF-Space bot: alloc from the configured slice.
cpu_alloc += float(b.cpu_cores or 0.0)
ram_alloc += int(b.ram_mb or 0)
storage_used += int(b.storage_used_mb or 0)
cpu_quota = float(user.cpu_quota_cores or 0.0)
ram_quota = int(user.ram_quota_mb or 0)
storage_quota = int(user.storage_quota_mb or 0)
# CPU percent: number of allocated cores / quota. Capped at 200%.
cpu_pct = _percent(cpu_alloc, cpu_quota) if cpu_quota > 0 else (100.0 if cpu_alloc > 0 else 0.0)
return {
"user": user,
"bots_total": len(bots),
"bots_running": sum(1 for b in bots if b.status == BotStatus.RUNNING),
"bots_by_status": bots_by_status,
"cpu_alloc_cores": round(cpu_alloc, 2),
"cpu_quota_cores": cpu_quota,
"cpu_percent": cpu_pct,
"ram_alloc_mb": int(ram_alloc),
"ram_quota_mb": ram_quota,
"ram_percent": _percent(ram_alloc, ram_quota),
"storage_used_mb": int(storage_used),
"storage_quota_mb": storage_quota,
"storage_percent": _percent(storage_used, storage_quota),
"bots_at_max": (user.plan_max_bots > 0 and len(bots) >= user.plan_max_bots),
}
async def collect_user_traffic(
db: AsyncSession,
user: User,
hours: int = 24,
) -> dict[str, Any]:
"""Recent proxy traffic for a user's bots."""
bot_ids = (await db.execute(
select(BotInstance.id).where(BotInstance.owner_id == user.id)
)).scalars().all()
if not bot_ids:
return {
"hours": hours,
"total_requests": 0,
"error_count": 0,
"avg_latency_ms": 0.0,
}
from datetime import datetime, timedelta
end = datetime.utcnow()
start = end - timedelta(hours=hours)
rows = (await db.execute(
select(AuditLog).where(
and_(
AuditLog.action == "proxy_request",
AuditLog.target_type == "bot",
AuditLog.target_id.in_(bot_ids),
AuditLog.created_at >= start,
AuditLog.created_at < end,
)
)
)).scalars().all()
total = errors = 0
latencies: list[float] = []
for r in rows:
total += 1
try:
details = json.loads(r.details_json) if r.details_json else {}
except (ValueError, TypeError):
details = {}
if int(details.get("status") or 0) >= 500:
errors += 1
lat = details.get("latency_ms")
if isinstance(lat, (int, float)):
latencies.append(float(lat))
return {
"hours": hours,
"total_requests": total,
"error_count": errors,
"avg_latency_ms": round(sum(latencies) / len(latencies), 1) if latencies else 0.0,
}
# ---------------------------------------------------------------------------
# Plan presets — quick-apply buttons on the users page
# ---------------------------------------------------------------------------
PLAN_PRESETS: dict[str, dict[str, Any]] = {
"free": {"cpu_cores": 0.4, "ram_mb": 512, "storage_mb": 2048, "max_bots": 1, "label": "Free"},
"trial": {"cpu_cores": 0.5, "ram_mb": 1024, "storage_mb": 4096, "max_bots": 2, "label": "Trial"},
"std": {"cpu_cores": 1.0, "ram_mb": 2 * 1024, "storage_mb": 10 * 1024, "max_bots": 5, "label": "Standard"},
"pro": {"cpu_cores": 4.0, "ram_mb": 8 * 1024, "storage_mb": 50 * 1024, "max_bots": 20, "label": "Pro"},
"biz": {"cpu_cores": 8.0, "ram_mb": 16 * 1024, "storage_mb": 200 * 1024, "max_bots": 50, "label": "Business"},
"unlim": {"cpu_cores": 16.0,"ram_mb": 64 * 1024, "storage_mb": 1024 * 1024,"max_bots": 1000,"label": "Unlimited"},
}
__all__ = [
"collect_user_usage",
"collect_user_traffic",
"PLAN_PRESETS",
] |