LLM推論における連続動的バッチ処理：Orca、vLLM、TGIのレイテンシベンチマーク

# vllm_server.py
import os
from vllm import LLM, SamplingParams
from fastapi import FastAPI, Request
from pydantic import BaseModel
import uvicorn
import time

# Configuration
MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
GPU_MEMORY_UTILIZATION = 0.9 # Adjust based on your GPU and model size

# Initialize LLM
print(f"Loading model: {MODEL_NAME}...")
llm = LLM(
    model=MODEL_NAME,
    tensor_parallel_size=1, # Single GPU
    gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
    trust_remote_code=True,
    dtype="bfloat16" # Use bfloat16 for better performance on A100
)
print("Model loaded.")

app = FastAPI()

class GenerateRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 256
    temperature: float = 0.7
    top_p: float = 0.95
    do_sample: bool = True

@app.post("/generate")
async def generate(request: GenerateRequest):
    sampling_params = SamplingParams(
        n=1,
        temperature=request.temperature,
        top_p=request.top_p,
        max_tokens=request.max_new_tokens,
        stop=["</s>"], # Llama-2 specific stop token
        do_sample=request.do_sample
    )

    start_time = time.time()
    outputs = await llm.generate_async(request.prompt, sampling_params)
    end_time = time.time()

    first_token_time = -1 # Placeholder, vLLM doesn't expose this directly in sync API
    # For accurate TTFT, you'd typically stream tokens and measure when the first one arrives.
    # For this benchmark, we'll approximate TTFT from the client side.

    generated_text = outputs[0].outputs[0].text
    num_output_tokens = len(outputs[0].outputs[0].token_ids)

    return {
        "generated_text": generated_text,
        "num_output_tokens": num_output_tokens,
        "total_time_s": end_time - start_time,
        "first_token_time_s": first_token_time # Will be calculated client-side
    }

if __name__ == "__main__":
    # To run: python vllm_server.py
    # Then in another terminal: uvicorn vllm_server:app --host 0.0.0.0 --port 8000 --workers 1
    uvicorn.run(app, host="0.0.0.0", port=8000, workers=1)

python vllm_server.py
# In a separate terminal:
uvicorn vllm_server:app --host 0.0.0.0 --port 8000 --workers 1

# TGI server command
# Ensure you have Docker and NVIDIA Container Toolkit installed
docker run --gpus all -p 8080:80 -v ~/.cache/huggingface:/data ghcr.io/huggingface/text-generation-inference:1.4 --model-id meta-llama/Llama-2-7b-chat-hf --dtype bfloat16 --max-input-length 1024 --max-total-tokens 2048

# benchmark_client.py
import asyncio
import httpx
import time
import random
import numpy as np
from typing import List, Dict

# Configuration
VLLM_ENDPOINT = "http://localhost:8000/generate"
TGI_ENDPOINT = "http://localhost:8080/generate" # TGI uses /generate for non-streaming
MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"

# Prompt templates for Llama-2
PROMPT_TEMPLATES = [
    "<s>[INST] {prompt} [/INST]",
    "<s>[INST] <<SYS>>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. If a question does not make any sense, or is not factually coherent, explain why instead of answering something incorrect. Do not share false information.\n<</SYS>>\n\n{prompt} [/INST]"
]

# Example prompts
BASE_PROMPTS = [
    "Explain the concept of quantum entanglement in simple terms.",
    "Write a short story about a detective solving a mystery in a futuristic city.",
    "Describe the economic impact of artificial intelligence on the job market.",
    "What are the main differences between classical and quantum computing?",
    "Provide a detailed recipe for authentic Italian lasagna.",
    "Discuss the ethical implications of autonomous vehicles.",
    "Summarize the plot of 'Dune' by Frank Herbert.",
    "Explain the process of photosynthesis.",
    "Write a poem about the beauty of the night sky.",
    "What is the significance of the 'butterfly effect' in chaos theory?"
]

def generate_random_prompt(min_len=50, max_len=200):
    base = random.choice(BASE_PROMPTS)
    # Pad with random words to reach desired length
    words = base.split()
    while len(" ".join(words)) < min_len:
        words.append(random.choice(BASE_PROMPTS).split()[0]) # Add a random word
    prompt = " ".join(words[:random.randint(len(words), len(words) + 50)]) # Add some variability
    return random.choice(PROMPT_TEMPLATES).format(prompt=prompt)

async def send_request(client: httpx.AsyncClient, endpoint: str, prompt: str, max_new_tokens: int):
    request_payload = {
        "prompt": prompt,
        "max_new_tokens": max_new_tokens,
        "temperature": 0.7,
        "top_p": 0.95,
        "do_sample": True
    }

    start_time = time.time()
    try:
        # For TGI, we need to stream to get TTFT accurately
        if "8080" in endpoint: # Heuristic for TGI
            async with client.stream("POST", endpoint, json=request_payload, timeout=60.0) as response:
                response.raise_for_status()
                first_token_received = False
                ttft = -1
                tokens = []
                token_gen_times = []

                async for chunk in response.aiter_bytes():
                    if not first_token_received:
                        ttft = time.time() - start_time
                        first_token_received = True

                    # TGI streaming format is Server-Sent Events (SSE)
                    # We need to parse it to extract tokens
                    try:
                        chunk_str = chunk.decode('utf-8')
                        for line in chunk_str.split('\n'):
                            if line.startswith('data:'):
                                data = line[len('data:'):].strip()
                                if data == '[DONE]':
                                    break
                                token_data = json.loads(data)
                                if 'token' in token_data and 'text' in token_data['token']:
                                    tokens.append(token_data['token']['text'])
                                    token_gen_times.append(time.time())
                    except json.JSONDecodeError:
                        # Handle incomplete JSON chunks
                        pass

                total_time = time.time() - start_time
                num_output_tokens = len(tokens)

                itl = -1
                if num_output_tokens > 1:
                    itl = np.mean(np.diff(token_gen_times))

                return {
                    "ttft": ttft,
                    "itl": itl,
                    "total_time": total_time,
                    "num_output_tokens": num_output_tokens,
                    "success": True
                }
        else: # vLLM non-streaming endpoint for simplicity in this benchmark
            response = await client.post(endpoint, json=request_payload, timeout=60.0)
            response.raise_for_status()
            data = response.json()

            # Approximate TTFT for vLLM as total_time / num_output_tokens for non-streaming
            # This is a simplification; a true streaming client would be needed for accurate TTFT.
            # For this benchmark, we'll use total_time as a proxy for TTFT for vLLM,
            # and focus on TGI's streaming TTFT.
            num_output_tokens = data.get("num_output_tokens", 0)
            total_time = time.time() - start_time

            # For vLLM, without streaming, ITL is hard to measure accurately from client
            # We'll report total time / tokens as an average token generation time.
            avg_token_gen_time = total_time / num_output_tokens if num_output_tokens > 0 else 0

            return {
                "ttft": total_time, # Proxy for TTFT for vLLM non-streaming
                "itl": avg_token_gen_time, # Proxy for ITL for vLLM non-streaming
                "total_time": total_time,
                "num_output_tokens": num_output_tokens,
                "success": True
            }
    except httpx.RequestError as e:
        print(f"Request failed: {e}")
        return {"ttft": -1, "itl": -1, "total_time": -1, "num_output_tokens": 0, "success": False}
    except httpx.HTTPStatusError as e:
        print(f"HTTP error: {e.response.status_code} - {e.response.text}")
        return {"ttft": -1, "itl": -1, "total_time": -1, "num_output_tokens": 0, "success": False}
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        return {"ttft": -1, "itl": -1, "total_time": -1, "num_output_tokens": 0, "success": False}


async def run_benchmark(endpoint: str, concurrency: int, num_requests: int = 100):
    print(f"\n--- Benchmarking {endpoint} with {concurrency} concurrent requests ---")
    results = []

    async with httpx.AsyncClient() as client:
        tasks = []
        for _ in range(num_requests):
            prompt = generate_random_prompt()
            max_new_tokens = random.randint(100, 500)
            tasks.append(send_request(client, endpoint, prompt, max_new_tokens))

        # Use a semaphore to limit concurrency
        semaphore = asyncio.Semaphore(concurrency)

        async def limited_task(task):
            async with semaphore:
                return await task

        start_benchmark_time = time.time()
        processed_results = await asyncio.gather(*[limited_task(t) for t in tasks])
        end_benchmark_time = time.time()

        for res in processed_results:
            if res["success"]:
                results.append(res)

    if not results:
        print("No successful requests to report.")
        return

    ttfts = [r["ttft"] for r in results if r["ttft"] > 0]
    itls = [r["itl"] for r in results if r["itl"] > 0]
    total_times = [r["total_time"] for r in results if r["total_time"] > 0]
    output_tokens = [r["num_output_tokens"] for r in results if r["num_output_tokens"] > 0]

    print(f"Total successful requests: {len(results)}")
    print(f"Overall benchmark duration: {end_benchmark_time - start_benchmark_time:.2f} s")
    print(f"Average TTFT: {np.mean(ttfts):.4f} s (Median: {np.median(ttfts):.4f} s)")
    print(f"Average ITL: {np.mean(itls):.4f} s (Median: {np.median(itls):.4f} s)")
    print(f"Average Total Request Time: {np.mean(total_times):.4f} s (Median: {np.median(total_times):.4f} s)")
    print(f"Average Output Tokens: {np.mean(output_tokens):.2f}")
    print(f"Total Output Tokens: {np.sum(output_tokens)}")
    print(f"Throughput (tokens/sec): {np.sum(output_tokens) / (end_benchmark_time - start_benchmark_time):.2f}")

async def main():
    concurrency_levels = [1, 4, 8, 16, 32]
    num_requests_per_level = 50 # Reduced for quicker run

    # Run vLLM benchmark
    # Note: For vLLM, the client-side TTFT/ITL will be less accurate without streaming.
    # The reported TTFT will be total request time, and ITL will be average token time.
    # This is to highlight the difference in how these metrics are typically measured for streaming vs non-streaming.
    # For a true comparison, vLLM's streaming API should be used.
    # for c in concurrency_levels:
    #     await run_benchmark(VLLM_ENDPOINT, c, num_requests_per_level)

    # Run TGI benchmark (streaming enabled for accurate TTFT/ITL)
    import json # Import json for TGI streaming parsing
    for c in concurrency_levels:
        await run_benchmark(TGI_ENDPOINT, c, num_requests_per_level)

if __name__ == "__main__":
    asyncio.run(main())