vLLM

# vLLM 완전 가이드: 서버 구축, 클라이언트, 벤치마킹
# pip install vllm openai

# ============================================
# 1. vLLM 서버 실행 (터미널에서)
# ============================================
# 기본 실행 (단일 GPU)
# python -m vllm.entrypoints.openai.api_server \
#     --model meta-llama/Llama-3.1-8B-Instruct \
#     --port 8000

# 멀티 GPU (Tensor Parallelism)
# python -m vllm.entrypoints.openai.api_server \
#     --model meta-llama/Llama-3.1-70B-Instruct \
#     --tensor-parallel-size 4 \
#     --port 8000

# 양자화 모델 서빙 (메모리 절약)
# python -m vllm.entrypoints.openai.api_server \
#     --model TheBloke/Llama-2-70B-Chat-AWQ \
#     --quantization awq \
#     --port 8000

# ============================================
# 2. Python에서 vLLM 직접 사용
# ============================================
from vllm import LLM, SamplingParams

def basic_vllm_inference():
    """vLLM으로 오프라인 배치 추론"""
    # 모델 로드
    llm = LLM(
        model="meta-llama/Llama-3.1-8B-Instruct",
        tensor_parallel_size=1,  # GPU 수
        gpu_memory_utilization=0.9,  # GPU 메모리 사용률
        max_model_len=4096  # 최대 컨텍스트 길이
    )

    # 샘플링 파라미터
    sampling_params = SamplingParams(
        temperature=0.7,
        top_p=0.95,
        max_tokens=512,
        stop=["", "[/INST]"]
    )

    # 배치 프롬프트
    prompts = [
        "[INST] Python으로 피보나치 함수를 작성해주세요. [/INST]",
        "[INST] 머신러닝과 딥러닝의 차이를 설명해주세요. [/INST]",
        "[INST] Docker의 장점을 3가지 알려주세요. [/INST]",
    ]

    # 추론 (자동 배칭)
    outputs = llm.generate(prompts, sampling_params)

    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt[:50]}...")
        print(f"Output: {generated_text[:200]}...")
        print("-" * 50)

    return outputs

# ============================================
# 3. OpenAI 호환 클라이언트 사용
# ============================================
from openai import OpenAI

def vllm_openai_client():
    """vLLM 서버를 OpenAI 클라이언트로 호출"""
    # vLLM 서버 주소 (로컬)
    client = OpenAI(
        base_url="http://localhost:8000/v1",
        api_key="dummy"  # vLLM은 API 키 검증 안 함
    )

    # Chat Completion
    response = client.chat.completions.create(
        model="meta-llama/Llama-3.1-8B-Instruct",
        messages=[
            {"role": "system", "content": "당신은 유능한 AI 어시스턴트입니다."},
            {"role": "user", "content": "vLLM의 장점을 설명해주세요."}
        ],
        temperature=0.7,
        max_tokens=500
    )

    print(response.choices[0].message.content)
    return response

# ============================================
# 4. 스트리밍 응답
# ============================================
def vllm_streaming():
    """스트리밍으로 토큰 단위 응답 받기"""
    client = OpenAI(
        base_url="http://localhost:8000/v1",
        api_key="dummy"
    )

    stream = client.chat.completions.create(
        model="meta-llama/Llama-3.1-8B-Instruct",
        messages=[
            {"role": "user", "content": "Python의 GIL에 대해 설명해주세요."}
        ],
        stream=True,
        max_tokens=300
    )

    full_response = ""
    for chunk in stream:
        if chunk.choices[0].delta.content:
            content = chunk.choices[0].delta.content
            print(content, end="", flush=True)
            full_response += content

    return full_response

# ============================================
# 5. 벤치마킹 및 성능 측정
# ============================================
import time
import asyncio
import aiohttp

async def benchmark_vllm(
    base_url: str = "http://localhost:8000",
    num_requests: int = 100,
    concurrency: int = 10
):
    """vLLM 서버 벤치마킹"""
    prompt = "What is the capital of France? Answer in one word."

    async def single_request(session):
        start = time.time()
        async with session.post(
            f"{base_url}/v1/completions",
            json={
                "model": "meta-llama/Llama-3.1-8B-Instruct",
                "prompt": prompt,
                "max_tokens": 50,
                "temperature": 0
            }
        ) as response:
            result = await response.json()
            latency = time.time() - start
            tokens = result.get("usage", {}).get("completion_tokens", 0)
            return latency, tokens

    # 동시 요청 실행
    async with aiohttp.ClientSession() as session:
        semaphore = asyncio.Semaphore(concurrency)

        async def bounded_request():
            async with semaphore:
                return await single_request(session)

        start_time = time.time()
        results = await asyncio.gather(*[
            bounded_request() for _ in range(num_requests)
        ])
        total_time = time.time() - start_time

    # 결과 분석
    latencies = [r[0] for r in results]
    total_tokens = sum(r[1] for r in results)

    print(f"=== vLLM Benchmark Results ===")
    print(f"Total requests: {num_requests}")
    print(f"Concurrency: {concurrency}")
    print(f"Total time: {total_time:.2f}s")
    print(f"Throughput: {num_requests / total_time:.2f} req/s")
    print(f"Token throughput: {total_tokens / total_time:.2f} tokens/s")
    print(f"Avg latency: {sum(latencies) / len(latencies) * 1000:.2f} ms")
    print(f"P50 latency: {sorted(latencies)[len(latencies)//2] * 1000:.2f} ms")
    print(f"P99 latency: {sorted(latencies)[int(len(latencies)*0.99)] * 1000:.2f} ms")

# ============================================
# 6. 프로덕션 설정 예시 (Docker Compose)
# ============================================
DOCKER_COMPOSE_EXAMPLE = """
# docker-compose.yml
version: '3.8'
services:
  vllm:
    image: vllm/vllm-openai:latest
    ports:
      - "8000:8000"
    volumes:
      - ~/.cache/huggingface:/root/.cache/huggingface
    environment:
      - HUGGING_FACE_HUB_TOKEN=${HF_TOKEN}
    command: >
      --model meta-llama/Llama-3.1-8B-Instruct
      --tensor-parallel-size 1
      --gpu-memory-utilization 0.9
      --max-model-len 8192
      --enable-prefix-caching
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
"""

# 실행 예시
if __name__ == "__main__":
    # 오프라인 추론 (vLLM 직접 사용)
    # basic_vllm_inference()

    # 서버 클라이언트 사용
    # vllm_openai_client()

    # 스트리밍
    # vllm_streaming()

    # 벤치마크 실행
    # asyncio.run(benchmark_vllm())

    print("vLLM 코드 예제 준비 완료!")
    print("서버 실행: python -m vllm.entrypoints.openai.api_server --model ")