From 451e05e53123f616b74bd655f2eba2d4927f863f Mon Sep 17 00:00:00 2001 From: shaofl <2899218482@qq.com> Date: Fri, 19 Jun 2026 21:24:14 +0800 Subject: [PATCH] feat(data): D5 opt-in bounded concurrency with a global rate limiter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Opt-in bounded concurrency for the data-update / cache-warm fetch stage, guarded by ONE shared global rate limiter. Default stays fully serial: every existing config validates and behaves byte-identically, default max_workers is 1, no live Tushare call is needed for acceptance. - scheduler: new data/feed/scheduler.py GlobalRateLimiter — a thread-safe ticket limiter that reserves the next evenly-spaced slot under a lock and sleeps outside it, so N workers funnel through ONE per-minute budget (the quota is never multiplied per thread). monotonic/sleep are injectable for fake-clock tests; it sees only an integer budget (no token / kwargs / exception payload). - throttle: request_with_retry gains a `scheduler` param — when set it acquires a global slot before EVERY attempt (including retries) and skips the per-call rate_limit sleep; when unset, behavior is the historical per-call throttle. - feeds: tushare_feed/flags/fina/covariates/intraday + index_feed take an optional `scheduler` (default None = unchanged) and forward it to request_with_retry. - cache: TushareCache gains `max_workers` (default 1). max_workers==1 keeps the byte-identical serial read-through; >1 (and >1 symbols) fans the planned dense per-symbol gap FETCH onto a bounded ThreadPoolExecutor, then performs store upserts + ledger writes on the MAIN thread in deterministic plan order. Each successful gap is made durable before the first failure is re-raised; a failed gap records no coverage and stays retryable. Store/ledger contents are independent of fetch completion order. Snapshot / index_weight paging / intraday remain serial in D5. - config: data_update.concurrency.max_workers (default 1; <1 rejected readably); config/data_update.yaml documents the block at max_workers: 1. - updater: run_data_update builds ONE GlobalRateLimiter (from rate_limit_per_min) only when max_workers>1, shared by all feeds via _build_feeds; the cache gets max_workers; UpdateResult + format_summary surface max_workers and the global rate_limit_per_min (non-secret). Invariants: cache public methods/return frames, ledger paths/columns/coverage semantics (ok/empty count; failed/not_ready do not), retry semantics, and factor/alpha/portfolio/backtest math all unchanged. phase0 anchor (0.9600/0.8408) unchanged. No endpoint schema registry, no PanelStore append/partition, no data-quality / research change. tests/test_scheduler.py + tests/test_data_update_concurrency.py (network-free): global spacing escalates under contention (fake clock); global-not-per-thread under real threads; retries acquire a slot per attempt; scheduler-mode skips the per-call sleep; secret-safe failure; feed forwards its scheduler; config validation; serial==concurrent frames/ledger/request-counts; failure leaves the gap uncovered+retryable while successes stay durable; empty/not_ready unchanged. --- config/data_update.yaml | 7 ++ data/cache/tushare_cache.py | 132 ++++++++++++++++++++ data/feed/index_feed.py | 5 + data/feed/scheduler.py | 73 +++++++++++ data/feed/throttle.py | 20 ++- data/feed/tushare_covariates.py | 6 +- data/feed/tushare_feed.py | 4 + data/feed/tushare_fina.py | 6 +- data/feed/tushare_flags.py | 6 +- data/feed/tushare_intraday.py | 4 + qt/config.py | 28 +++++ qt/data_updater.py | 58 +++++++-- tests/test_data_update_concurrency.py | 173 ++++++++++++++++++++++++++ tests/test_scheduler.py | 173 ++++++++++++++++++++++++++ 14 files changed, 677 insertions(+), 18 deletions(-) create mode 100644 data/feed/scheduler.py create mode 100644 tests/test_data_update_concurrency.py create mode 100644 tests/test_scheduler.py diff --git a/config/data_update.yaml b/config/data_update.yaml index efcbc56..822b188 100644 --- a/config/data_update.yaml +++ b/config/data_update.yaml @@ -85,6 +85,13 @@ data_update: - adj_factor - stk_mins_1min report_name: data_update_quality_report.md + # D5 opt-in bounded concurrency. max_workers=1 (default) keeps the warm fully + # serial — every existing config behaves exactly as before. >1 fans the dense + # per-symbol gap FETCH onto a bounded thread pool, all funneled through ONE + # shared global rate limiter (reusing rate_limit_per_min above), so the Tushare + # quota is never multiplied per thread; store + ledger writes stay deterministic. + concurrency: + max_workers: 1 # --- unused by data-update (schema-only stubs) ------------------------------- # factors: diff --git a/data/cache/tushare_cache.py b/data/cache/tushare_cache.py index 74aae54..7bccdbd 100644 --- a/data/cache/tushare_cache.py +++ b/data/cache/tushare_cache.py @@ -114,6 +114,28 @@ FetchSnapshot = Callable[..., "pd.DataFrame | None"] +class _FetchOk: + """A successful concurrent gap fetch (D5): carries the raw frame.""" + + __slots__ = ("raw",) + + def __init__(self, raw) -> None: + self.raw = raw + + +class _FetchError: + """A failed concurrent gap fetch (D5): carries the exception OBJECT only. + + Re-raised on the main thread after durable successes; the message is the + feed-level secret-safe ``RuntimeError`` (exception type, never the token). + """ + + __slots__ = ("exc",) + + def __init__(self, exc: Exception) -> None: + self.exc = exc + + class TushareCache: """Endpoint-level read-through cache (market bars + universe/tradability).""" @@ -130,6 +152,7 @@ def __init__( source_version: str | None = None, not_ready_days: int = 0, recent_tail_overrides: dict[str, int] | None = None, + max_workers: int = 1, ) -> None: self._store = store self._ledger = ledger @@ -150,6 +173,12 @@ def __init__( # REQUESTED range (not today-based) — e.g. fina_indicator refetches a long # trailing window of recent report periods to catch LATE disclosures. self._recent_tail_overrides = dict(recent_tail_overrides or {}) + # D5: bounded concurrency for the dense per-symbol gap FETCH stage. 1 + # (default) keeps the serial read-through path byte-identical; >1 fans the + # planned gap fetches onto a thread pool (the shared GlobalRateLimiter lives + # in the feeds, so the global quota holds), while store upserts + ledger + # writes stay on the main thread in deterministic order. + self._max_workers = max(1, int(max_workers)) # per-instance endpoint fetch counters (cache stats; one increment per # gap/window/snapshot actually sent to the API). A fully-covered repeat # run leaves these at zero — the read-through hit rate is observable from @@ -376,6 +405,12 @@ def _read_through( columns: list[str], key_cols: list[str], ) -> pd.DataFrame: + # D5: only the opted-in multi-worker, multi-symbol case takes the + # concurrent path; everything else stays on the byte-identical serial loop. + if self._max_workers > 1 and len(symbols) > 1: + return self._read_through_concurrent( + endpoint, symbols, start, end, fetch, parse, columns, key_cols + ) req_start = pd.Timestamp(start).normalize() req_end = pd.Timestamp(end).normalize() fields_hash = _fields_hash(columns) @@ -409,6 +444,103 @@ def _read_through( return pd.DataFrame(columns=columns) return pd.concat(out, ignore_index=True).reset_index(drop=True) + def _read_through_concurrent( + self, endpoint, symbols, start, end, fetch, parse, columns, key_cols + ) -> pd.DataFrame: + """D5 bounded-concurrency read-through: parallel fetch, serial durable writes. + + Phase 1 (main thread): plan every uncovered gap (ledger reads use the D4 + lookup cache). Phase 2: fetch the planned gaps on a bounded thread pool — + each fetch funnels through the feeds' shared ``GlobalRateLimiter``, so the + global quota holds. Phase 3 (main thread, deterministic plan order): + parse/upsert/record each SUCCESSFUL gap so it is durable, then re-raise the + first failure (a failed gap records no coverage and stays retryable, exactly + like the serial path). Phase 4: read the requested window back per symbol. + The store/ledger contents are independent of fetch completion order. + """ + req_start = pd.Timestamp(start).normalize() + req_end = pd.Timestamp(end).normalize() + fields_hash = _fields_hash(columns) + forced = endpoint in self._force_refresh + + # Phase 1 — plan (main thread). + plans: list[tuple[str, pd.Timestamp, pd.Timestamp]] = [] + n_covered = 0 + for symbol in symbols: + gaps = self._gaps_for(endpoint, symbol, req_start, req_end, forced) + if not gaps: + n_covered += 1 + for gap_start, gap_end in gaps: + plans.append((symbol, gap_start, gap_end)) + + # Phase 2 — concurrent fetch (bounded by max_workers). + raw_results = self._fetch_plans_concurrent(plans, fetch) + + # Phase 3 — deterministic durable writes (main thread, plan order). + first_exc: Exception | None = None + for (symbol, gap_start, gap_end), res in zip(plans, raw_results): + if isinstance(res, _FetchError): + if first_exc is None: + first_exc = res.exc + continue + self.fetch_counts[endpoint] = self.fetch_counts.get(endpoint, 0) + 1 + parsed = parse(res.raw) + if len(parsed): + self._store.upsert_symbol(endpoint, symbol, parsed, key_cols) + self.written_counts[endpoint] = ( + self.written_counts.get(endpoint, 0) + len(parsed) + ) + self._record_gap_coverage( + endpoint, symbol, gap_start, gap_end, parsed, fields_hash + ) + if first_exc is not None: + # earlier successful gaps are already durable; the failed gap is not + # recorded, so a later run retries it (serial-equivalent failure model). + raise first_exc + + # Phase 4 — read back the requested window (symbol order). + out: list[pd.DataFrame] = [] + for symbol in symbols: + cached = self._store.read_symbol(endpoint, symbol) + if not cached.empty: + mask = (cached["date"] >= req_start) & (cached["date"] <= req_end) + hit = cached.loc[mask, columns] + if not hit.empty: + out.append(hit) + _LOGGER.info( + "cache %s: %d symbols, %d gap-fetches, %d fully-covered " + "(api calls=%d, workers=%d)", + endpoint, len(symbols), len(plans), n_covered, + self.fetch_counts[endpoint], self._max_workers, + ) + if not out: + return pd.DataFrame(columns=columns) + return pd.concat(out, ignore_index=True).reset_index(drop=True) + + def _fetch_plans_concurrent(self, plans, fetch): + """Fetch every planned ``(symbol, start, end)`` gap on a bounded pool. + + Returns a list aligned to ``plans``: a ``_FetchOk(raw)`` or a + ``_FetchError(exc)`` per gap. Each task catches its OWN exception (so no + worker exception escapes the pool); the throttle/retry + the shared global + rate limiter live inside ``fetch`` (the feed closure). + """ + if not plans: + return [] + from concurrent.futures import ThreadPoolExecutor + + def task(i): + symbol, gap_start, gap_end = plans[i] + try: + raw = fetch(symbol, _compact(gap_start), _compact(gap_end)) + return _FetchOk(raw) + except Exception as exc: # secret-safe: only the object is kept + return _FetchError(exc) + + workers = min(self._max_workers, len(plans)) + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(task, range(len(plans)))) + def _gaps_for(self, endpoint, symbol, req_start, req_end, forced): """The uncovered sub-intervals to fetch (+ a forced trailing tail).""" if forced: diff --git a/data/feed/index_feed.py b/data/feed/index_feed.py index 3314b43..5e4d59e 100644 --- a/data/feed/index_feed.py +++ b/data/feed/index_feed.py @@ -33,6 +33,7 @@ def __init__( rate_limit: int | None = None, max_retries: int = 6, cache=None, + scheduler=None, ) -> None: self._secret_file = str(secret_file) self._token_key = token_key @@ -44,6 +45,8 @@ def __init__( # one. The cache stores RAW snapshots; the as-of membership stays # downstream, unchanged. self._cache = cache + # D5: optional shared GlobalRateLimiter; None keeps the per-call throttle. + self._scheduler = scheduler # -- secret handling (token never logged) ------------------------------- # def _client(self): @@ -90,6 +93,7 @@ def get_constituents(self, index_code: str, start: str, end: str) -> pd.DataFram pro.index_weight, max_retries=self._max_retries, rate_limit=self._rate_limit, + scheduler=self._scheduler, index_code=index_code, start_date=win_start.strftime("%Y%m%d"), end_date=win_end.strftime("%Y%m%d"), @@ -121,6 +125,7 @@ def fetch(index_code: str, start_compact: str, end_compact: str): self._client().index_weight, max_retries=self._max_retries, rate_limit=self._rate_limit, + scheduler=self._scheduler, index_code=index_code, start_date=start_compact, end_date=end_compact, diff --git a/data/feed/scheduler.py b/data/feed/scheduler.py new file mode 100644 index 0000000..a26fd52 --- /dev/null +++ b/data/feed/scheduler.py @@ -0,0 +1,73 @@ +"""Global request scheduler / rate limiter for bounded-concurrency warms (D5). + +When a cache warm runs with more than one worker, every worker must funnel its +API attempts through ONE shared budget — otherwise N threads each sleeping +``60 / rate_limit`` would multiply the global quota by N and blow the Tushare +ceiling. :class:`GlobalRateLimiter` reserves the next evenly-spaced time slot +under a lock, then sleeps OUTSIDE the lock until that slot, so concurrent callers +are serialized onto one global cadence. + +It is transport-agnostic and secret-free: it never sees a token, kwargs, an +endpoint, or an exception payload — only a per-minute integer budget. ``monotonic`` +and ``sleep`` are injectable so the spacing can be asserted with a fake clock (no +wall-clock timing in tests). +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable + + +class GlobalRateLimiter: + """A single shared per-minute request budget enforced ACROSS worker threads. + + ``rate_limit_per_min <= 0`` disables throttling (``acquire`` is a no-op). Each + ``acquire`` reserves the next slot spaced ``60 / rate_limit_per_min`` seconds + after the previous reservation, regardless of which thread calls it — so the + global rate holds no matter how many workers contend. + """ + + def __init__( + self, + rate_limit_per_min: int, + *, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, + ) -> None: + self._rate_limit_per_min = int(rate_limit_per_min) + self._interval = ( + 60.0 / self._rate_limit_per_min if self._rate_limit_per_min > 0 else 0.0 + ) + self._monotonic = monotonic + self._sleep = sleep + self._lock = threading.Lock() + self._next_allowed: float | None = None + + @property + def rate_limit_per_min(self) -> int: + """The configured global per-minute budget (0 == disabled).""" + return self._rate_limit_per_min + + def acquire(self) -> None: + """Block until this caller's globally-spaced slot; no-op if disabled. + + The slot is reserved under the lock (so two threads never claim the same + instant), then the wait happens outside the lock so other threads can keep + reserving their own later slots concurrently. + """ + if self._interval <= 0: + return + with self._lock: + now = self._monotonic() + if self._next_allowed is None or now >= self._next_allowed: + # the budget is idle: this caller goes now, next slot one interval out + self._next_allowed = now + self._interval + wait = 0.0 + else: + # the budget is busy: wait for the reserved slot, push the next one out + wait = self._next_allowed - now + self._next_allowed = self._next_allowed + self._interval + if wait > 0: + self._sleep(wait) diff --git a/data/feed/throttle.py b/data/feed/throttle.py index 35aa0a4..f991ad4 100644 --- a/data/feed/throttle.py +++ b/data/feed/throttle.py @@ -18,25 +18,35 @@ def request_with_retry( *, max_retries: int = 6, rate_limit: int | None = None, + scheduler: Any = None, **kwargs: Any, ) -> Any: - """Call ``fn(**kwargs)`` with retry-on-error + per-minute throttle. + """Call ``fn(**kwargs)`` with retry-on-error + a per-minute throttle. Retries transient exceptions with exponential backoff up to ``max_retries`` - attempts; on success sleeps ``60 / rate_limit`` seconds (no-op if unset). - Raises a readable ``RuntimeError`` (type only, never the token) if every - attempt fails. + attempts. Throttling has two mutually-exclusive modes: + + * ``scheduler`` set (D5 bounded concurrency): acquire one global slot BEFORE + each attempt, so every worker's attempts (including retries) share one budget + — no per-call ``rate_limit`` sleep is added (that would double the spacing). + * ``scheduler`` unset (the historical serial path): on success sleep + ``60 / rate_limit`` seconds (no-op if unset). + + Raises a readable ``RuntimeError`` (exception TYPE only, never the token / + payload) if every attempt fails. """ attempts = max(1, int(max_retries)) last_exc: Exception | None = None for attempt in range(attempts): + if scheduler is not None: + scheduler.acquire() # global slot before EVERY attempt (incl. retries) try: result = fn(**kwargs) except Exception as exc: # transient API / network error last_exc = exc time.sleep(min(2.0**attempt, 8.0)) continue - if rate_limit: + if scheduler is None and rate_limit: time.sleep(60.0 / rate_limit) return result raise RuntimeError( diff --git a/data/feed/tushare_covariates.py b/data/feed/tushare_covariates.py index 7fa7487..8c7d2d8 100644 --- a/data/feed/tushare_covariates.py +++ b/data/feed/tushare_covariates.py @@ -42,6 +42,7 @@ def __init__( rate_limit: int | None = None, max_retries: int = 6, cache=None, + scheduler=None, ) -> None: self._secret_file = str(secret_file) self._token_key = token_key @@ -53,6 +54,8 @@ def __init__( # ``pit_sw_intervals`` (index_member_all, P4-3) all read through it. # None keeps the historical direct fetch EXACTLY. self._cache = cache + # D5: optional shared GlobalRateLimiter; None keeps the per-call throttle. + self._scheduler = scheduler def _client(self): if self._pro is None: @@ -63,7 +66,8 @@ def _client(self): def _call(self, fn, **kwargs): return request_with_retry( - fn, max_retries=self._max_retries, rate_limit=self._rate_limit, **kwargs + fn, max_retries=self._max_retries, rate_limit=self._rate_limit, + scheduler=self._scheduler, **kwargs, ) def industry(self, symbols: list[str]) -> dict[str, str]: diff --git a/data/feed/tushare_feed.py b/data/feed/tushare_feed.py index dd9c29b..0e405be 100644 --- a/data/feed/tushare_feed.py +++ b/data/feed/tushare_feed.py @@ -47,6 +47,7 @@ def __init__( rate_limit: int | None = None, max_retries: int = 6, cache=None, + scheduler=None, ) -> None: self._secret_file = str(secret_file) self._token_key = token_key @@ -59,6 +60,8 @@ def __init__( # passes one in. The cache stores RAW rows only; front_adjust still runs # downstream, unchanged. self._cache = cache + # D5: optional shared GlobalRateLimiter; None keeps the per-call throttle. + self._scheduler = scheduler # -- secret handling ---------------------------------------------------- # def _read_token(self) -> str: @@ -82,6 +85,7 @@ def _call(self, fn, **kwargs): fn, max_retries=self._max_retries, rate_limit=self._rate_limit, + scheduler=self._scheduler, **kwargs, ) diff --git a/data/feed/tushare_fina.py b/data/feed/tushare_fina.py index ff78da2..0d053df 100644 --- a/data/feed/tushare_fina.py +++ b/data/feed/tushare_fina.py @@ -34,6 +34,7 @@ def __init__( rate_limit: int | None = None, max_retries: int = 6, cache=None, + scheduler=None, ) -> None: self._secret_file = str(secret_file) self._token_key = token_key @@ -45,6 +46,8 @@ def __init__( # rows (with ann_date) only — the ann_date<=trade_date as-of alignment # stays downstream, byte-identical. self._cache = cache + # D5: optional shared GlobalRateLimiter; None keeps the per-call throttle. + self._scheduler = scheduler def _client(self): if self._pro is None: @@ -55,7 +58,8 @@ def _client(self): def _call(self, fn, **kwargs): return request_with_retry( - fn, max_retries=self._max_retries, rate_limit=self._rate_limit, **kwargs + fn, max_retries=self._max_retries, rate_limit=self._rate_limit, + scheduler=self._scheduler, **kwargs, ) def get_fina_indicator( diff --git a/data/feed/tushare_flags.py b/data/feed/tushare_flags.py index db68602..e8064a2 100644 --- a/data/feed/tushare_flags.py +++ b/data/feed/tushare_flags.py @@ -32,6 +32,7 @@ def __init__( rate_limit: int | None = None, max_retries: int = 6, cache=None, + scheduler=None, ) -> None: self._secret_file = str(secret_file) self._token_key = token_key @@ -43,6 +44,8 @@ def __init__( # The cache stores RAW suspend_d / namechange / stk_limit rows — never a # derived flag as source of truth; the flag derivation stays here. self._cache = cache + # D5: optional shared GlobalRateLimiter; None keeps the per-call throttle. + self._scheduler = scheduler def _client(self): if self._pro is None: @@ -53,7 +56,8 @@ def _client(self): def _call(self, fn, **kwargs): return request_with_retry( - fn, max_retries=self._max_retries, rate_limit=self._rate_limit, **kwargs + fn, max_retries=self._max_retries, rate_limit=self._rate_limit, + scheduler=self._scheduler, **kwargs, ) # -- suspensions (停牌) -------------------------------------------------- # diff --git a/data/feed/tushare_intraday.py b/data/feed/tushare_intraday.py index 1be0987..ae106e5 100644 --- a/data/feed/tushare_intraday.py +++ b/data/feed/tushare_intraday.py @@ -57,6 +57,7 @@ def __init__( rate_limit: int | None = None, max_retries: int = 6, cache=None, + scheduler=None, ) -> None: self._secret_file = str(secret_file) self._token_key = token_key @@ -71,6 +72,8 @@ def __init__( # caller injects one. The cache stores RAW 1min bars only; normalization # still runs downstream, unchanged. self._cache = cache + # D5: optional shared GlobalRateLimiter; None keeps the per-call throttle. + self._scheduler = scheduler # -- secret handling / client ------------------------------------------ # def _client(self): @@ -89,6 +92,7 @@ def _call(self, fn, **kwargs): fn, max_retries=self._max_retries, rate_limit=self._rate_limit, + scheduler=self._scheduler, **kwargs, ) diff --git a/qt/config.py b/qt/config.py index 93e2f13..fe58360 100644 --- a/qt/config.py +++ b/qt/config.py @@ -653,6 +653,30 @@ def _check_report_name(cls, v: str) -> str: return name +class DataUpdateConcurrencyCfg(_Strict): + """D5 opt-in bounded concurrency for ``data-update`` cache warms (default serial). + + ``max_workers=1`` (the default) keeps the warm fully serial — byte/behavior + compatible with every existing config. ``max_workers>1`` fans the per-symbol gap + FETCH stage onto a bounded thread pool, all funneled through ONE shared global + rate limiter (reusing ``data_update.rate_limit_per_min``) so the Tushare quota is + never multiplied per thread; store upserts and ledger writes still happen in + deterministic main-thread order. + """ + + max_workers: int = 1 + + @field_validator("max_workers") + @classmethod + def _check_max_workers(cls, v: int) -> int: + if v < 1: + raise ValueError( + f"data_update.concurrency.max_workers must be >= 1 (1 == serial); " + f"got {v}." + ) + return v + + class DataUpdateCfg(_Strict): """Standalone data-updater section (P4-3) — consumed ONLY by ``data-update``. @@ -676,6 +700,10 @@ class DataUpdateCfg(_Strict): force_refresh: list[str] = Field(default_factory=list) # D3b report-only quality hook (default OFF; see DataUpdateQualityCfg). quality: DataUpdateQualityCfg = Field(default_factory=DataUpdateQualityCfg) + # D5 opt-in bounded concurrency (default max_workers=1 == serial). + concurrency: DataUpdateConcurrencyCfg = Field( + default_factory=DataUpdateConcurrencyCfg + ) @field_validator("endpoints", "force_refresh") @classmethod diff --git a/qt/data_updater.py b/qt/data_updater.py index f7bd878..d816b62 100644 --- a/qt/data_updater.py +++ b/qt/data_updater.py @@ -24,6 +24,7 @@ import pandas as pd +from data.feed.scheduler import GlobalRateLimiter from qt.config import RootConfig, load_config from qt.data_update_quality import collect_findings, write_quality_report @@ -59,6 +60,9 @@ class UpdateResult: quality_report_path: Path | None = None quality_findings_count: int = 0 quality_hard_count: int = 0 + # D5 bounded concurrency (max_workers=1 == serial; rate_limit is the global cap). + max_workers: int = 1 + rate_limit_per_min: int = 0 def update_endpoints( @@ -148,8 +152,15 @@ def _resolve_symbols(cfg: RootConfig, feeds: UpdateFeeds, start: str, end: str) return sorted(syms) -def _build_feeds(cfg: RootConfig, cache, intraday_cache, rate_limit: int) -> UpdateFeeds: - """Construct the real tushare feeds, all sharing the read-through caches.""" +def _build_feeds( + cfg: RootConfig, cache, intraday_cache, rate_limit: int, scheduler=None +) -> UpdateFeeds: + """Construct the real tushare feeds, all sharing the read-through caches. + + ``scheduler`` (D5) is the ONE shared global rate limiter handed to every feed + when bounded concurrency is on; ``None`` (serial mode) keeps the historical + per-call throttle on each feed, byte-identical to before. + """ from data.feed.index_feed import IndexConstituentsFeed from data.feed.tushare_covariates import TushareCovariatesFeed from data.feed.tushare_feed import TushareFeed @@ -160,17 +171,28 @@ def _build_feeds(cfg: RootConfig, cache, intraday_cache, rate_limit: int) -> Upd secret = cfg.data.external_secret_file key = cfg.data.tushare_token_key return UpdateFeeds( - market=TushareFeed(secret, token_key=key, rate_limit=rate_limit, cache=cache), - index=IndexConstituentsFeed(secret, token_key=key, cache=cache), - flags=TushareFlagsFeed(secret, token_key=key, rate_limit=rate_limit, cache=cache), + market=TushareFeed( + secret, token_key=key, rate_limit=rate_limit, cache=cache, + scheduler=scheduler, + ), + index=IndexConstituentsFeed( + secret, token_key=key, cache=cache, scheduler=scheduler + ), + flags=TushareFlagsFeed( + secret, token_key=key, rate_limit=rate_limit, cache=cache, + scheduler=scheduler, + ), covariates=TushareCovariatesFeed( - secret, token_key=key, rate_limit=rate_limit, cache=cache + secret, token_key=key, rate_limit=rate_limit, cache=cache, + scheduler=scheduler, ), fina=TushareFinancialFeed( - secret, token_key=key, rate_limit=rate_limit, cache=cache + secret, token_key=key, rate_limit=rate_limit, cache=cache, + scheduler=scheduler, ), intraday=TushareIntradayFeed( - secret, token_key=key, rate_limit=rate_limit, cache=intraday_cache + secret, token_key=key, rate_limit=rate_limit, cache=intraday_cache, + scheduler=scheduler, ), ) @@ -209,6 +231,15 @@ def run_data_update(config_path: str, *, today=None) -> UpdateResult: TushareIntradayCache, ) + # D5: bounded concurrency is opt-in. max_workers==1 (default) builds NO + # scheduler — every feed keeps its per-call throttle and the cache stays serial, + # byte-identical to before. >1 builds ONE global rate limiter shared by all + # feeds (so the quota is global, not per-thread) and a multi-worker cache. + max_workers = du.concurrency.max_workers + scheduler = ( + GlobalRateLimiter(du.rate_limit_per_min) if max_workers > 1 else None + ) + root = cfg.data.cache.root_dir cache = TushareCache( CacheParquetStore(root), @@ -219,11 +250,14 @@ def run_data_update(config_path: str, *, today=None) -> UpdateResult: today=today_ts, not_ready_days=du.not_ready_days, recent_tail_overrides={"fina_indicator": du.fina_tail_days}, + max_workers=max_workers, ) intraday_cache = TushareIntradayCache( IntradayParquetStore(root), IntradayCoverageLedger(root) ) - feeds = _build_feeds(cfg, cache, intraday_cache, du.rate_limit_per_min) + feeds = _build_feeds( + cfg, cache, intraday_cache, du.rate_limit_per_min, scheduler=scheduler + ) symbols = _resolve_symbols(cfg, feeds, start, end) intraday_window = ( @@ -253,6 +287,8 @@ def run_data_update(config_path: str, *, today=None) -> UpdateResult: quality_report_path=quality.report_path if quality is not None else None, quality_findings_count=quality.findings_count if quality is not None else 0, quality_hard_count=quality.hard_count if quality is not None else 0, + max_workers=max_workers, + rate_limit_per_min=du.rate_limit_per_min, ) @@ -284,9 +320,11 @@ def _maybe_run_quality(cfg: RootConfig, du, capture, symbols, start, end): def format_summary(result: UpdateResult) -> str: """One human line per endpoint: requests / rows_written / not_ready.""" + mode = "serial" if result.max_workers <= 1 else f"{result.max_workers} workers" lines = [ f"data-update window [{result.window_start.date()} .. " - f"{result.window_end.date()}], {len(result.symbols)} symbols" + f"{result.window_end.date()}], {len(result.symbols)} symbols " + f"({mode}, global rate_limit_per_min={result.rate_limit_per_min})" ] for ep in result.endpoints: s = result.summary.get(ep, {}) diff --git a/tests/test_data_update_concurrency.py b/tests/test_data_update_concurrency.py new file mode 100644 index 0000000..e84ee0b --- /dev/null +++ b/tests/test_data_update_concurrency.py @@ -0,0 +1,173 @@ +"""D5 bounded-concurrency cache warm: config + serial==concurrent + failure model. + +Network-free, synthetic fetch callables. Proves the opt-in concurrency config, +that ``max_workers>1`` produces the SAME frames / ledger rows / request counts as +serial, and that the failure / empty / not_ready semantics are unchanged +(successes durable, failures uncovered+retryable). +""" + +from __future__ import annotations + +import pandas as pd +import pytest +from pydantic import ValidationError + +from data.cache.coverage import CoverageLedger +from data.cache.parquet_store import CacheParquetStore +from data.cache.tushare_cache import TushareCache +from qt.config import DataUpdateCfg, DataUpdateConcurrencyCfg, load_config + +_CLK = lambda: pd.Timestamp("2026-06-13 21:00:00") # noqa: E731 + + +# --------------------------------------------------------------------------- # +# config +# --------------------------------------------------------------------------- # +def test_default_config_concurrency_is_serial(): + cfg = load_config("config/data_update.yaml") + assert cfg.data_update is not None + assert cfg.data_update.concurrency.max_workers == 1 + + +def test_data_update_cfg_without_concurrency_block_defaults_serial(): + assert DataUpdateCfg().concurrency.max_workers == 1 + + +@pytest.mark.parametrize("bad", [0, -1, -10]) +def test_max_workers_below_one_rejected(bad): + with pytest.raises(ValidationError): + DataUpdateConcurrencyCfg(max_workers=bad) + + +def test_max_workers_accepts_positive(): + assert DataUpdateConcurrencyCfg(max_workers=8).max_workers == 8 + + +# --------------------------------------------------------------------------- # +# synthetic dense fetch (daily_basic-shaped) +# --------------------------------------------------------------------------- # +_CAT = { + "000001.SZ": [("2024-01-02", 10.0, 1.0, 1000.0), ("2024-01-03", 11.0, 1.1, 1100.0)], + "000002.SZ": [("2024-01-02", 20.0, 2.0, 2000.0), ("2024-01-03", 21.0, 2.1, 2100.0)], + "000003.SZ": [("2024-01-02", 30.0, 3.0, 3000.0), ("2024-01-03", 31.0, 3.1, 3100.0)], + "000004.SZ": [("2024-01-02", 40.0, 4.0, 4000.0), ("2024-01-03", 41.0, 4.1, 4100.0)], +} +_SYMS = ["000001.SZ", "000002.SZ", "000003.SZ", "000004.SZ"] + + +class _DBFetch: + def __init__(self, catalog=_CAT, fail_on=None): + self.catalog = catalog + self.fail_on = fail_on or set() + self.calls: list[str] = [] + + def __call__(self, symbol, s, e): + self.calls.append(symbol) + if symbol in self.fail_on: + raise ConnectionError("transient boom") + S, E = pd.Timestamp(s), pd.Timestamp(e) + rows = [r for r in self.catalog.get(symbol, []) if S <= pd.Timestamp(r[0]) <= E] + if not rows: + return pd.DataFrame() + return pd.DataFrame({ + "ts_code": [symbol] * len(rows), + "trade_date": [r[0].replace("-", "") for r in rows], + "pe": [r[1] for r in rows], "pb": [r[2] for r in rows], + "total_mv": [r[3] for r in rows], + }) + + +def _cache(tmp_path, **kw): + root = str(tmp_path / "cache") + return TushareCache(CacheParquetStore(root), CoverageLedger(root), clock=_CLK, **kw) + + +def _sorted(df): + return df.sort_values(list(df.columns)).reset_index(drop=True) + + +def _ledger_sorted(led): + cols = [c for c in led.columns if c != "fetched_at"] + return led.sort_values(["key", "start_date"]).reset_index(drop=True)[cols] + + +# --------------------------------------------------------------------------- # +# serial == concurrent equivalence +# --------------------------------------------------------------------------- # +def test_concurrent_equals_serial_frames_and_ledger(tmp_path): + serial = _cache(tmp_path / "s", refresh_recent_days=0, max_workers=1) + concurrent = _cache(tmp_path / "c", refresh_recent_days=0, max_workers=4) + a = serial.daily_basic(_SYMS, "2024-01-02", "2024-01-03", _DBFetch()) + b = concurrent.daily_basic(_SYMS, "2024-01-02", "2024-01-03", _DBFetch()) + # identical output frames + pd.testing.assert_frame_equal(_sorted(a), _sorted(b)) + # identical ledger contents (fetched_at fixed by _CLK; order-independent) + pd.testing.assert_frame_equal( + _ledger_sorted(serial._ledger.read()), + _ledger_sorted(concurrent._ledger.read()), + ) + # identical request counts + assert serial.stats() == concurrent.stats() + + +def test_concurrent_warm_then_cold_warm_is_zero_calls(tmp_path): + cache = _cache(tmp_path, refresh_recent_days=0, max_workers=4) + f1 = _DBFetch() + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-03", f1) + assert sorted(f1.calls) == _SYMS # cold: each symbol fetched once + f2 = _DBFetch() + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-03", f2) + assert f2.calls == [] # warm: fully covered -> zero fetches + + +# --------------------------------------------------------------------------- # +# failure semantics: successes durable, failure uncovered + retryable +# --------------------------------------------------------------------------- # +def test_concurrent_failure_keeps_successes_durable_and_retries(tmp_path): + cache = _cache(tmp_path, refresh_recent_days=0, max_workers=4) + f = _DBFetch(fail_on={"000002.SZ"}) + with pytest.raises(ConnectionError): + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-03", f) + led = cache._ledger.read() + covered = set(led.loc[led["status"].isin(["ok", "empty"]), "key"].astype(str)) + # the three good symbols are durable; the failed one is NOT recorded + assert {"000001.SZ", "000003.SZ", "000004.SZ"} <= covered + assert "000002.SZ" not in covered + # a later run (now healthy) refetches ONLY the previously-failed symbol + f2 = _DBFetch() + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-03", f2) + assert f2.calls == ["000002.SZ"] + + +def test_concurrent_empty_return_counts_as_coverage(tmp_path): + cache = _cache(tmp_path, refresh_recent_days=0, max_workers=4) + # 000004 has no rows in this catalog -> empty return, still coverage + cat = {k: v for k, v in _CAT.items() if k != "000004.SZ"} + f = _DBFetch(catalog=cat) + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-03", f) + f2 = _DBFetch(catalog=cat) + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-03", f2) + assert f2.calls == [] # empty return recorded -> no needless refetch + + +def test_concurrent_not_ready_still_pending(tmp_path): + # today=2024-01-04, not_ready_days=1: 01-04 unpublished -> not_ready (not covered) + today = pd.Timestamp("2024-01-04") + cache = _cache( + tmp_path, refresh_recent_days=0, not_ready_days=1, today=today, max_workers=4 + ) + f = _DBFetch() # catalog has no 01-04 rows + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-04", f) + led = cache._ledger.read() + assert (led["status"] == "not_ready").any() + f2 = _DBFetch() + cache.daily_basic(_SYMS, "2024-01-02", "2024-01-04", f2) + assert f2.calls # the pending (not_ready) day is still a gap -> retried + + +def test_single_symbol_concurrent_uses_serial_path(tmp_path): + # one symbol -> len(symbols) == 1 -> serial loop even with max_workers>1 + cache = _cache(tmp_path, refresh_recent_days=0, max_workers=4) + f = _DBFetch() + out = cache.daily_basic(["000001.SZ"], "2024-01-02", "2024-01-03", f) + assert len(out) == 2 and f.calls == ["000001.SZ"] diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 0000000..df26505 --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,173 @@ +"""D5 global rate limiter + throttle scheduler wiring (network-free, fake clock). + +Proves the limiter spaces calls GLOBALLY (not per-thread), that retries acquire a +slot per attempt, and that a scheduler-mode failure stays secret-safe. No +wall-clock timing — ``monotonic`` / ``sleep`` are injected. +""" + +from __future__ import annotations + +import threading + +import pytest + +from data.feed import throttle +from data.feed.scheduler import GlobalRateLimiter +from data.feed.throttle import request_with_retry + + +class _FrozenClock: + """A monotonic frozen at 0.0 with a recording (non-advancing) sleep. + + Frozen time models all callers arriving 'at once' — so each reserved slot is + one interval further out and the recorded waits escalate, which only happens + if the budget is GLOBAL (a per-thread limiter would record a constant wait). + """ + + def __init__(self): + self.sleeps: list[float] = [] + + def monotonic(self): + return 0.0 + + def sleep(self, d): + self.sleeps.append(d) + + +def test_global_spacing_escalates_under_contention(): + clk = _FrozenClock() + rl = GlobalRateLimiter(120, monotonic=clk.monotonic, sleep=clk.sleep) # 0.5s + for _ in range(4): + rl.acquire() + # first goes immediately (no sleep); each later caller waits one more interval + assert clk.sleeps == [0.5, 1.0, 1.5] + + +def test_disabled_when_rate_non_positive(): + clk = _FrozenClock() + rl = GlobalRateLimiter(0, monotonic=clk.monotonic, sleep=clk.sleep) + for _ in range(5): + rl.acquire() + assert clk.sleeps == [] + assert rl.rate_limit_per_min == 0 + + +def test_advancing_clock_spaces_each_call(): + # a clock that advances by the slept amount -> sequential callers each wait one + # interval (the steady-state global cadence). + state = {"t": 0.0} + sleeps: list[float] = [] + + def mono(): + return state["t"] + + def slp(d): + sleeps.append(d) + state["t"] += d + + rl = GlobalRateLimiter(60, monotonic=mono, sleep=slp) # 1.0s + for _ in range(3): + rl.acquire() + assert sleeps == [1.0, 1.0] # first immediate, then 1s spacing each + + +def test_global_not_per_thread_under_real_threads(): + # Frozen clock + no-op-recording sleep; many threads each acquire once. The + # limiter must reserve exactly N global slots (next_allowed == N*interval), + # regardless of interleaving — proving the budget is shared, not per-thread. + clk = _FrozenClock() + lock = threading.Lock() + + def safe_sleep(d): + with lock: + clk.sleeps.append(d) + + rl = GlobalRateLimiter(120, monotonic=clk.monotonic, sleep=safe_sleep) # 0.5s + n = 8 + barrier = threading.Barrier(n) + + def worker(): + barrier.wait() + rl.acquire() + + threads = [threading.Thread(target=worker) for _ in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + # N reservations consumed N global slots (not N per-thread) + assert rl._next_allowed == pytest.approx(n * 0.5) + # cumulative wait is order-independent: 0 + 0.5 + 1.0 + ... + (n-1)*0.5 + assert sum(clk.sleeps) == pytest.approx(0.5 * n * (n - 1) / 2) + + +class _SpyScheduler: + def __init__(self): + self.acquires = 0 + + def acquire(self): + self.acquires += 1 + + +def test_retry_acquires_global_slot_per_attempt(monkeypatch): + monkeypatch.setattr(throttle.time, "sleep", lambda *_a: None) # no backoff wait + sched = _SpyScheduler() + calls = {"n": 0} + + def flaky(**_kw): + calls["n"] += 1 + if calls["n"] < 3: + raise RuntimeError("transient") + return "ok" + + assert request_with_retry(flaky, max_retries=5, scheduler=sched) == "ok" + assert calls["n"] == 3 + assert sched.acquires == 3 # one global slot per attempt, including retries + + +def test_scheduler_mode_skips_per_call_rate_sleep(monkeypatch): + # With a scheduler, the historical per-call ``rate_limit`` sleep must NOT also + # fire (that would double the spacing / multiply the quota). + recorded: list[float] = [] + monkeypatch.setattr(throttle.time, "sleep", recorded.append) + sched = _SpyScheduler() + request_with_retry(lambda **_kw: "ok", rate_limit=120, scheduler=sched) + assert recorded == [] # no 0.5s rate-limit sleep + assert sched.acquires == 1 + + +def test_scheduler_failure_is_secret_safe(monkeypatch): + monkeypatch.setattr(throttle.time, "sleep", lambda *_a: None) + sched = _SpyScheduler() + + def always_fail(**_kw): + raise ValueError("super-secret-token-leak") + + with pytest.raises(RuntimeError) as ei: + request_with_retry(always_fail, max_retries=2, scheduler=sched) + msg = str(ei.value) + assert "ValueError" in msg # exception TYPE only + assert "super-secret-token-leak" not in msg # never the payload + assert sched.acquires == 2 + + +def test_feed_call_forwards_its_scheduler(): + # the feed plumbing actually hands its scheduler to request_with_retry: a + # ``_call`` through a feed built with a spy scheduler acquires one slot. + from data.feed.tushare_feed import TushareFeed + + sched = _SpyScheduler() + feed = TushareFeed("x.json", rate_limit=120, scheduler=sched) + assert feed._call(lambda **_kw: "ok") == "ok" + assert sched.acquires == 1 + + +def test_feed_without_scheduler_is_unchanged(monkeypatch): + # default (no scheduler) keeps the per-call rate-limit sleep, byte-identical. + from data.feed.tushare_feed import TushareFeed + + recorded: list[float] = [] + monkeypatch.setattr(throttle.time, "sleep", recorded.append) + feed = TushareFeed("x.json", rate_limit=60) # 1.0s spacing, no scheduler + feed._call(lambda **_kw: "ok") + assert recorded == [1.0]