08 - 实战应用场景¶

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
import torch
from transformers import pipeline
from prometheus_client import Counter, Histogram, generate_latest
from fastapi.responses import Response
import time

# 创建FastAPI应用
app = FastAPI(
    title="AI Model API Service",
    description="AI模型推理API服务",
    version="1.0.0"
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Prometheus指标
request_counter = Counter('api_requests_total', 'Total API requests', ['endpoint', 'method'])
request_duration = Histogram('api_request_duration_seconds', 'API request duration', ['endpoint'])

# 加载模型
print("Loading AI models...")
classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
generator = pipeline("text-generation", model="gpt2")
print("Models loaded successfully!")

@app.get("/")
async def root():
    """根路径"""
    return {
        "message": "AI Model API Service",
        "version": "1.0.0",
        "endpoints": {
            "health": "/health",
            "classify": "/api/v1/classify",
            "generate": "/api/v1/generate",
            "metrics": "/metrics"
        }
    }

@app.get("/health")
async def health_check():
    """健康检查"""
    return {"status": "healthy"}

@app.post("/api/v1/classify")
async def classify_text(text: str):
    """文本分类"""
    start_time = time.time()
    request_counter.labels(endpoint="/api/v1/classify", method="POST").inc()

    try:
        result = classifier(text)
        duration = time.time() - start_time
        request_duration.labels(endpoint="/api/v1/classify").observe(duration)

        return {
            "text": text,
            "classification": result[0],
            "processing_time": duration
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/v1/generate")
async def generate_text(prompt: str, max_length: int = 100):
    """文本生成"""
    start_time = time.time()
    request_counter.labels(endpoint="/api/v1/generate", method="POST").inc()

    try:
        result = generator(prompt, max_length=max_length, num_return_sequences=1)
        duration = time.time() - start_time
        request_duration.labels(endpoint="/api/v1/generate").observe(duration)

        return {
            "prompt": prompt,
            "generated_text": result[0]["generated_text"],
            "processing_time": duration
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/metrics")
async def metrics():
    """Prometheus指标"""
    return Response(generate_latest(), media_type="text/plain")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

# 创建OSS配置文件
cat > app/oss_config.py << 'EOF'
import oss2

# OSS配置
OSS_ACCESS_KEY_ID = "your-access-key-id"
OSS_ACCESS_KEY_SECRET = "your-access-key-secret"
OSS_BUCKET_NAME = "ai-models-bucket"
OSS_ENDPOINT = "oss-cn-hangzhou.aliyuncs.com"

# 创建OSS客户端
auth = oss2.Auth(OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET)
bucket = oss2.Bucket(auth, OSS_ENDPOINT, OSS_BUCKET_NAME)

def upload_model(file_path, object_name):
    """上传模型文件到OSS"""
    bucket.put_object_from_file(object_name, file_path)
    print(f"模型文件上传成功: {object_name}")

def download_model(object_name, file_path):
    """从OSS下载模型文件"""
    bucket.get_object_to_file(object_name, file_path)
    print(f"模型文件下载成功: {object_name}")

def get_model_url(object_name, expires=3600):
    """获取模型文件的临时访问URL"""
    url = bucket.sign_url('GET', object_name, expires)
    return url
EOF

# 创建上传脚本
cat > upload_models.py << 'EOF'
from app.oss_config import upload_model
import os

# 模型文件列表
model_files = [
    ("models/text-classifier.pt", "models/text-classifier.pt"),
    ("models/text-generator.pt", "models/text-generator.pt"),
]

# 上传模型文件
for local_path, object_name in model_files:
    if os.path.exists(local_path):
        upload_model(local_path, object_name)
    else:
        print(f"文件不存在: {local_path}")
EOF

# 执行上传
python upload_models.py

# llm_service.py
from vllm import LLM, SamplingParams
from fastapi import FastAPI, HTTPException, Depends
from fastapi.responses import StreamingResponse
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel  # Pydantic数据验证模型
from typing import List, Optional, AsyncGenerator
import asyncio
import json
import redis
import time
from functools import wraps

# 初始化模型
print("Loading LLM model...")
llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",  # 或其他模型
    tensor_parallel_size=1,  # GPU数量
    gpu_memory_utilization=0.9,
    max_num_seqs=256,
)
print("Model loaded successfully!")

# FastAPI应用
app = FastAPI(title="LLM API Service")
security = HTTPBearer()

# Redis连接
redis_client = redis.Redis(host='localhost', port=6379, db=0)

# 数据模型
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    model: str = "llama-2-7b"
    messages: List[ChatMessage]
    temperature: float = 0.7
    max_tokens: int = 512
    stream: bool = False
    top_p: float = 1.0
    frequency_penalty: float = 0.0
    presence_penalty: float = 0.0

class ChatCompletionResponse(BaseModel):
    id: str
    object: str
    created: int
    model: str
    choices: List[dict]
    usage: dict

# 认证中间件
async def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """验证API Token"""
    token = credentials.credentials
    # 这里应该查询数据库或缓存验证token
    # 简化示例：检查token是否在Redis中
    user_id = redis_client.get(f"token:{token}")
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid token")
    return user_id.decode()

# 限流装饰器
def rate_limit(requests_per_minute: int = 60):
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):  # *args接收任意位置参数；**kwargs接收任意关键字参数
            # 获取用户ID（从依赖注入）
            user_id = kwargs.get('user_id')
            if not user_id:
                return await func(*args, **kwargs)

            # 检查限流
            key = f"rate_limit:{user_id}"
            current = redis_client.get(key)

            if current and int(current) >= requests_per_minute:
                raise HTTPException(status_code=429, detail="Rate limit exceeded")

            # 增加计数
            pipe = redis_client.pipeline()
            pipe.incr(key)
            pipe.expire(key, 60)
            pipe.execute()

            return await func(*args, **kwargs)
        return wrapper
    return decorator

# 格式化对话历史
def format_messages(messages: List[ChatMessage]) -> str:
    """将消息列表格式化为模型输入"""
    formatted = ""
    for msg in messages:
        if msg.role == "system":
            formatted += f"[INST] <<SYS>>\n{msg.content}\n<</SYS>>\n\n"
        elif msg.role == "user":
            formatted += f"{msg.content} [/INST] "
        elif msg.role == "assistant":
            formatted += f"{msg.content} </s><s>[INST] "
    return formatted

# 非流式生成
@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
@rate_limit(requests_per_minute=60)
async def chat_completion(
    request: ChatCompletionRequest,
    user_id: str = Depends(verify_token)
):
    """聊天补全接口（非流式）"""
    start_time = time.time()

    # 格式化输入
    prompt = format_messages(request.messages)

    # 设置采样参数
    sampling_params = SamplingParams(
        temperature=request.temperature,
        max_tokens=request.max_tokens,
        top_p=request.top_p,
        frequency_penalty=request.frequency_penalty,
        presence_penalty=request.presence_penalty,
    )

    # 生成回复
    outputs = llm.generate([prompt], sampling_params)
    generated_text = outputs[0].outputs[0].text

    # 计算token数（简化）
    prompt_tokens = len(prompt.split())
    completion_tokens = len(generated_text.split())

    return ChatCompletionResponse(
        id=f"chatcmpl-{int(time.time())}",
        object="chat.completion",
        created=int(time.time()),
        model=request.model,
        choices=[{
            "index": 0,
            "message": {
                "role": "assistant",
                "content": generated_text
            },
            "finish_reason": "stop"
        }],
        usage={
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": prompt_tokens + completion_tokens
        }
    )

# 流式生成
async def generate_stream(prompt: str, sampling_params: SamplingParams) -> AsyncGenerator[str, None]:
    """流式生成器"""
    # vLLM的流式生成
    stream = llm.generate([prompt], sampling_params, stream=True)

    for output in stream:
        text = output.outputs[0].text
        chunk = {
            "id": f"chatcmpl-{int(time.time())}",
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": "llama-2-7b",
            "choices": [{
                "index": 0,
                "delta": {"content": text},
                "finish_reason": None
            }]
        }
        yield f"data: {json.dumps(chunk)}\n\n"  # yield生成器：惰性产出值，节省内存
        await asyncio.sleep(0.01)  # 控制流速

    # 发送结束标记
    yield "data: [DONE]\n\n"

@app.post("/v1/chat/completions")
@rate_limit(requests_per_minute=60)
async def chat_completion_stream(
    request: ChatCompletionRequest,
    user_id: str = Depends(verify_token)
):
    """聊天补全接口（流式）"""
    if not request.stream:
        return await chat_completion(request, user_id)

    prompt = format_messages(request.messages)
    sampling_params = SamplingParams(
        temperature=request.temperature,
        max_tokens=request.max_tokens,
        top_p=request.top_p,
    )

    return StreamingResponse(
        generate_stream(prompt, sampling_params),
        media_type="text/event-stream"
    )

# 健康检查
@app.get("/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "model": "llama-2-7b",
        "gpu_available": True,
        "queue_length": 0
    }

# 模型信息
@app.get("/v1/models")
async def list_models():
    """列出可用模型"""
    return {
        "object": "list",
        "data": [
            {
                "id": "llama-2-7b",
                "object": "model",
                "created": 1677610602,
                "owned_by": "meta"
            }
        ]
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)