File size: 8,588 Bytes
f598ecc | 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 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 | """Admin-panel metrics: host load + per-user / per-bot breakdown.
Centralised here (instead of inline in routes.py) so:
* The HTML ``/admin/stats`` page and the JSON ``/api/admin/stats``
endpoint read from the same source — guaranteed parity.
* Tests can exercise the collector directly without spinning up
a full FastAPI client + DB session for what is fundamentally
a data-shape test.
Layout of the returned dict
---------------------------
::
{
"host": {cpu_*, ram_*, disk_*, boot_time, uptime_sec, ok},
"totals": {bots, running, cpu_live, ram_live_mb, bots_by_status: {...}},
"users": [ {user, bots_total, bots_running, cpu_live_percent,
ram_live_mb, cpu_alloc_cores, ram_alloc_mb,
storage_used_mb, bots_by_status: {...}} ... ],
"ts": float, # unix time at collection
}
Failure modes
-------------
``psutil`` is best-effort: if it's missing or the host denies the
read, host fields become ``None`` and ``host.ok`` flips to ``False``.
The rest of the dashboard still renders — a blank host card beats
a 500 on the whole page.
"""
from __future__ import annotations
import time
from datetime import datetime
from typing import Any
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
try:
import psutil # type: ignore
_HAS_PSUTIL = True
except Exception: # noqa: BLE001
_HAS_PSUTIL = False
# Prime psutil's CPU% baseline once at import so the very first
# ``cpu_percent(interval=None)`` call returns a real number instead
# of the always-zero first sample. Without this priming the
# dashboard would show "0%" for one render after every restart.
if _HAS_PSUTIL:
try:
psutil.cpu_percent(interval=None)
except Exception: # noqa: BLE001
pass
# ---------------------------------------------------------------------------
# Host metrics
# ---------------------------------------------------------------------------
def _collect_host() -> dict[str, Any]:
"""Snapshot of the host the panel itself runs on."""
out: dict[str, Any] = {
"ok": False,
"cpu_percent": None,
"cpu_count": None,
"ram_used_mb": None,
"ram_total_mb": None,
"ram_percent": None,
"disk_used_mb": None,
"disk_total_mb": None,
"disk_percent": None,
"boot_time": None,
"uptime_sec": None,
"error": None,
}
if not _HAS_PSUTIL:
out["error"] = "psutil unavailable"
return out
try:
out["cpu_percent"] = float(psutil.cpu_percent(interval=0.15))
out["cpu_count"] = psutil.cpu_count() or 1
vm = psutil.virtual_memory()
out["ram_total_mb"] = round(vm.total / (1024 * 1024), 1)
out["ram_used_mb"] = round((vm.total - vm.available) / (1024 * 1024), 1)
out["ram_percent"] = float(vm.percent)
try:
# The panel runs on Render / HF Spaces — data_dir is the
# most relevant mount, not necessarily "/".
from config import PANEL_CONFIG # type: ignore
du = psutil.disk_usage(str(PANEL_CONFIG.data_dir))
except Exception: # noqa: BLE001
du = psutil.disk_usage("/")
out["disk_total_mb"] = round(du.total / (1024 * 1024), 1)
out["disk_used_mb"] = round(du.used / (1024 * 1024), 1)
out["disk_percent"] = float(du.percent)
boot = psutil.boot_time()
out["boot_time"] = float(boot)
out["uptime_sec"] = int(time.time() - boot)
out["ok"] = True
return out
except Exception as exc: # noqa: BLE001
out["error"] = str(exc)
return out
# ---------------------------------------------------------------------------
# Bot-population metrics
# ---------------------------------------------------------------------------
async def _bot_status_counts(db: AsyncSession) -> dict[str, int]:
"""Count BotInstance rows grouped by status. Returns all known
statuses even when the count is 0, so the dashboard always renders
the same row order."""
# Imported lazily so this module can be imported standalone
# (e.g. by tests) without pulling the whole models graph.
from models import BotInstance, BotStatus
counts: dict[str, int] = {s.value: 0 for s in BotStatus}
rows = (await db.execute(
select(BotInstance.status, func.count(BotInstance.id))
.group_by(BotInstance.status)
)).all()
for status, count in rows:
key = status.value if hasattr(status, "value") else str(status)
counts[key] = int(count)
return counts
async def collect_admin_stats(db: AsyncSession) -> dict[str, Any]:
"""Build the full admin-stats snapshot.
The HTML page (``/admin/stats``) and the JSON endpoint
(``/api/admin/stats``) both call this — they only differ in how
they render the result.
"""
from models import (
BotInstance,
BotStatus,
DeploymentMode,
User,
)
host = _collect_host()
bots_by_status = await _bot_status_counts(db)
# One query, group by owner + status, so the per-user rollup
# is a single DB round-trip instead of N+1.
rows = (await db.execute(
select(
BotInstance.owner_id,
BotInstance.status,
func.count(BotInstance.id),
).group_by(BotInstance.owner_id, BotInstance.status)
)).all()
# owner_id -> {status -> count}
per_owner_status: dict[int, dict[str, int]] = {}
for owner_id, status, count in rows:
per_owner_status.setdefault(int(owner_id), {})[
status.value if hasattr(status, "value") else str(status)
] = int(count)
users_result = await db.execute(select(User).order_by(User.username.asc()))
users = list(users_result.scalars().all())
# Single query for all bots; per-user rows are O(1) lookups after.
all_bots = list((await db.execute(select(BotInstance))).scalars().all())
bots_by_owner: dict[int, list[BotInstance]] = {}
for b in all_bots:
bots_by_owner.setdefault(b.owner_id, []).append(b)
user_rows: list[dict[str, Any]] = []
totals = {
"bots": 0,
"running": 0,
"cpu_live": 0.0,
"ram_live_mb": 0.0,
"bots_by_status": dict(bots_by_status),
}
for u in users:
owned = bots_by_owner.get(u.id, [])
running_bots = [b for b in owned if b.status == BotStatus.RUNNING]
# Live process metrics only meaningful for MULTITENANT bots;
# LEGACY_SPACE bots run on their own HF Space and aren't
# visible to this process.
cpu_live = 0.0
ram_live_mb = 0.0
for b in owned:
if b.deployment_mode == DeploymentMode.MULTITENANT:
try:
import bot_runner # local import to avoid heavy cost on tests
st = bot_runner.bot_status(b)
if st.get("running"):
cpu_live += float(st.get("cpu_percent", 0.0))
ram_live_mb += float(st.get("rss_mb", 0.0))
except Exception: # noqa: BLE001
# If bot_runner can't probe (process exited, no
# permissions, …) we just don't add to the live
# totals — better than crashing the whole page.
pass
ram_alloc_mb = sum((b.ram_mb or 0) for b in running_bots)
cpu_alloc_cores = sum((b.cpu_cores or 0.0) for b in running_bots)
storage_mb = sum((b.storage_used_mb or 0) for b in owned)
row = {
"user": u,
"username": u.username,
"bots_total": len(owned),
"bots_running": len(running_bots),
"bots_by_status": dict(per_owner_status.get(u.id, {})),
"cpu_live_percent": round(cpu_live, 1),
"ram_live_mb": round(ram_live_mb, 1),
"cpu_alloc_cores": round(cpu_alloc_cores, 2),
"ram_alloc_mb": ram_alloc_mb,
"storage_used_mb": storage_mb,
}
user_rows.append(row)
totals["bots"] += len(owned)
totals["running"] += len(running_bots)
totals["cpu_live"] += cpu_live
totals["ram_live_mb"] += ram_live_mb
return {
"host": host,
"totals": totals,
"users": user_rows,
"ts": time.time(),
"ts_iso": datetime.utcnow().isoformat(timespec="seconds") + "Z",
}
__all__ = ["collect_admin_stats", "_collect_host", "_bot_status_counts"] |