diff --git a/README.md b/README.md index ec2fce37..55459969 100644 --- a/README.md +++ b/README.md @@ -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` = 完全串行(旧行为) | ### 跑起来后的第一次使用 diff --git a/backend/app/api/ext_data.py b/backend/app/api/ext_data.py index 47223eef..ba6da520 100644 --- a/backend/app/api/ext_data.py +++ b/backend/app/api/ext_data.py @@ -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: diff --git a/backend/app/api/overview.py b/backend/app/api/overview.py index 550bd7c7..395bc1df 100644 --- a/backend/app/api/overview.py +++ b/backend/app/api/overview.py @@ -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 diff --git a/backend/app/api/routes.py b/backend/app/api/routes.py index 6167a3d3..ba718b66 100644 --- a/backend/app/api/routes.py +++ b/backend/app/api/routes.py @@ -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: diff --git a/backend/app/indicators/pipeline.py b/backend/app/indicators/pipeline.py index e4e5ea29..d300c3e6 100644 --- a/backend/app/indicators/pipeline.py +++ b/backend/app/indicators/pipeline.py @@ -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: diff --git a/backend/app/services/heavy_job_limiter.py b/backend/app/services/heavy_job_limiter.py index d45c7069..c4d81931 100644 --- a/backend/app/services/heavy_job_limiter.py +++ b/backend/app/services/heavy_job_limiter.py @@ -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): diff --git a/backend/app/services/memory_budget.py b/backend/app/services/memory_budget.py new file mode 100644 index 00000000..ac319376 --- /dev/null +++ b/backend/app/services/memory_budget.py @@ -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))) diff --git a/backend/app/services/quote_service.py b/backend/app/services/quote_service.py index fd0aec47..a08afc4c 100644 --- a/backend/app/services/quote_service.py +++ b/backend/app/services/quote_service.py @@ -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 diff --git a/backend/scripts/benchmark_remaining_hotpaths.py b/backend/scripts/benchmark_remaining_hotpaths.py new file mode 100644 index 00000000..f3e65283 --- /dev/null +++ b/backend/scripts/benchmark_remaining_hotpaths.py @@ -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() diff --git a/backend/tests/test_ext_rows_projection.py b/backend/tests/test_ext_rows_projection.py new file mode 100644 index 00000000..abe2cad8 --- /dev/null +++ b/backend/tests/test_ext_rows_projection.py @@ -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}] diff --git a/backend/tests/test_memory_budget.py b/backend/tests/test_memory_budget.py new file mode 100644 index 00000000..9504dba5 --- /dev/null +++ b/backend/tests/test_memory_budget.py @@ -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) diff --git a/backend/tests/test_overview_singleflight.py b/backend/tests/test_overview_singleflight.py new file mode 100644 index 00000000..1f6cd0f1 --- /dev/null +++ b/backend/tests/test_overview_singleflight.py @@ -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} diff --git a/backend/tests/test_quote_etf_history_cache.py b/backend/tests/test_quote_etf_history_cache.py new file mode 100644 index 00000000..be360bd9 --- /dev/null +++ b/backend/tests/test_quote_etf_history_cache.py @@ -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)