dekun
291 lines
9.6 KiB
Python
291 lines
9.6 KiB
Python
"""随机任务下发与回传校验。
|
|
|
|
节点不暴露回调地址时，任务留在待拉取队列，客户端用
|
|
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, selectinload
|
|
|
|
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.alpha_queue import AUTO_CYCLE, ingest_alpha_factors
|
|
from app.services.common import get_node_by_sn
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ADMIN_TASK_TYPES = {
|
|
"alpha_mining": "量化因子挖掘",
|
|
"llm_inference": "大模型推理",
|
|
"lora_training": "QLoRA 微调",
|
|
"stress_test": "简易压测",
|
|
}
|
|
_BUSY = (TaskStatus.PENDING, TaskStatus.PROCESSING)
|
|
|
|
|
|
def create_admin_tasks(
|
|
session: Session,
|
|
task_type: str,
|
|
payload: dict | None,
|
|
target_node_sn: str | None = None,
|
|
now: datetime | None = None,
|
|
) -> list[ComputeTask]:
|
|
"""向指定节点或当前全部在线节点各放一条任务。payload 里补上一次性 nonce。"""
|
|
if task_type not in ADMIN_TASK_TYPES:
|
|
raise AppError("不支持的任务类型", 400)
|
|
moment = now or utcnow()
|
|
body = dict(payload or {})
|
|
if len(json.dumps(body, ensure_ascii=False)) > 8000:
|
|
raise AppError("任务参数过长", 400)
|
|
targets = _dispatch_targets(session, target_node_sn, moment)
|
|
deadline = moment + timedelta(minutes=settings.task_deadline_minutes)
|
|
tasks: list[ComputeTask] = []
|
|
for node in targets:
|
|
task_payload = dict(body)
|
|
task_payload["nonce"] = secrets.token_hex(16)
|
|
task = ComputeTask(
|
|
id=str(uuid.uuid4()),
|
|
node_id=node.id,
|
|
task_type=task_type,
|
|
payload=task_payload,
|
|
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(targets, tasks):
|
|
_push_task(node, task)
|
|
session.flush()
|
|
return tasks
|
|
|
|
|
|
def list_admin_tasks(session: Session) -> list[ComputeTask]:
|
|
return list(
|
|
session.scalars(
|
|
select(ComputeTask)
|
|
.options(selectinload(ComputeTask.node))
|
|
.order_by(ComputeTask.dispatched_at.desc(), ComputeTask.id.desc())
|
|
).all()
|
|
)
|
|
|
|
|
|
def result_preview(result: dict | None) -> str:
|
|
if not result:
|
|
return ""
|
|
for key in ("formula", "alpha", "output", "summary"):
|
|
value = result.get(key)
|
|
if isinstance(value, str) and value.strip():
|
|
return value.strip()[:120]
|
|
for key in ("patch_bytes", "patch_size", "adapter_bytes"):
|
|
if result.get(key) is not None:
|
|
return f"补丁 {result[key]}"
|
|
text = json.dumps(result, ensure_ascii=False)
|
|
return text if len(text) <= 120 else text[:117] + "..."
|
|
|
|
|
|
def _dispatch_targets(session: Session, target_node_sn: str | None, moment: datetime) -> list[ComputeNode]:
|
|
sn = (target_node_sn or "").strip()
|
|
if sn:
|
|
node = get_node_by_sn(session, sn)
|
|
if node.is_banned:
|
|
raise AppError("节点已封禁，不能下发任务", 409)
|
|
if node.license_expires_at <= moment:
|
|
raise AppError("节点授权已到期，不能下发任务", 409)
|
|
return [node]
|
|
nodes = session.scalars(
|
|
select(ComputeNode).where(
|
|
ComputeNode.status == NodeStatus.ONLINE,
|
|
ComputeNode.license_expires_at > moment,
|
|
ComputeNode.is_banned.is_(False),
|
|
)
|
|
).all()
|
|
if not nodes:
|
|
raise AppError("当前没有可下发的在线节点", 409)
|
|
return list(nodes)
|
|
|
|
|
|
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.node_id.is_not(None),
|
|
ComputeTask.status.in_(_BUSY),
|
|
ComputeTask.task_type.notin_(AUTO_CYCLE),
|
|
)
|
|
).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.in_(_BUSY))
|
|
.order_by(ComputeTask.dispatched_at, ComputeTask.id)
|
|
).all()
|
|
visible: list[ComputeTask] = []
|
|
for task in tasks:
|
|
if task.deadline_at < moment:
|
|
task.status = TaskStatus.EXPIRED
|
|
continue
|
|
if task.status == TaskStatus.PENDING:
|
|
task.status = TaskStatus.PROCESSING
|
|
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,
|
|
result: dict | 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 not in _BUSY:
|
|
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.result = result or None
|
|
task.status = TaskStatus.COMPLETED if success else TaskStatus.FAILED
|
|
session.flush()
|
|
if task.status == TaskStatus.COMPLETED and task.task_type == "alpha_mining":
|
|
ingest_alpha_factors(session, task)
|
|
return task
|
|
|
|
|
|
def expire_overdue_tasks(session: Session, now: datetime | None = None) -> int:
|
|
moment = now or utcnow()
|
|
tasks = session.scalars(
|
|
select(ComputeTask).where(
|
|
ComputeTask.node_id.is_not(None),
|
|
ComputeTask.status.in_(_BUSY),
|
|
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)
|