HyperNode-AI/app/services/heartbeat.py

"""心跳登记与离线巡检。"""

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.is_banned:
        raise AppError("节点已封禁", 403)
    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 set_node_banned(session: Session, node: ComputeNode, banned: bool, now: datetime | None = None) -> ComputeNode:
    """封禁立即结束在线区间。解除后要等下一次心跳才会重新上线。"""
    moment = now or utcnow()
    node.is_banned = banned
    node.updated_at = moment
    if banned and node.status == NodeStatus.ONLINE:
        _close_open_intervals(session, node.id, moment, "banned")
        node.status = NodeStatus.OFFLINE
    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