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:
@@ -0,0 +1,98 @@
from datetime import date, timedelta
import polars as pl
import pytest
from polars.testing import assert_frame_equal
from app.enriched_generation import get_enriched_generation
from app.indicators import pipeline
from app.services import preferences
@pytest.fixture
def sample(tmp_path, monkeypatch):
monkeypatch.setattr(preferences, "get_enriched_batch_size", lambda: 2)
monkeypatch.setattr(pipeline, "_custom_signal_exprs", {})
monkeypatch.setattr(pipeline, "_adaptive_sym_batch", lambda default, rows: 2)
symbols = [f"{600000 + i}.SH" for i in range(5)]
rows = []
for day in range(12):
for i, symbol in enumerate(symbols):
close = 10.0 + i + day * 0.1
rows.append({
"symbol": symbol, "date": date(2026, 8, 1) + timedelta(days=day),
"open": close, "high": close + 0.1, "low": close - 0.1,
"close": close, "volume": 0.0 if i == 4 else 1000.0 + day,
"amount": 0.0 if i == 4 else 10000.0, "quote_ts": 0,
})
raw = pl.DataFrame(rows)
for frame in raw.partition_by("date"):
out = tmp_path / "kline_daily" / f"date={frame['date'][0]}" / "part.parquet"
out.parent.mkdir(parents=True)
frame.write_parquet(out)
instruments = pl.DataFrame({
"symbol": symbols, "name": ["stock"] * 5, "float_shares": [1000000.0] * 5,
})
factors = pl.DataFrame({
"symbol": symbols[:3], "trade_date": [date(2026, 8, 8)] * 3, "ex_factor": [1.1] * 3,
})
shares = pl.DataFrame({
"symbol": symbols[:2], "period_end": [date(2026, 6, 30)] * 2,
"announce_date": [date(2026, 8, 5)] * 2, "float_shares": [800000.0] * 2,
})
for name, frame in (("instruments/all.parquet", instruments), ("adj_factor/all.parquet", factors),
("financials/shares/part.parquet", shares)):
out = tmp_path / name
out.parent.mkdir(parents=True, exist_ok=True)
frame.write_parquet(out)
return raw, instruments, factors, shares
@pytest.mark.parametrize("adjustment_only", [False, True])
def test_incremental_batches_preserve_values_and_generation(sample, tmp_path, monkeypatch, adjustment_only):
raw, instruments, factors, shares = sample
expected = pipeline._select_storage_cols(pipeline.compute_enriched(
raw, instruments=instruments, factors=factors, historical_shares=shares,
)).sort("symbol", "date")
if adjustment_only:
for frame in raw.partition_by("date"):
out = tmp_path / "kline_daily_enriched" / f"date={frame['date'][0]}" / "part.parquet"
out.parent.mkdir(parents=True)
frame.write_parquet(out)
monkeypatch.setattr(pipeline, "_load_recent_history", lambda *args, **kwargs: pl.DataFrame())
original_compute = pipeline.compute_enriched
def bounded_compute(frame, **kwargs):
assert frame["symbol"].n_unique() <= 2, "incremental wide compute must be batched"
return original_compute(frame, **kwargs)
monkeypatch.setattr(pipeline, "compute_enriched", bounded_compute)
generation = get_enriched_generation(tmp_path)
symbols = raw["symbol"].unique().to_list() if adjustment_only else None
written = pipeline.run_pipeline(tmp_path, symbols=symbols, new_dates_only=True)
actual = pl.read_parquet(str(tmp_path / "kline_daily_enriched/**/*.parquet"))
if adjustment_only:
actual = actual.select(expected.columns)
assert written > 0
assert_frame_equal(actual.sort("symbol", "date"), expected, check_exact=True)
assert get_enriched_generation(tmp_path) != generation
def test_failed_compute_batch_publishes_no_partial_new_dates(sample, tmp_path, monkeypatch):
original_compute = pipeline.compute_enriched
calls = 0
def fail_second(frame, **kwargs):
nonlocal calls
calls += 1
if calls == 2:
raise RuntimeError("batch failed")
return original_compute(frame, **kwargs)
monkeypatch.setattr(pipeline, "compute_enriched", fail_second)
monkeypatch.setattr(pipeline, "_load_recent_history", lambda *args, **kwargs: pl.DataFrame())
generation = get_enriched_generation(tmp_path)
with pytest.raises(RuntimeError, match="batch failed"):
pipeline.run_pipeline(tmp_path, new_dates_only=True)
assert not list((tmp_path / "kline_daily_enriched").rglob("*.parquet"))
assert get_enriched_generation(tmp_path) == generation
@@ -9,6 +9,7 @@ from __future__ import annotations
from datetime import date
import polars as pl
import pytest
from app.tickflow.repository import KlineRepository
@@ -54,3 +55,39 @@ def test_get_enriched_range_rebuilds_when_cold_and_not_warming():
assert result is not None
assert result.height == 2
assert result["symbol"].unique().to_list() == ["600000.SH"]
@pytest.mark.parametrize("symbols", [None, ["600000.SH"], ["missing"], []])
def test_range_projects_before_filter_and_preserves_empty_schema(monkeypatch, symbols):
repo = _bare_repo()
cache = pl.DataFrame({
"symbol": ["600000.SH", "600001.SH"],
"date": [date(2026, 1, 1), date(2026, 1, 2)],
"close": [10.0, 11.0], "unused": [100.0, 200.0],
})
repo._enriched_history_cache = cache
original_filter = pl.DataFrame.filter
def narrow_filter(frame, *args, **kwargs):
assert "unused" not in frame.columns
return original_filter(frame, *args, **kwargs)
monkeypatch.setattr(pl.DataFrame, "filter", narrow_filter)
result = repo.get_enriched_range(date(2026, 1, 1), date(2026, 1, 2), symbols, ["close", "absent"])
assert result is not None
if symbols is not None and "600000.SH" not in symbols:
assert result.is_empty()
assert result.schema == cache.schema
else:
assert result.columns == ["symbol", "date", "close"]
assert result.height == (2 if symbols is None else 1)
def test_range_duplicate_columns_keep_legacy_empty_behavior():
repo = _bare_repo()
cache = pl.DataFrame({"symbol": ["600000.SH"], "date": [date(2026, 1, 1)], "close": [10.0]})
repo._enriched_history_cache = cache
empty = repo.get_enriched_range(date(2026, 1, 1), date(2026, 1, 1), [], ["close", "close"])
assert empty.schema == cache.schema
with pytest.raises(pl.exceptions.DuplicateError):
repo.get_enriched_range(date(2026, 1, 1), date(2026, 1, 1), None, ["close", "close"])
@@ -0,0 +1,80 @@
from datetime import date, timedelta
import polars as pl
import pytest
from polars.testing import assert_frame_equal
from app.indicators import pipeline
@pytest.fixture
def history(monkeypatch):
monkeypatch.setattr(pipeline, "_custom_signal_exprs", {
"signal_test_previous": (pl.col("close") > pl.col("close").shift(1).over("symbol")),
})
symbols = ["600000.SH", "300001.SZ", "688001.SH", "000001.SZ", "920001.BJ"]
rows = []
for i, symbol in enumerate(symbols):
for day in range(160):
close = 10 + i + day * 0.01 + (day % 7) * 0.2
rows.append({
"symbol": symbol, "date": date(2025, 1, 1) + timedelta(days=day),
"open": close - 0.1, "high": close + 0.3, "low": close - 0.3,
"close": close, "volume": float(1000 + day * 7), "amount": close * 1000,
"raw_close": close, "raw_high": close + 0.3, "raw_low": close - 0.3,
})
return pl.DataFrame(rows).sample(fraction=1, shuffle=True, seed=17)
@pytest.mark.parametrize("with_shares", [False, True])
def test_history_metadata_matches_original_refresh(history, tmp_path, with_shares):
symbols = history["symbol"].unique().sort().to_list()
instruments = pl.DataFrame({
"symbol": symbols[:-1], "name": ["stock", "ST stock", "stock", "stock"],
"total_shares": [2000000.0] * 4, "float_shares": [1000000.0] * 4,
})
shares = pl.DataFrame({
"symbol": [symbols[0]], "period_end": [date(2025, 2, 1)],
"announce_date": [date(2025, 3, 1)], "float_shares": [500000.0],
}) if with_shares else None
benchmark = history.filter(pl.col("symbol") == symbols[0]).with_columns(
pl.lit("000001.SH").alias("symbol"),
)
index_path = tmp_path / "kline_index_daily" / "part.parquet"
index_path.parent.mkdir()
benchmark.write_parquet(index_path)
expected = pipeline.compute_limit_signals(
pipeline.compute_signals(pipeline.attach_deviation_columns(
pipeline.compute_indicators(history.sort("symbol", "date")), tmp_path,
)),
instruments, historical_shares=shares,
)
missing = [c for c in ("name", "total_shares", "float_shares") if c not in expected.columns]
expected = expected.join(instruments.select("symbol", *missing), on="symbol", how="left")
actual = pipeline.compute_enriched_history_window(
history, tmp_path, instruments=instruments, historical_shares=shares, sym_batch=2,
include_instrument_metadata=True,
)
assert_frame_equal(actual, expected.sort("symbol", "date"), check_exact=True)
@pytest.mark.parametrize("instruments", [None, pl.DataFrame(), pl.DataFrame({"symbol": ["600000.SH"]})])
def test_history_optional_metadata_keeps_legacy_inputs(history, tmp_path, instruments):
expected = pipeline.compute_enriched_history_window(history, tmp_path, instruments, sym_batch=2)
actual = pipeline.compute_enriched_history_window(
history, tmp_path, instruments, sym_batch=2, include_instrument_metadata=True,
)
assert_frame_equal(actual, expected, check_exact=True)
def test_history_wide_sort_is_bounded_by_batch(history, tmp_path, monkeypatch):
original_sort = pl.DataFrame.sort
def bounded_sort(frame, *args, **kwargs):
if frame.width > history.width:
assert frame["symbol"].n_unique() <= 2, "full history wide sort copies the cache"
return original_sort(frame, *args, **kwargs)
monkeypatch.setattr(pl.DataFrame, "sort", bounded_sort)
result = pipeline.compute_enriched_history_window(history, tmp_path, sym_batch=2)
assert result.select("symbol", "date").equals(history.select("symbol", "date").sort("symbol", "date"))
+64
View File
@@ -94,3 +94,67 @@ def test_invalid_release_does_not_overfill_capacity() -> None:
def test_module_aliases_share_the_capacity_two_singleton() -> None:
assert heavy_job_limiter is shared_heavy_job_limiter
assert shared_heavy_job_limiter.capacity == 2
def test_exclusive_job_waits_for_all_normal_jobs() -> None:
limiter = HeavyJobLimiter(capacity=2)
assert limiter.acquire("normal", timeout=0)
assert not limiter.acquire("exclusive", timeout=0)
limiter.release("normal")
with limiter.slot("exclusive", timeout=0):
assert limiter.in_use == 2
assert not limiter.acquire("normal", timeout=0)
assert limiter.in_use == 0
def test_nested_cache_refresh_reuses_exclusive_reservation() -> None:
limiter = HeavyJobLimiter(capacity=2)
with limiter.slot("exclusive", timeout=0):
with limiter.slot("exclusive", timeout=0), limiter.slot("normal", timeout=0):
assert limiter.in_use == 2
assert limiter.in_use == 2
assert limiter.in_use == 0
def test_nested_upgrade_fails_instead_of_deadlocking() -> None:
limiter = HeavyJobLimiter(capacity=2)
with (
limiter.slot("normal", timeout=0),
pytest.raises(RuntimeError, match="upgrade"),
limiter.slot("exclusive", timeout=0),
):
pytest.fail("unreachable")
assert limiter.in_use == 0
def test_waiting_exclusive_job_cannot_be_overtaken() -> None:
limiter = HeavyJobLimiter(capacity=2)
assert limiter.acquire("normal", timeout=0)
entered = threading.Event()
finish = threading.Event()
def exclusive():
with limiter.slot("exclusive", timeout=2):
entered.set()
assert finish.wait(2)
waiter = threading.Thread(target=exclusive)
waiter.start()
try:
deadline = time.monotonic() + 1
while time.monotonic() < deadline:
with limiter._condition:
if limiter._waiters:
break
time.sleep(0.005)
else:
pytest.fail("exclusive job did not queue")
# One normal slot is free, but belongs to the exclusive job ahead.
assert not limiter.acquire("normal", timeout=0)
limiter.release("normal")
assert entered.wait(1)
finally:
finish.set()
waiter.join(2)
assert not waiter.is_alive()
assert limiter.in_use == 0
+286
View File
@@ -0,0 +1,286 @@
from __future__ import annotations
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, date, datetime, timedelta
from types import SimpleNamespace
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from app.services import heavy_job_limiter, pipeline_jobs, preferences
from app.services.heavy_job_limiter import HeavyJobLimiter
from app.services.pipeline_jobs import JobCancelledError, JobStore, run_with_capacity
from app.tickflow.repository import KlineRepository
def wait_for(predicate, timeout=2):
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if predicate():
return
time.sleep(0.005)
raise AssertionError("condition did not become true")
def waiting(store, jid):
log = store.get(jid)["log"]
return bool(log) and log[-1]["stage"] == "init"
@pytest.fixture
def capacity(monkeypatch, tmp_path):
monkeypatch.setattr(preferences, "load", lambda: {})
monkeypatch.setattr(pipeline_jobs, "_CANCEL_FLAGS", {})
monkeypatch.setattr(pipeline_jobs, "_run_slot_owner", None)
store = JobStore(store_dir=tmp_path / "jobs")
limiter = HeavyJobLimiter(capacity=2, cancel_poll_interval=0.005)
monkeypatch.setattr(pipeline_jobs, "job_store", store)
monkeypatch.setattr(heavy_job_limiter, "shared_heavy_job_limiter", limiter)
return store, limiter
def test_queued_pipeline_is_cancellable_and_not_reaped(capacity):
store, limiter = capacity
jid, _ = store.create(timeout_s=1)
called = threading.Event()
assert limiter.acquire("normal", timeout=0)
with ThreadPoolExecutor(max_workers=1) as pool:
future = pool.submit(run_with_capacity, jid, called.set)
try:
wait_for(lambda: waiting(store, jid))
assert store.get(jid)["status"] == "pending"
stale = (datetime.now(UTC) - timedelta(hours=2)).isoformat()
store._active_jobs[jid]["started_at"] = stale
store._active_jobs[jid]["last_progress_at"] = stale
store.reap_stale()
assert store.get(jid)["status"] == "pending"
store.terminate(jid, "cancelled in queue")
with pytest.raises(JobCancelledError):
future.result(timeout=1)
assert not called.is_set()
assert limiter.in_use == 1
finally:
limiter.release("normal")
assert limiter.in_use == 0
def test_cancel_does_not_release_running_worker_capacity(capacity):
store, limiter = capacity
jid, _ = store.create()
entered, finish = threading.Event(), threading.Event()
def work():
entered.set()
assert finish.wait(2)
raise JobCancelledError(jid)
with ThreadPoolExecutor(max_workers=1) as pool:
future = pool.submit(run_with_capacity, jid, work)
try:
assert entered.wait(1)
store.terminate(jid, "cancelled while computing")
assert store.get(jid)["status"] == "failed"
assert limiter.in_use == 2
assert not limiter.acquire("normal", timeout=0)
finally:
finish.set()
with pytest.raises(JobCancelledError):
future.result(timeout=1)
assert limiter.in_use == 0
def test_cache_refresh_queues_and_can_nest_in_pipeline(capacity, monkeypatch):
store, limiter = capacity
repo = object.__new__(KlineRepository)
calls = []
monkeypatch.setattr(repo, "_refresh_enriched_impl", lambda: calls.append(limiter.in_use))
assert limiter.acquire("normal", timeout=0)
with ThreadPoolExecutor(max_workers=1) as pool:
future = pool.submit(repo._refresh_enriched)
try:
wait_for(lambda: bool(limiter._waiters))
assert calls == []
finally:
limiter.release("normal")
future.result(timeout=1)
jid, _ = store.create()
run_with_capacity(jid, repo._refresh_enriched)
assert calls == [2, 2]
assert limiter.in_use == 0
@pytest.mark.parametrize("kind", ["signal", "factor", "factor_batch"])
def test_in_process_backtests_share_capacity(capacity, monkeypatch, kind):
from app.backtest.factor import FactorBacktestService
from app.services.backtest import BacktestService
_, limiter = capacity
cls = BacktestService if kind == "signal" else FactorBacktestService
service = object.__new__(cls)
method = "run_batch" if kind == "factor_batch" else "run"
received = []
config = object()
def compute(actual, **kwargs):
assert actual is config
assert limiter.in_use == 2
received.append(True)
return "unchanged result"
monkeypatch.setattr(service, "_" + method, compute)
assert limiter.acquire("normal", timeout=0)
with ThreadPoolExecutor(max_workers=1) as pool:
future = pool.submit(getattr(service, method), config)
try:
wait_for(lambda: bool(limiter._waiters))
assert received == []
finally:
limiter.release("normal")
assert future.result(timeout=1) == "unchanged result"
# Auto-mining evaluates factors within its existing two-slot reservation.
with limiter.slot("mining"):
assert getattr(service, method)(config) == "unchanged result"
assert limiter.in_use == 0
@pytest.mark.parametrize("endpoint,body,module_name,function_name,result", [
("/api/pipeline/run", {}, "app.jobs.daily_pipeline", "run_now", {}),
("/api/kline/extend_history", {"value": 1, "unit": "month"},
"app.services.extend_history", "run_extend_history", {}),
("/api/kline/repair_daily", {"start_date": "2026-01-01"},
"app.services.repair_daily", "run_repair_daily", {}),
("/api/kline/rebuild_enriched", {}, "app.indicators.pipeline", "run_pipeline", 3),
("/api/kline/sync_minute", {"days": 1},
"app.services.kline_sync", "sync_and_persist_minute", 3),
])
def test_api_jobs_wait_before_computing(
capacity, monkeypatch, tmp_path, endpoint, body, module_name, function_name, result,
):
import importlib
from app.api import kline, pipeline
from app.jobs import daily_pipeline
store, limiter = capacity
entered, finish = threading.Event(), threading.Event()
def compute(*args, **kwargs):
assert limiter.in_use == 2
entered.set()
assert finish.wait(3)
return result
monkeypatch.setattr(importlib.import_module(module_name), function_name, compute)
monkeypatch.setattr(pipeline, "job_store", store)
monkeypatch.setattr("app.tickflow.pools.get_pool", lambda *args: [])
monkeypatch.setattr(daily_pipeline, "_refresh_single_view", lambda *args: None)
app = FastAPI()
app.include_router(pipeline.router)
app.include_router(kline.router)
app.state.repo = SimpleNamespace(
store=SimpleNamespace(data_dir=tmp_path),
db=SimpleNamespace(execute=lambda *args: None),
refresh_cache=lambda: None,
get_index_symbol_set=lambda: set(),
)
app.state.capabilities = SimpleNamespace(has=lambda _: True)
assert limiter.acquire("normal", timeout=0)
released = False
with TestClient(app) as client:
try:
response = client.post(endpoint, json=body)
assert response.status_code == 200, response.text
jid = response.json()["job_id"]
wait_for(lambda: waiting(store, jid))
assert not entered.is_set()
assert client.get(f"/api/pipeline/jobs/{jid}").json()["status"] == "pending"
limiter.release("normal")
released = True
assert entered.wait(2)
assert not limiter.acquire("normal", timeout=0)
finish.set()
wait_for(lambda: store.get(jid)["status"] in ("succeeded", "failed"))
assert store.get(jid)["status"] == "succeeded", store.get(jid)
finally:
finish.set()
if not released:
limiter.release("normal")
wait_for(lambda: limiter.in_use == 0)
@pytest.mark.parametrize("fail", [False, True])
def test_manual_pipeline_holds_capacity_through_refresh_off_event_loop(capacity, monkeypatch, fail):
from app.api import pipeline
store, limiter = capacity
refreshing, finish = threading.Event(), threading.Event()
def compute(*args, **kwargs):
if fail:
raise ValueError("partial pipeline failure")
return {"rows": 3}
def refresh():
assert limiter.in_use == 2
refreshing.set()
assert finish.wait(3)
monkeypatch.setattr(pipeline, "job_store", store)
monkeypatch.setattr(pipeline.daily_pipeline, "run_now", compute)
app = FastAPI()
app.include_router(pipeline.router)
app.state.repo = SimpleNamespace(refresh_cache=refresh)
app.state.capabilities = object()
with TestClient(app) as client:
try:
jid = client.post("/api/pipeline/run").json()["job_id"]
assert refreshing.wait(1)
assert not limiter.acquire("normal", timeout=0)
# This request must complete while the worker is blocked in refresh.
assert client.get(f"/api/pipeline/jobs/{jid}").json()["status"] == "running"
finish.set()
wait_for(lambda: store.get(jid)["status"] in ("succeeded", "failed"))
assert store.get(jid)["status"] == ("failed" if fail else "succeeded")
finally:
finish.set()
assert limiter.in_use == 0
@pytest.mark.parametrize("entry", ["scheduled", "integrity"])
def test_background_pipeline_entries_share_capacity(capacity, monkeypatch, entry):
from app.jobs import daily_pipeline
from app.services import data_integrity, repair_daily
store, limiter = capacity
called = threading.Event()
def compute(*args, **kwargs):
assert limiter.in_use == 2
called.set()
return {}
monkeypatch.setattr(repair_daily, "run_repair_daily", compute)
assert limiter.acquire("normal", timeout=0)
with ThreadPoolExecutor(max_workers=1) as pool:
if entry == "scheduled":
future = pool.submit(daily_pipeline._run_tracked, compute, "test")
else:
state = SimpleNamespace(repo=object(), capabilities=SimpleNamespace(has=lambda _: True))
data_integrity.launch_integrity_repair(state, date(2026, 1, 1), "test")
future = None
try:
wait_for(lambda: store.active_id() is not None)
jid = store.active_id()
wait_for(lambda: waiting(store, jid))
assert store.get(jid)["status"] == "pending"
assert not called.is_set()
finally:
limiter.release("normal")
wait_for(lambda: store.get(jid)["status"] == "succeeded")
if future is not None:
assert future.result(timeout=1)
assert called.is_set()
assert limiter.in_use == 0