mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 14:24:15 +08:00
fix(pipeline): enriched 全量重建流式化, 峰值内存与历史长度解耦 (#208/#174)
小内存机器 (2C2G) 延长历史后全量重建必然 OOM 的两个根因:
1. 全量模式把所有批次结果累积在内存 date_buffers 直到统一写盘,
5 年 x 5500 只 (~800 万行) 全表驻留 0.5-1GB;
2. 批次大小按「每只 ~244 行」的年代设定, 历史延长后单批宽表
(指标全列) 瞬时 ~600MB。
修复 (双核模拟实测: 275 万行峰值 RSS 958MB -> 415MB, 耗时 +37%):
- 流式暂存发布: 每批结果立即写暂存文件 (enriched 树外隐藏目录,
不被 **/*.parquet 业务 glob 扫到; polars/duckdb 的 ** 均匹配
点目录), 最后按日期分块流式合并、逐分区原子替换; 日期覆盖
校验与 EnrichedPublication 崩溃语义保持不变, finally 保证
暂存清理, 历史残留按 mtime (24h) 启动清扫;
- 自适应批次: 检测内存 <8GB 时按单批目标行数 15 万收缩,
>=8GB 保持用户设置 (大机器行为不变, 20 核实测耗时持平);
- 刷新路径同修: 启动/盘后的 300 天窗口计算 (repository) 改走
compute_enriched_history_window 分批执行, 与整帧顺序等价
(各步骤均 over(symbol) 分组), 峰值与标的总量解耦。
验证: 新增 5 项回归测试 (流式输出与整帧直算逐列相等/暂存清理/
残留清扫/自适应批次/窗口分批等价), 全套 backend 1498 项通过,
ruff 基线零新增。
This commit is contained in:
@@ -16,6 +16,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import shutil
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
@@ -1287,6 +1290,105 @@ def attach_deviation_columns_today(
|
||||
)
|
||||
|
||||
|
||||
# ================================================================
|
||||
# 全量重建流式暂存 + 自适应批次 (#208/#174)
|
||||
#
|
||||
# 旧全量模式把所有批次结果累积在内存 date_buffers 直到统一写盘:
|
||||
# 延长历史后 (5年 × 5500 只 ≈ 800 万行) 「全表驻留 + 单批宽表」双双
|
||||
# 超出小内存机器上限, 重建必然 OOM。现改为:
|
||||
# - 每批结果立即写暂存文件 (enriched 树外的隐藏目录 —— polars/duckdb
|
||||
# 的 **/*.parquet glob 均会匹配点目录, 树内暂存会被业务读取扫到),
|
||||
# 最后按日期分块流式合并、逐分区原子替换;
|
||||
# - 批次大小按单批目标行数自适应收缩 (指标/信号全部 over("symbol")
|
||||
# 分组, symbol 级分批不改变计算结果, 只约束单批宽表峰值)。
|
||||
# 任一时刻峰值内存 = 单批计算 + 单个日期块合并, 与总历史长度无关。
|
||||
# ================================================================
|
||||
|
||||
_STAGING_ROOT = Path(".staging") / "enriched_rebuild"
|
||||
_STALE_STAGING_MAX_AGE_S = 24 * 3600
|
||||
_RAM_LARGE_BYTES = 8 * 1024 ** 3 # ≥8GB 视为内存充裕, 批次保持用户设置
|
||||
_BATCH_TARGET_ROWS = 150_000 # 小内存单批目标行数 (宽表 ~60-80MB)
|
||||
_BATCH_MIN_SYMBOLS = 50
|
||||
_MERGE_DATE_CHUNKS = 15 # 最终合并按日期切 15 块流式执行
|
||||
|
||||
_ram_bytes_cache: int | None | bool = False # False = 未探测
|
||||
|
||||
|
||||
def _total_ram_bytes() -> int | None:
|
||||
global _ram_bytes_cache
|
||||
if _ram_bytes_cache is False:
|
||||
try:
|
||||
import psutil
|
||||
_ram_bytes_cache = psutil.virtual_memory().total
|
||||
except Exception:
|
||||
_ram_bytes_cache = None
|
||||
return _ram_bytes_cache # type: ignore[return-value]
|
||||
|
||||
|
||||
def _adaptive_sym_batch(default_batch: int, rows_per_symbol: int) -> int:
|
||||
"""小内存机器按单批目标行数收缩批次; 大内存机器保持原值 (#208)。"""
|
||||
if (_total_ram_bytes() or 0) >= _RAM_LARGE_BYTES:
|
||||
return default_batch
|
||||
return max(
|
||||
_BATCH_MIN_SYMBOLS,
|
||||
min(default_batch, _BATCH_TARGET_ROWS // max(rows_per_symbol, 1)),
|
||||
)
|
||||
|
||||
|
||||
def _sweep_stale_staging(data_dir: Path) -> None:
|
||||
"""清理崩溃/取消运行残留的暂存目录 (按 mtime 判定, 不碰活跃目录)。"""
|
||||
root = data_dir / _STAGING_ROOT
|
||||
if not root.exists():
|
||||
return
|
||||
cutoff = time.time() - _STALE_STAGING_MAX_AGE_S
|
||||
for run_dir in root.iterdir():
|
||||
try:
|
||||
if run_dir.is_dir() and run_dir.stat().st_mtime < cutoff:
|
||||
shutil.rmtree(run_dir, ignore_errors=True)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def compute_enriched_history_window(
|
||||
df_hist: pl.DataFrame,
|
||||
data_dir: Path,
|
||||
instruments: pl.DataFrame | None = None,
|
||||
historical_shares: pl.DataFrame | None = None,
|
||||
sym_batch: int | None = None,
|
||||
) -> 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:
|
||||
return df_hist
|
||||
symbols = df_hist["symbol"].unique().sort().to_list()
|
||||
if sym_batch is None:
|
||||
rows_per_sym = max(1, df_hist.height // max(len(symbols), 1))
|
||||
sym_batch = _adaptive_sym_batch(2000, rows_per_sym)
|
||||
parts: list[pl.DataFrame] = []
|
||||
for bs in range(0, len(symbols), sym_batch):
|
||||
batch = symbols[bs:bs + sym_batch]
|
||||
part = df_hist.filter(pl.col("symbol").is_in(batch)).sort(["symbol", "date"])
|
||||
part = compute_indicators(part)
|
||||
part = attach_deviation_columns(part, data_dir)
|
||||
part = compute_signals(part)
|
||||
if instruments is not None and not instruments.is_empty():
|
||||
inst_batch = instruments.filter(pl.col("symbol").is_in(batch))
|
||||
shares_batch = (
|
||||
historical_shares.filter(pl.col("symbol").is_in(batch))
|
||||
if historical_shares is not None and not historical_shares.is_empty()
|
||||
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"])
|
||||
|
||||
|
||||
def run_pipeline(data_dir: Path | None = None,
|
||||
symbols: list[str] | None = None,
|
||||
new_dates_only: bool = False,
|
||||
@@ -1460,9 +1562,10 @@ def run_pipeline(data_dir: Path | None = None,
|
||||
|
||||
import gc
|
||||
|
||||
# ── 按 symbol 分批处理: 每只股只有 ~244 行, 无冗余计算 ──
|
||||
# 先获取全部 symbol 列表
|
||||
lf_all = scan_daily_parquet(daily_glob, cast_options=_cast)
|
||||
# ── 按 symbol 分批处理: 指标全部 over("symbol") 分组, 分批不改变结果 ──
|
||||
# 文件列表只收集一次, 批间复用 (避免每批重新展开 glob)
|
||||
daily_files = sorted(str(p) for p in daily_dir.rglob("*.parquet"))
|
||||
lf_all = scan_daily_parquet(daily_files, cast_options=_cast)
|
||||
if symbols:
|
||||
sym_set = set(symbols)
|
||||
lf_all = lf_all.filter(pl.col("symbol").is_in(list(sym_set)))
|
||||
@@ -1476,7 +1579,6 @@ def run_pipeline(data_dir: Path | None = None,
|
||||
return 0
|
||||
|
||||
total_syms = len(all_symbols)
|
||||
logger.info("全量计算: %d 只标的, 按 symbol 分批 [%s]", total_syms, mode)
|
||||
|
||||
if not factors.is_empty() and symbols:
|
||||
factors = factors.filter(pl.col("symbol").is_in(list(sym_set)))
|
||||
@@ -1486,106 +1588,139 @@ def run_pipeline(data_dir: Path | None = None,
|
||||
inst_use = instruments.filter(pl.col("symbol").is_in(list(sym_set)))
|
||||
|
||||
from app.services import preferences as prefs_mod
|
||||
SYM_BATCH = prefs_mod.get_enriched_batch_size() # 每批 N 只 × ~244 天, 可在设置中调整
|
||||
# 自适应批次 (#208): 单批体积按目标行数恒定, 与总历史长度解耦;
|
||||
# 小内存机器自动收缩, 大内存机器保持用户设置
|
||||
total_rows = lf_all.select(pl.len()).collect(streaming=True).item()
|
||||
rows_per_sym = max(1, -(-int(total_rows) // total_syms))
|
||||
SYM_BATCH = _adaptive_sym_batch(prefs_mod.get_enriched_batch_size(), rows_per_sym)
|
||||
total_batches = (total_syms + SYM_BATCH - 1) // SYM_BATCH
|
||||
logger.info("全量计算: %d 只标的 (%d 行, ~%d 行/只), symbol 分批 %d 只/批, %d 批 [%s]",
|
||||
total_syms, total_rows, rows_per_sym, SYM_BATCH, total_batches, mode)
|
||||
|
||||
# 全量模式: 收集所有批次结果, 最后按日期分区覆盖写入
|
||||
from collections import defaultdict
|
||||
date_buffers: dict[str, list[pl.DataFrame]] = defaultdict(list)
|
||||
# 全量模式: 流式暂存发布 (#208) —— 每批落盘暂存文件, 不再内存累积;
|
||||
# 暂存目录在 enriched 树外, 不会被任何 **/*.parquet 业务 glob 扫到
|
||||
staging_dir: Path | None = None
|
||||
staging_files: list[str] = []
|
||||
if not symbols:
|
||||
_sweep_stale_staging(d)
|
||||
staging_dir = d / _STAGING_ROOT / uuid.uuid4().hex
|
||||
staging_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for batch_start in range(0, total_syms, SYM_BATCH):
|
||||
batch_end = min(batch_start + SYM_BATCH, total_syms)
|
||||
batch_syms = all_symbols[batch_start:batch_end]
|
||||
try:
|
||||
for batch_start in range(0, total_syms, SYM_BATCH):
|
||||
batch_end = min(batch_start + SYM_BATCH, total_syms)
|
||||
batch_syms = all_symbols[batch_start:batch_end]
|
||||
|
||||
# 只读取本批 symbol 的数据
|
||||
lf_batch = scan_daily_parquet(daily_glob, cast_options=_cast)
|
||||
lf_batch = lf_batch.filter(pl.col("symbol").is_in(batch_syms))
|
||||
raw = lf_batch.sort(["symbol", "date"]).collect(streaming=True)
|
||||
# 只读取本批 symbol 的数据
|
||||
lf_batch = scan_daily_parquet(daily_files, cast_options=_cast)
|
||||
lf_batch = lf_batch.filter(pl.col("symbol").is_in(batch_syms))
|
||||
raw = lf_batch.sort(["symbol", "date"]).collect(streaming=True)
|
||||
|
||||
if raw.is_empty():
|
||||
continue
|
||||
if raw.is_empty():
|
||||
continue
|
||||
|
||||
# 本批的 factors / instruments
|
||||
batch_factors = (
|
||||
factors.filter(pl.col("symbol").is_in(batch_syms))
|
||||
if not factors.is_empty() else factors
|
||||
)
|
||||
batch_inst = (
|
||||
inst_use.filter(pl.col("symbol").is_in(batch_syms))
|
||||
if not inst_use.is_empty() else inst_use
|
||||
)
|
||||
batch_shares = (
|
||||
historical_shares.filter(pl.col("symbol").is_in(batch_syms))
|
||||
if not historical_shares.is_empty() else historical_shares
|
||||
)
|
||||
# 本批的 factors / instruments
|
||||
batch_factors = (
|
||||
factors.filter(pl.col("symbol").is_in(batch_syms))
|
||||
if not factors.is_empty() else factors
|
||||
)
|
||||
batch_inst = (
|
||||
inst_use.filter(pl.col("symbol").is_in(batch_syms))
|
||||
if not inst_use.is_empty() else inst_use
|
||||
)
|
||||
batch_shares = (
|
||||
historical_shares.filter(pl.col("symbol").is_in(batch_syms))
|
||||
if not historical_shares.is_empty() else historical_shares
|
||||
)
|
||||
|
||||
# 计算
|
||||
enriched = compute_enriched(
|
||||
raw,
|
||||
factors=batch_factors,
|
||||
instruments=batch_inst,
|
||||
historical_shares=batch_shares,
|
||||
)
|
||||
# 计算
|
||||
enriched = compute_enriched(
|
||||
raw,
|
||||
factors=batch_factors,
|
||||
instruments=batch_inst,
|
||||
historical_shares=batch_shares,
|
||||
)
|
||||
|
||||
if not enriched.is_empty():
|
||||
if symbols:
|
||||
# 局部模式: 直接按日期合并写入
|
||||
for date_df in 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"
|
||||
if not enriched.is_empty():
|
||||
if symbols:
|
||||
# 局部模式: 直接按日期合并写入
|
||||
for date_df in 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"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
date_df_storage = _select_storage_cols(date_df)
|
||||
if out.exists():
|
||||
existing = pl.read_parquet(out)
|
||||
existing = existing.filter(~pl.col("symbol").is_in(batch_syms))
|
||||
date_df_storage = pl.concat([existing, date_df_storage], how="diagonal_relaxed")
|
||||
date_df_storage = date_df_storage.sort(["symbol"])
|
||||
publication.write_parquet(date_df_storage, out)
|
||||
written += date_df_storage.height
|
||||
else:
|
||||
# 全量模式: 写单批暂存文件 (按 date,symbol 排序 →
|
||||
# 合并期 parquet 行组统计可按日期裁剪), 随即释放本批内存
|
||||
out = staging_dir / f"batch-{batch_start // SYM_BATCH:04d}.parquet"
|
||||
_select_storage_cols(enriched).sort(["date", "symbol"]).write_parquet(out)
|
||||
staging_files.append(str(out))
|
||||
written += enriched.height
|
||||
|
||||
del raw, enriched, batch_factors, batch_inst, batch_shares
|
||||
gc.collect()
|
||||
|
||||
logger.info("symbol 批次 %d/%d (%s ~ %s), 已处理 %d 行",
|
||||
batch_start // SYM_BATCH + 1,
|
||||
total_batches,
|
||||
batch_syms[0], batch_syms[-1], written)
|
||||
|
||||
# 通知进度
|
||||
if on_batch_done:
|
||||
on_batch_done(batch_start // SYM_BATCH + 1, total_batches)
|
||||
|
||||
# 全量模式: 日期覆盖校验 → 按日期分块流式合并 → 逐分区原子替换
|
||||
if not symbols and staging_files:
|
||||
existing_dates = {
|
||||
p.name.removeprefix("date=")
|
||||
for p in base.glob("date=*")
|
||||
if p.is_dir()
|
||||
}
|
||||
unique_dates = sorted(
|
||||
scan_enriched_parquet(staging_files).select("date").unique()
|
||||
.collect()["date"].to_list()
|
||||
)
|
||||
rebuilt_dates = {
|
||||
ds.isoformat() if hasattr(ds, "isoformat") else str(ds)
|
||||
for ds in unique_dates
|
||||
}
|
||||
missing_dates = existing_dates - rebuilt_dates
|
||||
if missing_dates:
|
||||
sample = ", ".join(sorted(missing_dates)[:5])
|
||||
raise RuntimeError(f"全量重建结果缺少已有日期分区,拒绝覆盖: {sample}")
|
||||
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
chunk = max(1, -(-len(unique_dates) // _MERGE_DATE_CHUNKS))
|
||||
for ci in range(0, len(unique_dates), chunk):
|
||||
lo = unique_dates[ci]
|
||||
hi = unique_dates[min(ci + chunk, len(unique_dates)) - 1]
|
||||
block = (
|
||||
scan_enriched_parquet(staging_files)
|
||||
.filter((pl.col("date") >= lo) & (pl.col("date") <= hi))
|
||||
.sort(["date", "symbol"])
|
||||
.collect(streaming=True)
|
||||
)
|
||||
for date_df in block.partition_by("date"):
|
||||
ds = date_df["date"][0]
|
||||
ds_str = ds.isoformat() if hasattr(ds, "isoformat") else str(ds)
|
||||
out = base / f"date={ds_str}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
date_df_storage = _select_storage_cols(date_df)
|
||||
if out.exists():
|
||||
existing = pl.read_parquet(out)
|
||||
existing = existing.filter(~pl.col("symbol").is_in(batch_syms))
|
||||
date_df_storage = pl.concat([existing, date_df_storage], how="diagonal_relaxed")
|
||||
date_df_storage = date_df_storage.sort(["symbol"])
|
||||
publication.write_parquet(date_df_storage, out)
|
||||
written += date_df_storage.height
|
||||
else:
|
||||
# 全量模式: 缓冲到 date_buffers, 最后一次性写入
|
||||
for date_df in enriched.partition_by("date"):
|
||||
dt = date_df["date"][0]
|
||||
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
|
||||
date_buffers[ds].append(_select_storage_cols(date_df).sort(["symbol"]))
|
||||
written += date_df.height
|
||||
|
||||
del raw, enriched, batch_factors, batch_inst, batch_shares
|
||||
gc.collect()
|
||||
|
||||
logger.info("symbol 批次 %d/%d (%s ~ %s), 已处理 %d 行",
|
||||
batch_start // SYM_BATCH + 1,
|
||||
total_batches,
|
||||
batch_syms[0], batch_syms[-1], written)
|
||||
|
||||
# 通知进度
|
||||
if on_batch_done:
|
||||
on_batch_done(batch_start // SYM_BATCH + 1, total_batches)
|
||||
|
||||
# 全量模式: 按日期分区写入
|
||||
if not symbols and date_buffers:
|
||||
existing_dates = {
|
||||
p.name.removeprefix("date=")
|
||||
for p in base.glob("date=*")
|
||||
if p.is_dir()
|
||||
}
|
||||
rebuilt_dates = set(date_buffers)
|
||||
missing_dates = existing_dates - rebuilt_dates
|
||||
if missing_dates:
|
||||
sample = ", ".join(sorted(missing_dates)[:5])
|
||||
raise RuntimeError(f"全量重建结果缺少已有日期分区,拒绝覆盖: {sample}")
|
||||
|
||||
base.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for ds, dfs in date_buffers.items():
|
||||
out = base / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
merged = pl.concat(dfs, how="diagonal_relaxed").sort(["symbol"])
|
||||
publication.write_parquet(merged, out)
|
||||
|
||||
date_buffers.clear()
|
||||
gc.collect()
|
||||
publication.write_parquet(date_df.sort(["symbol"]), out)
|
||||
gc.collect()
|
||||
logger.info("全量暂存合并完成: %d 个日期分区", len(unique_dates))
|
||||
finally:
|
||||
# 无论成功/失败/取消都清掉本次暂存 (历史残留由 _sweep_stale_staging 兜底)
|
||||
if staging_dir is not None:
|
||||
shutil.rmtree(staging_dir, ignore_errors=True)
|
||||
|
||||
publication.commit()
|
||||
t_done = _t.perf_counter()
|
||||
|
||||
@@ -582,7 +582,7 @@ class KlineRepository:
|
||||
# 300 日历天 ≈ 210 交易日, 覆盖 filter_history 最大 lookback(90) + warmup(60)
|
||||
try:
|
||||
from datetime import timedelta
|
||||
from app.indicators.pipeline import compute_indicators, compute_signals, compute_limit_signals
|
||||
from app.indicators.pipeline import compute_enriched_history_window
|
||||
start_full = latest - timedelta(days=300)
|
||||
read_cols = [c for c in ["symbol", "date", "open", "high", "low", "close",
|
||||
"volume", "amount", "raw_close", "raw_high", "raw_low"]
|
||||
@@ -600,28 +600,22 @@ 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 indicators")
|
||||
df_full = compute_indicators(df_hist)
|
||||
logger.info("enriched refresh step done: compute indicators rows=%d (%.2fs)", len(df_full), time.perf_counter() - step)
|
||||
|
||||
# 异动偏离列 (deviate_Nd = 个股动量 - 基准指数动量), 运行时附着
|
||||
from app.indicators.pipeline import attach_deviation_columns
|
||||
df_full = attach_deviation_columns(df_full, self.store.data_dir)
|
||||
|
||||
step = time.perf_counter()
|
||||
logger.info("enriched refresh step start: compute signals")
|
||||
df_full = compute_signals(df_full)
|
||||
logger.info("enriched refresh step done: compute signals (%.2fs)", time.perf_counter() - step)
|
||||
if instruments is not None and not instruments.is_empty():
|
||||
step = time.perf_counter()
|
||||
logger.info("enriched refresh step start: compute limit signals")
|
||||
df_full = compute_limit_signals(
|
||||
df_full,
|
||||
instruments,
|
||||
historical_shares=self.get_historical_shares(),
|
||||
)
|
||||
logger.info("enriched refresh step done: compute limit signals (%.2fs)", time.perf_counter() - step)
|
||||
logger.info("enriched refresh step start: compute window (batched)")
|
||||
df_full = compute_enriched_history_window(
|
||||
df_hist,
|
||||
self.store.data_dir,
|
||||
instruments=instruments,
|
||||
historical_shares=(
|
||||
self.get_historical_shares()
|
||||
if instruments is not None and not instruments.is_empty()
|
||||
else None
|
||||
),
|
||||
)
|
||||
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():
|
||||
|
||||
@@ -78,3 +78,120 @@ def test_full_rebuild_rejects_missing_existing_dates_before_writing(tmp_path, mo
|
||||
tmp_path / "kline_daily_enriched" / "date=2026-07-15" / "part.parquet"
|
||||
)
|
||||
assert existing["close"].to_list() == [1.0]
|
||||
|
||||
|
||||
# ================================================================
|
||||
# 流式暂存发布 + 自适应批次 (#208/#174)
|
||||
# ================================================================
|
||||
|
||||
def _write_daily_multi(data_dir, symbols: list[str], dates: list[str]) -> None:
|
||||
for ds in dates:
|
||||
rows = {
|
||||
"symbol": symbols,
|
||||
"date": [date.fromisoformat(ds)] * len(symbols),
|
||||
"open": [10.0] * len(symbols),
|
||||
"high": [11.0] * len(symbols),
|
||||
"low": [9.5] * len(symbols),
|
||||
"close": [10.5] * len(symbols),
|
||||
"volume": [100.0] * len(symbols),
|
||||
"amount": [1000.0] * len(symbols),
|
||||
}
|
||||
out = data_dir / "kline_daily" / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
pl.DataFrame(rows).write_parquet(out)
|
||||
|
||||
|
||||
def test_full_rebuild_streaming_output_matches_direct_compute(tmp_path, monkeypatch):
|
||||
"""多批流式暂存 + 分块合并的输出与整帧直算完全一致 (#208)。"""
|
||||
symbols = [f"{600000 + i}.SH" for i in range(4)]
|
||||
dates = [f"2026-07-{d:02d}" for d in range(1, 6)]
|
||||
_write_daily_multi(tmp_path, symbols, dates)
|
||||
# 强制 2 只/批 → 2 批暂存 + 合并, 覆盖流式路径
|
||||
monkeypatch.setattr(pipeline, "_adaptive_sym_batch", lambda default, rows: 2)
|
||||
|
||||
written = pipeline.run_pipeline(data_dir=tmp_path)
|
||||
assert written == len(symbols) * len(dates)
|
||||
|
||||
from app.parquet import scan_daily_parquet
|
||||
raw = (scan_daily_parquet(str(tmp_path / "kline_daily" / "**" / "*.parquet"))
|
||||
.sort(["symbol", "date"]).collect())
|
||||
expected = pipeline._select_storage_cols(pipeline.compute_enriched(raw)).sort(["symbol", "date"])
|
||||
|
||||
got = pl.read_parquet(str(tmp_path / "kline_daily_enriched" / "**" / "*.parquet"))
|
||||
got_cols = [c for c in expected.columns if c in got.columns]
|
||||
assert got.select(got_cols).sort(["symbol", "date"]).equals(expected.select(got_cols))
|
||||
|
||||
|
||||
def test_full_rebuild_cleans_staging_and_keeps_it_outside_globs(tmp_path, monkeypatch):
|
||||
"""重建完成后暂存目录被清理, 业务 glob 不会扫到暂存文件 (#208)。"""
|
||||
_write_daily(tmp_path, "2026-07-14", 14.0)
|
||||
_write_daily(tmp_path, "2026-07-15", 15.0)
|
||||
monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
|
||||
|
||||
pipeline.run_pipeline(data_dir=tmp_path)
|
||||
|
||||
staging_root = tmp_path / ".staging" / "enriched_rebuild"
|
||||
assert not staging_root.exists() or not any(staging_root.iterdir())
|
||||
# 暂存位于 enriched 树外: enriched glob 只见 date=* 分区
|
||||
parts = list((tmp_path / "kline_daily_enriched").glob("date=*"))
|
||||
assert sorted(p.name for p in parts) == ["date=2026-07-14", "date=2026-07-15"]
|
||||
|
||||
|
||||
def test_stale_staging_swept_on_full_rebuild(tmp_path, monkeypatch):
|
||||
"""崩溃/取消残留的暂存目录按 mtime 被清扫 (#208)。"""
|
||||
import os
|
||||
import time as _time
|
||||
stale = tmp_path / ".staging" / "enriched_rebuild" / "dead-run"
|
||||
stale.mkdir(parents=True)
|
||||
junk = stale / "batch-0000.parquet"
|
||||
junk.write_bytes(b"junk")
|
||||
old = _time.time() - 2 * 24 * 3600
|
||||
os.utime(stale, (old, old))
|
||||
|
||||
_write_daily(tmp_path, "2026-07-14", 14.0)
|
||||
monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
|
||||
pipeline.run_pipeline(data_dir=tmp_path)
|
||||
|
||||
assert not stale.exists()
|
||||
|
||||
|
||||
def test_adaptive_sym_batch_shrinks_only_on_small_ram(monkeypatch):
|
||||
"""小内存: 按目标行数收缩; 大内存: 保持用户设置 (#208)。"""
|
||||
monkeypatch.setattr(pipeline, "_total_ram_bytes", lambda: 2 * 1024 ** 3)
|
||||
got = pipeline._adaptive_sym_batch(1000, 1500)
|
||||
assert got == max(pipeline._BATCH_MIN_SYMBOLS, min(1000, 150_000 // 1500))
|
||||
# 行数极少时不低于下限
|
||||
assert pipeline._adaptive_sym_batch(1000, 1) == 1000
|
||||
|
||||
monkeypatch.setattr(pipeline, "_total_ram_bytes", lambda: 16 * 1024 ** 3)
|
||||
assert pipeline._adaptive_sym_batch(1000, 1500) == 1000
|
||||
|
||||
|
||||
def test_history_window_batched_equals_direct(tmp_path):
|
||||
"""刷新窗口分批计算与整帧顺序执行等价 (#208)。"""
|
||||
import numpy as np
|
||||
rng = np.random.default_rng(11)
|
||||
n_syms, n_days = 6, 30
|
||||
close = 10.0 * np.cumprod(1 + rng.normal(0, 0.02, (n_syms, n_days)), axis=1)
|
||||
df_hist = pl.DataFrame({
|
||||
"symbol": np.repeat([f"{600000 + i}.SH" for i in range(n_syms)], n_days),
|
||||
"_day": np.tile(np.arange(n_days), n_syms),
|
||||
"open": (close * 0.99).reshape(-1),
|
||||
"high": (close * 1.01).reshape(-1),
|
||||
"low": (close * 0.98).reshape(-1),
|
||||
"close": close.reshape(-1),
|
||||
"volume": rng.integers(1000, 9000, n_syms * n_days).astype(float),
|
||||
"amount": rng.integers(1, 99, n_syms * n_days).astype(float),
|
||||
}).with_columns(
|
||||
(pl.lit(date(2026, 6, 1)) + pl.duration(days=pl.col("_day"))).alias("date")
|
||||
).drop("_day")
|
||||
|
||||
direct = pipeline.compute_signals(
|
||||
pipeline.attach_deviation_columns(
|
||||
pipeline.compute_indicators(df_hist.clone()), tmp_path
|
||||
)
|
||||
)
|
||||
batched = pipeline.compute_enriched_history_window(
|
||||
df_hist.clone(), tmp_path, sym_batch=2
|
||||
)
|
||||
assert batched.sort(["symbol", "date"]).equals(direct.sort(["symbol", "date"]))
|
||||
|
||||
Reference in New Issue
Block a user