feat: concurrent scanning with ProcessPoolExecutor

- Add workers param to SignalScanner.scan() (default=0 for serial)
- workers=2+ uses ProcessPoolExecutor for parallel .day file processing
- Extract _scan_one_file as top-level function for pickle compatibility
- Results identical between serial and parallel modes
- Add 4 tests with synthetic .day file fixtures
This commit is contained in:
Justin Gu
2026-06-11 02:09:10 +08:00
parent af005d9fe4
commit b7e0f17842
2 changed files with 236 additions and 2 deletions
+121 -2
View File
@@ -87,6 +87,7 @@ class SignalScanner:
self,
universe: str = "all",
progress_callback: Any = None,
workers: int = 0,
) -> list[ScanResult]:
"""扫描全市场,返回触发买入信号的股票列表。
@@ -97,6 +98,10 @@ class SignalScanner:
- "sz": 仅深圳
- 文件路径: 每行一个 "市场 代码"(如 "SZ 000001"
progress_callback: 进度回调函数(current, total, filename)
workers: 并发工作进程数
- 0: 串行模式(默认,向后兼容)
- 1: 串行但使用 executor 基础设施
- 2+: ProcessPoolExecutor 并发执行
Returns:
触发买入信号的 ScanResult 列表
@@ -107,9 +112,24 @@ class SignalScanner:
if not files:
return []
results: list[ScanResult] = []
total = len(files)
# 串行模式(workers=0,向后兼容)
if workers <= 0:
return self._scan_serial(files, total, progress_callback)
# 并发模式
return self._scan_parallel(files, total, workers, progress_callback)
def _scan_serial(
self,
files: list[tuple[Path, str, str]],
total: int,
progress_callback: Any,
) -> list[ScanResult]:
"""串行扫描(原有逻辑)。"""
results: list[ScanResult] = []
for idx, (filepath, market, code) in enumerate(files):
if progress_callback:
progress_callback(idx, total, filepath.name)
@@ -119,7 +139,6 @@ class SignalScanner:
if result is not None:
results.append(result)
except Exception:
# 单个文件出错不中断整体扫描
continue
if progress_callback:
@@ -127,6 +146,49 @@ class SignalScanner:
return results
def _scan_parallel(
self,
files: list[tuple[Path, str, str]],
total: int,
workers: int,
progress_callback: Any,
) -> list[ScanResult]:
"""并发扫描(ProcessPoolExecutor)。"""
import concurrent.futures
# 构建参数:每个任务需要的独立数据
tasks = [
(str(filepath), market, code, self._strategy_cls, self._cash, self._commission)
for filepath, market, code in files
]
results: list[ScanResult] = []
completed = 0
with concurrent.futures.ProcessPoolExecutor(max_workers=workers) as executor:
future_to_idx = {
executor.submit(_scan_one_file, *task): i for i, task in enumerate(tasks)
}
for future in concurrent.futures.as_completed(future_to_idx):
completed += 1
idx = future_to_idx[future]
if progress_callback:
progress_callback(completed, total, files[idx][0].name)
try:
result = future.result()
if result is not None:
results.append(result)
except Exception:
continue
if progress_callback:
progress_callback(total, total, "done")
return results
def _collect_files(self, universe: str) -> list[tuple[Path, str, str]]:
"""收集需要扫描的 .day 文件列表。
@@ -318,3 +380,60 @@ def _bars_to_df(bars: list[Any]) -> pd.DataFrame:
)
return pd.DataFrame(rows)
def _scan_one_file(
filepath: str,
market: str,
code: str,
strategy_cls: type[Strategy],
cash: float,
commission: float,
) -> ScanResult | None:
"""顶层扫描函数(供 ProcessPoolExecutor 调用)。
将实例方法逻辑提取为独立函数,避免 pickle 实例方法的问题。
逻辑与 SignalScanner._scan_one 完全一致。
Args:
filepath: .day 文件路径字符串
market: 市场代码(SZ/SH
code: 6 位股票代码
strategy_cls: Strategy 子类
cash: 初始资金
commission: 佣金率
Returns:
ScanResult 如果触发信号,否则 None
"""
bars = read_daily_bars(filepath)
if len(bars) < 30:
return None
df = _bars_to_df(bars)
if df.empty:
return None
try:
factor_signals = extract_factor_signals(
strategy_cls,
df,
cash=cash,
commission=commission,
)
except Exception:
return None
if not factor_signals.buy_mask[-1]:
return None
last_bar = bars[-1]
signal_date = last_bar.year * 10000 + last_bar.month * 100 + last_bar.day
last_close = last_bar.close
return ScanResult(
code=code,
market=market,
signal_date=signal_date,
last_close=last_close,
)
+115
View File
@@ -0,0 +1,115 @@
"""单元测试:信号扫描引擎。
测试 SignalScanner 的并发扫描和增量扫描功能。
使用临时目录构造 .day 文件 fixture,无需真实数据。
"""
from __future__ import annotations
import struct
from pathlib import Path
import pandas as pd
import pytest
from easy_tdx.backtest.strategy import Strategy
from easy_tdx.screen.scanner import SignalScanner
class AlwaysBuyStrategy(Strategy):
"""策略:每个 bar 都产生买入信号(用于扫描测试)。"""
def init(self) -> None:
pass
def next(self) -> None:
self.buy(size=0)
def _write_day_file(
path: Path,
n_bars: int = 50,
base_price: float = 10.0,
) -> None:
"""写一个最小的 .day 文件(通达信日线格式)。
格式: date(I) open(I) high(I) low(I) close(I) amount(f) vol(I) reserved(I)
每条 32 字节, 小端序. 价格以 0.01 为系数存储.
"""
dates = pd.date_range("2024-01-01", periods=n_bars, freq="D")
with open(path, "wb") as f:
for i in range(n_bars):
dt = dates[i]
day = dt.year * 10000 + dt.month * 100 + dt.day
price = base_price + i * 0.01
f.write(
struct.pack(
"<IIIIIfII",
day,
int(price * 100),
int((price + 0.5) * 100),
int((price - 0.5) * 100),
int(price * 100),
float(1000000 + i * 100),
10000 + i * 10,
0,
)
)
@pytest.fixture
def vipdoc(tmp_path: Path) -> Path:
"""创建包含 .day 文件的临时 vipdoc 目录."""
sz_lday = tmp_path / "sz" / "lday"
sz_lday.mkdir(parents=True)
for code in ("000001", "000002", "000003"):
_write_day_file(sz_lday / f"sz{code}.day", n_bars=50)
# 指数文件 (应被过滤)
_write_day_file(sz_lday / "sz399001.day", n_bars=50)
return tmp_path
class TestConcurrentScan:
"""测试并发扫描."""
def test_scan_produces_results(self, vipdoc: Path) -> None:
"""基本扫描应返回触发信号的股票."""
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
results = scanner.scan(universe="all")
assert len(results) >= 1, f"Expected >= 1 result, got {len(results)}"
def test_concurrent_same_as_serial(self, vipdoc: Path) -> None:
"""并发扫描结果应与串行扫描一致."""
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
serial = scanner.scan(universe="all", workers=1)
parallel = scanner.scan(universe="all", workers=2)
serial_codes = sorted(r.code for r in serial)
parallel_codes = sorted(r.code for r in parallel)
assert serial_codes == parallel_codes
def test_scan_with_zero_workers_uses_serial(self, vipdoc: Path) -> None:
"""workers=0 应退回串行模式."""
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
results = scanner.scan(universe="all", workers=0)
assert len(results) >= 1
def test_progress_callback(self, vipdoc: Path) -> None:
"""进度回调应被正确调用."""
scanner = SignalScanner(AlwaysBuyStrategy, vipdoc_path=vipdoc)
progress: list[tuple[int, int, str]] = []
def on_progress(current: int, total: int, name: str) -> None:
progress.append((current, total, name))
scanner.scan(universe="all", progress_callback=on_progress)
assert len(progress) >= 2
assert progress[-1][2] == "done"