kis_bot/kis_trader/backtest/optuna_common.py

#!/usr/bin/env python3
"""
kis_trader/backtest/optuna_common.py — Optuna storage·DB 공통 (MariaDB 141)
=========================================================================
TradeDB(database.py) 와 동일 호스트·계정, 전용 DB kis_optuna 에 study 저장.
Win11·VM 양쪽에서 같은 storage 로 trial 공유·재개 가능.
"""
from __future__ import annotations

import logging
import os
from typing import Optional
from urllib.parse import quote_plus

from kis_trader.utils.env import get_env_from_db

logger = logging.getLogger("optuna_common")

# Optuna 전용 MariaDB (매매 DB kis_quant_db 와 분리)
DEFAULT_OPTUNA_DB_NAME = "kis_optuna"
OPTUNA_STRATEGIES = ("tail", "momentum", "breakout")


def mariadb_creds() -> dict:
    """TradeDB(database.py) 와 동일 우선순위 — env > 기본 141."""
    return {
        "host": os.environ.get("DB_HOST", "192.168.0.141"),
        "port": int(os.environ.get("DB_PORT", "3306")),
        "user": os.environ.get("DB_USER", "jae"),
        "password": os.environ.get("DB_PASS", "1234"),
    }


def resolve_optuna_db_name() -> str:
    """
    Optuna storage DB — 기본 kis_optuna (매매 kis_quant_db 와 분리).
    env OPTUNA_DB_NAME 로 오버라이드 가능.
    """
    raw = get_env_from_db("OPTUNA_DB_NAME", "")
    if raw and str(raw).strip() not in ("", "None"):
        return str(raw).strip()
    env = os.environ.get("OPTUNA_DB_NAME", "")
    if env and str(env).strip():
        return str(env).strip()
    return DEFAULT_OPTUNA_DB_NAME


def build_mariadb_storage_url(db_name: Optional[str] = None) -> str:
    """mysql+pymysql://…@141/optuna 형식 storage URL."""
    creds = mariadb_creds()
    name = (db_name or resolve_optuna_db_name()).strip()
    user = quote_plus(creds["user"])
    passwd = quote_plus(creds["password"])
    return (
        f"mysql+pymysql://{user}:{passwd}@{creds['host']}:{creds['port']}/{name}"
        f"?charset=utf8mb4"
    )


def ensure_optuna_database(db_name: Optional[str] = None) -> str:
    """
    MariaDB 141 — kis_optuna 존재 확인 (없으면 CREATE 시도).
    """
    name = (db_name or resolve_optuna_db_name()).strip()
    creds = mariadb_creds()

    try:
        import pymysql
    except ImportError as exc:
        raise ImportError(
            "Optuna MariaDB storage 는 pymysql 필요: pip install PyMySQL"
        ) from exc

    # DB 존재 여부만 확인 (이미 있으면 CREATE 생략)
    conn = pymysql.connect(
        host=creds["host"],
        port=creds["port"],
        user=creds["user"],
        password=creds["password"],
        charset="utf8mb4",
        autocommit=True,
        connect_timeout=10,
    )
    try:
        with conn.cursor() as cur:
            cur.execute("SHOW DATABASES LIKE %s", (name,))
            exists = cur.fetchone() is not None
            if not exists:
                cur.execute(
                    f"CREATE DATABASE IF NOT EXISTS `{name}` "
                    "DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
                )
                logger.info(
                    "📦 Optuna DB 생성: %s@%s:%s/%s",
                    creds["user"], creds["host"], creds["port"], name,
                )
            else:
                logger.info(
                    "📦 Optuna storage DB: %s@%s:%s/%s",
                    creds["user"], creds["host"], creds["port"], name,
                )
    except Exception as exc:
        logger.error("❌ Optuna DB '%s' 접속/확인 실패: %s", name, exc)
        raise
    finally:
        conn.close()
    return name


def resolve_optuna_storage_url(cli_override: Optional[str] = None) -> str:
    """
    Storage URL 우선순위:
      1) CLI --storage
      2) OPTUNA_STORAGE_URL (DB/env)
      3) MariaDB 141 / kis_optuna (TradeDB 동일 계정)
    """
    if cli_override and str(cli_override).strip():
        return str(cli_override).strip()
    from_db = get_env_from_db("OPTUNA_STORAGE_URL", "")
    if from_db and str(from_db).strip() not in ("", "None"):
        return str(from_db).strip()
    db_name = ensure_optuna_database()
    return build_mariadb_storage_url(db_name)


def resolve_study_name(
    *,
    strategy: str,
    mode: str,
    start: str,
    end: str,
    cli_override: Optional[str] = None,
) -> str:
    """Study 이름 — 전략·기간·모드 포함."""
    if cli_override and str(cli_override).strip():
        return str(cli_override).strip()
    env_key = f"OPTUNA_{strategy.upper()}_STUDY_NAME"
    from_db = get_env_from_db(env_key, "")
    if from_db and str(from_db).strip() not in ("", "None"):
        return str(from_db).strip()
    legacy = get_env_from_db("OPTUNA_TAIL_STUDY_NAME", "")
    if strategy == "tail" and legacy and str(legacy).strip() not in ("", "None"):
        return str(legacy).strip()
    return f"{strategy}_{mode}_{start}_{end}"


def optuna_run_lock_name(strategy: str) -> str:
    return f"{strategy}_param_search_optuna"