123 lines
5.1 KiB
Python
123 lines
5.1 KiB
Python
"""
|
|
TensorFlow/Keras 기반 주가 예측 서비스
|
|
|
|
- LSTM 모델로 단기 주가 방향(상승/하락) 예측
|
|
- KIS API 차트 데이터에서 피처 추출 → 정규화 → 모델 추론
|
|
- 모델 파일 없으면 더미 예측(0.5) 반환 (서비스 무중단)
|
|
- ML_TF_ENABLED=false 시 비활성화 가능
|
|
"""
|
|
import asyncio
|
|
import functools
|
|
import logging
|
|
from typing import List, Dict, Any, Optional
|
|
|
|
import numpy as np
|
|
|
|
from app.services.korea_investment_api import korea_investment_api
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TensorFlowNotAvailable(Exception):
|
|
pass
|
|
|
|
|
|
class TensorFlowPredictor:
|
|
"""Lightweight wrapper around a Keras model + scaler.
|
|
|
|
- Lazily attempts to import TensorFlow/Keras and load model files.
|
|
- If TensorFlow or model/scaler files are missing, raises TensorFlowNotAvailable
|
|
so callers can handle gracefully (e.g., return 503 or a friendly message).
|
|
- Fetches recent daily closes from KIS API.
|
|
|
|
이슈 7: model.predict()는 CPU-bound 동기 호출이므로 predict_stock() 내에서
|
|
asyncio.to_thread()로 워커 스레드에 위임해 이벤트 루프 블로킹을 방지한다.
|
|
"""
|
|
|
|
def __init__(self, model_path: str = "models/lstm_model.h5", scaler_path: str = "models/scaler.pkl", time_steps: int = 60):
|
|
self.model_path = model_path
|
|
self.scaler_path = scaler_path
|
|
self.time_steps = time_steps
|
|
|
|
try:
|
|
import tensorflow as tf # noqa: F401
|
|
from tensorflow import keras # noqa: F401
|
|
import joblib # noqa: F401
|
|
except Exception as e:
|
|
logger.warning(f"TensorFlow or joblib not available: {e}")
|
|
raise TensorFlowNotAvailable("TensorFlow/joblib not installed in backend environment")
|
|
|
|
try:
|
|
from tensorflow import keras
|
|
import joblib
|
|
|
|
self.model = keras.models.load_model(self.model_path)
|
|
self.scaler = joblib.load(self.scaler_path)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to load TF model/scaler: {e}")
|
|
raise TensorFlowNotAvailable(f"Model or scaler not found: {e}")
|
|
|
|
async def fetch_kis_closes(self, stock_code: str, days: int = 60) -> np.ndarray:
|
|
"""Fetch recent daily closes from KIS. Returns array oldest→latest."""
|
|
try:
|
|
resp = await korea_investment_api.get_daily_chart(stock_code, "D")
|
|
output = resp.get("output") or []
|
|
if not output:
|
|
raise RuntimeError("empty_chart")
|
|
closes_desc = [float(item.get("stck_clpr", 0) or 0) for item in output]
|
|
closes = list(reversed(closes_desc))
|
|
arr = np.array(closes[-days:], dtype=float)
|
|
if arr.size < days:
|
|
logger.info(f"{stock_code}: only {arr.size} closes available (requested {days})")
|
|
return arr
|
|
except Exception as e:
|
|
logger.error(f"Failed to fetch KIS closes for {stock_code}: {e}")
|
|
raise
|
|
|
|
async def predict_stock(self, stock_code: str) -> Dict[str, Any]:
|
|
"""Return prediction dict with expected_return and simple signal.
|
|
|
|
model.predict()를 asyncio.to_thread()로 워커 스레드에서 실행해
|
|
Keras/TF의 GIL-holding 연산이 이벤트 루프를 막지 않게 한다.
|
|
"""
|
|
try:
|
|
prices = await self.fetch_kis_closes(stock_code, days=self.time_steps)
|
|
if prices.size == 0:
|
|
raise RuntimeError("no_price_data")
|
|
|
|
scaled = self.scaler.transform(prices.reshape(-1, 1))
|
|
if scaled.shape[0] < self.time_steps:
|
|
pad = np.repeat(scaled[0:1], self.time_steps - scaled.shape[0], axis=0)
|
|
scaled = np.vstack([pad, scaled])
|
|
X_input = scaled[-self.time_steps:].reshape(1, self.time_steps, 1)
|
|
|
|
# ── 이슈 7: CPU-bound model.predict → 워커 스레드로 분리 ──────
|
|
predict_fn = functools.partial(self.model.predict, X_input, verbose=0)
|
|
predicted_scaled = await asyncio.to_thread(predict_fn)
|
|
# ─────────────────────────────────────────────────────────────
|
|
|
|
predicted_price = float(self.scaler.inverse_transform(predicted_scaled)[0][0])
|
|
current_price = float(prices[-1])
|
|
|
|
expected_return = (
|
|
(predicted_price - current_price) / current_price * 100.0
|
|
if current_price > 0 else 0.0
|
|
)
|
|
signal = "BUY" if predicted_price > current_price * 1.02 else "HOLD"
|
|
|
|
return {
|
|
"stock_code": stock_code,
|
|
"current_price": current_price,
|
|
"predicted_price": predicted_price,
|
|
"expected_return": expected_return,
|
|
"signal": signal,
|
|
}
|
|
except TensorFlowNotAvailable:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Prediction failed for {stock_code}: {e}")
|
|
return {
|
|
"stock_code": stock_code,
|
|
"error": str(e),
|
|
}
|