ai-trade/main_server.py

import os
import signal
import threading
import uvicorn
import multiprocessing
from fastapi import FastAPI, Request
from pydantic import BaseModel
from contextlib import asynccontextmanager

from modules.config import Config
from modules.services.ollama import OllamaManager
from modules.services.kis import KISClient
from modules.services.news import NewsCollector
from modules.services.telegram import TelegramMessenger
from modules.bot import AutoTradingBot
from modules.utils.process_tracker import ProcessTracker, ProcessWatchdog
from modules.services.telegram_bot.runner import run_telegram_bot_standalone

# 전역 객체
messenger = TelegramMessenger()
ai_agent = None
kis_client = None
news_collector = None
watchdog = None


def run_trading_bot(ipc_lock, command_queue, shutdown_event, eod_event=None):
    """트레이딩 봇 실행 래퍼"""
    ProcessTracker.register("Trading Bot Main")
    bot = AutoTradingBot(ipc_lock=ipc_lock, command_queue=command_queue,
                         shutdown_event=shutdown_event, eod_event=eod_event)
    bot.loop()


@asynccontextmanager
async def lifespan(app: FastAPI):
    global ai_agent, kis_client, news_collector, watchdog

    # 1. 설정 검증
    Config.validate()

    # 2. 좀비 프로세스 정리
    try:
        ProcessTracker.check_and_kill_zombies()
        ProcessTracker.clear()
        ProcessTracker.register("Main Server (Uvicorn Worker)")
    except Exception:
        pass

    # 3. 전역 객체 초기화
    ai_agent = OllamaManager()
    kis_client = KISClient()
    news_collector = NewsCollector()

    # 4. 공유 리소스 생성
    ipc_lock = multiprocessing.Lock()
    command_queue = multiprocessing.Queue()
    shutdown_event = multiprocessing.Event()
    eod_event = multiprocessing.Event()     # [v3.1] EOD 셧다운 시그널

    print("[Server] Starting AI Trading Bot & Telegram Bot...")

    # 5. 자식 프로세스 생성
    bot_args = (ipc_lock, command_queue, shutdown_event, eod_event)
    telegram_args = (ipc_lock, command_queue, shutdown_event)

    bot_process = multiprocessing.Process(
        target=run_trading_bot, args=bot_args)
    bot_process.start()

    telegram_process = multiprocessing.Process(
        target=run_telegram_bot_standalone, args=telegram_args)
    telegram_process.start()

    # 6. Watchdog 시작
    watchdog = ProcessWatchdog(shutdown_event=shutdown_event)
    watchdog.watch("Trading Bot", bot_process, run_trading_bot, bot_args)
    watchdog.watch("Telegram Bot", telegram_process,
                   run_telegram_bot_standalone, telegram_args)
    watchdog.start()

    messenger.send_message("[Server Started] Windows AI Server Online.")

    # [v3.1] EOD 모니터 스레드: 봇이 EOD 시그널을 보내면 서버 프로세스 자동 종료
    _server_pid = os.getpid()

    def _eod_monitor():
        """eod_event 감지 시 SIGTERM으로 uvicorn 우아하게 종료"""
        while not shutdown_event.is_set():
            if eod_event.is_set():
                print("[Server] EOD 시그널 수신 — 서버 종료 중 (15초 후)...")
                import time as _time
                _time.sleep(15)  # 자식 프로세스 정리 시간
                print(f"[Server] SIGTERM → PID {_server_pid}")
                os.kill(_server_pid, signal.SIGTERM)
                return
            import time as _time
            _time.sleep(5)

    _eod_thread = threading.Thread(target=_eod_monitor, daemon=True, name="EODMonitor")
    _eod_thread.start()

    yield

    # [Shutdown]
    print("[Server] Shutting down...")
    shutdown_event.set()

    if watchdog:
        watchdog.stop()

    # 자식 프로세스 종료
    for name in ["Trading Bot", "Telegram Bot"]:
        proc = watchdog.get_process(name) if watchdog else None
        if proc and proc.is_alive():
            print(f"   - Stopping {name}...")
            proc.join(timeout=5)
            if proc.is_alive():
                proc.terminate()
                proc.join(timeout=3)

    # SharedMemory 정리
    try:
        from multiprocessing.shared_memory import SharedMemory
        shm = SharedMemory(name=Config.SHM_NAME, create=False)
        shm.close()
        shm.unlink()
    except Exception:
        pass

    messenger.send_message("[Server Stopped] Server Shutting Down.")


app = FastAPI(title="Windows AI Stock Server", lifespan=lifespan)


@app.middleware("http")
async def log_requests(request: Request, call_next):
    print(f"[HTTP] {request.method} {request.url}")
    response = await call_next(request)
    return response


class ManualOrderRequest(BaseModel):
    ticker: str
    action: str
    quantity: int


@app.get("/")
def index():
    vram = 0
    if ai_agent:
        vram = ai_agent.check_vram()
    return {
        "status": "online",
        "gpu_vram": round(vram, 2),
        "service": "Windows AI Server (Refactored)"
    }


@app.get("/trade/balance")
@app.get("/api/trade/balance")
async def get_balance():
    if not kis_client:
        return {"error": "Server not initialized"}
    return kis_client.get_balance()


@app.post("/trade/order")
@app.post("/api/trade/order")
async def manual_order(req: ManualOrderRequest):
    ticker = req.ticker
    qty = req.quantity
    action = req.action.upper()

    result = "No Action"
    if action == "BUY":
        result = kis_client.buy_stock(ticker, qty)
    elif action == "SELL":
        result = kis_client.sell_stock(ticker, qty)

    return {"status": "executed", "kis_result": result}


@app.post("/analyze/portfolio")
@app.post("/api/analyze/portfolio")
async def analyze_portfolio():
    balance = kis_client.get_balance()
    news = news_collector.get_market_news()

    prompt = f"""
    Analyze this portfolio with recent news:
    Portfolio: {balance}
    News: {news}
    Response in Korean.
    """
    analysis = ai_agent.request_inference(prompt)
    return {"analysis": analysis}


if __name__ == "__main__":
    # 서버 시작 시 좀비 프로세스 정리
    try:
        ProcessTracker.check_and_kill_zombies()
    except Exception:
        pass

    print("[Server] Starting Windows AI Server...")
    uvicorn.run("main_server:app", host="0.0.0.0", port=8000, reload=False)