mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 14:24:15 +08:00
fix(oom): 降低补算内存峰值并统一重任务并发控制
This commit is contained in:
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
# 拉取线程在分块回调处自行退出, 这里无需(也无法)再写状态。
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
@@ -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(
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"))
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user