main.py

"""
vLLM 서버 애플리케이션
OpenAI 호환 API를 제공하는 vLLM 서버를 실행합니다.
"""
import os
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.async_llm_engine import AsyncLLMEngine
from vllm.entrypoints.openai.api_server import (
    init_app_state,
    create_app,
)
import uvicorn


def main():
    """vLLM 서버를 시작합니다."""
    # 환경 변수에서 모델 이름 가져오기
    model_name = os.getenv("MODEL_NAME", "google/gemma-2-2b-it")
    
    # Hugging Face 토큰 가져오기 (Secret Manager에서 주입됨)
    hf_token = os.getenv("HF_TOKEN")
    
    # vLLM 엔진 인수 설정
    engine_args = AsyncEngineArgs(
        model=model_name,
        trust_remote_code=True,
        gpu_memory_utilization=0.9,
        tensor_parallel_size=1,  # Cloud Run은 GPU 1개만 제공
        max_model_len=4096,
        download_dir="/tmp/models",  # 모델 다운로드 경로
    )
    
    # Hugging Face 토큰이 있으면 환경 변수로 설정
    if hf_token:
        os.environ["HF_TOKEN"] = hf_token
    
    # 비동기 엔진 생성
    engine = AsyncLLMEngine.from_engine_args(engine_args)
    
    # 애플리케이션 상태 초기화
    app_state = init_app_state(engine)
    
    # FastAPI 앱 생성
    app = create_app(app_state)
    
    # 서버 시작
    port = int(os.getenv("PORT", 8080))
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=port,
        log_level="info",
    )


if __name__ == "__main__":
    main()