"""心跳登记与离线巡检。""" 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