inpaintServer/main.py

278 lines
9.4 KiB
Python

#!/usr/bin/env python3
"""
인페인팅 서버 메인 애플리케이션
iopaint와 호환되는 API를 제공합니다.
"""
import time
import logging
import json
import asyncio
from contextlib import asynccontextmanager
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import uvicorn
from app.core.config import settings
from app.core.worker_manager import worker_manager
from app.core.session_pool import session_pool
from app.api.endpoints import router
from app.monitoring.dashboard import monitor_app
# 로깅 설정
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
# 서버 시작 시간 기록
start_time = time.time()
async def save_status_periodically():
"""주기적으로 워커와 세션 상태를 파일에 저장합니다."""
while True:
try:
status = {
"worker_status": worker_manager.get_status(),
"session_status": session_pool.get_status(),
"timestamp": time.time()
}
with open("status.json", "w") as f:
json.dump(status, f)
except Exception as e:
logger.warning(f"상태 저장 실패: {e}")
await asyncio.sleep(1)
@asynccontextmanager
async def lifespan(app: FastAPI):
"""애플리케이션 생명주기 관리"""
# 시작 시
logger.info("🚀 인페인팅 서버 시작 중...")
# app.state에 공유 객체 저장
app.state.worker_manager = worker_manager
app.state.session_pool = session_pool
logger.info("✅ 공유 객체를 app.state에 저장 완료")
# 상태 저장 백그라운드 작업 시작
status_task = asyncio.create_task(save_status_periodically())
try:
# 세션 풀 초기화
await session_pool.initialize()
logger.info("✅ 세션 풀 초기화 완료")
# 워커 매니저 시작
await worker_manager.start()
logger.info("✅ 워커 매니저 시작 완료")
logger.info("🎉 인페인팅 서버 시작 완료!")
except Exception as e:
logger.error(f"❌ 서버 시작 실패: {e}")
raise
yield
# 종료 시
logger.info("🛑 인페인팅 서버 종료 중...")
# 상태 저장 백그라운드 작업 취소
status_task.cancel()
try:
# 워커 매니저 중지
await worker_manager.stop()
logger.info("✅ 워커 매니저 중지 완료")
logger.info("👋 인페인팅 서버 종료 완료")
except Exception as e:
logger.error(f"❌ 서버 종료 중 오류: {e}")
# 메인 애플리케이션 생성
app = FastAPI(
title="인페인팅 서버",
description="Simple LAMA, MIGAN, REMBG를 활용한 병렬 처리 인페인팅 서버 (iopaint 호환)",
version="1.0.0",
lifespan=lifespan
)
# CORS 미들웨어 추가
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# API 라우터 포함
app.include_router(router)
# 모니터링은 start_server.sh를 통해 독립적으로 실행됩니다.
# app.mount("/monitoring", monitor_app, name="monitoring")
# 모니터링 데이터 직접 제공 (완전 통합)
@app.get("/monitoring/api/status")
async def get_monitoring_status():
"""모니터링 상태를 직접 반환합니다."""
try:
import psutil
from app.utils.gpu_monitor import GPUMonitor
# 시스템 정보 수집
cpu_percent = psutil.cpu_percent()
memory = psutil.virtual_memory()
process_count = len(psutil.pids())
# GPU 정보 수집
gpu_monitor = GPUMonitor()
gpu_info = await gpu_monitor.get_gpu_info()
# 워커 매니저 상태 (안전한 방식)
try:
worker_status = {
"total_workers": getattr(worker_manager, 'workers', {}).__len__() if hasattr(worker_manager, 'workers') else 0,
"queue_size": getattr(worker_manager, 'queue', None).qsize() if hasattr(worker_manager, 'queue') else 0,
"workers_by_status": {"idle": [], "busy": [], "starting": [], "stopping": [], "error": []},
"running": getattr(worker_manager, 'running', False)
}
except Exception as e:
logger.warning(f"워커 상태 수집 실패: {e}")
worker_status = {
"total_workers": 0,
"queue_size": 0,
"workers_by_status": {"idle": [], "busy": [], "starting": [], "stopping": [], "error": []},
"running": False
}
# 세션 풀 상태 (안전한 방식)
try:
sessions = getattr(session_pool, 'sessions', {})
total_sessions = len(sessions) if sessions else 0
available_sessions = len([s for s in sessions.values() if getattr(s, 'available', False)]) if sessions else 0
session_status = {
"total_sessions": total_sessions,
"available_sessions": available_sessions,
"model_distribution": {"simple-lama": 0, "migan": 0, "rembg": 0}
}
except Exception as e:
logger.warning(f"세션 상태 수집 실패: {e}")
session_status = {
"total_sessions": 0,
"available_sessions": 0,
"model_distribution": {"simple-lama": 0, "migan": 0, "rembg": 0}
}
# API 통계 (간단한 버전)
api_stats = {
"total_requests": 0, # 실제로는 카운터가 필요
"success_rate": 100.0,
"average_response_time": 0.0,
"error_count": 0
}
return {
"timestamp": time.time(),
"system_type": "Jetson Xavier" if settings.IS_JETSON else "x86_64",
"system": {
"cpu_percent": cpu_percent,
"memory_percent": memory.percent,
"process_count": process_count
},
"gpu": gpu_info,
"worker_status": worker_status,
"session_status": session_status,
"api_stats": api_stats
}
except Exception as e:
logger.error(f"모니터링 데이터 수집 실패: {e}")
return {"error": f"모니터링 데이터 수집 실패: {str(e)}"}
@app.get("/monitoring/api/simple")
async def get_simple_monitoring():
"""간단한 모니터링 상태를 반환합니다."""
try:
import psutil
return {
"timestamp": time.time(),
"system_type": "Jetson Xavier" if settings.IS_JETSON else "x86_64",
"cpu_percent": psutil.cpu_percent(),
"memory_percent": psutil.virtual_memory().percent,
"status": "running"
}
except Exception as e:
return {"error": f"간단한 상태 수집 실패: {str(e)}"}
@app.get("/monitoring/api/worker-status")
async def get_worker_status():
"""워커 상태를 반환합니다."""
try:
workers = getattr(worker_manager, 'workers', {})
queue = getattr(worker_manager, 'queue', None)
running = getattr(worker_manager, 'running', False)
return {
"total_workers": len(workers) if workers else 0,
"queue_size": queue.qsize() if queue else 0,
"running": running,
"status": "active" if running else "stopped"
}
except Exception as e:
logger.warning(f"워커 상태 조회 실패: {e}")
return {"error": f"워커 상태 조회 실패: {str(e)}"}
@app.get("/monitoring/api/session-status")
async def get_session_status():
"""세션 풀 상태를 반환합니다."""
try:
sessions = getattr(session_pool, 'sessions', {})
total_sessions = len(sessions) if sessions else 0
available_sessions = len([s for s in sessions.values() if getattr(s, 'available', False)]) if sessions else 0
return {
"total_sessions": total_sessions,
"available_sessions": available_sessions,
"model_distribution": {"simple-lama": 0, "migan": 0, "rembg": 0}
}
except Exception as e:
logger.warning(f"세션 상태 조회 실패: {e}")
return {"error": f"세션 상태 조회 실패: {str(e)}"}
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="인페인팅 서버")
parser.add_argument("--dev", action="store_true", help="개발 모드로 실행")
parser.add_argument("--host", default=settings.HOST, help="호스트 주소")
parser.add_argument("--port", type=int, default=settings.PORT, help="포트 번호")
parser.add_argument("--workers", type=int, default=settings.WORKERS, help="워커 수")
args = parser.parse_args()
if args.dev:
logger.info("🔧 개발 모드로 실행합니다")
uvicorn.run(
"main:app",
host=args.host,
port=args.port,
reload=True,
log_level="info"
)
else:
logger.info("🚀 프로덕션 모드로 실행합니다")
uvicorn.run(
"main:app",
host=args.host,
port=args.port,
workers=args.workers,
log_level="info"
)