vLLM 高性能推理系列——入门篇

from vllm import LLM, SamplingParams
import time

# 初始化模型（第一次会从 HuggingFace 下载）
print("Loading model...")
start = time.time()
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    trust_remote_code=True,  # Qwen 模型需要这个
    dtype="float16",         # 半精度，平衡速度和精度
)
print(f"Model loaded in {time.time() - start:.1f}s")

# 采样参数
sampling_params = SamplingParams(
    temperature=0.7,    # 控制随机性，0=确定性，1=高随机
    top_p=0.9,          # 核采样，只从概率累积达到 90% 的 token 中采样
    max_tokens=512,     # 最大生成长度
    stop=["<|im_end|>"] # 停止符
)

# 准备一批 prompt
prompts = [
    "用简单的语言解释量子计算",
    "写一首关于编程的俳句",
    "Python 的 GIL 是什么？为什么它是个问题？",
    "解释为什么天空是蓝色的",
    "如何向一个 5 岁小孩解释人工智能",
]

# 批量推理
print(f"\nProcessing {len(prompts)} prompts...")
start = time.time()
outputs = llm.generate(prompts, sampling_params)
elapsed = time.time() - start

# 统计
total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
print(f"\nGenerated {total_tokens} tokens in {elapsed:.1f}s")
print(f"Throughput: {total_tokens / elapsed:.1f} tokens/s")

# 打印结果
for output in outputs:
    print(f"\n{'='*60}")
    print(f"Prompt: {output.prompt[:50]}...")
    print(f"Response: {output.outputs[0].text[:200]}...")

from openai import OpenAI

# 只需要改 base_url，其他代码完全不变！
client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="not-needed"  # vLLM 默认不验证 key
)

# 和调用 GPT-4 一模一样的代码
response = client.chat.completions.create(
    model="Qwen/Qwen2.5-7B-Instruct",
    messages=[
        {"role": "system", "content": "你是一个资深的 Python 工程师"},
        {"role": "user", "content": "如何用 Python 实现一个高性能的 Web 爬虫？"}
    ],
    temperature=0.7,
    max_tokens=1024,
    stream=True  # 支持流式输出！
)

# 流式打印
for chunk in response:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

# benchmark.py
import asyncio
import aiohttp
import time

async def send_request(session, prompt):
    start = time.time()
    async with session.post(
        "http://localhost:8000/v1/chat/completions",
        json={
            "model": "Qwen/Qwen2.5-7B-Instruct",
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": 256
        }
    ) as response:
        result = await response.json()
        return time.time() - start

async def benchmark(concurrency=10, total_requests=100):
    prompts = ["用100字介绍Python语言"] * total_requests
    
    async with aiohttp.ClientSession() as session:
        start = time.time()
        tasks = [send_request(session, p) for p in prompts[:concurrency]]
        
        completed = 0
        latencies = []
        
        while tasks or completed < total_requests:
            done, tasks = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for task in done:
                latencies.append(task.result())
                completed += 1
                if completed < total_requests:
                    tasks.add(asyncio.create_task(
                        send_request(session, prompts[completed])
                    ))
        
        elapsed = time.time() - start
        
    print(f"Concurrency: {concurrency}")
    print(f"Total requests: {total_requests}")
    print(f"Total time: {elapsed:.1f}s")
    print(f"Throughput: {total_requests / elapsed:.1f} req/s")
    print(f"Avg latency: {sum(latencies) / len(latencies) * 1000:.0f}ms")
    print(f"P99 latency: {sorted(latencies)[int(len(latencies) * 0.99)] * 1000:.0f}ms")

asyncio.run(benchmark(concurrency=20, total_requests=200))

from langchain_openai import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings

# 加载向量数据库（假设已经构建好）
embeddings = HuggingFaceEmbeddings(
    model_name="BAAI/bge-small-zh-v1.5",
    model_kwargs={'device': 'cuda'}  # embedding 也用 GPU 加速
)
vectorstore = FAISS.load_local(
    "faiss_index", 
    embeddings,
    allow_dangerous_deserialization=True
)

# 连接 vLLM 后端
llm = ChatOpenAI(
    base_url="http://localhost:8000/v1",
    api_key="not-needed",
    model="Qwen/Qwen2.5-7B-Instruct",
    temperature=0.7,
    max_tokens=1024
)

# 自定义 prompt 模板
prompt_template = """你是一个专业的技术顾问。请根据以下参考资料回答用户的问题。
如果参考资料中没有相关信息，请诚实地说"我不知道"，不要编造答案。

参考资料：
{context}

用户问题：{question}

请用清晰、专业的语言回答："""

PROMPT = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "question"]
)

# 构建 RAG Chain
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=vectorstore.as_retriever(
        search_type="mmr",       # 最大边际相关性，避免重复
        search_kwargs={"k": 5}   # 检索 5 个相关片段
    ),
    chain_type_kwargs={"prompt": PROMPT},
    return_source_documents=True
)

# 查询
result = qa_chain.invoke({"query": "vLLM 的 PagedAttention 是如何工作的？"})

print("Answer:", result["result"])
print("\nSources:")
for doc in result["source_documents"]:
    print(f"  - {doc.metadata.get('source', 'unknown')}: {doc.page_content[:100]}...")

# monitor.py
import requests
import time
import smtplib
from email.mime.text import MIMEText

class VLLMMonitor:
    def __init__(self, endpoints, alert_email):
        self.endpoints = endpoints
        self.alert_email = alert_email
        self.healthy = {ep: True for ep in endpoints}
    
    def check_health(self, endpoint):
        try:
            resp = requests.get(f"{endpoint}/health", timeout=10)
            return resp.status_code == 200
        except:
            return False
    
    def check_metrics(self, endpoint):
        """检查关键指标"""
        try:
            resp = requests.get(f"{endpoint}/metrics", timeout=10)
            # 解析 Prometheus 格式的指标
            metrics = {}
            for line in resp.text.split('\n'):
                if line and not line.startswith('#'):
                    parts = line.split(' ')
                    if len(parts) == 2:
                        metrics[parts[0]] = float(parts[1])
            return metrics
        except:
            return None
    
    def send_alert(self, message):
        print(f"ALERT: {message}")
        # 这里可以接入钉钉、飞书、PagerDuty 等
    
    def run(self, interval=30):
        print(f"Monitoring {len(self.endpoints)} endpoints...")
        while True:
            for endpoint in self.endpoints:
                is_healthy = self.check_health(endpoint)
                
                # 状态变化时告警
                if is_healthy != self.healthy[endpoint]:
                    if is_healthy:
                        self.send_alert(f"✅ {endpoint} recovered")
                    else:
                        self.send_alert(f"🚨 {endpoint} is DOWN!")
                    self.healthy[endpoint] = is_healthy
                
                # 检查指标
                if is_healthy:
                    metrics = self.check_metrics(endpoint)
                    if metrics:
                        # GPU 显存使用率过高
                        gpu_util = metrics.get('vllm:gpu_cache_usage_perc', 0)
                        if gpu_util > 95:
                            self.send_alert(f"⚠️ {endpoint} GPU memory usage: {gpu_util:.1f}%")
                        
                        # 队列积压
                        pending = metrics.get('vllm:num_requests_waiting', 0)
                        if pending > 100:
                            self.send_alert(f"⚠️ {endpoint} request queue: {pending}")
            
            time.sleep(interval)

if __name__ == "__main__":
    monitor = VLLMMonitor(
        endpoints=[
            "http://localhost:8000",
            "http://localhost:8001",
        ],
        alert_email="[email protected]"
    )
    monitor.run()

SamplingParams(
    temperature=0.7,      # 太高会胡说八道，太低会重复
    top_p=0.9,            # 和 temperature 配合使用
    repetition_penalty=1.1,  # 避免重复
    presence_penalty=0.1,    # 鼓励新话题
)