HyperNode-AI/app/services/tasks.py

"""随机任务下发与回传校验。

节点不暴露回调地址时，任务留在待拉取队列，客户端用
GET /api/v1/node/{sn}/tasks 领取。配置了 callback_url 时，调度器会额外 POST 一次。
回传必须带上任务里的 nonce，避免未收到任务的客户端直接上报成功。
"""

from __future__ import annotations

import hmac
import json
import logging
import random
import secrets
import uuid
from datetime import datetime, timedelta
from urllib.error import URLError
from urllib.request import Request, urlopen

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, ComputeTask, NodeStatus, PushStatus, TaskStatus
from app.services.common import get_node_by_sn

logger = logging.getLogger(__name__)


def dispatch_random_tasks(session: Session, now: datetime | None = None) -> list[ComputeTask]:
    """向当前在线、授权仍有效、且没有未完成任务的节点随机下发一轮任务。"""
    moment = now or utcnow()
    expire_overdue_tasks(session, moment)
    pending_node_ids = set(
        session.scalars(select(ComputeTask.node_id).where(ComputeTask.status == TaskStatus.PENDING)).all()
    )
    candidates = [
        node
        for node in session.scalars(
            select(ComputeNode).where(
                ComputeNode.status == NodeStatus.ONLINE,
                ComputeNode.license_expires_at > moment,
            )
        ).all()
        if node.id not in pending_node_ids
    ]
    if not candidates:
        return []

    batch = min(len(candidates), settings.dispatch_batch_size)
    chosen = random.sample(candidates, batch)
    deadline = moment + timedelta(minutes=settings.task_deadline_minutes)
    tasks: list[ComputeTask] = []
    for node in chosen:
        task = ComputeTask(
            id=str(uuid.uuid4()),
            node_id=node.id,
            task_type="simulated_compute",
            payload={
                "op": "simulated_compute",
                "nonce": secrets.token_hex(16),
                "seed": secrets.token_hex(8),
            },
            status=TaskStatus.PENDING,
            push_status=PushStatus.QUEUED,
            dispatched_at=moment,
            deadline_at=deadline,
        )
        session.add(task)
        tasks.append(task)
    session.flush()
    for node, task in zip(chosen, tasks):
        _push_task(node, task)
    session.flush()
    return tasks


def list_pending_tasks(session: Session, sn: str, now: datetime | None = None) -> list[ComputeTask]:
    moment = now or utcnow()
    node = get_node_by_sn(session, sn.strip())
    tasks = session.scalars(
        select(ComputeTask)
        .where(ComputeTask.node_id == node.id, ComputeTask.status == TaskStatus.PENDING)
        .order_by(ComputeTask.dispatched_at)
    ).all()
    visible: list[ComputeTask] = []
    for task in tasks:
        if task.deadline_at < moment:
            task.status = TaskStatus.EXPIRED
            continue
        visible.append(task)
    session.flush()
    return visible


def report_task(
    session: Session,
    sn: str,
    task_id: str,
    success: bool,
    nonce: str,
    now: datetime | None = None,
) -> ComputeTask:
    """校验任务归属、时限和 nonce。只有校验通过且 success 为真才记为有效响应。"""
    moment = now or utcnow()
    node = get_node_by_sn(session, sn.strip())
    task = session.get(ComputeTask, task_id)
    if task is None:
        raise AppError("任务不存在", 404)
    if task.node_id != node.id:
        raise AppError("任务不属于该节点", 403)
    if task.status != TaskStatus.PENDING:
        raise AppError("任务已结束，不能重复回传", 409)
    if task.deadline_at < moment:
        task.status = TaskStatus.EXPIRED
        session.flush()
        raise AppError("任务已超过回传时限", 409)

    expected = str((task.payload or {}).get("nonce", ""))
    if not _nonce_matches(nonce, expected):
        raise AppError("任务响应校验失败")

    task.responded_at = moment
    task.success = success
    task.status = TaskStatus.COMPLETED if success else TaskStatus.FAILED
    session.flush()
    return task


def expire_overdue_tasks(session: Session, now: datetime | None = None) -> int:
    moment = now or utcnow()
    tasks = session.scalars(
        select(ComputeTask).where(
            ComputeTask.status == TaskStatus.PENDING,
            ComputeTask.deadline_at < moment,
        )
    ).all()
    for task in tasks:
        task.status = TaskStatus.EXPIRED
    if tasks:
        session.flush()
    return len(tasks)


def _nonce_matches(provided: str, expected: str) -> bool:
    if not expected or len(provided) != len(expected):
        return False
    return hmac.compare_digest(provided, expected)


def _push_task(node: ComputeNode, task: ComputeTask) -> None:
    if not node.callback_url:
        task.push_status = PushStatus.QUEUED
        return
    body = json.dumps(
        {
            "task_id": task.id,
            "sn": node.sn,
            "task_type": task.task_type,
            "payload": task.payload,
            "deadline_at": task.deadline_at.isoformat(),
        }
    ).encode()
    request = Request(
        node.callback_url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urlopen(request, timeout=2) as response:
            if 200 <= response.status < 300:
                task.push_status = PushStatus.PUSHED
                task.push_error = None
                return
            task.push_status = PushStatus.FAILED
            task.push_error = f"HTTP {response.status}"
    except (URLError, TimeoutError, OSError) as exc:
        task.push_status = PushStatus.FAILED
        task.push_error = str(exc)[:255]
        logger.warning("task push failed sn=%s task=%s error=%s", node.sn, task.id, exc)