从零搭建 LLM 推理服务

# src/config.py
"""
配置管理
"""
from pydantic_settings import BaseSettings
from typing import Optional, List
from functools import lru_cache

class Settings(BaseSettings):
    """应用配置"""

    # 服务配置
    app_name: str = "LLM Inference Service"
    host: str = "0.0.0.0"
    port: int = 8000
    workers: int = 1

    # 模型配置
    model_name: str = "meta-llama/Llama-2-7b-chat-hf"
    model_path: Optional[str] = None
    dtype: str = "float16"
    quantization: Optional[str] = None  # "awq", "gptq", None

    # vLLM 配置
    tensor_parallel_size: int = 1
    max_model_len: int = 4096
    gpu_memory_utilization: float = 0.9
    enforce_eager: bool = False

    # 推理配置
    max_tokens: int = 2048
    temperature: float = 0.7
    top_p: float = 0.9
    top_k: int = 50

    # 并发配置
    max_concurrent_requests: int = 100
    max_batch_size: int = 32
    request_timeout: float = 60.0

    # 认证配置
    api_key_enabled: bool = False
    api_keys: List[str] = []

    # 日志配置
    log_level: str = "INFO"
    log_format: str = "json"

    class Config:
        env_file = ".env"
        env_prefix = "LLM_"


@lru_cache()
def get_settings() -> Settings:
    return Settings()

# src/engine.py
"""
vLLM 推理引擎封装
"""
import asyncio
from typing import AsyncGenerator, List, Optional, Dict, Any
from dataclasses import dataclass
import time
import logging

from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
from vllm.outputs import RequestOutput

from .config import get_settings

logger = logging.getLogger(__name__)


@dataclass
class GenerationResult:
    """生成结果"""
    request_id: str
    text: str
    tokens: List[int]
    finish_reason: str
    prompt_tokens: int
    completion_tokens: int
    latency_ms: float


@dataclass
class StreamChunk:
    """流式输出块"""
    request_id: str
    text: str
    token_id: int
    is_finished: bool
    finish_reason: Optional[str] = None


class InferenceEngine:
    """推理引擎"""

    def __init__(self):
        self.settings = get_settings()
        self.engine: Optional[AsyncLLMEngine] = None
        self._request_counter = 0
        self._lock = asyncio.Lock()

        # 统计信息
        self.stats = {
            "total_requests": 0,
            "total_tokens_generated": 0,
            "total_latency_ms": 0,
            "active_requests": 0
        }

    async def initialize(self):
        """初始化引擎"""
        logger.info(f"Initializing inference engine with model: {self.settings.model_name}")

        engine_args = AsyncEngineArgs(
            model=self.settings.model_path or self.settings.model_name,
            dtype=self.settings.dtype,
            tensor_parallel_size=self.settings.tensor_parallel_size,
            max_model_len=self.settings.max_model_len,
            gpu_memory_utilization=self.settings.gpu_memory_utilization,
            enforce_eager=self.settings.enforce_eager,
            quantization=self.settings.quantization,
            trust_remote_code=True,
        )

        self.engine = AsyncLLMEngine.from_engine_args(engine_args)
        logger.info("Inference engine initialized successfully")

    async def shutdown(self):
        """关闭引擎"""
        if self.engine:
            # vLLM 清理
            pass
        logger.info("Inference engine shut down")

    async def _get_request_id(self) -> str:
        """生成请求 ID"""
        async with self._lock:
            self._request_counter += 1
            return f"req_{self._request_counter}_{int(time.time()*1000)}"

    async def generate(
        self,
        prompt: str,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        top_k: Optional[int] = None,
        stop: Optional[List[str]] = None,
        **kwargs
    ) -> GenerationResult:
        """非流式生成"""
        request_id = await self._get_request_id()
        start_time = time.time()

        sampling_params = SamplingParams(
            max_tokens=max_tokens or self.settings.max_tokens,
            temperature=temperature or self.settings.temperature,
            top_p=top_p or self.settings.top_p,
            top_k=top_k or self.settings.top_k,
            stop=stop,
        )

        self.stats["active_requests"] += 1

        try:
            results_generator = self.engine.generate(
                prompt=prompt,
                sampling_params=sampling_params,
                request_id=request_id
            )

            final_output = None
            async for output in results_generator:
                final_output = output

            if final_output is None:
                raise RuntimeError("No output generated")

            output = final_output.outputs[0]
            latency_ms = (time.time() - start_time) * 1000

            # 更新统计
            self.stats["total_requests"] += 1
            self.stats["total_tokens_generated"] += len(output.token_ids)
            self.stats["total_latency_ms"] += latency_ms

            return GenerationResult(
                request_id=request_id,
                text=output.text,
                tokens=output.token_ids,
                finish_reason=output.finish_reason,
                prompt_tokens=len(final_output.prompt_token_ids),
                completion_tokens=len(output.token_ids),
                latency_ms=latency_ms
            )

        finally:
            self.stats["active_requests"] -= 1

    async def generate_stream(
        self,
        prompt: str,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        top_p: Optional[float] = None,
        top_k: Optional[int] = None,
        stop: Optional[List[str]] = None,
        **kwargs
    ) -> AsyncGenerator[StreamChunk, None]:
        """流式生成"""
        request_id = await self._get_request_id()

        sampling_params = SamplingParams(
            max_tokens=max_tokens or self.settings.max_tokens,
            temperature=temperature or self.settings.temperature,
            top_p=top_p or self.settings.top_p,
            top_k=top_k or self.settings.top_k,
            stop=stop,
        )

        self.stats["active_requests"] += 1
        previous_text = ""

        try:
            results_generator = self.engine.generate(
                prompt=prompt,
                sampling_params=sampling_params,
                request_id=request_id
            )

            async for output in results_generator:
                if output.outputs:
                    current_output = output.outputs[0]
                    current_text = current_output.text

                    # 计算增量文本
                    delta_text = current_text[len(previous_text):]
                    previous_text = current_text

                    if delta_text:
                        # 获取最新的 token id
                        token_id = current_output.token_ids[-1] if current_output.token_ids else 0

                        is_finished = current_output.finish_reason is not None

                        yield StreamChunk(
                            request_id=request_id,
                            text=delta_text,
                            token_id=token_id,
                            is_finished=is_finished,
                            finish_reason=current_output.finish_reason
                        )

                        if is_finished:
                            self.stats["total_tokens_generated"] += len(current_output.token_ids)

        finally:
            self.stats["active_requests"] -= 1
            self.stats["total_requests"] += 1

    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        stats = self.stats.copy()

        if stats["total_requests"] > 0:
            stats["avg_latency_ms"] = stats["total_latency_ms"] / stats["total_requests"]
            stats["avg_tokens_per_request"] = stats["total_tokens_generated"] / stats["total_requests"]
        else:
            stats["avg_latency_ms"] = 0
            stats["avg_tokens_per_request"] = 0

        return stats


# 全局引擎实例
_engine: Optional[InferenceEngine] = None


async def get_engine() -> InferenceEngine:
    """获取引擎实例"""
    global _engine
    if _engine is None:
        _engine = InferenceEngine()
        await _engine.initialize()
    return _engine

# src/schemas.py
"""
请求/响应模型
"""
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
from enum import Enum


class Role(str, Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"


class Message(BaseModel):
    """聊天消息"""
    role: Role
    content: str


class ChatCompletionRequest(BaseModel):
    """聊天补全请求 (OpenAI 兼容)"""
    model: str = Field(default="default")
    messages: List[Message]
    max_tokens: Optional[int] = Field(default=None, ge=1, le=4096)
    temperature: Optional[float] = Field(default=None, ge=0, le=2)
    top_p: Optional[float] = Field(default=None, ge=0, le=1)
    top_k: Optional[int] = Field(default=None, ge=1)
    stop: Optional[List[str]] = None
    stream: bool = False
    user: Optional[str] = None


class CompletionRequest(BaseModel):
    """文本补全请求"""
    model: str = Field(default="default")
    prompt: str
    max_tokens: Optional[int] = Field(default=None, ge=1, le=4096)
    temperature: Optional[float] = Field(default=None, ge=0, le=2)
    top_p: Optional[float] = Field(default=None, ge=0, le=1)
    stop: Optional[List[str]] = None
    stream: bool = False


class Usage(BaseModel):
    """使用量"""
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


class Choice(BaseModel):
    """生成选项"""
    index: int
    message: Message
    finish_reason: str


class ChatCompletionResponse(BaseModel):
    """聊天补全响应"""
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[Choice]
    usage: Usage


class DeltaMessage(BaseModel):
    """流式消息增量"""
    role: Optional[str] = None
    content: Optional[str] = None


class StreamChoice(BaseModel):
    """流式选项"""
    index: int
    delta: DeltaMessage
    finish_reason: Optional[str] = None


class ChatCompletionStreamResponse(BaseModel):
    """流式响应"""
    id: str
    object: str = "chat.completion.chunk"
    created: int
    model: str
    choices: List[StreamChoice]


class HealthResponse(BaseModel):
    """健康检查响应"""
    status: str
    model: str
    gpu_memory_used: Optional[float] = None
    active_requests: int

# src/api.py
"""
API 路由
"""
import time
import asyncio
from typing import AsyncGenerator
from fastapi import FastAPI, HTTPException, Depends, Header
from fastapi.responses import StreamingResponse
from sse_starlette.sse import EventSourceResponse
import json

from .schemas import (
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatCompletionStreamResponse,
    CompletionRequest,
    Message,
    Choice,
    StreamChoice,
    DeltaMessage,
    Usage,
    HealthResponse,
    Role
)
from .engine import get_engine, InferenceEngine
from .config import get_settings

app = FastAPI(title="LLM Inference Service", version="1.0.0")


def format_chat_prompt(messages: list[Message]) -> str:
    """格式化聊天消息为 prompt"""
    # Llama 2 Chat 格式
    prompt = ""
    system_message = None

    for msg in messages:
        if msg.role == Role.SYSTEM:
            system_message = msg.content
        elif msg.role == Role.USER:
            if system_message:
                prompt += f"<s>[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n{msg.content} [/INST]"
                system_message = None
            else:
                prompt += f"<s>[INST] {msg.content} [/INST]"
        elif msg.role == Role.ASSISTANT:
            prompt += f" {msg.content} </s>"

    return prompt


async def verify_api_key(authorization: str = Header(None)):
    """API Key 验证"""
    settings = get_settings()

    if not settings.api_key_enabled:
        return True

    if not authorization:
        raise HTTPException(status_code=401, detail="Missing Authorization header")

    # Bearer token 格式
    if not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Invalid Authorization format")

    api_key = authorization[7:]
    if api_key not in settings.api_keys:
        raise HTTPException(status_code=401, detail="Invalid API key")

    return True


@app.get("/health", response_model=HealthResponse)
async def health_check(engine: InferenceEngine = Depends(get_engine)):
    """健康检查"""
    stats = engine.get_stats()

    return HealthResponse(
        status="healthy",
        model=get_settings().model_name,
        active_requests=stats["active_requests"]
    )


@app.get("/v1/models")
async def list_models():
    """列出可用模型"""
    settings = get_settings()
    return {
        "object": "list",
        "data": [
            {
                "id": settings.model_name,
                "object": "model",
                "owned_by": "organization",
                "permission": []
            }
        ]
    }


@app.post("/v1/chat/completions")
async def chat_completions(
    request: ChatCompletionRequest,
    engine: InferenceEngine = Depends(get_engine),
    _: bool = Depends(verify_api_key)
):
    """聊天补全接口 (OpenAI 兼容)"""
    # 格式化 prompt
    prompt = format_chat_prompt(request.messages)

    if request.stream:
        return EventSourceResponse(
            stream_chat_response(request, prompt, engine),
            media_type="text/event-stream"
        )

    # 非流式响应
    result = await engine.generate(
        prompt=prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        top_k=request.top_k,
        stop=request.stop
    )

    return ChatCompletionResponse(
        id=result.request_id,
        created=int(time.time()),
        model=request.model,
        choices=[
            Choice(
                index=0,
                message=Message(role=Role.ASSISTANT, content=result.text),
                finish_reason=result.finish_reason
            )
        ],
        usage=Usage(
            prompt_tokens=result.prompt_tokens,
            completion_tokens=result.completion_tokens,
            total_tokens=result.prompt_tokens + result.completion_tokens
        )
    )


async def stream_chat_response(
    request: ChatCompletionRequest,
    prompt: str,
    engine: InferenceEngine
) -> AsyncGenerator[str, None]:
    """流式聊天响应"""
    request_id = f"chatcmpl-{int(time.time()*1000)}"

    # 发送角色
    first_chunk = ChatCompletionStreamResponse(
        id=request_id,
        created=int(time.time()),
        model=request.model,
        choices=[
            StreamChoice(
                index=0,
                delta=DeltaMessage(role="assistant"),
                finish_reason=None
            )
        ]
    )
    yield f"data: {first_chunk.model_dump_json()}\n\n"

    # 流式生成
    async for chunk in engine.generate_stream(
        prompt=prompt,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        top_p=request.top_p,
        top_k=request.top_k,
        stop=request.stop
    ):
        response = ChatCompletionStreamResponse(
            id=request_id,
            created=int(time.time()),
            model=request.model,
            choices=[
                StreamChoice(
                    index=0,
                    delta=DeltaMessage(content=chunk.text),
                    finish_reason=chunk.finish_reason if chunk.is_finished else None
                )
            ]
        )
        yield f"data: {response.model_dump_json()}\n\n"

        if chunk.is_finished:
            break

    yield "data: [DONE]\n\n"


@app.get("/metrics")
async def metrics(engine: InferenceEngine = Depends(get_engine)):
    """Prometheus 格式指标"""
    stats = engine.get_stats()

    metrics_text = f"""# HELP llm_requests_total Total number of requests
# TYPE llm_requests_total counter
llm_requests_total {stats['total_requests']}

# HELP llm_tokens_generated_total Total tokens generated
# TYPE llm_tokens_generated_total counter
llm_tokens_generated_total {stats['total_tokens_generated']}

# HELP llm_active_requests Current active requests
# TYPE llm_active_requests gauge
llm_active_requests {stats['active_requests']}

# HELP llm_avg_latency_ms Average latency in milliseconds
# TYPE llm_avg_latency_ms gauge
llm_avg_latency_ms {stats['avg_latency_ms']:.2f}
"""

    return metrics_text

# src/main.py
"""
应用入口
"""
import uvicorn
import asyncio
from contextlib import asynccontextmanager
from fastapi import FastAPI

from .api import app
from .engine import get_engine
from .config import get_settings


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    # 启动时初始化引擎
    engine = await get_engine()
    yield
    # 关闭时清理
    await engine.shutdown()


app.router.lifespan_context = lifespan


def main():
    settings = get_settings()

    uvicorn.run(
        "src.main:app",
        host=settings.host,
        port=settings.port,
        workers=settings.workers,
        log_level=settings.log_level.lower()
    )


if __name__ == "__main__":
    main()

# tests/test_api.py
"""
API 测试
"""
import pytest
import httpx
import asyncio

BASE_URL = "http://localhost:8000"


@pytest.mark.asyncio
async def test_health_check():
    """测试健康检查"""
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{BASE_URL}/health")
        assert response.status_code == 200
        data = response.json()
        assert data["status"] == "healthy"


@pytest.mark.asyncio
async def test_chat_completion():
    """测试聊天补全"""
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(
            f"{BASE_URL}/v1/chat/completions",
            json={
                "messages": [
                    {"role": "user", "content": "Hello, how are you?"}
                ],
                "max_tokens": 100
            }
        )
        assert response.status_code == 200
        data = response.json()
        assert "choices" in data
        assert len(data["choices"]) > 0


@pytest.mark.asyncio
async def test_streaming():
    """测试流式输出"""
    async with httpx.AsyncClient(timeout=60.0) as client:
        async with client.stream(
            "POST",
            f"{BASE_URL}/v1/chat/completions",
            json={
                "messages": [
                    {"role": "user", "content": "Count from 1 to 5."}
                ],
                "stream": True,
                "max_tokens": 50
            }
        ) as response:
            assert response.status_code == 200
            chunks = []
            async for line in response.aiter_lines():
                if line.startswith("data: "):
                    chunks.append(line)
            assert len(chunks) > 0

# tests/benchmark.py
"""
性能测试
"""
import asyncio
import httpx
import time
from statistics import mean, stdev

BASE_URL = "http://localhost:8000"
CONCURRENT_REQUESTS = 10
TOTAL_REQUESTS = 100


async def single_request(client: httpx.AsyncClient) -> dict:
    """单个请求"""
    start_time = time.time()

    response = await client.post(
        f"{BASE_URL}/v1/chat/completions",
        json={
            "messages": [
                {"role": "user", "content": "What is 2+2?"}
            ],
            "max_tokens": 50
        }
    )

    latency = (time.time() - start_time) * 1000
    data = response.json()

    return {
        "latency_ms": latency,
        "tokens": data["usage"]["completion_tokens"],
        "status": response.status_code
    }


async def run_benchmark():
    """运行基准测试"""
    async with httpx.AsyncClient(timeout=60.0) as client:
        semaphore = asyncio.Semaphore(CONCURRENT_REQUESTS)

        async def limited_request():
            async with semaphore:
                return await single_request(client)

        # 预热
        print("Warming up...")
        await asyncio.gather(*[single_request(client) for _ in range(5)])

        # 正式测试
        print(f"Running {TOTAL_REQUESTS} requests with {CONCURRENT_REQUESTS} concurrency...")
        start_time = time.time()

        results = await asyncio.gather(*[
            limited_request() for _ in range(TOTAL_REQUESTS)
        ])

        total_time = time.time() - start_time

        # 统计结果
        latencies = [r["latency_ms"] for r in results]
        tokens = sum(r["tokens"] for r in results)
        success = sum(1 for r in results if r["status"] == 200)

        print("\n=== Benchmark Results ===")
        print(f"Total Requests: {TOTAL_REQUESTS}")
        print(f"Successful: {success}")
        print(f"Total Time: {total_time:.2f}s")
        print(f"Throughput: {TOTAL_REQUESTS/total_time:.2f} req/s")
        print(f"Tokens Generated: {tokens}")
        print(f"Token Throughput: {tokens/total_time:.2f} tokens/s")
        print(f"\nLatency (ms):")
        print(f"  Mean: {mean(latencies):.2f}")
        print(f"  Std: {stdev(latencies):.2f}")
        print(f"  Min: {min(latencies):.2f}")
        print(f"  Max: {max(latencies):.2f}")
        print(f"  P50: {sorted(latencies)[len(latencies)//2]:.2f}")
        print(f"  P99: {sorted(latencies)[int(len(latencies)*0.99)]:.2f}")


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