413 lines
15 KiB
Python
413 lines
15 KiB
Python
"""
|
||
kis_trader/engine/indicator_cache.py — RSI/EMA materialized 캐시 (파라서치 가속)
|
||
================================================================================
|
||
- 종목별 IndicatorCache: RSI·EMA period 별 시리즈 1회 계산 → 조합 루프 재사용
|
||
- ws_candles DB 컬럼 materialize (RSI·EMA period 별 rsi_N / ema_N)
|
||
|
||
[ws_candles 컬럼]
|
||
- 실시간 WS 확정봉: rsi_2, rsi_3, rsi_5 (CandleAggregator 가 INSERT)
|
||
- materialize 배치: rsi_7, rsi_14, rsi_21 + ema_5,9,12,15,21,34 (env 로 period 확장)
|
||
|
||
파라서치 ④ 구간(조합×RSI/EMA 재계산) 가속용 — 실매 TRIGGER 로직 변경 없음.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||
|
||
from kis_trader.engine.ema_trend_filter import compute_ema_series
|
||
from kis_trader.utils.env import get_env_bool, get_env_float, get_env_from_db, get_env_int
|
||
|
||
|
||
def _parse_period_list(env_key: str, fallback: str) -> List[int]:
|
||
raw = str(get_env_from_db(env_key, fallback) or fallback).strip()
|
||
out: List[int] = []
|
||
for part in raw.split(","):
|
||
part = part.strip()
|
||
if not part:
|
||
continue
|
||
try:
|
||
p = int(float(part))
|
||
if p > 0:
|
||
out.append(p)
|
||
except (TypeError, ValueError):
|
||
continue
|
||
return sorted(set(out)) if out else [int(x) for x in fallback.split(",") if x.strip()]
|
||
|
||
|
||
def default_rsi_periods() -> List[int]:
|
||
return _parse_period_list("PARAM_SEARCH_RSI_PERIODS", "2,3,5,7,14,21")
|
||
|
||
|
||
def default_ema_periods() -> List[int]:
|
||
return _parse_period_list("PARAM_SEARCH_EMA_PERIODS", "5,9,12,15,21,34")
|
||
|
||
|
||
def ws_candles_materialize_rsi_periods() -> List[int]:
|
||
"""DB materialize 대상 RSI period — env ``WS_CANDLE_MATERIALIZE_RSI_PERIODS`` 우선."""
|
||
raw = _parse_period_list(
|
||
"WS_CANDLE_MATERIALIZE_RSI_PERIODS",
|
||
",".join(str(p) for p in default_rsi_periods() if p not in (2, 3, 5)),
|
||
)
|
||
# rsi_2/3/5 는 CandleAggregator 가 실시간 기록 — 배치는 7+ 위주
|
||
return sorted(set(raw))
|
||
|
||
|
||
def ws_candles_materialize_ema_periods() -> List[int]:
|
||
"""DB materialize 대상 EMA period."""
|
||
return _parse_period_list(
|
||
"WS_CANDLE_MATERIALIZE_EMA_PERIODS",
|
||
",".join(str(p) for p in default_ema_periods()),
|
||
)
|
||
|
||
|
||
def ws_candles_all_rsi_column_names() -> List[str]:
|
||
"""ws_candles SELECT/마이그레이션용 RSI 컬럼명 (rsi_2 … env period)."""
|
||
periods = sorted(set(default_rsi_periods()) | {2, 3, 5})
|
||
return [f"rsi_{p}" for p in periods]
|
||
|
||
|
||
def ws_candles_all_ema_column_names() -> List[str]:
|
||
periods = ws_candles_materialize_ema_periods()
|
||
return [f"ema_{p}" for p in periods]
|
||
|
||
|
||
def ws_candles_indicator_column_migrations() -> List[Tuple[str, str]]:
|
||
"""(컬럼명, ALTER DDL suffix) — TradeDB._migrate_add_columns 에서 사용.
|
||
|
||
※ TradeDB 기동 중 호출되므로 get_env_from_db(DB) 를 쓰지 않는다.
|
||
(env 조회 → TradeDB() 재진입 → _tables_lock 데드락 → journal 로그 0줄 hang)
|
||
"""
|
||
out: List[Tuple[str, str]] = []
|
||
# env 기본값과 동일한 정적 period — 런타임 확장 period 는 ensure 시 information_schema 로 ADD
|
||
migration_rsi = (7, 14, 21)
|
||
migration_ema = (5, 9, 12, 15, 21, 34)
|
||
for period in migration_rsi:
|
||
out.append((
|
||
f"rsi_{period}",
|
||
f"DOUBLE NULL COMMENT 'RSI({period}) materialized'",
|
||
))
|
||
for period in migration_ema:
|
||
out.append((
|
||
f"ema_{period}",
|
||
f"DOUBLE NULL COMMENT 'EMA({period}) materialized'",
|
||
))
|
||
return out
|
||
|
||
|
||
def ensure_ws_candles_indicator_columns(db) -> int:
|
||
"""ws_candles 에 누락된 rsi_N / ema_N 컬럼 ADD (1회 마이그레이션). Returns 추가된 컬럼 수."""
|
||
added = 0
|
||
try:
|
||
cols = set(db.conn.get_columns("ws_candles"))
|
||
except Exception:
|
||
return 0
|
||
for col, ddl in ws_candles_indicator_column_migrations():
|
||
if col in cols:
|
||
continue
|
||
try:
|
||
db.conn.execute(f"ALTER TABLE ws_candles ADD COLUMN {col} {ddl}")
|
||
cols.add(col)
|
||
added += 1
|
||
except Exception:
|
||
pass
|
||
return added
|
||
|
||
|
||
def indicator_cache_enabled() -> bool:
|
||
return get_env_bool("PARAM_SEARCH_INDICATOR_CACHE_ENABLED", True)
|
||
|
||
|
||
def materialize_db_on_load() -> bool:
|
||
return get_env_bool("WS_CANDLE_MATERIALIZE_ON_LOAD", False)
|
||
|
||
|
||
class IndicatorCache:
|
||
"""단일 종목 캔들에 대한 RSI/EMA 시리즈 메모이제이션."""
|
||
|
||
# _rsi_final / _ema_final: 해당 period 시리즈를 (DB seed + 1회 계산 병합 후)
|
||
# '확정'했음을 표시. 확정된 period 는 재계산 없이 캐시를 그대로 반환한다.
|
||
# (기존엔 rsi_at/ema_at 호출마다 compute_*_series 를 다시 돌려 O(N²) 였음 — 결과는
|
||
# 동일했으나 파라서치 단일 백테가 봉당 전체 재계산으로 수십 초 소요되던 병목)
|
||
__slots__ = ("closes", "_rsi", "_ema", "_rsi_final", "_ema_final")
|
||
|
||
def __init__(self, candles: List[Dict[str, Any]]) -> None:
|
||
self.closes: List[float] = [float(c.get("close") or 0) for c in candles]
|
||
self._rsi: Dict[int, List[Optional[float]]] = {}
|
||
self._ema: Dict[int, List[Optional[float]]] = {}
|
||
self._rsi_final: set = set()
|
||
self._ema_final: set = set()
|
||
self._seed_from_candle_fields(candles)
|
||
|
||
def _seed_from_candle_fields(self, candles: List[Dict[str, Any]]) -> None:
|
||
"""DB materialized 컬럼(rsi_14, ema_9 등)이 있으면 해당 period 시리즈에 반영."""
|
||
for i, c in enumerate(candles):
|
||
for period in default_rsi_periods():
|
||
key = f"rsi_{period}"
|
||
if key in c and c[key] is not None:
|
||
ser = self._rsi.setdefault(period, [None] * len(candles))
|
||
if i < len(ser):
|
||
try:
|
||
ser[i] = float(c[key])
|
||
except (TypeError, ValueError):
|
||
pass
|
||
legacy = {2: "rsi_2", 3: "rsi_3", 5: "rsi_5"}.get(period)
|
||
if legacy and legacy in c and c[legacy] is not None:
|
||
ser = self._rsi.setdefault(period, [None] * len(candles))
|
||
if i < len(ser) and ser[i] is None:
|
||
try:
|
||
ser[i] = float(c[legacy])
|
||
except (TypeError, ValueError):
|
||
pass
|
||
for period in default_ema_periods():
|
||
key = f"ema_{period}"
|
||
if key in c and c[key] is not None:
|
||
ser = self._ema.setdefault(period, [None] * len(candles))
|
||
if i < len(ser):
|
||
try:
|
||
ser[i] = float(c[key])
|
||
except (TypeError, ValueError):
|
||
pass
|
||
|
||
def rsi_series(self, period: int) -> List[Optional[float]]:
|
||
p = max(1, int(period))
|
||
# 이미 확정된 period 면 재계산 없이 캐시 반환 (핫패스 가속, 결과 불변)
|
||
if p in self._rsi_final:
|
||
return self._rsi[p]
|
||
from kis_trader.engine.momentum_engine import compute_rsi_series
|
||
computed = compute_rsi_series(self.closes, p)
|
||
seeded = self._rsi.get(p)
|
||
if seeded is None:
|
||
self._rsi[p] = computed
|
||
self._rsi_final.add(p)
|
||
return computed
|
||
if len(seeded) != len(computed):
|
||
seeded = (list(seeded) + [None] * len(computed))[: len(computed)]
|
||
# warm-up prepend 등으로 DB rsi_N 이 앞구간만 있으면 뒤 구간은 계산값으로 보완
|
||
merged = [
|
||
s if s is not None else b
|
||
for s, b in zip(seeded, computed)
|
||
]
|
||
self._rsi[p] = merged
|
||
self._rsi_final.add(p)
|
||
return merged
|
||
|
||
def rsi_at(self, i: int, period: int) -> Optional[float]:
|
||
ser = self.rsi_series(period)
|
||
if 0 <= i < len(ser):
|
||
return ser[i]
|
||
return None
|
||
|
||
def ema_series(self, period: int) -> List[Optional[float]]:
|
||
p = max(1, int(period))
|
||
# 이미 확정된 period 면 재계산 없이 캐시 반환 (핫패스 가속, 결과 불변)
|
||
if p in self._ema_final:
|
||
return self._ema[p]
|
||
computed = compute_ema_series(self.closes, p)
|
||
seeded = self._ema.get(p)
|
||
if seeded is None:
|
||
self._ema[p] = computed
|
||
self._ema_final.add(p)
|
||
return computed
|
||
if len(seeded) != len(computed):
|
||
seeded = (list(seeded) + [None] * len(computed))[: len(computed)]
|
||
merged = [
|
||
s if s is not None else b
|
||
for s, b in zip(seeded, computed)
|
||
]
|
||
self._ema[p] = merged
|
||
self._ema_final.add(p)
|
||
return merged
|
||
|
||
def ema_at(self, i: int, period: int) -> Optional[float]:
|
||
ser = self.ema_series(period)
|
||
if 0 <= i < len(ser):
|
||
return ser[i]
|
||
return None
|
||
|
||
|
||
def build_indicator_cache(candles: List[Dict[str, Any]]) -> IndicatorCache:
|
||
return IndicatorCache(candles)
|
||
|
||
|
||
def build_indicator_cache_by_code(
|
||
codes_candles: Dict[str, List[Dict[str, Any]]],
|
||
*,
|
||
rsi_periods: Optional[Sequence[int]] = None,
|
||
ema_periods: Optional[Sequence[int]] = None,
|
||
) -> Dict[str, IndicatorCache]:
|
||
"""종목별 캐시 — rsi/ema period 는 warm-up 용 (첫 접근 시 lazy 계산)."""
|
||
if not indicator_cache_enabled():
|
||
return {}
|
||
out: Dict[str, IndicatorCache] = {}
|
||
rsi_p = list(rsi_periods or default_rsi_periods())
|
||
ema_p = list(ema_periods or default_ema_periods())
|
||
for code, rows in (codes_candles or {}).items():
|
||
if not rows:
|
||
continue
|
||
ic = IndicatorCache([dict(r) for r in rows])
|
||
for p in rsi_p:
|
||
ic.rsi_series(p)
|
||
for p in ema_p:
|
||
ic.ema_series(p)
|
||
out[str(code)] = ic
|
||
return out
|
||
|
||
|
||
def attach_indicator_caches_to_params(
|
||
params: Dict[str, Any],
|
||
codes_candles: Dict[str, List[Dict[str, Any]]],
|
||
) -> None:
|
||
"""params['_indicator_cache_by_code'] 에 종목별 캐시 주입 (in-place)."""
|
||
if not indicator_cache_enabled():
|
||
return
|
||
existing = params.get("_indicator_cache_by_code")
|
||
if isinstance(existing, dict) and existing:
|
||
return
|
||
params["_indicator_cache_by_code"] = build_indicator_cache_by_code(codes_candles)
|
||
|
||
|
||
def get_indicator_cache_from_params(
|
||
params: Dict[str, Any],
|
||
code: str,
|
||
) -> Optional[IndicatorCache]:
|
||
by_code = params.get("_indicator_cache_by_code")
|
||
if isinstance(by_code, dict):
|
||
return by_code.get(str(code))
|
||
ic = params.get("_indicator_cache")
|
||
return ic if isinstance(ic, IndicatorCache) else None
|
||
|
||
|
||
def enrich_candles_with_materialized_fields(
|
||
candles: List[Dict[str, Any]],
|
||
ic: IndicatorCache,
|
||
*,
|
||
rsi_periods: Optional[Sequence[int]] = None,
|
||
ema_periods: Optional[Sequence[int]] = None,
|
||
) -> None:
|
||
"""캔들 dict 에 rsi_N / ema_N 필드 부착 (DB 저장·재로드용)."""
|
||
rsi_p = list(rsi_periods or default_rsi_periods())
|
||
ema_p = list(ema_periods or default_ema_periods())
|
||
for i, c in enumerate(candles):
|
||
for p in rsi_p:
|
||
v = ic.rsi_at(i, p)
|
||
if v is not None:
|
||
c[f"rsi_{p}"] = round(v, 4)
|
||
for p in ema_p:
|
||
v = ic.ema_at(i, p)
|
||
if v is not None:
|
||
c[f"ema_{p}"] = round(v, 4)
|
||
|
||
|
||
def materialize_ws_candles_batch(
|
||
db,
|
||
candles_by_code: Dict[str, List[Dict[str, Any]]],
|
||
timeframe: int,
|
||
) -> int:
|
||
"""
|
||
ws_candles 행에 rsi_N·ema_N UPDATE (배치).
|
||
env ``WS_CANDLE_MATERIALIZE_ON_LOAD=true`` 일 때만 실행.
|
||
Returns: 업데이트 시도 행 수.
|
||
"""
|
||
if not materialize_db_on_load() or db is None:
|
||
return 0
|
||
ensure_ws_candles_indicator_columns(db)
|
||
try:
|
||
db_cols = set(db.conn.get_columns("ws_candles"))
|
||
except Exception:
|
||
return 0
|
||
|
||
rsi_p = [p for p in ws_candles_materialize_rsi_periods() if f"rsi_{p}" in db_cols]
|
||
ema_p = [p for p in ws_candles_materialize_ema_periods() if f"ema_{p}" in db_cols]
|
||
if not rsi_p and not ema_p:
|
||
return 0
|
||
|
||
set_parts = []
|
||
for p in rsi_p:
|
||
set_parts.append(f"rsi_{p} = ?")
|
||
for p in ema_p:
|
||
set_parts.append(f"ema_{p} = ?")
|
||
set_sql = ", ".join(set_parts)
|
||
if not set_sql:
|
||
return 0
|
||
|
||
batch: List[Tuple[Any, ...]] = []
|
||
for code, rows in (candles_by_code or {}).items():
|
||
if not rows:
|
||
continue
|
||
ic = build_indicator_cache([dict(r) for r in rows])
|
||
enrich_candles_with_materialized_fields(
|
||
rows, ic,
|
||
rsi_periods=rsi_p or None,
|
||
ema_periods=ema_p or None,
|
||
)
|
||
for c in rows:
|
||
ct = str(c.get("candle_time") or "")
|
||
if not ct:
|
||
continue
|
||
vals: List[Any] = []
|
||
for p in rsi_p:
|
||
vals.append(c.get(f"rsi_{p}"))
|
||
for p in ema_p:
|
||
vals.append(c.get(f"ema_{p}"))
|
||
vals.extend([str(code), int(timeframe), ct])
|
||
batch.append(tuple(vals))
|
||
|
||
if not batch:
|
||
return 0
|
||
|
||
where = "WHERE code = ? AND timeframe = ? AND candle_time = ? AND source = 'kis'"
|
||
sql_sqlite = f"""
|
||
UPDATE ws_candles
|
||
SET {set_sql}, updated_at = datetime('now','localtime')
|
||
{where}
|
||
"""
|
||
set_mysql = ", ".join(
|
||
[f"rsi_{p} = %s" for p in rsi_p]
|
||
+ [f"ema_{p} = %s" for p in ema_p]
|
||
)
|
||
sql_mysql = f"""
|
||
UPDATE ws_candles
|
||
SET {set_mysql}
|
||
WHERE code = %s AND timeframe = %s AND candle_time = %s AND source = 'kis'
|
||
"""
|
||
try:
|
||
with db.conn:
|
||
db.conn.executemany(sql_sqlite, batch)
|
||
except Exception:
|
||
try:
|
||
db.conn.executemany(sql_mysql, batch)
|
||
except Exception:
|
||
return 0
|
||
return len(batch)
|
||
|
||
|
||
def ws_candles_select_indicator_cols(db=None) -> str:
|
||
"""SELECT 절 추가 컬럼 — DB 에 실제 존재하는 rsi_N / ema_N 만."""
|
||
names: List[str] = []
|
||
if db is not None:
|
||
try:
|
||
cols = set(db.conn.get_columns("ws_candles"))
|
||
for n in ws_candles_all_rsi_column_names() + ws_candles_all_ema_column_names():
|
||
if n in cols:
|
||
names.append(n)
|
||
except Exception:
|
||
pass
|
||
if not names:
|
||
names = ws_candles_all_rsi_column_names() + ws_candles_all_ema_column_names()
|
||
return (", " + ", ".join(names)) if names else ""
|
||
|
||
|
||
def ws_candles_has_materialized_cols(db) -> bool:
|
||
"""materialize 컬럼(ema_9 또는 ema_34 등)이 1개 이상 있으면 True."""
|
||
try:
|
||
cols = set(db.conn.get_columns("ws_candles"))
|
||
for n in ws_candles_all_ema_column_names():
|
||
if n in cols:
|
||
return True
|
||
for n in ("rsi_7", "rsi_14", "rsi_21"):
|
||
if n in cols:
|
||
return True
|
||
return False
|
||
except Exception:
|
||
return False
|