fix(oom): 降低补算内存峰值并统一重任务并发控制

This commit is contained in:
0112020179
2026-09-09 11:09:48 +08:00
parent f325a9676c
commit 771a95a6d1
16 changed files with 757 additions and 81 deletions
+8 -12
View File
@@ -1130,7 +1130,7 @@ async def sync_minute(request: Request):
"""
import asyncio
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, try_acquire_run_slot
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, run_with_capacity, try_acquire_run_slot
from app.api.data import invalidate_storage_cache
from app.services.preferences import get_minute_sync_days
from app.tickflow.capabilities import Cap
@@ -1167,7 +1167,6 @@ async def sync_minute(request: Request):
job_store.progress(job_id, stage, pct, msg)
try:
job_store.start(job_id)
progress("sync_minute", 5, "解析标的池…")
universe = sorted(set(get_pool("watchlist")) | set(get_pool("CN_Equity_A")))
# 补充 instruments 全量标的,覆盖北交所、新股等
@@ -1200,7 +1199,7 @@ async def sync_minute(request: Request):
on_chunk_done=_on_chunk,
)
written = await loop.run_in_executor(_long_task_executor, _run)
written = await loop.run_in_executor(_long_task_executor, run_with_capacity, job_id, _run)
# 刷新视图
from app.jobs.daily_pipeline import _refresh_single_view
@@ -1336,7 +1335,7 @@ async def extend_history(request: Request):
raise HTTPException(status_code=403, detail="需要 Pro+ 权限 (batch K-line)")
from app.services.extend_history import run_extend_history
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, try_acquire_run_slot
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, run_with_capacity, try_acquire_run_slot
from app.api.data import invalidate_storage_cache
job_id, is_new = job_store.create()
@@ -1355,9 +1354,8 @@ async def extend_history(request: Request):
stage_pct=stage_pct, skip_log=skip_log)
try:
job_store.start(job_id)
result = await loop.run_in_executor(
_long_task_executor,
_long_task_executor, run_with_capacity, job_id,
lambda: run_extend_history(repo, capset, value, unit, on_progress=progress),
)
if "error" in result:
@@ -1418,7 +1416,7 @@ async def repair_daily(request: Request):
raise HTTPException(status_code=403, detail="需要 Pro+ 权限 (batch K-line)")
from app.services.repair_daily import run_repair_daily
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, try_acquire_run_slot
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, run_with_capacity, try_acquire_run_slot
from app.api.data import invalidate_storage_cache
job_id, is_new = job_store.create()
@@ -1445,8 +1443,7 @@ async def repair_daily(request: Request):
return run_repair_daily(repo, capset, start_date, on_progress=progress)
try:
job_store.start(job_id)
result = await loop.run_in_executor(_long_task_executor, _run)
result = await loop.run_in_executor(_long_task_executor, run_with_capacity, job_id, _run)
if "error" in result:
job_store.fail(job_id, result["error"])
else:
@@ -1481,7 +1478,7 @@ async def rebuild_enriched(request: Request):
try:
repo = request.app.state.repo
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, try_acquire_run_slot
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, run_with_capacity, try_acquire_run_slot
from app.api.data import invalidate_storage_cache
job_id, is_new = job_store.create()
@@ -1500,7 +1497,6 @@ async def rebuild_enriched(request: Request):
stage_pct=stage_pct, skip_log=skip_log)
try:
job_store.start(job_id)
progress("rebuild_enriched", 10, "全量计算 enriched…")
from app.indicators.pipeline import run_pipeline
@@ -1511,7 +1507,7 @@ async def rebuild_enriched(request: Request):
stage_pct=int(100 * cur / tot), skip_log=True)
written = await loop.run_in_executor(
_long_task_executor,
_long_task_executor, run_with_capacity, job_id,
lambda: run_pipeline(on_batch_done=_batch_progress),
)
+9 -7
View File
@@ -12,6 +12,7 @@ from app.services.pipeline_jobs import (
JobCancelledError,
job_store,
release_run_slot,
run_with_capacity,
try_acquire_run_slot,
)
from app.api.data import invalidate_storage_cache
@@ -52,7 +53,6 @@ async def run_now(request: Request) -> dict:
# 管道运行期间暂停实时行情取数, 防止覆写同一批 parquet 竞态
qs = getattr(request.app.state, "quote_service", None)
try:
job_store.start(job_id)
loop = asyncio.get_event_loop()
def progress(stage: str, pct: int, msg: str, stage_pct: int | None = None,
@@ -60,15 +60,17 @@ async def run_now(request: Request) -> dict:
job_store.progress(job_id, stage, pct, msg, stage_pct=stage_pct, skip_log=skip_log)
def _run() -> dict:
if qs:
with qs.paused():
return daily_pipeline.run_now(repo, capset, on_progress=progress)
return daily_pipeline.run_now(repo, capset, on_progress=progress)
try:
if qs:
with qs.paused():
return daily_pipeline.run_now(repo, capset, on_progress=progress)
return daily_pipeline.run_now(repo, capset, on_progress=progress)
finally:
repo.refresh_cache()
result = await loop.run_in_executor(_long_task_executor, _run)
result = await loop.run_in_executor(_long_task_executor, run_with_capacity, job_id, _run)
job_store.succeed(job_id, result)
invalidate_storage_cache()
repo.refresh_cache() # 刷新 Polars 缓存
except JobCancelledError:
# 已被 reap/手动取消终止: job 状态已由 terminate() 写为 failed,
# 拉取线程在分块回调处自行退出, 这里无需(也无法)再写状态。
+22
View File
@@ -166,6 +166,17 @@ class FactorBacktestService:
config: FactorConfig,
*,
regime_by_date: Mapping[object, Any] | None = None,
) -> FactorResult:
from app.services.heavy_job_limiter import shared_heavy_job_limiter
with shared_heavy_job_limiter.slot("exclusive"):
return self._run(config, regime_by_date=regime_by_date)
def _run(
self,
config: FactorConfig,
*,
regime_by_date: Mapping[object, Any] | None = None,
) -> FactorResult:
t0 = time.perf_counter()
run_id = uuid.uuid4().hex[:10]
@@ -210,6 +221,17 @@ class FactorBacktestService:
config: FactorBatchConfig,
*,
regime_by_date: Mapping[object, Any] | None = None,
) -> FactorBatchResult:
from app.services.heavy_job_limiter import shared_heavy_job_limiter
with shared_heavy_job_limiter.slot("exclusive"):
return self._run_batch(config, regime_by_date=regime_by_date)
def _run_batch(
self,
config: FactorBatchConfig,
*,
regime_by_date: Mapping[object, Any] | None = None,
) -> FactorBatchResult:
"""在同一份 Panel 上依次评估多个因子, 避免重复读取和计算指标。"""
t0 = time.perf_counter()
+54 -9
View File
@@ -1391,12 +1391,13 @@ def compute_enriched_history_window(
instruments: pl.DataFrame | None = None,
historical_shares: pl.DataFrame | None = None,
sym_batch: int | None = None,
*,
include_instrument_metadata: bool = False,
) -> pl.DataFrame:
"""按 symbol 分批执行历史窗口计算: 指标 → 偏离列 → 信号 → 涨跌停。
与整帧顺序执行完全等价 (各步骤均 over("symbol") 分组), 分批约束
峰值内存: repository._refresh_enriched 的 300 天窗口在 5500 只、
210 交易日下整帧宽表 ~1.2GB, 小内存机器启动即 OOM (#208)。
分批约束计算中间表, 最终完整历史仍常驻内存。排序和可选元数据关联
均在单批完成, 避免合并后再复制整张宽表。
sym_batch 显式传入时跳过自适应 (测试用)。
"""
if df_hist.is_empty() or "symbol" not in df_hist.columns:
@@ -1420,9 +1421,52 @@ def compute_enriched_history_window(
else historical_shares
)
part = compute_limit_signals(part, inst_batch, historical_shares=shares_batch)
parts.append(part)
out = parts[0] if len(parts) == 1 else pl.concat(parts, how="diagonal_relaxed")
return out.sort(["symbol", "date"])
if include_instrument_metadata:
inst_cols = [c for c in ("name", "total_shares", "float_shares")
if c in inst_batch.columns and c not in part.columns]
if inst_cols:
part = part.join(
inst_batch.select("symbol", *inst_cols).unique(subset=["symbol"]),
on="symbol", how="left",
)
# 连续、互不重叠的已排序 symbol 批次, 拼接后天然有序。
parts.append(part.sort(["symbol", "date"]))
return parts[0] if len(parts) == 1 else pl.concat(parts, how="diagonal_relaxed", rechunk=False)
def _compute_storage_batches(
raw: pl.DataFrame,
*,
factors: pl.DataFrame,
instruments: pl.DataFrame,
historical_shares: pl.DataFrame,
) -> pl.DataFrame:
"""保留完整标的历史输入, 单批计算宽表后仅累积落盘窄表。"""
from app.services import preferences
if raw.is_empty():
return _select_storage_cols(raw)
symbols = raw["symbol"].unique().sort().to_list()
rows_per_sym = max(1, -(-raw.height // len(symbols)))
batch_size = _adaptive_sym_batch(preferences.get_enriched_batch_size(), rows_per_sym)
parts = []
for start in range(0, len(symbols), batch_size):
batch = symbols[start:start + batch_size]
part = compute_enriched(
raw.filter(pl.col("symbol").is_in(batch)),
factors=factors.filter(pl.col("symbol").is_in(batch)) if not factors.is_empty() else factors,
instruments=(instruments.filter(pl.col("symbol").is_in(batch))
if not instruments.is_empty() else instruments),
historical_shares=(historical_shares.filter(pl.col("symbol").is_in(batch))
if not historical_shares.is_empty() else historical_shares),
)
# 下一批开始前释放宽表; 分区发布仍在所有计算批次成功之后。
if not part.is_empty():
parts.append(_select_storage_cols(part))
del part
if not parts:
return _select_storage_cols(raw.head(0))
return pl.concat(parts, how="diagonal_relaxed", rechunk=False)
def run_pipeline(data_dir: Path | None = None,
@@ -1518,7 +1562,7 @@ def run_pipeline(data_dir: Path | None = None,
else:
raw_full = raw_new
enriched_new = compute_enriched(
enriched_new = _compute_storage_batches(
raw_full,
factors=factors,
instruments=instruments,
@@ -1549,6 +1593,7 @@ def run_pipeline(data_dir: Path | None = None,
written += date_df.height
t_write_new = _t.perf_counter()
logger.info("增量写入: %.2fs, %d", t_write_new - t_new, written)
del raw_new, hist_df, raw_full, enriched_new
# 3. 受除权因子影响的个股: 重算全部已有日期 (累积因子链变了)
if symbols:
@@ -1560,7 +1605,7 @@ def run_pipeline(data_dir: Path | None = None,
factors_sym = factors.filter(pl.col("symbol").is_in(list(sym_set))) if not factors.is_empty() else factors
inst_sym = instruments.filter(pl.col("symbol").is_in(list(sym_set))) if not instruments.is_empty() else instruments
shares_sym = historical_shares.filter(pl.col("symbol").is_in(list(sym_set))) if not historical_shares.is_empty() else historical_shares
enriched_sym = compute_enriched(
enriched_sym = _compute_storage_batches(
raw_sym,
factors=factors_sym,
instruments=inst_sym,
@@ -1680,7 +1725,7 @@ def run_pipeline(data_dir: Path | None = None,
if not enriched.is_empty():
if symbols:
# 局部模式: 直接按日期合并写入
for date_df in enriched.partition_by("date"):
for date_df in _select_storage_cols(enriched).partition_by("date"):
dt = date_df["date"][0]
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
out = base / f"date={ds}" / "part.parquet"
+2 -3
View File
@@ -846,7 +846,7 @@ def _run_tracked(fn, job_label: str) -> bool:
重任务执行槽: 再挡一层僵尸并发(reap 后线程仍活时不得并行写 parquet)。
返回 True 仅表示任务已成功并且执行槽已释放。
"""
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, try_acquire_run_slot
from app.services.pipeline_jobs import JobCancelledError, job_store, release_run_slot, run_with_capacity, try_acquire_run_slot
job_id, is_new = job_store.create()
if not is_new:
@@ -863,8 +863,7 @@ def _run_tracked(fn, job_label: str) -> bool:
succeeded = False
try:
job_store.start(job_id)
result = fn(on_progress=progress)
result = run_with_capacity(job_id, lambda: fn(on_progress=progress))
job_store.succeed(job_id, result)
succeeded = True
logger.info("scheduled %s completed: job_id=%s", job_label, job_id)
+1 -1
View File
@@ -288,7 +288,7 @@ async def _application_lifespan(app: FastAPI):
return
with shared_heavy_job_limiter.slot(
"normal",
"exclusive",
cancel_event=matrix_prewarm_owner.cancel_event,
):
result = prewarm_matrix_cache(
+6
View File
@@ -241,6 +241,12 @@ class BacktestService:
return result if result is not None else pd.DataFrame()
def run(self, config: BacktestConfig) -> BacktestResult:
from app.services.heavy_job_limiter import shared_heavy_job_limiter
with shared_heavy_job_limiter.slot("exclusive"):
return self._run(config)
def _run(self, config: BacktestConfig) -> BacktestResult:
vbt = _get_vbt()
run_id = uuid.uuid4().hex[:10]
+2 -2
View File
@@ -313,6 +313,7 @@ def launch_integrity_repair(app_state, start_date: date, reason: str) -> tuple[s
JobCancelledError,
job_store,
release_run_slot,
run_with_capacity,
try_acquire_run_slot,
)
from app.services.repair_daily import run_repair_daily
@@ -339,8 +340,7 @@ def launch_integrity_repair(app_state, start_date: date, reason: str) -> tuple[s
if not try_acquire_run_slot(job_id):
job_store.fail(job_id, "已有数据任务在运行(或上一次任务卡死未结束),请稍后再试")
return
job_store.start(job_id)
result = _run()
result = run_with_capacity(job_id, _run)
if isinstance(result, dict) and "error" in result:
job_store.fail(job_id, str(result["error"]))
else:
+44 -23
View File
@@ -4,11 +4,12 @@ from __future__ import annotations
import threading
import time
from collections import deque
from collections.abc import Iterator
from contextlib import contextmanager
from typing import ClassVar, Literal
HeavyJobKind = Literal["normal", "mining"]
HeavyJobKind = Literal["normal", "mining", "exclusive"]
class HeavyJobLimitTimeoutError(TimeoutError):
@@ -20,7 +21,7 @@ class HeavyJobCancelledError(RuntimeError):
class HeavyJobLimiter:
"""A weighted limiter where normal jobs cost one slot and mining costs two."""
"""FIFO weighted capacity; exclusive jobs reserve the entire process budget."""
_WEIGHTS: ClassVar[dict[HeavyJobKind, int]] = {"normal": 1, "mining": 2}
@@ -32,8 +33,10 @@ class HeavyJobLimiter:
self.capacity = capacity
self._cancel_poll_interval = cancel_poll_interval
self._used = 0
self._acquired = {"normal": 0, "mining": 0}
self._acquired = {"normal": 0, "mining": 0, "exclusive": 0}
self._condition = threading.Condition()
self._waiters: deque[object] = deque()
self._local = threading.local()
@property
def in_use(self) -> int:
@@ -61,23 +64,29 @@ class HeavyJobLimiter:
deadline = None if timeout is None else time.monotonic() + timeout
with self._condition:
while True:
if cancel_event is not None and cancel_event.is_set():
return False
if self._used + weight <= self.capacity:
self._used += weight
self._acquired[kind] += 1
return True
ticket = object()
self._waiters.append(ticket)
try:
while True:
if cancel_event is not None and cancel_event.is_set():
return False
if self._waiters[0] is ticket and self._used + weight <= self.capacity:
self._used += weight
self._acquired[kind] += 1
return True
remaining = None if deadline is None else deadline - time.monotonic()
if remaining is not None and remaining <= 0:
return False
wait_for = remaining
if cancel_event is not None:
wait_for = self._cancel_poll_interval
if remaining is not None:
wait_for = min(wait_for, remaining)
self._condition.wait(wait_for)
remaining = None if deadline is None else deadline - time.monotonic()
if remaining is not None and remaining <= 0:
return False
wait_for = remaining
if cancel_event is not None:
wait_for = self._cancel_poll_interval
if remaining is not None:
wait_for = min(wait_for, remaining)
self._condition.wait(wait_for)
finally:
self._waiters.remove(ticket)
self._condition.notify_all()
def release(self, kind: HeavyJobKind = "normal") -> None:
"""Return capacity previously acquired for ``kind``."""
@@ -97,21 +106,33 @@ class HeavyJobLimiter:
timeout: float | None = None,
cancel_event: threading.Event | None = None,
) -> Iterator[HeavyJobLimiter]:
"""Acquire weighted capacity for the duration of a ``with`` block."""
"""Reserve capacity in the executing thread, reusing an outer reservation."""
weight = self._weight(kind)
held = getattr(self._local, "weight", 0)
if held:
if weight > held:
raise RuntimeError("cannot upgrade a held heavy-job reservation")
if cancel_event is not None and cancel_event.is_set():
raise HeavyJobCancelledError(f"{kind} job was cancelled")
yield self
return
acquired = self.acquire(kind, timeout=timeout, cancel_event=cancel_event)
if not acquired:
if cancel_event is not None and cancel_event.is_set():
raise HeavyJobCancelledError(f"{kind} job was cancelled while waiting")
raise HeavyJobLimitTimeoutError(f"timed out waiting for {kind} job capacity")
try:
self._local.weight = weight
yield self
finally:
self._local.weight = 0
self.release(kind)
@classmethod
def _weight(cls, kind: HeavyJobKind) -> int:
def _weight(self, kind: HeavyJobKind) -> int:
if kind == "exclusive":
return self.capacity
try:
return cls._WEIGHTS[kind]
return self._WEIGHTS[kind]
except KeyError as exc:
raise ValueError(f"unsupported heavy job kind: {kind!r}") from exc
+24 -1
View File
@@ -19,9 +19,10 @@ import logging
import os
import threading
import uuid
from collections.abc import Callable
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Literal
from typing import Any, Literal, TypeVar
logger = logging.getLogger(__name__)
@@ -479,6 +480,28 @@ def _duration_s(j: dict[str, Any]) -> float | None:
# 进程内单例
job_store = JobStore()
_Result = TypeVar("_Result")
def run_with_capacity(job_id: str, fn: Callable[[], _Result]) -> _Result:
"""Wait in the worker, keeping the reservation until its real execution ends."""
from app.services.heavy_job_limiter import (
HeavyJobCancelledError,
shared_heavy_job_limiter,
)
job_store.progress(job_id, "init", 0, "等待其他计算任务完成…")
with _CANCEL_FLAGS_LOCK:
cancel_event = _CANCEL_FLAGS.get(job_id)
try:
with shared_heavy_job_limiter.slot("exclusive", cancel_event=cancel_event):
if is_cancelled(job_id):
raise JobCancelledError(job_id)
job_store.start(job_id)
return fn()
except HeavyJobCancelledError as exc:
raise JobCancelledError(job_id) from exc
# ================================================================
# 重任务互斥执行槽 — 防「僵尸并发」, 带所有权 token
+20 -23
View File
@@ -539,6 +539,12 @@ class KlineRepository:
self._index_enriched_cache_date = None
def _refresh_enriched(self) -> None:
from app.services.heavy_job_limiter import shared_heavy_job_limiter
with shared_heavy_job_limiter.slot("exclusive"):
self._refresh_enriched_impl()
def _refresh_enriched_impl(self) -> None:
"""从 parquet 加载 enriched 最新日到内存 + 构建聚合表。
enriched parquet 仅存 14 列基础数据。启动时读入历史数据并即时计算完整指标,
@@ -601,8 +607,7 @@ class KlineRepository:
if not df_hist.is_empty():
instruments = self._instruments_cache if self._instruments_cache is not None else pl.DataFrame()
# 分批执行 指标→偏离→信号→涨跌停 (与整帧顺序等价, 各步骤均
# over("symbol") 分组), 峰值内存与标的总量解耦 (#208)
# 分批计算并关联元数据, 保留完整历史, 限制宽表临时副本。
step = time.perf_counter()
logger.info("enriched refresh step start: compute window (batched)")
df_full = compute_enriched_history_window(
@@ -614,24 +619,12 @@ class KlineRepository:
if instruments is not None and not instruments.is_empty()
else None
),
include_instrument_metadata=True,
)
del df_hist
logger.info("enriched refresh step done: compute window rows=%d (%.2fs)",
len(df_full), time.perf_counter() - step)
# JOIN instruments 到完整历史 (filter_history/basic_filter 需要 name/股本等列)
if instruments is not None and not instruments.is_empty():
inst_cols = [c for c in ["name", "total_shares", "float_shares"]
if c in instruments.columns and c not in df_full.columns]
if inst_cols:
step = time.perf_counter()
logger.info("enriched refresh step start: join instruments")
df_full = df_full.join(
instruments.select(["symbol", *inst_cols]).unique(subset=["symbol"]),
on="symbol",
how="left",
)
logger.info("enriched refresh step done: join instruments (%.2fs)", time.perf_counter() - step)
# 缓存完整历史 (含指标+必要基础信息) 供 filter_history/backtest 直接复用
if self.get_matrix_data_generation("stock") != refresh_generation:
raise EnrichedGenerationUnavailableError(
@@ -772,9 +765,9 @@ class KlineRepository:
needed = [c for c in base_cols if c in hist_all.columns]
step = time.perf_counter()
logger.info("live agg step start: slice history cache")
df_hist = hist_all.filter(
df_hist = hist_all.select(needed).filter(
(pl.col("date") >= start_60d) & (pl.col("date") <= latest)
).select(needed).sort(["symbol", "date"])
).sort(["symbol", "date"])
logger.info("live agg step done: slice history cache rows=%d (%.2fs)", len(df_hist), time.perf_counter() - step)
state_cols = [
@@ -1236,16 +1229,20 @@ class KlineRepository:
if cache_min > start or cache_max < end:
return None
df = cache.filter((pl.col("date") >= start) & (pl.col("date") <= end))
if symbols is not None:
df = df.filter(pl.col("symbol").is_in(symbols))
if columns and not df.is_empty():
df = cache
if columns:
existing = [c for c in columns if c in df.columns]
if "symbol" not in existing and "symbol" in df.columns:
existing.insert(0, "symbol")
if "date" not in existing and "date" in df.columns:
existing.insert(1, "date")
df = df.select(existing)
df = df.select(list(dict.fromkeys(existing)))
df = df.filter((pl.col("date") >= start) & (pl.col("date") <= end))
if symbols is not None:
df = df.filter(pl.col("symbol").is_in(symbols))
if columns:
# 保持旧接口空结果的完整 schema, 非空时沿用请求列校验。
df = cache.clear() if df.is_empty() else df.select(existing)
return df.sort(["symbol", "date"])
def get_live_agg(self) -> pl.DataFrame: