dekun
324 lines
9.6 KiB
Python
324 lines
9.6 KiB
Python
"""自动任务排队。
|
|
|
|
大模型推理、量化因子、微调按这个顺序轮流补一批。
|
|
没有待领取的自动任务时立刻补下一种；距上一批已满 2 小时也会再补一批。
|
|
切片先不绑定节点，心跳到来时再改成 processing。高 IC 因子写入 alpha_factors。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import random
|
|
import secrets
|
|
import uuid
|
|
from datetime import datetime, timedelta, timezone
|
|
|
|
from sqlalchemy import 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
|
|
AUTO_CYCLE = ("llm_inference", "alpha_mining", "lora_training")
|
|
LLM_MODEL = "qwen2.5-3b-instruct-q4_k_m"
|
|
LORA_BASE = "qwen2.5-3b"
|
|
PROMPTS = (
|
|
"用三句话说明去中心化算力网络做什么",
|
|
"列出显卡推理时的三个资源限制",
|
|
"解释量化因子的 IC 表示什么",
|
|
)
|
|
OPERATORS = (
|
|
"ts_mean",
|
|
"ts_std",
|
|
"rank",
|
|
"delta",
|
|
"decay_linear",
|
|
"correlation",
|
|
"ts_rank",
|
|
"scale",
|
|
)
|
|
SYMBOLS = ("BTCUSDT", "ETHUSDT", "BNBUSDT", "SOLUSDT")
|
|
|
|
|
|
def maybe_generate_auto_batch(
|
|
session: Session,
|
|
now: datetime | None = None,
|
|
force: bool = False,
|
|
) -> list[ComputeTask]:
|
|
"""没有待领取的自动任务，或距上一批已满 2 小时，就按轮询生成下一种任务。"""
|
|
moment = now or utcnow()
|
|
pending, latest = _auto_queue_state(session)
|
|
due = latest is None or moment - latest >= timedelta(hours=ALPHA_BATCH_HOURS)
|
|
if pending and not due and not force:
|
|
return []
|
|
kind = _next_cycle_type(session)
|
|
if kind == "llm_inference":
|
|
return _create_llm_batch(session, moment)
|
|
if kind == "lora_training":
|
|
return _create_lora_batch(session, moment)
|
|
return _create_alpha_batch(session, moment)
|
|
|
|
|
|
def maybe_generate_alpha_batch(
|
|
session: Session,
|
|
now: datetime | None = None,
|
|
force: bool = False,
|
|
) -> list[ComputeTask]:
|
|
return maybe_generate_auto_batch(session, now=now, force=force)
|
|
|
|
|
|
def assign_queued_auto_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.in_(AUTO_CYCLE),
|
|
ComputeTask.status == TaskStatus.PROCESSING,
|
|
)
|
|
)
|
|
if holding is not None:
|
|
return None
|
|
task = session.scalar(
|
|
select(ComputeTask)
|
|
.where(
|
|
ComputeTask.task_type.in_(AUTO_CYCLE),
|
|
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 assign_queued_alpha_task(
|
|
session: Session,
|
|
node: ComputeNode,
|
|
now: datetime | None = None,
|
|
) -> ComputeTask | None:
|
|
return assign_queued_auto_task(session, node, now=now)
|
|
|
|
|
|
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_llm_batch(session: Session, moment: datetime) -> list[ComputeTask]:
|
|
seed, rng, batch_id = _batch_identity()
|
|
prompts = rng.sample(PROMPTS, k=SLICE_COUNT)
|
|
tasks = [
|
|
_queue_slice(
|
|
session,
|
|
moment,
|
|
"llm_inference",
|
|
{
|
|
"auto": True,
|
|
"batch_id": batch_id,
|
|
"seed": seed,
|
|
"model": LLM_MODEL,
|
|
"prompt": prompt,
|
|
"max_tokens": 256,
|
|
"nonce": secrets.token_hex(16),
|
|
},
|
|
)
|
|
for prompt in prompts
|
|
]
|
|
session.flush()
|
|
return tasks
|
|
|
|
|
|
def _create_lora_batch(session: Session, moment: datetime) -> list[ComputeTask]:
|
|
seed, _rng, batch_id = _batch_identity()
|
|
tasks = [
|
|
_queue_slice(
|
|
session,
|
|
moment,
|
|
"lora_training",
|
|
{
|
|
"auto": True,
|
|
"batch_id": batch_id,
|
|
"seed": seed,
|
|
"base_model": LORA_BASE,
|
|
"rank": rank,
|
|
"steps": 100,
|
|
"nonce": secrets.token_hex(16),
|
|
},
|
|
)
|
|
for rank in (8, 16, 32)
|
|
]
|
|
session.flush()
|
|
return tasks
|
|
|
|
|
|
def _create_alpha_batch(session: Session, moment: datetime) -> list[ComputeTask]:
|
|
seed, rng, batch_id = _batch_identity()
|
|
combos = _operator_combos(rng, SLICE_COUNT)
|
|
symbols = rng.sample(SYMBOLS, k=SLICE_COUNT)
|
|
tasks = [
|
|
_queue_slice(
|
|
session,
|
|
moment,
|
|
"alpha_mining",
|
|
{
|
|
"auto": True,
|
|
"batch_id": batch_id,
|
|
"seed": seed,
|
|
"operators": combo,
|
|
"symbol": symbol,
|
|
"nonce": secrets.token_hex(16),
|
|
},
|
|
)
|
|
for combo, symbol in zip(combos, symbols)
|
|
]
|
|
session.flush()
|
|
return tasks
|
|
|
|
|
|
def _queue_slice(session: Session, moment: datetime, task_type: str, payload: dict) -> ComputeTask:
|
|
task = ComputeTask(
|
|
id=str(uuid.uuid4()),
|
|
node_id=None,
|
|
task_type=task_type,
|
|
payload=payload,
|
|
status=TaskStatus.PENDING,
|
|
push_status=PushStatus.QUEUED,
|
|
dispatched_at=moment,
|
|
deadline_at=moment + timedelta(days=QUEUE_HOLD_DAYS),
|
|
)
|
|
session.add(task)
|
|
return task
|
|
|
|
|
|
def _batch_identity() -> tuple[int, random.Random, str]:
|
|
seed = secrets.randbelow(2**31)
|
|
return seed, random.Random(seed), uuid.uuid4().hex[:12]
|
|
|
|
|
|
def _auto_queue_state(session: Session) -> tuple[int, datetime | None]:
|
|
rows = session.scalars(
|
|
select(ComputeTask).where(
|
|
ComputeTask.task_type.in_(AUTO_CYCLE),
|
|
ComputeTask.node_id.is_(None),
|
|
ComputeTask.status == TaskStatus.PENDING,
|
|
)
|
|
).all()
|
|
pending = sum(1 for row in rows if (row.payload or {}).get("auto"))
|
|
latest = _latest_auto_dispatched(session)
|
|
return pending, latest
|
|
|
|
|
|
def _latest_auto_dispatched(session: Session) -> datetime | None:
|
|
rows = session.scalars(
|
|
select(ComputeTask)
|
|
.where(ComputeTask.task_type.in_(AUTO_CYCLE))
|
|
.order_by(ComputeTask.dispatched_at.desc(), ComputeTask.id.desc())
|
|
.limit(100)
|
|
).all()
|
|
for row in rows:
|
|
if (row.payload or {}).get("auto"):
|
|
return _as_utc(row.dispatched_at)
|
|
return None
|
|
|
|
|
|
def _next_cycle_type(session: Session) -> str:
|
|
rows = session.scalars(
|
|
select(ComputeTask)
|
|
.where(ComputeTask.task_type.in_(AUTO_CYCLE))
|
|
.order_by(ComputeTask.dispatched_at.desc(), ComputeTask.id.desc())
|
|
.limit(100)
|
|
).all()
|
|
latest = None
|
|
for row in rows:
|
|
if (row.payload or {}).get("auto"):
|
|
latest = row.task_type
|
|
break
|
|
if latest not in AUTO_CYCLE:
|
|
return AUTO_CYCLE[0]
|
|
return AUTO_CYCLE[(AUTO_CYCLE.index(latest) + 1) % len(AUTO_CYCLE)]
|
|
|
|
|
|
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
|