当前位置：首页 > news >正文

vLLM-v0.17.1保姆级教程：vLLM + Airflow构建定时批量推理工作流

news 2026/4/29 21:05:43

vLLM-v0.17.1保姆级教程：vLLM + Airflow构建定时批量推理工作流

1. vLLM框架简介

vLLM是一个专为大型语言模型(LLM)设计的高性能推理和服务库，以其出色的吞吐量和易用性著称。这个项目最初由加州大学伯克利分校的天空计算实验室开发，现在已经发展成为一个活跃的社区驱动项目。

vLLM的核心优势在于其创新的内存管理技术PagedAttention，这项技术能够高效地管理注意力机制中的键值对内存，显著提升了推理效率。以下是vLLM的主要功能特点：

高效推理：采用连续批处理技术处理并发请求，结合优化的CUDA/HIP图实现快速模型执行
量化支持：全面支持GPTQ、AWQ、INT4、INT8和FP8等多种量化方式
性能优化：集成FlashAttention和FlashInfer等先进技术，提供优化的CUDA内核
高级功能：支持推测性解码和分块预填充等前沿技术

在易用性方面，vLLM提供了：

与HuggingFace模型的完美兼容
多种解码算法支持，包括并行采样和束搜索
分布式推理能力，支持张量并行和流水线并行
流式输出和OpenAI兼容的API服务器
广泛的硬件支持，涵盖NVIDIA/AMD/Intel等多种平台
前缀缓存和多LoRA支持等实用功能

2. 环境准备与安装

2.1 系统要求

在开始之前，请确保您的系统满足以下基本要求：

操作系统：Linux (推荐Ubuntu 20.04或更高版本)
Python版本：3.8或更高
GPU：NVIDIA GPU (推荐显存16GB以上)
CUDA版本：11.8或更高

2.2 安装vLLM

使用pip安装最新版vLLM非常简单：

pip install vllm==0.17.1

如果您需要使用特定功能，可以选择安装额外依赖：

# 安装带有AWQ支持的版本 pip install vllm[awq]==0.17.1 # 安装带有GPTQ支持的版本 pip install vllm[gptq]==0.17.1

2.3 安装Airflow

Apache Airflow是一个强大的工作流管理平台，我们将用它来调度批量推理任务：

pip install apache-airflow==2.7.1

初始化Airflow数据库：

airflow db init

3. 基础vLLM服务搭建

3.1 启动基础推理服务

首先，让我们启动一个基本的vLLM服务：

from vllm import LLM, SamplingParams # 初始化模型 llm = LLM(model="meta-llama/Llama-2-7b-chat-hf") # 定义采样参数 sampling_params = SamplingParams(temperature=0.8, top_p=0.95) # 执行推理 outputs = llm.generate(["AI的未来发展方向是什么？"], sampling_params) # 输出结果 print(outputs[0].outputs[0].text)

3.2 创建OpenAI兼容API

vLLM提供了与OpenAI兼容的API服务，可以这样启动：

python -m vllm.entrypoints.openai.api_server \ --model meta-llama/Llama-2-7b-chat-hf \ --port 8000

启动后，您可以通过curl测试API：

curl http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ -d '{ "model": "meta-llama/Llama-2-7b-chat-hf", "prompt": "AI的未来发展方向是什么？", "max_tokens": 100, "temperature": 0.8 }'

4. 构建Airflow定时批量推理工作流

4.1 设计工作流架构

我们将构建一个完整的批量推理系统，包含以下组件：

任务调度器：Airflow DAG定时触发批量推理任务
任务队列：Redis或RabbitMQ管理待处理请求
推理工作节点：vLLM实例处理实际推理请求
结果存储：数据库或文件系统保存推理结果

4.2 创建Airflow DAG

在Airflow的dags目录下创建vllm_batch_inference.py：

from datetime import datetime, timedelta from airflow import DAG from airflow.operators.python import PythonOperator default_args = { 'owner': 'airflow', 'depends_on_past': False, 'email_on_failure': False, 'email_on_retry': False, 'retries': 1, 'retry_delay': timedelta(minutes=5), } def prepare_batch(**context): # 这里可以从数据库或文件中读取批量推理任务 prompts = [ "解释深度学习的基本概念", "写一篇关于机器学习未来发展的短文", "用简单的语言说明神经网络如何工作" ] context['ti'].xcom_push(key='prompts', value=prompts) def run_vllm_inference(**context): from vllm import LLM, SamplingParams prompts = context['ti'].xcom_pull(key='prompts') llm = LLM(model="meta-llama/Llama-2-7b-chat-hf") sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=256) outputs = llm.generate(prompts, sampling_params) results = [output.outputs[0].text for output in outputs] context['ti'].xcom_push(key='results', value=results) return results def save_results(**context): results = context['ti'].xcom_pull(key='results') # 这里可以将结果保存到数据库或文件系统 for i, result in enumerate(results): print(f"Result {i+1}: {result}") with DAG( 'vllm_batch_inference', default_args=default_args, description='定时批量LLM推理工作流', schedule_interval=timedelta(hours=6), start_date=datetime(2023, 1, 1), catchup=False, ) as dag: prepare_task = PythonOperator( task_id='prepare_batch', python_callable=prepare_batch, provide_context=True, ) inference_task = PythonOperator( task_id='run_vllm_inference', python_callable=run_vllm_inference, provide_context=True, ) save_task = PythonOperator( task_id='save_results', python_callable=save_results, provide_context=True, ) prepare_task >> inference_task >> save_task

4.3 扩展为分布式工作流

对于大规模批量推理，我们可以将工作流扩展为分布式架构：

任务分发：使用Airflow的CeleryExecutor分发任务到多个工作节点
并行处理：每个工作节点运行独立的vLLM实例
结果聚合：将所有节点的结果汇总到中心存储

修改后的DAG可能包含并行任务分支：

from airflow.utils.task_group import TaskGroup # ... (前面的代码保持不变) with DAG( # ... (DAG参数保持不变) ) as dag: prepare_task = PythonOperator( task_id='prepare_batch', python_callable=prepare_batch, provide_context=True, ) with TaskGroup('parallel_inference') as inference_group: # 创建多个并行推理任务 for i in range(4): # 假设有4个工作节点 inference_task = PythonOperator( task_id=f'run_vllm_inference_{i}', python_callable=run_vllm_inference, provide_context=True, op_kwargs={'worker_id': i}, ) aggregate_task = PythonOperator( task_id='aggregate_results', python_callable=aggregate_results, provide_context=True, ) save_task = PythonOperator( task_id='save_results', python_callable=save_results, provide_context=True, ) prepare_task >> inference_group >> aggregate_task >> save_task

5. 性能优化与最佳实践

5.1 vLLM性能调优

为了获得最佳性能，可以考虑以下优化措施：

批处理大小：调整max_num_seqs参数找到最佳批处理大小
内存管理：合理设置gpu_memory_utilization(默认0.9)
量化模型：使用GPTQ或AWQ量化减少内存占用
张量并行：对于大模型，使用tensor_parallel_size参数实现多GPU并行

示例优化配置：

llm = LLM( model="meta-llama/Llama-2-7b-chat-hf", tensor_parallel_size=2, # 使用2个GPU gpu_memory_utilization=0.85, max_num_seqs=64, quantization="awq", )