diff --git a/backend/app/strategy/engine.py b/backend/app/strategy/engine.py index 2c28ae7..3082133 100644 --- a/backend/app/strategy/engine.py +++ b/backend/app/strategy/engine.py @@ -12,6 +12,7 @@ import sys import threading import time from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field, replace from datetime import date from pathlib import Path @@ -20,6 +21,7 @@ from typing import Any import numpy as np import polars as pl +from app.config import settings from app.strategy.scoring import ( SCORING_DIRECTION_LOW, effective_scoring, @@ -1119,8 +1121,15 @@ class StrategyEngine: overrides_map: dict | None = None, *, strategy_ids: list[str] | None = None, + parallel: bool = True, ) -> dict[str, StrategyResult]: - """批量执行策略;当前数据、历史和矩阵均来自同一个调用上下文。""" + """批量执行策略;当前数据、历史和矩阵均来自同一个调用上下文。 + + parallel=True 时用有界线程池并发执行: 策略对 context 是只读纯函数 + (polars 计算释放 GIL), 并发不改变结果, 逐策略耗时日志不变。composite + 子策略的递归 run_all 以 parallel=False 调用, 保证嵌套时线程总数仍 + 不超过 worker 上限, 不随叠加层数放大。 + """ if context.current is None: raise ValueError("strategy run_all context requires current data") df = context.current @@ -1164,15 +1173,22 @@ class StrategyEngine: params_map.get(sid), ) ) + matrix_t0 = time.perf_counter() shared_matrix = build_market_data_matrix( shared_history, field_columns=field_columns, ) + logger.info( + "run_all: shared matrix built in %.0fms (fields=%d)", + (time.perf_counter() - matrix_t0) * 1000, + len(field_columns), + ) results: dict[str, StrategyResult] = {} - for sid, _ in selected: - results[sid] = self.run( + def _execute(sid: str) -> tuple[str, StrategyResult]: + started = time.perf_counter() + result = self.run( sid, replace( context, @@ -1183,6 +1199,31 @@ class StrategyEngine: params=params_map.get(sid), overrides=overrides_map.get(sid), ) + elapsed_ms = (time.perf_counter() - started) * 1000 + # >=1s 打 INFO 供热点归因 (哪些策略吃掉了 run_all 的大头), 其余 DEBUG 防噪。 + log_fn = logger.info if elapsed_ms >= 1000 else logger.debug + log_fn( + "run_all: strategy %s took %.0fms (total=%d)", + sid, + elapsed_ms, + result.total, + ) + return sid, result + + workers = min(settings.strategy_run_all_workers, len(selected)) + if parallel and workers > 1: + with ThreadPoolExecutor( + max_workers=workers, thread_name_prefix="strategy-run" + ) as pool: + futures = [pool.submit(_execute, sid) for sid, _ in selected] + # 按原顺序收集: 首个失败策略的异常语义与串行执行一致。 + for future in futures: + sid, result = future.result() + results[sid] = result + else: + for sid, _ in selected: + sid, result = _execute(sid) + results[sid] = result return results @@ -1420,6 +1461,9 @@ class StrategyEngine: params_map={}, overrides_map=overrides_map, strategy_ids=child_ids, + # 嵌套调用串行: 父级 worker 已并发, 子级再开池会使线程总数随叠加 + # 层数放大 (4×4×...), 超出并发闸与核数的合理范围。 + parallel=False, ) ordered_results = [child_results[cid] for cid in child_ids] diff --git a/backend/tests/test_strategy_run_all_parallel.py b/backend/tests/test_strategy_run_all_parallel.py new file mode 100644 index 0000000..e6e10ec --- /dev/null +++ b/backend/tests/test_strategy_run_all_parallel.py @@ -0,0 +1,112 @@ +"""run_all 并行执行的等价性回归测试。 + +engine.run_all 支持有界线程池并发 (策略对共享 context 只读纯函数)。此处验证: +- parallel=True 与 parallel=False 对同一批策略产出完全一致的结果 (总数 + 标的集); +- 失败策略的异常语义一致 (原顺序首个失败抛出); +- composite 子策略的嵌套 run_all 不受影响 (递归恒串行)。 +""" +from __future__ import annotations + +from datetime import date +from pathlib import Path + +import polars as pl +import pytest + +from app.strategy.engine import StrategyDataContext, StrategyEngine + +_FILTER_TEMPLATE = ''' +import polars as pl + +META = {{ + "id": "{sid}", + "name": "{sid}", + "timeframes": ["1d"], + "asset_types": ["stock"], +}} +def filter(df, params): + return pl.col("close") > {threshold} +''' + + +def _write_strategy(directory: Path, code: str) -> None: + directory.mkdir(parents=True, exist_ok=True) + (directory / f"strategy_{abs(hash(code)) % 10**8}.py").write_text(code, encoding="utf-8") + + +def _make_engine(tmp_path: Path) -> StrategyEngine: + d = tmp_path / "strategies" + _write_strategy(d, _FILTER_TEMPLATE.format(sid="cheap_a", threshold=10.0)) + _write_strategy(d, _FILTER_TEMPLATE.format(sid="cheap_b", threshold=15.0)) + _write_strategy(d, _FILTER_TEMPLATE.format(sid="cheap_c", threshold=20.0)) + return StrategyEngine(strategy_dirs=[d]) + + +def _context() -> StrategyDataContext: + n = 30 + current = pl.DataFrame({ + "symbol": [f"{i:06d}.SZ" for i in range(n)], + "name": [f"股票{i}" for i in range(n)], + "open": [5.0 + i for i in range(n)], + "high": [5.5 + i for i in range(n)], + "low": [4.5 + i for i in range(n)], + "close": [5.0 + i for i in range(n)], + "volume": [1000.0 * (i + 1) for i in range(n)], + "amount": [5000.0 * (i + 1) for i in range(n)], + "turnover_rate": [1.0 + i * 0.1 for i in range(n)], + "total_shares": [1e8 for _ in range(n)], + "float_shares": [5e7 for _ in range(n)], + }) + return StrategyDataContext( + asset_type="stock", + timeframe="1d", + as_of=date(2026, 9, 7), + current=current, + ) + + +def _signature(results: dict) -> dict: + return { + sid: (r.total, tuple(sorted(row["symbol"] for row in r.rows))) + for sid, r in results.items() + } + + +def test_parallel_run_all_matches_sequential_results(tmp_path: Path) -> None: + engine = _make_engine(tmp_path) + context = _context() + # 关闭默认基础过滤, 让结果只取决于策略谓词本身 + overrides = { + meta["id"]: {"basic_filter": {"enabled": False}} + for meta in engine.list_strategies() + } + + sequential = engine.run_all(context, overrides_map=overrides, parallel=False) + parallel = engine.run_all(context, overrides_map=overrides, parallel=True) + + assert list(parallel) == list(sequential) # 结果键序一致 + assert _signature(parallel) == _signature(sequential) + # close 序列 5..34: >10 → 11..34 共 24 只; >15 → 19 只; >20 → 14 只 + assert sequential["cheap_a"].total == 24 + assert sequential["cheap_b"].total == 19 + assert sequential["cheap_c"].total == 14 + + +def test_parallel_run_all_preserves_failure_semantics(tmp_path: Path) -> None: + d = tmp_path / "strategies" + _write_strategy(d, _FILTER_TEMPLATE.format(sid="ok_first", threshold=10.0)) + _write_strategy( + d, + ''' +META = {"id": "boom", "name": "boom"} +def filter(df, params): + raise ValueError("injected strategy failure") +''', + ) + _write_strategy(d, _FILTER_TEMPLATE.format(sid="ok_last", threshold=15.0)) + engine = StrategyEngine(strategy_dirs=[d]) + context = _context() + + for parallel in (False, True): + with pytest.raises(ValueError, match="injected strategy failure"): + engine.run_all(context, parallel=parallel)