mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 16:44: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:
|
||||
|
||||
Reference in New Issue
Block a user