"""Ollama와 vLLM의 동시 요청 처리 성능을 같은 조건에서 측정·비교한다.

vllm.md '측정 결과' 절의 수치를 재현하는 스크립트다.
두 서버를 OpenAI/Ollama 채팅 엔드포인트로 호출해 생성 경로를 정렬하고,
동시성(1,4,8,...)을 올려가며 처리량(req/s)·지연시간(p50/p95)을 잰다.
정밀도가 다르면(Ollama Q4 vs vLLM bf16) 절대 수치보다 '동시성 1 대비 정규화 처리량'을 본다.

표준 라이브러리만 쓴다(추가 설치 불필요). 파이썬 3.8+.

준비 — 두 서버를 같은 GPU에서 순차로 띄운다(동시 적재 시 VRAM 부족):

  # vLLM (네이티브 리눅스는 최신 이미지 사용. WSL2는 'UVA is not available'를 피해 v0.8.5로 고정)
  docker run --gpus all --ipc=host -p 8000:8000 \
    -v ~/.cache/huggingface:/root/.cache/huggingface \
    vllm/vllm-openai:latest Qwen/Qwen3-4B \
    --dtype bfloat16 --max-model-len 2048 --gpu-memory-utilization 0.9 --seed 42

  # Ollama (컨텍스트를 제한해야 모델 전체가 GPU에 올라간다)
  docker run -d --gpus all -p 11434:11434 \
    -e OLLAMA_NUM_PARALLEL=8 -e OLLAMA_CONTEXT_LENGTH=2048 \
    ollama/ollama:latest
  docker exec <컨테이너> ollama pull qwen3:4b

동시 처리 상한은 vLLM은 `--max-num-seqs`(기본값 큼), Ollama는 `OLLAMA_NUM_PARALLEL`로
정해진다 — 이 값보다 큰 동시성은 큐에 쌓여 처리량이 포화하고 지연이 늘어난다.

실행:
  python3 scripts/benchmark-serving.py --target vllm   --model Qwen/Qwen3-4B
  python3 scripts/benchmark-serving.py --target ollama --model qwen3:4b
"""
import argparse
import json
import statistics
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

프롬프트 = (
    "Explain the difference between continuous batching and static batching "
    "in LLM serving, in detail."
)


def 한_요청(설정):
    대상, 주소, 모델, 최대토큰 = 설정
    if 대상 == "vllm":
        url = f"{주소}/v1/chat/completions"
        본문 = {
            "model": 모델,
            "messages": [{"role": "user", "content": 프롬프트}],
            "temperature": 0.0,
            "max_tokens": 최대토큰,
        }
    else:
        url = f"{주소}/api/chat"
        본문 = {
            "model": 모델,
            "messages": [{"role": "user", "content": 프롬프트}],
            "stream": False,
            "options": {"temperature": 0.0, "num_predict": 최대토큰, "num_ctx": 2048},
        }
    요청 = urllib.request.Request(
        url, data=json.dumps(본문).encode(), headers={"Content-Type": "application/json"}
    )
    시작 = time.perf_counter()
    응답 = json.load(urllib.request.urlopen(요청, timeout=300))
    지연 = (time.perf_counter() - 시작) * 1000.0
    토큰 = (
        응답.get("usage", {}).get("completion_tokens", 0)
        if 대상 == "vllm"
        else 응답.get("eval_count", 0)
    )
    return 지연, 토큰


def 스윕(설정, 동시성, 총요청):
    인자목록 = [설정] * 총요청
    with ThreadPoolExecutor(max_workers=동시성) as 풀:
        시작 = time.perf_counter()
        결과 = list(풀.map(한_요청, 인자목록))
        경과 = time.perf_counter() - 시작
    지연들 = [r[0] for r in 결과]
    토큰들 = [r[1] for r in 결과]
    초당요청 = 총요청 / 경과
    초당토큰 = sum(토큰들) / 경과
    p50 = statistics.median(지연들)
    p95 = statistics.quantiles(지연들, n=100)[94] if len(지연들) >= 20 else max(지연들)
    평균토큰 = sum(토큰들) / len(토큰들)
    return 초당요청, 초당토큰, p50, p95, 평균토큰


def main():
    파서 = argparse.ArgumentParser(description="Ollama/vLLM 동시 요청 처리 벤치마크")
    파서.add_argument("--target", choices=["ollama", "vllm"], required=True)
    파서.add_argument("--model", required=True, help="예: qwen3:4b 또는 Qwen/Qwen3-4B")
    파서.add_argument(
        "--base-url",
        default=None,
        help="기본값: vllm=http://127.0.0.1:8000, ollama=http://127.0.0.1:11434",
    )
    파서.add_argument("--concurrency", type=int, nargs="+", default=[1, 4, 8, 16, 32])
    파서.add_argument("--total-requests", type=int, default=64)
    파서.add_argument("--max-tokens", type=int, default=128)
    인자 = 파서.parse_args()

    주소 = 인자.base_url or (
        "http://127.0.0.1:8000" if 인자.target == "vllm" else "http://127.0.0.1:11434"
    )
    설정 = (인자.target, 주소, 인자.model, 인자.max_tokens)

    print(
        f"# {인자.target}  model={인자.model}  base={주소}  "
        f"max_tokens={인자.max_tokens}  요청수/동시성={인자.total_requests}"
    )
    행들 = []
    for 동시성 in 인자.concurrency:
        초당요청, 초당토큰, p50, p95, 평균토큰 = 스윕(설정, 동시성, 인자.total_requests)
        행들.append((동시성, 초당요청))
        # 평균 생성 토큰이 목표보다 크게 적으면 조기 종료(길이 아티팩트)를 의심한다.
        경고 = "  ⚠ 조기종료?" if 평균토큰 < 인자.max_tokens * 0.9 else ""
        print(
            f"동시성={동시성:>3}  req/s={초당요청:6.2f}  out_tok/s={초당토큰:7.1f}  "
            f"p50={p50:7.0f}ms  p95={p95:7.0f}ms  평균토큰={평균토큰:.0f}{경고}",
            flush=True,
        )
    기준 = 행들[0][1]
    print(
        "\n정규화 처리량(동시성 1 대비): "
        + "  ".join(f"c{동시성}={초당요청 / 기준:.2f}x" for 동시성, 초당요청 in 행들)
    )


if __name__ == "__main__":
    main()