dekun
132 lines
5.1 KiB
Python
132 lines
5.1 KiB
Python
"""心跳登记与离线巡检。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hmac
|
|
from datetime import datetime, timedelta
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.clock import utcnow
|
|
from app.config import settings
|
|
from app.errors import AppError
|
|
from app.models import ComputeNode, NodeStatus, UptimeInterval
|
|
from app.services.common import get_node_by_sn
|
|
from app.services.datacenter import assert_fingerprint, assert_subnet_gpu
|
|
from app.services.ip_guard import assert_ip_available
|
|
from app.services.hardware import GENESIS_RECOVERY_DAYS, HardwareReport, classify_generation, is_genesis, score_hardware
|
|
|
|
|
|
def assert_node_token(node: ComputeNode, token: str | None) -> None:
|
|
"""旧节点没有 Token 时仍按序列号接入。审核开通的节点必须带上专属 Token。"""
|
|
if not node.node_token:
|
|
return
|
|
supplied = (token or "").strip()
|
|
if not supplied or not hmac.compare_digest(node.node_token, supplied):
|
|
raise AppError("节点 Token 不正确", 403)
|
|
|
|
|
|
def record_heartbeat(
|
|
session: Session,
|
|
sn: str,
|
|
token: str | None = None,
|
|
now: datetime | None = None,
|
|
hardware: HardwareReport | None = None,
|
|
client_ip: str | None = None,
|
|
hardware_fingerprint: str | None = None,
|
|
) -> ComputeNode:
|
|
"""刷新最后活跃时间。离线节点收到心跳后重新打开一段在线区间。"""
|
|
moment = now or utcnow()
|
|
node = get_node_by_sn(session, sn.strip())
|
|
assert_node_token(node, token)
|
|
if node.license_expires_at <= moment:
|
|
raise AppError("节点授权已到期", 403)
|
|
bound_ip = (client_ip or "").strip() or None
|
|
if bound_ip:
|
|
assert_ip_available(session, bound_ip, node=node, now=moment)
|
|
reported_gpu = hardware.gpu_model if hardware is not None and hardware.gpu_model else node.gpu_model
|
|
assert_subnet_gpu(session, bound_ip, reported_gpu)
|
|
node.client_ip = bound_ip
|
|
fingerprint = assert_fingerprint(session, hardware_fingerprint, node.id)
|
|
if fingerprint:
|
|
node.hardware_fingerprint = fingerprint
|
|
|
|
if hardware is not None:
|
|
_apply_hardware(node, hardware)
|
|
node.last_heartbeat_at = moment
|
|
node.updated_at = moment
|
|
if node.status != NodeStatus.ONLINE:
|
|
node.status = NodeStatus.ONLINE
|
|
session.add(UptimeInterval(node_id=node.id, started_at=moment))
|
|
else:
|
|
_ensure_open_interval(session, node.id, moment)
|
|
session.flush()
|
|
return node
|
|
|
|
|
|
def mark_stale_nodes(session: Session, now: datetime | None = None) -> list[str]:
|
|
"""把超过心跳时限或授权已到期的在线节点置为 offline，并关闭在线区间。"""
|
|
moment = now or utcnow()
|
|
cutoff = moment - timedelta(minutes=settings.heartbeat_timeout_minutes)
|
|
nodes = session.scalars(select(ComputeNode).where(ComputeNode.status == NodeStatus.ONLINE)).all()
|
|
changed: list[str] = []
|
|
for node in nodes:
|
|
reasons: list[tuple[datetime, str]] = []
|
|
if node.last_heartbeat_at is None or node.last_heartbeat_at < cutoff:
|
|
reasons.append((node.last_heartbeat_at or moment, "heartbeat_timeout"))
|
|
if node.license_expires_at <= moment:
|
|
reasons.append((node.license_expires_at, "license_expired"))
|
|
if not reasons:
|
|
continue
|
|
ended_at, reason = min(reasons, key=lambda item: item[0])
|
|
_close_open_intervals(session, node.id, ended_at, reason)
|
|
node.status = NodeStatus.OFFLINE
|
|
node.updated_at = moment
|
|
changed.append(node.sn)
|
|
session.flush()
|
|
return changed
|
|
|
|
|
|
def _apply_hardware(node: ComputeNode, hardware: HardwareReport) -> None:
|
|
node.gpu_model = (hardware.gpu_model or "").strip()[:128] or None
|
|
node.vram_gb = hardware.vram_gb
|
|
node.cpu_count = hardware.cpu_count
|
|
node.memory_gb = hardware.memory_gb
|
|
node.hardware_weight = float(score_hardware(node.gpu_model, node.vram_gb))
|
|
if node.generation_locked or not node.gpu_model:
|
|
return
|
|
generation, recovery_days = classify_generation(node.gpu_model)
|
|
node.gpu_generation = generation
|
|
node.generation_locked = True
|
|
if is_genesis(node.node_index):
|
|
node.recovery_days = GENESIS_RECOVERY_DAYS
|
|
return
|
|
node.recovery_days = recovery_days
|
|
horizon = node.licensed_at + timedelta(days=recovery_days)
|
|
if horizon > node.license_expires_at:
|
|
node.license_expires_at = horizon
|
|
|
|
|
|
def _ensure_open_interval(session: Session, node_id: int, started_at: datetime) -> None:
|
|
opened = session.scalar(
|
|
select(UptimeInterval).where(
|
|
UptimeInterval.node_id == node_id,
|
|
UptimeInterval.ended_at.is_(None),
|
|
)
|
|
)
|
|
if opened is None:
|
|
session.add(UptimeInterval(node_id=node_id, started_at=started_at))
|
|
|
|
|
|
def _close_open_intervals(session: Session, node_id: int, ended_at: datetime, reason: str) -> None:
|
|
rows = session.scalars(
|
|
select(UptimeInterval).where(
|
|
UptimeInterval.node_id == node_id,
|
|
UptimeInterval.ended_at.is_(None),
|
|
)
|
|
).all()
|
|
for row in rows:
|
|
row.ended_at = ended_at if ended_at >= row.started_at else row.started_at
|
|
row.close_reason = reason
|