自动排队生成量化因子挖掘任务，心跳派发后将高 IC 因子入库。

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-10-11 01:49:16 +08:00
parent f0b35e6b34
commit 598cecad5f
11 changed files with 529 additions and 23 deletions
@@ -0,0 +1,200 @@
+"""量化因子挖掘的排队与自动调度。
+
+没有待领取的 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