perf: 优化扩展查询、概览并发与 ETF 历史读取

按需裁剪 Parquet 列,合并概览并发构建并防止旧结果回填。为 ETF 历史输入增加有界缓存,按容器内存预算进行任务准入及批次控制,保留原计算口径。补充回归测试与可复现性能基准。

Signed-off-by: shy3130 <shy3130@qq.com>
This commit is contained in:
shy3130
2026-11-02 00:23:06 +08:00
parent ab2b82da37
commit 7bd290289a
13 changed files with 861 additions and 38 deletions
+1 -1
View File
@@ -560,7 +560,7 @@ POLARS_MAX_THREADS=2 # 主要收益项: 用约 40% 的构建耗时换 500
| 变量 | 默认 | 作用 |
| :-- | :-- | :-- |
| `TICK_WORKER_RSS_LIMIT_BYTES` | `6442450944` (6GiB) | 回测/优化独立 worker 的 RSS 上限, 超限终止并返回明确错误(而非拖垮整机); `0` = 不限制 |
| `TICK_HEAVY_MIN_AVAILABLE_BYTES` | `1073741824` (1GiB) | 系统可用内存低于该值时, 新的重任务(回测/因子批量/挖掘)挂起等待内存恢复; `0` = 不检查 |
| `TICK_HEAVY_MIN_AVAILABLE_BYTES` | `1073741824` (1GiB) | 有效可用内存低于该值时, 新的重任务(回测/因子批量/挖掘)挂起等待内存恢复; 有效预算取宿主机与容器 cgroup v1/v2 的较小值, 包含父级限额和 `memory.high`; `0` = 不检查 |
| `TICK_DUCKDB_READ_CONCURRENCY` | `3` | DuckDB 并发读上限。读查询走独立连接(MVCC), 同步期的长扫描不再阻塞交互读; 并发数越高峰值内存越大, `1` = 完全串行(旧行为) |
### 跑起来后的第一次使用
+34 -10
View File
@@ -265,6 +265,9 @@ def _read_ext_dataframe(
snapshot_date: str | None = None,
start_date: str | None = None,
end_date: str | None = None,
*,
columns: list[str] | None = None,
required_columns: list[str] | None = None,
) -> tuple[pl.DataFrame, str | None]:
"""读取扩展表数据。
@@ -275,13 +278,24 @@ def _read_ext_dataframe(
"""
cfg_dir = data_dir / "ext_data" / config.id
def collect(lf: pl.LazyFrame) -> pl.DataFrame:
if columns:
available = lf.collect_schema().names()
keep = [c for c in columns if c in available]
required = required_columns or []
# Preserve full-row fallback for unknown output fields, and retain
# the original validation/errors for unknown filter/sort fields.
if keep and all(c in available for c in required):
lf = lf.select(list(dict.fromkeys([*keep, *required])))
return lf.collect()
if config.mode == "snapshot":
if start_date or end_date:
raise HTTPException(400, "快照模式无历史分区, 不支持日期范围查询")
path = cfg_dir / "part.parquet"
if not path.exists():
return pl.DataFrame(), None
return pl.read_parquet(path), _latest_sync_date(config, data_dir)
return collect(pl.scan_parquet(path)), _latest_sync_date(config, data_dir)
base = cfg_dir / "timeseries"
if not base.exists():
@@ -292,7 +306,7 @@ def _read_ext_dataframe(
path = base / f"date={day}" / "part.parquet"
if not path.exists():
return pl.DataFrame(), day
return pl.read_parquet(path), day
return collect(pl.scan_parquet(path)), day
partitions = sorted(
d for d in base.iterdir()
@@ -317,21 +331,22 @@ def _read_ext_dataframe(
# 没有它, 范围合并后消费方无法区分行来自哪一天。
frames = []
for d in selected:
f = pl.read_parquet(d / "part.parquet")
f = pl.scan_parquet(d / "part.parquet")
names = f.collect_schema().names()
day = d.name[5:]
if "date" not in f.columns:
if "date" not in names:
f = f.with_columns(pl.lit(day).alias("date"))
elif "data_date" not in f.columns:
elif "data_date" not in names:
f = f.with_columns(pl.lit(day).alias("data_date"))
frames.append(f)
# diagonal: 各分区 schema 可能不一致 (配置中途加字段 → 早期分区少列),
# 缺列自动补 null 并取超类型; 单分区时无需 concat
df = pl.concat(frames, how="diagonal") if len(frames) > 1 else frames[0]
return df, f"{selected[0].name[5:]}..{selected[-1].name[5:]}"
return collect(df), f"{selected[0].name[5:]}..{selected[-1].name[5:]}"
latest = partitions[-1]
latest_date = latest.name[5:]
return pl.read_parquet(latest / "part.parquet"), latest_date
return collect(pl.scan_parquet(latest / "part.parquet")), latest_date
def _with_instrument_name(df: pl.DataFrame, data_dir: Path) -> pl.DataFrame:
@@ -636,17 +651,26 @@ def list_rows(
raise HTTPException(404, f"配置 '{config_id}' 不存在")
data_dir = _data_dir(request)
df, active_date = _read_ext_dataframe(config, data_dir, snapshot_date, start_date, end_date)
requested = [c.strip() for c in (columns or "").split(",") if c.strip()]
tf = config.pull.time_field if config.pull else None
required = [entry.partition(":")[0].strip() for entry in (filter or [])]
if sort:
required.append(sort.partition(":")[0].strip())
if tf:
required.append(tf)
df, active_date = _read_ext_dataframe(
config, data_dir, snapshot_date, start_date, end_date,
columns=["symbol", "code", "name", *requested] if requested else None,
required_columns=required,
)
df = _with_instrument_name(df, data_dir)
# 日内序列表: 按时间列升序, 保证分页/截断取到的是最早的盘
tf = config.pull.time_field if config.pull else None
if tf and tf in df.columns:
df = df.sort(tf)
df = _apply_row_filters(df, filter)
df = _apply_row_sort(df, sort)
total = len(df)
df = df.slice(offset, limit)
requested = [c.strip() for c in (columns or "").split(",") if c.strip()]
if requested:
keep = [c for c in ["symbol", "code", "name", *requested] if c in df.columns]
if keep:
+30 -7
View File
@@ -5,6 +5,7 @@ import math
import re
import threading
import time
from concurrent.futures import Future
from datetime import date
from typing import Any
@@ -24,6 +25,8 @@ _cache_ts: float = 0.0
# 缓存跨线程读写锁: market_overview 在 FastAPI 线程池读, invalidate 在数据刷新线程清,
# 无锁会读到撕裂/过期状态。用模块级 Lock 守护 check-then-set 与 clear。
_cache_lock = threading.Lock()
_cache_generation = 0
_inflight: dict[tuple[int, str], Future] = {}
def invalidate_overview_cache() -> None:
@@ -31,11 +34,12 @@ def invalidate_overview_cache() -> None:
清除数据后调用, 避免看板在 TTL 窗口内继续返回旧的聚合结果。
"""
global _cache, _cache_key, _cache_ts
global _cache, _cache_key, _cache_ts, _cache_generation
with _cache_lock:
_cache = None
_cache_key = None
_cache_ts = 0.0
_cache_generation += 1
_DIMENSION_SEP = re.compile(r"[、,,;;|/\s]+")
@@ -363,16 +367,35 @@ def _build_overview(request: Request, as_of: date | None = None) -> dict:
def market_overview(request: Request, as_of: date | None = None):
"""总览页单次请求聚合数据,避免前端拉全市场明细后再计算。"""
global _cache, _cache_key, _cache_ts
now = time.time()
now = time.monotonic()
cache_key = as_of.isoformat() if as_of else "latest"
# 读缓存持锁, 避免与 invalidate 的 clear 竞态读到撕裂状态
with _cache_lock:
if _cache is not None and _cache_key == cache_key and (now - _cache_ts) < _CACHE_TTL:
return _cache
# 装配在锁外进行 (耗时), 允许并发未命中时各自构建, 不长时间持锁串行化请求
data = _build_overview(request, as_of)
generation = _cache_generation
flight_key = (generation, cache_key)
flight = _inflight.get(flight_key)
leader = flight is None
if leader:
flight = Future()
_inflight[flight_key] = flight
if not leader:
return flight.result()
# Both computation and follower waits stay outside the cache lock.
try:
data = _build_overview(request, as_of)
except BaseException as exc:
with _cache_lock:
_inflight.pop(flight_key, None)
flight.set_exception(exc)
raise
with _cache_lock:
_cache = data
_cache_key = cache_key
_cache_ts = now
# A refresh may have invalidated this build while it was running.
if generation == _cache_generation:
_cache = data
_cache_key = cache_key
_cache_ts = now
_inflight.pop(flight_key, None)
flight.set_result(data)
return data
+5
View File
@@ -34,11 +34,16 @@ def _memory_snapshot() -> dict:
import psutil
vm = psutil.virtual_memory()
from app.services.memory_budget import memory_budget
budget = memory_budget(vm)
snap: dict[str, object] = {
"rss_bytes": psutil.Process().memory_info().rss,
"system_total_bytes": vm.total,
"system_available_bytes": vm.available,
"system_used_pct": vm.percent,
"effective_total_bytes": budget.total_bytes,
"effective_available_bytes": budget.available_bytes,
}
cg = Path("/sys/fs/cgroup")
if _read_cgroup_int(cg / "memory.current") is not None:
+5 -10
View File
@@ -1409,18 +1409,13 @@ _BATCH_TARGET_ROWS = 150_000 # 小内存单批目标行数 (宽表 ~60-
_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]
try:
from app.services.memory_budget import memory_budget
return memory_budget().total_bytes
except Exception:
return None
def _adaptive_sym_batch(default_batch: int, rows_per_symbol: int) -> int:
+2 -2
View File
@@ -33,9 +33,9 @@ _MEMORY_POLL_S = 5.0
def _memory_admission_ok() -> bool:
if _HEAVY_MIN_AVAILABLE_BYTES <= 0:
return True
import psutil
from app.services.memory_budget import memory_budget
return psutil.virtual_memory().available >= _HEAVY_MIN_AVAILABLE_BYTES
return memory_budget().available_bytes >= _HEAVY_MIN_AVAILABLE_BYTES
class HeavyJobLimitTimeoutError(TimeoutError):
+91
View File
@@ -0,0 +1,91 @@
"""Process-visible memory budget, including Linux cgroup ancestors.
System counters keep their psutil semantics; admission and batch sizing use the
smaller container budget. No cgroup files are written by this module.
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from pathlib import Path
import psutil
_PROC_SELF = Path("/proc/self")
@dataclass(frozen=True)
class MemoryBudget:
total_bytes: int
available_bytes: int
def _read_int(path: Path) -> int | None:
try:
value = int(path.read_text().strip())
return value if value >= 0 else None
except (OSError, ValueError):
return None
def _unescape(value: str) -> str:
return re.sub(r"\\([0-7]{3})", lambda m: chr(int(m[1], 8)), value)
def _memory_groups() -> list[tuple[Path, Path, bool]]:
try:
membership = _PROC_SELF.joinpath("cgroup").read_text().splitlines()
mounts = _PROC_SELF.joinpath("mountinfo").read_text().splitlines()
except OSError:
return []
groups = []
for line in membership:
parts = line.split(":", 2)
if len(parts) != 3:
continue
unified = parts[0] == "0" and not parts[1]
if not unified and "memory" not in parts[1].split(","):
continue
member = Path(_unescape(parts[2]))
if not member.is_absolute() or ".." in member.parts:
continue
for mount in mounts:
before, sep, after = mount.partition(" - ")
fields, filesystem = before.split(), after.split()
if not sep or len(fields) < 5 or len(filesystem) < 3:
continue
if filesystem[0] != ("cgroup2" if unified else "cgroup"):
continue
if not unified and "memory" not in filesystem[2].split(","):
continue
root, point = Path(_unescape(fields[3])), Path(_unescape(fields[4]))
try:
relative = member.relative_to(root)
except ValueError:
continue
groups.append((point / relative, point, unified))
return groups
def memory_budget(vm=None) -> MemoryBudget:
vm = psutil.virtual_memory() if vm is None else vm
total = int(getattr(vm, "total", vm.available))
available = int(vm.available)
for leaf, mount, unified in _memory_groups():
current = leaf
while True:
usage = _read_int(current / ("memory.current" if unified else "memory.usage_in_bytes"))
names = ("memory.max", "memory.high") if unified else ("memory.limit_in_bytes",)
for name in names:
limit = _read_int(current / name)
# 'max', unreadable controls and v1's huge unlimited sentinel
# cannot tighten the host budget. Include ancestor usage because
# sibling processes consume the same parent allocation.
if limit is not None and limit < int(getattr(vm, "total", vm.available)):
total = min(total, limit)
if usage is not None:
available = min(available, max(0, limit - usage))
if current == mount:
break
current = current.parent
return MemoryBudget(total, min(total, max(0, available)))
+54 -8
View File
@@ -29,8 +29,9 @@ import threading
import time
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from datetime import date, datetime
from datetime import date, datetime, timedelta
from datetime import time as dt_time
from pathlib import Path
import polars as pl
@@ -308,6 +309,7 @@ class QuoteService:
# prev 不早于本时段开盘时计算
self._volume_delta: dict[str, tuple[float, float]] = {}
self._volume_delta_span_s: float = 0.0
self._etf_history_cache: tuple[tuple, pl.DataFrame] | None = None
# ================================================================
# 生命周期
@@ -339,6 +341,7 @@ class QuoteService:
self._thread = None
if thread:
thread.join(timeout=10)
self._etf_history_cache = None
# 此处不持久化关闭: lifespan shutdown (容器停止/重启) 也调用 stop,
# 持久化 False 会让每次重启后实时行情都变关闭、需手动再开。
# 持久化语义归 disable() (用户主动关闭)。
@@ -431,6 +434,7 @@ class QuoteService:
def set_repo(self, repo) -> None:
"""注入 KlineRepository, 用于实时落盘。"""
self._repo = repo
self._etf_history_cache = None
def set_app_state(self, app_state) -> None:
"""注入 FastAPI app.state, 用于获取 strategy_monitor 等单例。"""
@@ -1889,6 +1893,45 @@ class QuoteService:
# enriched 增量计算
# ================================================================
@staticmethod
def _etf_history_files(root: Path, today: date) -> tuple:
# Daily storage uses date partitions. Probe only the existing 150-day
# warmup horizon, so old years and today's repeated writes cost no scan.
files = []
for offset in range(150, 0, -1):
folder = root / f"date={today - timedelta(days=offset)}"
for path in sorted(folder.glob("*.parquet")):
try:
stat = path.stat()
except FileNotFoundError:
continue
files.append((str(path), stat.st_mtime_ns, stat.st_ctime_ns, stat.st_size, stat.st_ino))
return tuple(files)
def _etf_live_history(self, today: date, columns: list[str]) -> pl.DataFrame | None:
root = self._repo.store.data_dir / "kline_etf_daily"
files = self._etf_history_files(root, today)
if not files:
self._etf_history_cache = None
return None # Preserve the original empty-history fallback.
key = (str(root), today, files)
cached = self._etf_history_cache
if cached is not None and cached[0] == key:
return cached[1]
frame = guarded_collect(
scan_daily_parquet([entry[0] for entry in files])
.select(columns)
.filter((pl.col("date") >= today - timedelta(days=150)) & (pl.col("date") != today))
.sort(["symbol", "date"]),
priority="background",
)
# A bounded immutable input cache only: formulas, adjustment factors and
# signals are still evaluated on each tick by the original pipeline.
self._etf_history_cache = None
if frame.estimated_size() <= 64 * 1024**2 and files == self._etf_history_files(root, today):
self._etf_history_cache = (key, frame)
return frame
def _flush_live_enriched(self, daily_df: pl.DataFrame, quote_extra: pl.DataFrame = None, asset_type: str = "stock", merge: bool = False) -> None:
"""增量计算今天的 enriched: 用昨天的递推状态 + 今天 OHLCV → 只算今天 5500 行。
@@ -1955,13 +1998,16 @@ class QuoteService:
table = {"etf": "kline_etf_daily", "index": "kline_index_daily"}.get(asset_type, "kline_daily")
daily_glob = str(self._repo.store.data_dir / table / "**" / "*.parquet")
ohlcv_cols = ["symbol", "date", "open", "high", "low", "close", "volume", "amount", "quote_ts"]
hist_df = guarded_collect(
scan_daily_parquet(daily_glob)
# 多读一段自然日, 再按交易日计数确定窗口起点 (长假前后 90 个自然日不足 60 个交易日)
.filter(pl.col("date") >= cutoff - timedelta(days=60))
.sort(["symbol", "date"]),
priority="background",
)
hist_df = self._etf_live_history(today, ohlcv_cols) if asset_type == "etf" else None
if hist_df is None:
hist_df = guarded_collect(
scan_daily_parquet(daily_glob)
# 多读一段自然日, 再按交易日计数确定窗口起点 (长假前后 90 个自然日不足 60 个交易日)
.filter(pl.col("date") >= cutoff - timedelta(days=60))
.select(ohlcv_cols)
.sort(["symbol", "date"]),
priority="background",
)
if hist_df.is_empty():
return
@@ -0,0 +1,203 @@
"""Offline, isolated benchmark for projection, overview and ETF history caching.
Run with the backend on PYTHONPATH; use the same interpreter/script against an
archived baseline to compare. Each run creates and removes its own synthetic
data. No provider calls or real user data are used. Timings exclude fixture
generation and ETF disk output; output digests must match across revisions.
"""
from __future__ import annotations
import hashlib
import json
import logging
import os
import statistics
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import date, timedelta
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
import polars as pl
def digest(value) -> str:
return hashlib.sha256(json.dumps(value, sort_keys=True, default=str).encode()).hexdigest()
def timed(fn, runs=7) -> dict:
samples = []
for _ in range(runs):
start = time.perf_counter()
fn()
samples.append((time.perf_counter() - start) * 1000)
return {"median_ms": round(statistics.median(samples), 3),
"runs_ms": [round(x, 3) for x in samples]}
def ext_benchmark(root: Path) -> dict:
from app.api import ext_data
from app.services.ext_data import ExtConfig, ExtField
rng = np.random.default_rng(7)
days = [date(2026, 1, 1) + timedelta(days=i) for i in range(120)]
for day in days:
path = root / "ext_data/bench/timeseries" / f"date={day}" / "part.parquet"
path.parent.mkdir(parents=True)
pl.DataFrame({"symbol": [f"{i:06}.SZ" for i in range(2000)], **{
f"metric_{i}": rng.random(2000) for i in range(30)
}}).write_parquet(path)
cfg = ExtConfig(id="bench", label="bench", mode="timeseries", fields=[ExtField("symbol")])
materialized = []
original = ext_data._with_instrument_name
def inspect(frame, data_dir):
materialized.append(frame.estimated_size())
return original(frame, data_dir)
def query():
return ext_data.list_rows(None, "bench", snapshot_date=None, columns="metric_0",
limit=100, offset=0, filter=None, sort=None,
start_date=str(days[0]), end_date=str(days[-1]))
with (patch.object(ext_data, "_store", lambda _: SimpleNamespace(get=lambda _: cfg)),
patch.object(ext_data, "_data_dir", lambda _: root),
patch.object(ext_data, "_with_instrument_name", inspect)):
result = query()
timing = timed(query)
assert result["total"] == 240000 and len(result["rows"]) == 100
return {"partitions": 120, "rows": 240000, "stored_columns": 31,
"materialized_bytes": materialized[-1], "output_digest": digest(result), **timing}
def overview_benchmark() -> dict:
from app.api import overview
counts, durations = [], []
for _ in range(7):
overview.invalidate_overview_cache()
barrier = threading.Barrier(8)
builds = []
def build(*args, builds=builds):
builds.append(1)
# Controlled miss window: only duplicate-build counts are evidence
# about production work; this sleep is not a latency speedup claim.
time.sleep(0.04)
return {"market": {"symbols": 5500, "up": 2100, "down": 3100}}
def request(_, barrier=barrier):
barrier.wait(timeout=5)
return overview.market_overview(None)
with patch.object(overview, "_build_overview", build), ThreadPoolExecutor(8) as pool:
start = time.perf_counter()
result = list(pool.map(request, range(8)))
durations.append((time.perf_counter() - start) * 1000)
assert all(row == result[0] for row in result)
counts.append(len(builds))
return {"concurrent_requests": 8, "builds_per_burst": counts,
"synthetic_build_delay_ms": 40, "output_digest": digest(result),
"median_burst_ms": round(statistics.median(durations), 3)}
def etf_benchmark(root: Path) -> dict:
from app.indicators import pipeline
from app.services import quote_service
today, n = date(2026, 10, 29), 500
days = sorted(today - timedelta(days=i) for i in range(1, 145)
if (today - timedelta(days=i)).weekday() < 5)
for j, day in enumerate(days):
frame = pl.DataFrame({
"symbol": [f"{510000+i}.SH" for i in range(n)], "date": [day] * n,
"open": [10. + i % 30 * .1 for i in range(n)],
"high": [12. + i % 30 * .1 for i in range(n)],
"low": [9. + i % 30 * .1 for i in range(n)],
"close": [10.5 + i % 30 * .1 + j * .001 for i in range(n)],
"volume": [100000. + i + j for i in range(n)],
"amount": [1000000. + i + j for i in range(n)],
})
path = root / "kline_etf_daily" / f"date={day}" / "part.parquet"
path.parent.mkdir(parents=True)
frame.write_parquet(path)
current = frame.with_columns(pl.lit(today).alias("date"))
outputs, scans, computes = [], [], []
repo = SimpleNamespace(store=SimpleNamespace(data_dir=root),
get_enriched_latest_asset=lambda _: (current, today),
flush_live_enriched_asset=lambda asset, df: outputs.append(df))
service = quote_service.QuoteService()
service.set_repo(repo)
collect, compute = quote_service.guarded_collect, pipeline.compute_enriched
def counted_scan(*args, **kwargs):
scans.append(1)
return collect(*args, **kwargs)
def counted_compute(frame, *args, **kwargs):
computes.append(frame.height)
return compute(frame, *args, **kwargs)
with (patch.object(quote_service, "cn_today", lambda: today),
patch.object(quote_service, "guarded_collect", counted_scan),
patch.object(pipeline, "compute_enriched", counted_compute)):
cold = timed(lambda: service._flush_live_enriched(current, asset_type="etf"), runs=1)
for _ in range(2):
service._flush_live_enriched(current, asset_type="etf")
scans.clear()
warm = timed(lambda: service._flush_live_enriched(current, asset_type="etf"), runs=21)
assert len(outputs) == 24 and all(df.height == n for df in outputs)
hashes = [digest(df.sort("symbol").to_dicts()) for df in outputs]
assert len(set(hashes)) == 1
cache = getattr(service, "_etf_history_cache", None)
return {"symbols": n, "history_partitions": len(days), "warm_history_scans": len(scans),
"warm_updates": 21, "full_compute_rows": computes[-1],
"cache_bytes": cache[1].estimated_size() if cache else 0,
"cold_ms": cold["median_ms"], "output_digest": hashes[0], **warm}
def memory_benchmark(root: Path) -> dict:
import psutil
from app.indicators import pipeline
from app.services import heavy_job_limiter
gib = 1024**3
with patch.object(psutil, "virtual_memory", lambda: SimpleNamespace(total=32*gib, available=16*gib)):
try:
from app.services import memory_budget
except ImportError:
memory_budget = None
if memory_budget is not None:
proc, mount = root / "proc", root / "cgroup"
proc.mkdir()
mount.mkdir()
(proc / "cgroup").write_text("0::/\n")
(proc / "mountinfo").write_text(f"1 0 0:1 / {mount} rw - cgroup2 cgroup rw\n")
(mount / "memory.max").write_text(str(2*gib))
(mount / "memory.current").write_text(str(3*gib//2))
memory_budget._PROC_SELF = proc
with patch.object(heavy_job_limiter, "_HEAVY_MIN_AVAILABLE_BYTES", gib):
return {"simulated_host_gib": 32, "container_limit_gib": 2,
"container_available_gib": 0.5,
"heavy_job_admitted": heavy_job_limiter._memory_admission_ok(),
"batch_symbols_for_1500_rows_each": pipeline._adaptive_sym_batch(1000, 1500)}
def main():
logging.basicConfig(level=logging.ERROR)
with tempfile.TemporaryDirectory(prefix="tsp-hotpaths-") as temp:
root = Path(temp)
os.environ["DATA_DIR"] = str(root / "state")
os.environ["TICKFLOW_ENV_FILE"] = str(root / "absent.env")
result = {"ext_query": ext_benchmark(root), "overview": overview_benchmark(),
"etf": etf_benchmark(root), "memory": memory_benchmark(root)}
print(json.dumps(result, indent=2))
if __name__ == "__main__":
main()
+101
View File
@@ -0,0 +1,101 @@
from types import SimpleNamespace
import polars as pl
import pytest
from fastapi import HTTPException
from app.api import ext_data
from app.services.ext_data import ExtConfig, ExtField, PullConfig
@pytest.fixture
def query(tmp_path, monkeypatch):
cfg = ExtConfig(id="bench", label="测试", mode="timeseries", fields=[ExtField("symbol")])
cfg.pull = PullConfig(time_field="time")
for i, day in enumerate(["2026-10-28", "2026-10-29"]):
path = tmp_path / "ext_data/bench/timeseries" / f"date={day}" / "part.parquet"
path.parent.mkdir(parents=True)
frame = pl.DataFrame({
"symbol": ["600000.SH", "000001.SZ", "600519.SH"],
"score": [float("nan"), 1.0 + i, None],
"group": ["A", "B", "A"], "time": ["09:32", "09:31", "09:33"],
"unused": ["large payload"] * 3,
"date": [day] * 3,
})
if i:
frame = frame.with_columns(pl.lit("added").alias("extra"))
frame.write_parquet(path)
inst = tmp_path / "instruments/instruments.parquet"
inst.parent.mkdir()
pl.DataFrame({"symbol": ["000001.SZ"], "name": ["平安银行"]}).write_parquet(inst)
monkeypatch.setattr(ext_data, "_store", lambda _: SimpleNamespace(get=lambda _: cfg))
monkeypatch.setattr(ext_data, "_data_dir", lambda _: tmp_path)
def run(**overrides):
args = dict(snapshot_date=None, columns="score", limit=3, offset=1, filter=None,
sort=None, start_date="2026-10-28", end_date="2026-10-29")
args.update(overrides)
return ext_data.list_rows(None, "bench", **args)
return run, cfg, tmp_path
@pytest.mark.parametrize("kwargs", [
{},
{"filter": ["group:A"], "sort": "time:desc"},
{"filter": ["name:平安银行"], "sort": "score:desc"},
{"columns": "extra,data_date,date,score,score", "limit": 10, "offset": 0},
{"columns": "unknown"},
{"columns": " , "},
{"snapshot_date": "2026-10-28"},
{"start_date": None, "end_date": None},
{"start_date": "2026-11-01", "end_date": "2026-11-02"},
])
def test_projected_response_matches_full_read(query, monkeypatch, kwargs):
run, _, _ = query
actual = run(**kwargs)
original = ext_data._read_ext_dataframe
def full_read(*args, **ignored):
return original(*args)
monkeypatch.setattr(ext_data, "_read_ext_dataframe", full_read)
assert actual == run(**kwargs)
@pytest.mark.parametrize("kwargs", [
{"filter": ["absent:1"]}, {"filter": ["group"]}, {"sort": "absent:desc"},
])
def test_projection_preserves_error_details(query, monkeypatch, kwargs):
run, _, _ = query
with pytest.raises(HTTPException) as projected:
run(**kwargs)
original = ext_data._read_ext_dataframe
monkeypatch.setattr(ext_data, "_read_ext_dataframe", lambda *a, **kw: original(*a))
with pytest.raises(HTTPException) as full:
run(**kwargs)
assert (projected.value.status_code, projected.value.detail) == (full.value.status_code, full.value.detail)
def test_unused_columns_are_not_materialized(query, monkeypatch):
run, _, _ = query
original = ext_data._with_instrument_name
materialized = []
def inspect(frame, data_dir):
materialized.append(frame.columns)
return original(frame, data_dir)
monkeypatch.setattr(ext_data, "_with_instrument_name", inspect)
result = run(filter=["group:B"], sort="time:desc")
assert set(materialized[0]) == {"symbol", "score", "group", "time"}
assert result["total"] == 2
assert list(result["rows"][0]) == ["symbol", "name", "score"]
def test_snapshot_unknown_projection_keeps_full_row_fallback(query):
run, cfg, root = query
cfg.mode, cfg.pull = "snapshot", None
pl.DataFrame({"value": [1], "other": [2]}).write_parquet(root / "ext_data/bench/part.parquet")
result = run(columns="missing", start_date=None, end_date=None, offset=0)
assert result["rows"] == [{"value": 1, "other": 2}]
+103
View File
@@ -0,0 +1,103 @@
from types import SimpleNamespace
import psutil
import pytest
from app.indicators import pipeline
from app.services import heavy_job_limiter, memory_budget
GIB = 1024**3
@pytest.fixture
def cgroup(tmp_path, monkeypatch):
proc = tmp_path / "proc"
proc.mkdir()
mount = tmp_path / "cgroup"
mount.mkdir()
(proc / "cgroup").write_text("0::/parent/job\n")
(proc / "mountinfo").write_text(f"1 0 0:1 / {mount} rw - cgroup2 cgroup rw\n")
(mount / "parent/job").mkdir(parents=True)
monkeypatch.setattr(memory_budget, "_PROC_SELF", proc)
monkeypatch.setattr(psutil, "virtual_memory", lambda: SimpleNamespace(total=32*GIB, available=16*GIB))
return proc, mount
def write(group, **controls):
for name, value in controls.items():
(group / name.replace("_", ".")).write_text(str(value))
def test_parent_limit_accounts_for_sibling_usage_and_recovers(cgroup):
_, mount = cgroup
write(mount / "parent", memory_max=4*GIB, memory_current=3*GIB)
write(mount / "parent/job", memory_max="max", memory_current=GIB)
assert memory_budget.memory_budget() == memory_budget.MemoryBudget(4*GIB, GIB)
write(mount / "parent", memory_current=GIB)
assert memory_budget.memory_budget().available_bytes == 3*GIB
def test_memory_high_is_an_admission_budget(cgroup):
_, mount = cgroup
write(mount / "parent/job", memory_max=4*GIB, memory_high=2*GIB, memory_current=GIB)
assert memory_budget.memory_budget() == memory_budget.MemoryBudget(2*GIB, GIB)
def test_heavy_gate_and_batch_size_use_container_budget(cgroup, monkeypatch):
_, mount = cgroup
write(mount / "parent/job", memory_max=2*GIB, memory_current=3*GIB//2)
monkeypatch.setattr(heavy_job_limiter, "_HEAVY_MIN_AVAILABLE_BYTES", GIB)
assert not heavy_job_limiter._memory_admission_ok()
assert pipeline._adaptive_sym_batch(1000, 1500) == 100
write(mount / "parent/job", memory_current=GIB//2, memory_max=16*GIB)
assert heavy_job_limiter._memory_admission_ok()
assert pipeline._adaptive_sym_batch(1000, 1500) == 1000
def test_v1_bind_mount_and_unlimited_sentinel(cgroup):
proc, mount = cgroup
(proc / "cgroup").write_text("2:cpu:/other\n4:memory:/docker/app\n")
(proc / "mountinfo").write_text(f"1 0 0:1 /docker/app {mount} rw - cgroup cgroup rw,memory\n")
(mount / "memory.limit_in_bytes").write_text(str(2*GIB))
(mount / "memory.usage_in_bytes").write_text(str(GIB))
assert memory_budget.memory_budget() == memory_budget.MemoryBudget(2*GIB, GIB)
(mount / "memory.limit_in_bytes").write_text(str(2**63-4096))
assert memory_budget.memory_budget() == memory_budget.MemoryBudget(32*GIB, 16*GIB)
def test_missing_or_malformed_controls_preserve_host_fallback(cgroup):
proc, mount = cgroup
write(mount / "parent/job", memory_max="bad", memory_current=GIB)
assert memory_budget.memory_budget() == memory_budget.MemoryBudget(32*GIB, 16*GIB)
(proc / "mountinfo").unlink()
assert memory_budget.memory_budget().available_bytes == 16*GIB
def test_exhausted_limit_never_returns_negative_available(cgroup):
_, mount = cgroup
write(mount / "parent/job", memory_max=GIB, memory_current=2*GIB)
assert memory_budget.memory_budget().available_bytes == 0
def test_health_keeps_host_counters_and_exposes_effective_budget(cgroup, monkeypatch):
from app.api.routes import _memory_snapshot
_, mount = cgroup
write(mount / "parent/job", memory_max=2*GIB, memory_current=GIB)
monkeypatch.setattr(psutil, "virtual_memory", lambda: SimpleNamespace(total=32*GIB, available=16*GIB, percent=50))
snapshot = _memory_snapshot()
assert snapshot["system_total_bytes"] == 32*GIB
assert snapshot["system_available_bytes"] == 16*GIB
assert snapshot["effective_total_bytes"] == 2*GIB
assert snapshot["effective_available_bytes"] == GIB
def test_escaped_mount_path_and_host_pressure(cgroup, monkeypatch):
proc, mount = cgroup
moved = mount.with_name("group with space")
mount.rename(moved)
escaped = str(moved).replace(" ", r"\040")
(proc / "mountinfo").write_text(f"1 0 0:1 / {escaped} rw - cgroup2 cgroup rw\n")
write(moved / "parent/job", memory_max=4*GIB, memory_current=GIB)
monkeypatch.setattr(psutil, "virtual_memory", lambda: SimpleNamespace(total=32*GIB, available=GIB//2))
assert memory_budget.memory_budget() == memory_budget.MemoryBudget(4*GIB, GIB//2)
+110
View File
@@ -0,0 +1,110 @@
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import date
import pytest
from app.api import overview
@pytest.fixture(autouse=True)
def reset_cache():
overview.invalidate_overview_cache()
yield
overview.invalidate_overview_cache()
def test_concurrent_cold_requests_share_one_build(monkeypatch):
barrier = threading.Barrier(8)
calls = []
def build(*args):
calls.append(1)
time.sleep(0.08)
return {"rows": [{"symbol": "600000.SH", "change_pct": 0.05}]}
def request():
barrier.wait(timeout=2)
return overview.market_overview(None)
monkeypatch.setattr(overview, "_build_overview", build)
with ThreadPoolExecutor(8) as pool:
results = list(pool.map(lambda _: request(), range(8)))
assert len(calls) == 1
assert all(result == results[0] for result in results)
def test_invalidation_during_build_never_repopulates_old_cache(monkeypatch):
entered, release = threading.Event(), threading.Event()
calls = []
def build(*args):
calls.append(1)
if len(calls) == 1:
entered.set()
assert release.wait(2)
return {"value": "old"}
return {"value": "new"}
monkeypatch.setattr(overview, "_build_overview", build)
with ThreadPoolExecutor(2) as pool:
pending = pool.submit(overview.market_overview, None)
assert entered.wait(2)
overview.invalidate_overview_cache()
assert overview.market_overview(None) == {"value": "new"}
release.set()
pending.result(timeout=2)
assert overview.market_overview(None) == {"value": "new"}
def test_failed_build_releases_followers_and_can_retry(monkeypatch):
barrier = threading.Barrier(4)
calls = []
def fail(*args):
calls.append(1)
time.sleep(0.08)
raise ValueError("unavailable")
def request():
barrier.wait(timeout=2)
with pytest.raises(ValueError, match="unavailable"):
overview.market_overview(None)
monkeypatch.setattr(overview, "_build_overview", fail)
with ThreadPoolExecutor(4) as pool:
list(pool.map(lambda _: request(), range(4)))
assert len(calls) == 1
monkeypatch.setattr(overview, "_build_overview", lambda *args: {"ok": True})
assert overview.market_overview(None) == {"ok": True}
def test_different_dates_build_independently(monkeypatch):
barrier = threading.Barrier(2)
def build(request, as_of):
barrier.wait(timeout=2)
return {"date": as_of}
monkeypatch.setattr(overview, "_build_overview", build)
days = [date(2026, 10, 28), date(2026, 10, 29)]
with ThreadPoolExecutor(2) as pool:
assert list(pool.map(lambda day: overview.market_overview(None, day), days)) == [{"date": d} for d in days]
def test_slow_build_does_not_extend_existing_freshness_window(monkeypatch):
now, calls = [100.0], []
monkeypatch.setattr(overview.time, "monotonic", lambda: now[0])
def build(*args):
calls.append(1)
now[0] += 2.0
return {"build": len(calls)}
monkeypatch.setattr(overview, "_build_overview", build)
assert overview.market_overview(None) == {"build": 1}
now[0] = 104.0
assert overview.market_overview(None) == {"build": 1}
now[0] = 105.0
assert overview.market_overview(None) == {"build": 2}
@@ -0,0 +1,122 @@
from datetime import date, timedelta
from types import SimpleNamespace
import polars as pl
import pytest
from polars.testing import assert_frame_equal
from app.services import quote_service
TODAY = date(2026, 10, 29)
@pytest.fixture
def sample(tmp_path, monkeypatch):
monkeypatch.setattr(quote_service, "cn_today", lambda: TODAY)
days = [TODAY - timedelta(days=i) for i in range(1, 145)]
days = sorted(d for d in days if d.weekday() < 5)
for i, day in enumerate(days):
frame = pl.DataFrame({
"symbol": ["510300.SH", "159915.SZ"], "date": [day] * 2,
"open": [10.0, 20.0], "high": [11.0+i/100, 21.0+i/100],
"low": [9.0, 19.0], "close": [10.0+i/100, 20.0+i/100],
"volume": [10000.0, 20000.0], "amount": [100000.0, 400000.0],
})
path = tmp_path / "kline_etf_daily" / f"date={day}" / "part.parquet"
path.parent.mkdir(parents=True)
frame.write_parquet(path)
today = frame.with_columns(pl.lit(TODAY).alias("date"))
output = []
repo = SimpleNamespace(
store=SimpleNamespace(data_dir=tmp_path),
get_enriched_latest_asset=lambda _: (today, TODAY),
flush_live_enriched_asset=lambda asset, df: output.append(df),
merge_live_enriched_asset=lambda asset, df: output.append(df),
)
service = quote_service.QuoteService()
service.set_repo(repo)
return service, repo, today, output, days
def test_warm_etf_updates_do_not_rescan_history(sample, monkeypatch):
service, _repo, today, output, _ = sample
original = quote_service.guarded_collect
scans = []
def collect(*args, **kwargs):
scans.append(1)
return original(*args, **kwargs)
monkeypatch.setattr(quote_service, "guarded_collect", collect)
service._flush_live_enriched(today, asset_type="etf")
service._flush_live_enriched(today, asset_type="etf", merge=True)
assert len(output) == 2
assert_frame_equal(output[0], output[1], check_exact=True)
assert len(scans) == 1
@pytest.mark.parametrize("change", ["history", "factor", "rollover", "new_symbol", "delete"])
def test_warm_etf_output_matches_cold_recompute_after_changes(sample, monkeypatch, change):
service, repo, today, output, days = sample
service._flush_live_enriched(today, asset_type="etf")
root = repo.store.data_dir
if change == "history":
path = root / "kline_etf_daily" / f"date={days[-1]}" / "part.parquet"
pl.read_parquet(path).with_columns(pl.lit(15.0).alias("close")).write_parquet(path)
elif change == "delete":
(root / "kline_etf_daily" / f"date={days[-1]}" / "part.parquet").unlink()
elif change == "factor":
path = root / "adj_factor_etf/all.parquet"
path.parent.mkdir(parents=True)
pl.DataFrame({"symbol": ["510300.SH"], "trade_date": [days[-5]], "ex_factor": [1.1]}).write_parquet(path)
elif change == "rollover":
path = root / "kline_etf_daily" / f"date={TODAY}" / "part.parquet"
path.parent.mkdir(parents=True)
today.write_parquet(path)
monkeypatch.setattr(quote_service, "cn_today", lambda: TODAY + timedelta(days=1))
today = today.with_columns(pl.lit(TODAY + timedelta(days=1)).alias("date"))
else:
today = pl.concat([today, today.head(1).with_columns(pl.lit("510500.SH").alias("symbol"))])
today = today.with_columns((pl.col("close") + 0.25).alias("close"))
service._flush_live_enriched(today, asset_type="etf")
cold = quote_service.QuoteService()
cold.set_repo(repo)
# Compare with the original uncached glob-read path, including all numeric
# columns, nulls, adjusted prices and indicator warmup boundaries.
monkeypatch.setattr(cold, "_etf_live_history", lambda *args: None, raising=False)
cold._flush_live_enriched(today, asset_type="etf")
assert len(output) == 3
assert_frame_equal(output[-2], output[-1], check_exact=True)
def test_today_disk_writes_do_not_invalidate_history(sample, monkeypatch):
service, repo, today, output, _ = sample
original = quote_service.guarded_collect
scans = []
def collect(*args, **kwargs):
scans.append(1)
return original(*args, **kwargs)
monkeypatch.setattr(quote_service, "guarded_collect", collect)
service._flush_live_enriched(today, asset_type="etf")
path = repo.store.data_dir / "kline_etf_daily" / f"date={TODAY}" / "part.parquet"
path.parent.mkdir(parents=True)
today = today.with_columns((pl.col("close") + 0.5).alias("close"))
today.write_parquet(path)
service._flush_live_enriched(today, asset_type="etf")
assert len(scans) == 1
assert_frame_equal(output[-1].select("symbol", "raw_close").sort("symbol"),
today.select("symbol", pl.col("close").alias("raw_close")).sort("symbol"))
def test_oversized_history_is_not_retained_and_repo_change_clears_cache(sample, monkeypatch):
service, repo, today, output, _ = sample
service._flush_live_enriched(today, asset_type="etf")
assert service._etf_history_cache is not None
service.set_repo(repo)
assert service._etf_history_cache is None
monkeypatch.setattr(pl.DataFrame, "estimated_size", lambda *args: 65 * 1024**2)
service._flush_live_enriched(today, asset_type="etf")
assert service._etf_history_cache is None
assert_frame_equal(output[-2], output[-1], check_exact=True)