diff --git a/backend/app/api/kline.py b/backend/app/api/kline.py index 5fb6356..ac934e7 100644 --- a/backend/app/api/kline.py +++ b/backend/app/api/kline.py @@ -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), ) diff --git a/backend/app/api/pipeline.py b/backend/app/api/pipeline.py index a0d4672..2d97379 100644 --- a/backend/app/api/pipeline.py +++ b/backend/app/api/pipeline.py @@ -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, # 拉取线程在分块回调处自行退出, 这里无需(也无法)再写状态。 diff --git a/backend/app/backtest/factor.py b/backend/app/backtest/factor.py index 2a7518c..fd9acc7 100644 --- a/backend/app/backtest/factor.py +++ b/backend/app/backtest/factor.py @@ -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() diff --git a/backend/app/indicators/pipeline.py b/backend/app/indicators/pipeline.py index 480fa06..25917d0 100644 --- a/backend/app/indicators/pipeline.py +++ b/backend/app/indicators/pipeline.py @@ -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" diff --git a/backend/app/jobs/daily_pipeline.py b/backend/app/jobs/daily_pipeline.py index 3c3e70f..b10233f 100644 --- a/backend/app/jobs/daily_pipeline.py +++ b/backend/app/jobs/daily_pipeline.py @@ -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) diff --git a/backend/app/main.py b/backend/app/main.py index f54e3c4..e9beea5 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -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( diff --git a/backend/app/services/backtest.py b/backend/app/services/backtest.py index 769b17a..c58171f 100644 --- a/backend/app/services/backtest.py +++ b/backend/app/services/backtest.py @@ -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] diff --git a/backend/app/services/data_integrity.py b/backend/app/services/data_integrity.py index d0b0cae..34df520 100644 --- a/backend/app/services/data_integrity.py +++ b/backend/app/services/data_integrity.py @@ -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: diff --git a/backend/app/services/heavy_job_limiter.py b/backend/app/services/heavy_job_limiter.py index 65c7e88..ca37f43 100644 --- a/backend/app/services/heavy_job_limiter.py +++ b/backend/app/services/heavy_job_limiter.py @@ -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 diff --git a/backend/app/services/pipeline_jobs.py b/backend/app/services/pipeline_jobs.py index 08160e5..91d2641 100644 --- a/backend/app/services/pipeline_jobs.py +++ b/backend/app/services/pipeline_jobs.py @@ -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 diff --git a/backend/app/tickflow/repository.py b/backend/app/tickflow/repository.py index aa42645..0098bbb 100644 --- a/backend/app/tickflow/repository.py +++ b/backend/app/tickflow/repository.py @@ -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: diff --git a/backend/tests/test_enriched_incremental_memory.py b/backend/tests/test_enriched_incremental_memory.py new file mode 100644 index 0000000..0043c49 --- /dev/null +++ b/backend/tests/test_enriched_incremental_memory.py @@ -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 diff --git a/backend/tests/test_enriched_range_warming_guard.py b/backend/tests/test_enriched_range_warming_guard.py index ece7f2b..94e5d45 100644 --- a/backend/tests/test_enriched_range_warming_guard.py +++ b/backend/tests/test_enriched_range_warming_guard.py @@ -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"]) diff --git a/backend/tests/test_enriched_refresh_memory.py b/backend/tests/test_enriched_refresh_memory.py new file mode 100644 index 0000000..e7f1cb4 --- /dev/null +++ b/backend/tests/test_enriched_refresh_memory.py @@ -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")) diff --git a/backend/tests/test_heavy_job_limiter.py b/backend/tests/test_heavy_job_limiter.py index 5ebc56d..a62ff3c 100644 --- a/backend/tests/test_heavy_job_limiter.py +++ b/backend/tests/test_heavy_job_limiter.py @@ -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 diff --git a/backend/tests/test_pipeline_capacity.py b/backend/tests/test_pipeline_capacity.py new file mode 100644 index 0000000..b10ee0d --- /dev/null +++ b/backend/tests/test_pipeline_capacity.py @@ -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