from fastapi import APIRouter, BackgroundTasks, HTTPException
from fastapi.responses import HTMLResponse
from sse_starlette.sse import EventSourceResponse
from loguru import logger
import asyncio
import json
import config
import train.trainer as trainer
from train.schemas import TrainRequest, TrainStatus
from utils.cache import load_cache
router = APIRouter()
@router.post("/start", summary="학습 실행")
async def start_train(
background_tasks: BackgroundTasks,
req: TrainRequest = None,
):
if req is None:
req = TrainRequest()
if trainer.train_status.running:
raise HTTPException(status_code=409, detail="이미 학습이 실행 중입니다.")
samples = load_cache()
if not samples:
raise HTTPException(
status_code=400,
detail="학습 데이터가 없습니다. /dataset/build 먼저 실행하세요.",
)
logger.info(f"학습 요청: {req.dict()}")
background_tasks.add_task(trainer.run_train, req)
return {"message": "학습 시작", "config": req.dict()}
@router.post("/stop", summary="학습 중단")
def stop_train():
if not trainer.train_status.running:
raise HTTPException(status_code=409, detail="학습이 실행 중이 아닙니다.")
trainer.stop_requested = True
trainer.train_status.message = "중단 요청됨..."
return {"message": "중단 요청 완료"}
@router.get("/status", summary="학습 진행 상태 조회", response_model=TrainStatus)
def get_train_status():
return trainer.train_status
@router.get("/stream", summary="학습 진행 상태 SSE 스트리밍")
async def stream_status():
"""
SSE(Server-Sent Events)로 학습 상태 실시간 전송
학습 완료 시 자동 종료
"""
async def event_generator():
prev_msg = ""
while True:
status = trainer.train_status
data = json.dumps({
"running": status.running,
"epoch": status.epoch,
"total": status.total,
"loss": status.loss,
"val_acc": status.val_acc,
"message": status.message,
# "progress": round(status.epoch / status.total * 100, 1) if status.total > 0 else 0,
"progress": round(status.batch / status.total_batches * 100, 1) if status.total_batches > 0 else 0, # ← 배치 기준으로 변경
"batch_message": status.message, # 배치 단위 메시지
}, ensure_ascii=False)
if data != prev_msg:
yield {"event": "status", "data": data}
prev_msg = data
if not status.running and status.epoch > 0:
yield {"event": "done", "data": json.dumps({"message": status.message})}
break
await asyncio.sleep(0.5) # 0.5초마다 polling → 배치 진행 반영
return EventSourceResponse(event_generator())
@router.get("/monitor", summary="학습 모니터 HTML", response_class=HTMLResponse)
def train_monitor():
"""브라우저에서 학습 진행 상태 실시간 확인"""
html = """
학습 모니터
대기 중
대기 중...
"""
return HTMLResponse(content=html)