mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 17:54:15 +08:00
shy3130
a9af2bce1c
小内存机器 (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 基线零新增。
198 lines
7.8 KiB
Python
198 lines
7.8 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import date
|
|
|
|
import polars as pl
|
|
import pytest
|
|
|
|
from app.indicators import pipeline
|
|
|
|
|
|
def _write_daily(data_dir, ds: str, close: float) -> None:
|
|
out = data_dir / "kline_daily" / f"date={ds}" / "part.parquet"
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
pl.DataFrame({
|
|
"symbol": ["600000.SH"],
|
|
"date": [date.fromisoformat(ds)],
|
|
"open": [close],
|
|
"high": [close],
|
|
"low": [close],
|
|
"close": [close],
|
|
"volume": [100.0],
|
|
"amount": [1000.0],
|
|
"quote_ts": [0],
|
|
}).write_parquet(out)
|
|
|
|
|
|
def _write_existing(data_dir, ds: str, close: float) -> None:
|
|
out = data_dir / "kline_daily_enriched" / f"date={ds}" / "part.parquet"
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
pl.DataFrame({
|
|
"symbol": ["600000.SH"],
|
|
"date": [date.fromisoformat(ds)],
|
|
"close": [close],
|
|
}).write_parquet(out)
|
|
|
|
|
|
def _fake_compute_enriched(raw: pl.DataFrame, **_kwargs) -> pl.DataFrame:
|
|
return raw.with_columns(
|
|
pl.col("close").alias("raw_close"),
|
|
pl.col("high").alias("raw_high"),
|
|
pl.col("low").alias("raw_low"),
|
|
pl.lit(None, dtype=pl.Float64).alias("turnover_rate"),
|
|
pl.lit(0, dtype=pl.UInt32).alias("consecutive_limit_ups"),
|
|
pl.lit(0, dtype=pl.UInt32).alias("consecutive_limit_downs"),
|
|
)
|
|
|
|
|
|
def test_full_rebuild_overwrites_existing_partitions_without_deleting_base(tmp_path, monkeypatch):
|
|
_write_daily(tmp_path, "2026-07-14", 14.0)
|
|
_write_daily(tmp_path, "2026-07-15", 15.0)
|
|
_write_existing(tmp_path, "2026-07-15", 1.0)
|
|
marker = tmp_path / "kline_daily_enriched" / "keep.txt"
|
|
marker.write_text("keep", encoding="utf-8")
|
|
monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
|
|
|
|
written = pipeline.run_pipeline(data_dir=tmp_path)
|
|
|
|
assert written == 2
|
|
assert marker.read_text(encoding="utf-8") == "keep"
|
|
assert pl.read_parquet(
|
|
tmp_path / "kline_daily_enriched" / "date=2026-07-14" / "part.parquet"
|
|
)["close"].to_list() == [14.0]
|
|
assert pl.read_parquet(
|
|
tmp_path / "kline_daily_enriched" / "date=2026-07-15" / "part.parquet"
|
|
)["close"].to_list() == [15.0]
|
|
|
|
|
|
def test_full_rebuild_rejects_missing_existing_dates_before_writing(tmp_path, monkeypatch):
|
|
_write_daily(tmp_path, "2026-07-15", 15.0)
|
|
_write_existing(tmp_path, "2026-07-14", 14.0)
|
|
_write_existing(tmp_path, "2026-07-15", 1.0)
|
|
monkeypatch.setattr(pipeline, "compute_enriched", _fake_compute_enriched)
|
|
|
|
with pytest.raises(RuntimeError, match="缺少已有日期分区"):
|
|
pipeline.run_pipeline(data_dir=tmp_path)
|
|
|
|
existing = pl.read_parquet(
|
|
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"]))
|