HyperNode-AI/app/services/alpha_queue.py

"""量化因子挖掘的排队与自动调度。

没有待领取的 alpha_mining 任务时补一批；距上一批已满 2 小时时再补一批。
切片先不绑定节点，心跳到来时再改成 processing。高 IC 因子写入 alpha_factors。
"""

from __future__ import annotations

import random
import secrets
import uuid
from datetime import datetime, timedelta, timezone

from sqlalchemy import func, select
from sqlalchemy.orm import Session

from app.clock import utcnow
from app.config import settings
from app.models import AlphaFactor, ComputeNode, ComputeTask, PushStatus, TaskStatus

ALPHA_BATCH_HOURS = 2
MIN_ABS_IC = 0.03
SLICE_COUNT = 3
QUEUE_HOLD_DAYS = 7
OPERATORS = (
    "ts_mean",
    "ts_std",
    "rank",
    "delta",
    "decay_linear",
    "correlation",
    "ts_rank",
    "scale",
)
SYMBOLS = ("BTCUSDT", "ETHUSDT", "BNBUSDT", "SOLUSDT")


def maybe_generate_alpha_batch(
    session: Session,
    now: datetime | None = None,
    force: bool = False,
) -> list[ComputeTask]:
    """队列里没有 pending 的挖掘任务，或距上一批已满 2 小时，就生成新的一批切片。"""
    moment = now or utcnow()
    pending = session.scalar(
        select(func.count())
        .select_from(ComputeTask)
        .where(
            ComputeTask.task_type == "alpha_mining",
            ComputeTask.status == TaskStatus.PENDING,
        )
    ) or 0
    latest = _as_utc(
        session.scalar(select(func.max(ComputeTask.dispatched_at)).where(ComputeTask.task_type == "alpha_mining"))
    )
    due = latest is None or moment - latest >= timedelta(hours=ALPHA_BATCH_HOURS)
    if pending and not due and not force:
        return []
    return _create_alpha_batch(session, moment)


def assign_queued_alpha_task(
    session: Session,
    node: ComputeNode,
    now: datetime | None = None,
) -> ComputeTask | None:
    """把最早的一条未绑定切片派给当前节点。节点已有执行中的挖掘任务时不再加派。"""
    moment = now or utcnow()
    holding = session.scalar(
        select(ComputeTask.id).where(
            ComputeTask.node_id == node.id,
            ComputeTask.task_type == "alpha_mining",
            ComputeTask.status == TaskStatus.PROCESSING,
        )
    )
    if holding is not None:
        return None
    task = session.scalar(
        select(ComputeTask)
        .where(
            ComputeTask.task_type == "alpha_mining",
            ComputeTask.status == TaskStatus.PENDING,
            ComputeTask.node_id.is_(None),
        )
        .order_by(ComputeTask.dispatched_at, ComputeTask.id)
        .limit(1)
    )
    if task is None:
        return None
    task.node_id = node.id
    task.status = TaskStatus.PROCESSING
    task.deadline_at = moment + timedelta(minutes=settings.task_deadline_minutes)
    session.flush()
    return task


def ingest_alpha_factors(session: Session, task: ComputeTask) -> list[AlphaFactor]:
    """成功回传的挖掘结果里，只保存公式非空且绝对 IC 达到门槛的因子。"""
    if task.task_type != "alpha_mining" or task.status != TaskStatus.COMPLETED or not task.success:
        return []
    if task.node_id is None:
        return []
    payload = task.payload or {}
    saved: list[AlphaFactor] = []
    for item in _factor_items(task.result or {}):
        formula = str(item.get("formula") or item.get("alpha") or "").strip()
        ic = _as_float(item.get("ic"))
        if not formula or ic is None or abs(ic) < MIN_ABS_IC:
            continue
        symbol = item.get("symbol") or payload.get("symbol")
        seed = payload.get("seed")
        row = AlphaFactor(
            task_id=task.id,
            node_id=task.node_id,
            formula=formula,
            ic=ic,
            symbol=str(symbol) if symbol else None,
            seed=None if seed is None else str(seed),
            created_at=task.responded_at or utcnow(),
        )
        session.add(row)
        saved.append(row)
    if saved:
        session.flush()
    return saved


def _create_alpha_batch(session: Session, moment: datetime) -> list[ComputeTask]:
    seed = secrets.randbelow(2**31)
    rng = random.Random(seed)
    batch_id = uuid.uuid4().hex[:12]
    combos = _operator_combos(rng, SLICE_COUNT)
    symbols = rng.sample(SYMBOLS, k=SLICE_COUNT)
    deadline = moment + timedelta(days=QUEUE_HOLD_DAYS)
    tasks: list[ComputeTask] = []
    for combo, symbol in zip(combos, symbols):
        task = ComputeTask(
            id=str(uuid.uuid4()),
            node_id=None,
            task_type="alpha_mining",
            payload={
                "batch_id": batch_id,
                "seed": seed,
                "operators": combo,
                "symbol": symbol,
                "nonce": secrets.token_hex(16),
            },
            status=TaskStatus.PENDING,
            push_status=PushStatus.QUEUED,
            dispatched_at=moment,
            deadline_at=deadline,
        )
        session.add(task)
        tasks.append(task)
    session.flush()
    return tasks


def _operator_combos(rng: random.Random, count: int) -> list[list[str]]:
    combos: list[list[str]] = []
    seen: set[tuple[str, ...]] = set()
    while len(combos) < count:
        size = 2 if rng.random() < 0.7 else 3
        combo = tuple(sorted(rng.sample(OPERATORS, k=size)))
        if combo in seen:
            continue
        seen.add(combo)
        combos.append(list(combo))
    return combos


def _factor_items(result: dict) -> list[dict]:
    factors = result.get("factors")
    if isinstance(factors, list):
        return [item for item in factors if isinstance(item, dict)]
    if any(key in result for key in ("formula", "alpha", "ic")):
        return [result]
    return []


def _as_utc(value: datetime | str | None) -> datetime | None:
    if value is None:
        return None
    parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) if isinstance(value, str) else value
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed


def _as_float(value: object) -> float | None:
    if isinstance(value, bool) or value is None:
        return None
    if isinstance(value, (int, float)):
        return float(value)
    if isinstance(value, str):
        try:
            return float(value)
        except ValueError:
            return None
    return None