305 lines
12 KiB
Python
305 lines
12 KiB
Python
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 = """
|
|
<!DOCTYPE html>
|
|
<html>
|
|
<head>
|
|
<meta charset="utf-8">
|
|
<title>학습 모니터</title>
|
|
<style>
|
|
* { box-sizing: border-box; }
|
|
body { font-family: sans-serif; margin: 0; background: #f5f5f5;
|
|
display: flex; flex-direction: column; height: 100vh; }
|
|
.header { padding: 16px 24px; background: #1976d2; color: #fff; flex-shrink: 0; }
|
|
.header h2 { margin: 0; font-size: 18px; }
|
|
.container { padding: 20px 24px; flex: 1; overflow-y: auto; }
|
|
|
|
/* 상태 카드 */
|
|
.cards { display: flex; gap: 16px; margin-bottom: 20px; flex-wrap: wrap; }
|
|
.card { background: #fff; border-radius: 8px; padding: 16px 20px;
|
|
flex: 1; min-width: 140px; box-shadow: 0 1px 4px rgba(0,0,0,0.1); }
|
|
.card .label { font-size: 12px; color: #888; margin-bottom: 6px; }
|
|
.card .value { font-size: 24px; font-weight: bold; color: #1976d2; }
|
|
|
|
/* 진행바 */
|
|
.progress-wrap { background: #fff; border-radius: 8px; padding: 16px 20px;
|
|
margin-bottom: 20px; box-shadow: 0 1px 4px rgba(0,0,0,0.1); }
|
|
.progress-label { display: flex; justify-content: space-between;
|
|
font-size: 13px; color: #555; margin-bottom: 8px; }
|
|
.progress-bar { height: 18px; background: #e0e0e0; border-radius: 9px; overflow: hidden; }
|
|
.progress-fill { height: 100%; background: #1976d2; border-radius: 9px;
|
|
transition: width 0.5s ease; width: 0%; }
|
|
|
|
/* 배치 진행 메시지 */
|
|
.batch-msg { background: #fff; border-radius: 8px; padding: 12px 20px;
|
|
margin-bottom: 20px; box-shadow: 0 1px 4px rgba(0,0,0,0.1);
|
|
font-size: 13px; color: #555; font-family: monospace; }
|
|
.log-wrap { background: #fff; border-radius: 8px; padding: 16px 20px;
|
|
box-shadow: 0 1px 4px rgba(0,0,0,0.1); }
|
|
.log-wrap h3 { margin: 0 0 10px; font-size: 14px; color: #333; }
|
|
.log-list { list-style: none; margin: 0; padding: 0;
|
|
max-height: 450px; overflow-y: auto; }
|
|
.log-list li { padding: 5px 8px; font-size: 13px; border-bottom: 1px solid #f0f0f0;
|
|
display: flex; gap: 10px; }
|
|
.log-list li .time { color: #aaa; flex-shrink: 0; }
|
|
.log-list li.done { color: #2e7d32; font-weight: bold; }
|
|
.log-list li.error { color: #c62828; }
|
|
|
|
/* 상태 뱃지 */
|
|
.badge { display: inline-block; padding: 3px 10px; border-radius: 12px;
|
|
font-size: 12px; font-weight: bold; }
|
|
.badge.running { background: #e3f2fd; color: #1565c0; }
|
|
.badge.done { background: #e8f5e9; color: #2e7d32; }
|
|
.badge.idle { background: #f5f5f5; color: #757575; }
|
|
|
|
/* 버튼 */
|
|
.btn { padding: 8px 20px; border: none; border-radius: 6px; cursor: pointer;
|
|
font-size: 13px; font-weight: bold; }
|
|
.btn-start { background: #1976d2; color: #fff; }
|
|
.btn-start:hover { background: #1565c0; }
|
|
.btn-start:disabled { background: #90caf9; cursor: not-allowed; }
|
|
.btn-stop { background: #d32f2f; color: #fff; }
|
|
.btn-stop:hover { background: #b71c1c; }
|
|
.actions { display: flex; gap: 10px; align-items: center; margin-bottom: 20px; }
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class="header">
|
|
<h2>🤖 LayoutLMv3 학습 모니터</h2>
|
|
</div>
|
|
<div class="container">
|
|
<div class="actions">
|
|
<button class="btn btn-start" id="btnStart" onclick="startTrain()">▶ 학습 시작</button>
|
|
<button class="btn btn-stop" id="btnStop" onclick="stopTrain()" style="display:none">■ 학습 중단</button>
|
|
<span id="badge" class="badge idle">대기 중</span>
|
|
</div>
|
|
|
|
<!-- 진행바 -->
|
|
<div class="progress-wrap">
|
|
<div class="progress-label">
|
|
<span id="epochLabel">Epoch 0 / 0</span>
|
|
<span id="pctLabel">0%</span>
|
|
</div>
|
|
<div class="progress-bar">
|
|
<div class="progress-fill" id="progressFill"></div>
|
|
</div>
|
|
</div>
|
|
|
|
<!-- 카드 -->
|
|
<div class="cards">
|
|
<div class="card">
|
|
<div class="label">현재 Loss</div>
|
|
<div class="value" id="cardLoss">-</div>
|
|
</div>
|
|
<div class="card">
|
|
<div class="label">Val Accuracy</div>
|
|
<div class="value" id="cardAcc">-</div>
|
|
</div>
|
|
<div class="card">
|
|
<div class="label">진행률</div>
|
|
<div class="value" id="cardPct">0%</div>
|
|
</div>
|
|
</div>
|
|
|
|
<!-- 배치 진행 메시지 -->
|
|
<div class="batch-msg" id="batchMsg">대기 중...</div>
|
|
|
|
<!-- 메시지 로그 -->
|
|
<div class="log-wrap">
|
|
<h3>📋 진행 로그</h3>
|
|
<ul class="log-list" id="logList"></ul>
|
|
</div>
|
|
</div>
|
|
|
|
<script>
|
|
let es = null;
|
|
|
|
function now() {
|
|
return new Date().toLocaleTimeString("ko-KR");
|
|
}
|
|
|
|
function addLog(msg, cls = "") {
|
|
const li = document.createElement("li");
|
|
li.className = cls;
|
|
li.innerHTML = `<span class="time">${now()}</span><span>${msg}</span>`;
|
|
const list = document.getElementById("logList");
|
|
list.prepend(li);
|
|
}
|
|
|
|
function updateUI(d) {
|
|
// 배치 단위 실시간 메시지
|
|
document.getElementById("batchMsg").textContent = d.message || "";
|
|
|
|
// 진행바 (에포크 기준)
|
|
document.getElementById("progressFill").style.width = d.progress + "%";
|
|
document.getElementById("epochLabel").textContent =
|
|
`Epoch ${d.epoch} / ${d.total}`;
|
|
document.getElementById("pctLabel").textContent = d.progress + "%";
|
|
|
|
// 카드
|
|
document.getElementById("cardLoss").textContent = d.loss || "-";
|
|
document.getElementById("cardAcc").textContent =
|
|
d.val_acc ? (d.val_acc * 100).toFixed(1) + "%" : "-";
|
|
document.getElementById("cardPct").textContent = d.progress + "%";
|
|
|
|
// 뱃지
|
|
const badge = document.getElementById("badge");
|
|
if (d.running) {
|
|
badge.textContent = "학습 중";
|
|
badge.className = "badge running";
|
|
}
|
|
}
|
|
|
|
function startTrain() {
|
|
fetch("/train/start", { method: "POST" })
|
|
.then(r => r.json())
|
|
.then(data => {
|
|
addLog("학습 시작 요청: " + JSON.stringify(data.config));
|
|
document.getElementById("btnStart").disabled = true;
|
|
document.getElementById("btnStop").style.display = "inline-block";
|
|
subscribeSSE();
|
|
})
|
|
.catch(e => addLog("오류: " + e, "error"));
|
|
}
|
|
|
|
function stopTrain() {
|
|
fetch("/train/stop", { method: "POST" })
|
|
.then(r => r.json())
|
|
.then(d => addLog("⏹ " + d.message))
|
|
.catch(e => addLog("오류: " + e, "error"));
|
|
}
|
|
|
|
function subscribeSSE() {
|
|
if (es) es.close();
|
|
es = new EventSource("/train/stream");
|
|
|
|
es.addEventListener("status", e => {
|
|
const d = JSON.parse(e.data);
|
|
updateUI(d);
|
|
if (d.message) addLog(d.message);
|
|
});
|
|
|
|
es.addEventListener("done", e => {
|
|
const d = JSON.parse(e.data);
|
|
addLog("✅ " + d.message, "done");
|
|
|
|
const badge = document.getElementById("badge");
|
|
badge.textContent = "완료";
|
|
badge.className = "badge done";
|
|
document.getElementById("btnStart").disabled = false;
|
|
document.getElementById("btnStop").style.display = "none";
|
|
es.close();
|
|
});
|
|
|
|
es.onerror = () => {
|
|
addLog("SSE 연결 끊김", "error");
|
|
es.close();
|
|
document.getElementById("btnStart").disabled = false;
|
|
};
|
|
}
|
|
|
|
// 페이지 로드 시 이미 학습 중이면 자동 구독
|
|
fetch("/train/status")
|
|
.then(r => r.json())
|
|
.then(d => {
|
|
updateUI({...d, progress: d.total > 0 ? d.epoch / d.total * 100 : 0});
|
|
if (d.running) {
|
|
document.getElementById("btnStart").disabled = true;
|
|
addLog("학습 진행 중 - 자동 연결");
|
|
subscribeSSE();
|
|
}
|
|
});
|
|
</script>
|
|
</body>
|
|
</html>
|
|
"""
|
|
return HTMLResponse(content=html) |