AIBrix OpenAI Batch API 端到端测试实战:从环境搭建到完整流程验证
【免费下载链接】aibrixCost-efficient and pluggable Infrastructure components for GenAI inference项目地址: https://gitcode.com/GitHub_Trending/ai/aibrix
本指南以 AIBrix 仓库中的端到端(E2E)测试文档为主体,讲解如何针对真实运行的 AIBrix 元数据服务(Metadata Service)验证 OpenAI Batch API 的完整批量推理工作流。文章覆盖测试环境准备(服务端口转发、S3/TOS 对象存储凭据生成)、pytest 运行方式、测试用例结构与断言逻辑,并结合 python/aibrix/aibrix/metadata/api/v1/batch.py 与 python/aibrix/aibrix/metadata/api/v1/files.py 等源码,深入剖析测试背后对应的服务端实现。读完本文,你将能够独立搭建 E2E 测试环境、运行并解读 Batch API 测试结果,并理解"上传文件 → 创建批量任务 → 轮询状态 → 下载结果"这一 OpenAI 兼容批量推理闭环的服务端状态机。
一、E2E 测试是什么
AIBrix 的端到端测试位于 python/aibrix/tests/e2e/,其特点是针对真实运行的服务实例进行验证,而非使用 mock 对象。当前该目录下的测试主体是 python/aibrix/tests/e2e/test_batch_api.py:
test_batch_api.py—— OpenAI Batch API 端点的端到端测试
该测试文件会向一个真实启动的 AIBrix 元数据服务发送 HTTP 请求,依次执行文件上传、批量任务创建、状态轮询、结果下载与批量列表校验,从而验证整套批量推理链路在真实存储后端与真实推理引擎下的可用性。这是文档明确说明的测试定位:"This directory contains end-to-end tests for Aibrix services that run against real service instances."
说明:原 README 中提及的
test_batch_api_service_availability()与test_batch_api_error_handling_real_service()属于文档描述的历史测试形态;从当前源码结构看,test_batch_api.py 实际包含test_batch_api_e2e_real_service(httpx 直连版)与test_openai_batch_api(OpenAI Python SDK 版)两个用例,下文均以实际源码为准展开。
二、前置条件:一个可访问的服务与一套存储凭据
2.1 运行中的 AIBrix 服务
测试默认连接http://localhost:8888/。如果你在 Kubernetes 集群中部署了元数据服务,最简单的方式是使用kubectl port-forward将服务端口暴露到本机:
kubectl -n envoy-gateway-system port-forward service/envoy-aibrix-system-aibrix-eg-903790dc 8888:80该命令把集群内envoy-gateway-system命名空间下的网关服务映射到本机8888端口。作为参考,AIBrix 元数据服务的实际部署清单位于 config/metadata/metadata.yaml,其中 Service 暴露端口为8090,Deployment 以--host 0.0.0.0 --port 8090启动aibrix_metadata进程,并通过--enable-k8s-support开启 Kubernetes 任务执行能力——生产环境通常还需要在网关层做路由与端口转发。
2.2 生成对象存储凭据
批量推理的输入文件与输出文件都需要持久化到对象存储(S3、TOS 或本地存储),因此测试前必须确保对象存储可访问、且元数据服务能读取相应凭据。以 S3 为例,在仓库的 Python 包根目录执行:
cd /path/to/aibrix/python/aibrix python -m scripts.generate_secrets s3 --bucket <your-bucket-name>该命令会读取你通过aws configure配置好的 AWS 凭据(访问密钥、区域),并在 Kubernetes 集群中创建对应的 Secret。其 CLI 实现在 python/aibrix/scripts/generate_secrets.py,支持四个子命令:
| 子命令 | 作用 | 关键参数 |
|---|---|---|
s3 | 创建 S3 凭据 Secret | --bucket/-b(必填)、--name、--namespace/-n |
tos | 创建 TOS 凭据 Secret | --bucket/-b(必填)、--name、--namespace/-n |
delete | 删除指定 Secret | secret_name(位置参数)、--namespace |
list | 列出命名空间内所有 Secret | --namespace |
命令示例(摘自 generate_secrets.py 的 help 文本):
# 使用默认名称创建 S3 Secret python -m scripts.generate_secrets s3 --bucket my-bucket # 自定义 Secret 名称 python -m scripts.generate_secrets s3 --bucket my-bucket --name my-s3-creds # 创建 TOS Secret(需要 TOS_* 环境变量) python -m scripts.generate_secrets tos --bucket my-tos-bucket # 删除与列出 python -m scripts.generate_secrets delete my-secret-name python -m scripts.generate_secrets list # --namespace 既可以放在子命令前,也可以放在子命令后 python -m scripts.generate_secrets --namespace my-namespace s3 --bucket my-bucket python -m scripts.generate_secrets s3 --bucket my-bucket --namespace my-namespace底层逻辑位于 python/aibrix/aibrix/metadata/secret_gen.py,SecretGenerator类通过boto3.Session().get_credentials()获取 AWS 凭据(区域缺省为us-east-1),将access_key、secret_key、region、bucket-name写入基于模板的 Kubernetes Secret:
- S3 模板 python/aibrix/aibrix/metadata/setting/s3_secret_template.yaml(
type: Opaque,默认名称aibrix-s3-credentials); - TOS 模板 python/aibrix/aibrix/metadata/setting/tos_secret_template.yaml(需要
TOS_ACCESS_KEY、TOS_SECRET_KEY、TOS_ENDPOINT、TOS_REGION四个环境变量)。
_encode_data()会将全部值做 Base64 编码后写入data字段,符合 Kubernetes Secret 的存储规范;创建前若同名 Secret 已存在会先删除再重建,避免残留冲突。
三、运行测试:四种 pytest 姿势
所有命令都需在python/aibrix目录下执行(即pyproject.toml所在目录)。测试依赖httpx、pytest、pytest-asyncio、openai等包,均已在 python/aibrix/pyproject.toml 中声明(如pytest = "^8.3.2"、pytest-asyncio = "^1.1.0"、openai = "^2.32.0")。
3.1 运行全部 E2E 测试
cd /path/to/aibrix/python/aibrix pytest tests/e2e/ -v3.2 仅运行 Batch API 测试
cd /path/to/aibrix/python/aibrix pytest tests/e2e/test_batch_api.py -v3.3 运行单个指定测试
cd /path/to/aibrix/python/aibrix pytest tests/e2e/test_batch_api.py::test_batch_api_e2e_real_service -v3.4 输出详细日志
cd /path/to/aibrix/python/aibrix pytest tests/e2e/test_batch_api.py -v -s-s会关闭 pytest 对 stdout 的捕获,测试中用print()输出的进度信息(如Step 1: Uploading batch input file...、Attempt 3: Status = in_progress)会实时显示,便于观察批量任务的状态迁移过程。
四、测试结构拆解:Fixture、用例与断言
4.1service_health:会话级健康检查 Fixture
test_batch_api.py 定义了一个scope="session"的 fixture:
@pytest.fixture(scope="session") def service_health(): """Fixture to check service health and skip tests if service is not available.""" base_url = "http://localhost:8888" print(f"🔍 Checking service health at {base_url}...") is_healthy = asyncio.run(check_service_health(base_url)) if not is_healthy: pytest.skip(f"Service at {base_url} is not available or healthy") print(f"✅ Service at {base_url} is healthy") return base_url要点:
- 会话级(session-scoped):整个测试会话只检查一次服务健康状态,避免每个用例都重复探测;
- 自动跳过:若服务不可用,直接
pytest.skip(),全部依赖该 fixture 的用例都会被标记为 SKIPPED 而非 FAILED,这正是"针对真实服务"的测试应有的容错行为; - 健康检查方式:
check_service_health()通过httpx.AsyncClient(timeout=10.0)发起GET {base_url}/v1/batches,只要返回200即认为服务健康(见 test_batch_api.py)。这里借用了"列出所有批量任务"这个读接口做通用可用性探测——README 的 API Endpoints 一节对此有明确说明:/v1/batches用于"General service availability check by list all batches"。
注意:README 描述 fixture 时提到"Tests
/healthzendpoint",但实际代码以GET /v1/batches作为健康检查路径,请以源码行为为准。
4.2 完整工作流测试:test_batch_api_e2e_real_service
这是核心用例(test_batch_api.py),用httpx.AsyncClient(timeout=60.0)直连服务,完整走一遍 OpenAI Batch API 的五个步骤:
- 上传输入文件(Files API):构造包含 3 条请求的 JSONL 数据,以
multipart/form-data形式POST /v1/files,purpose固定为batch。断言返回的object == "file"、purpose == "batch"、status == "uploaded",并取出id作为input_file_id; - 创建批量任务(Batch API):
POST /v1/batches,请求体为{"input_file_id": ..., "endpoint": "/v1/chat/completions", "completion_window": "24h"},断言返回的object == "batch"且input_file_id、endpoint回显一致,取得batch_id; - 轮询任务状态:每 5 秒
GET /v1/batches/{batch_id}一次,最多轮询 60 次(即最长等待 300 秒)。根据状态分流处理:completed→ 取出output_file_id,校验request_counts为total == 3、completed == 3、failed == 0;failed→ 读取errors字段并pytest.fail;cancelled/expired→ 直接判失败;scheduling、validating、in_progress、finalizing→ 预期的中间状态,继续等待;- 未知状态 → 打印告警继续轮询;
- 下载并校验输出(Files API):
GET /v1/files/{output_file_id}/content,将响应内容按 UTF-8 解码后交给verify_batch_output_content()校验; - 验证批量列表接口:
GET /v1/batches,断言object == "list"、data非空,且能在列表中找回本测试创建的batch_id且状态为completed。
4.3 输出内容校验:verify_batch_output_content
该函数(test_batch_api.py)逐行解析输出 JSONL,验证 OpenAI Batch 输出格式:
- 输出行数必须等于预期的请求数(此处为 3);
- 每条输出必须包含顶层字段
id、custom_id、response; custom_id必须按序匹配request-1、request-2、request-3;response内必须包含status_code、request_id、body,且status_code == 200;body内必须包含model、choices。
这一断言与 OpenAI Batch API 的官方输出契约保持一致:批量输出文件中的每一行都携带原始请求的custom_id,便于用户将结果与输入一一对应。
4.4 OpenAI SDK 版本:test_openai_batch_api
test_batch_api.py 提供了使用官方openaiPython SDK 的等价实现,逻辑与 httpx 版完全一致,但调用方式更贴近真实用户:
with OpenAI(base_url=f"{base_url}/v1", api_key="aibrix") as client: upload_result = client.files.create(file=f, purpose="batch") batch_result = client.batches.create( input_file_id=input_file_id, endpoint="/v1/chat/completions", completion_window="24h", ) status_result = client.batches.retrieve(batch_id=batch_id) output = client.files.content(output_file_id) list_result = client.batches.list()注意 SDK 客户端的base_url指向http://localhost:8888/v1,api_key任意填写(此处为"aibrix")。该用例验证了 AIBrix 服务与 OpenAI SDK 的兼容性——用户可以直接用官方 SDK 驱动 AIBrix 的批量推理。
4.5 测试输入数据的生成
generate_batch_input_data()(test_batch_api.py)用于构造 JSONL 批量请求,其中ENDPOINT_SAMPLE_BODIES定义了四类端点的样例请求体:
| 端点 | 样例模型 | 关键字段 |
|---|---|---|
/v1/chat/completions | gpt-3.5-turbo-0125 | messages、max_tokens: 1000 |
/v1/completions | gpt-3.5-turbo-0125 | prompt、max_tokens: 100 |
/v1/embeddings | text-embedding-ada-002 | input |
/v1/rerank | reranker-v1 | query、documents |
生成的数据格式为 OpenAI Batch 标准的 JSONL:
{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo-0125", "messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello world!"}], "max_tokens": 1000}}每条记录包含custom_id、method、url、body四个字段——这正是批量任务中每条独立请求的完整描述。
五、服务端实现:测试背后的 Batch API 与 Files API
E2E 测试并非孤立存在,它与元数据服务的实现一一对应。以下是从源码角度对测试覆盖面的印证。
5.1 Files API:/v1/files系列端点
实现在 python/aibrix/aibrix/metadata/api/v1/files.py:
POST /v1/files(create_file):校验文件扩展名(仅支持json、jsonl),通过Reader包裹上传内容并受settings.MAX_FILE_SIZE大小限制,生成 UUID 作为file_id,调用request.app.state.storage.put_object()写入对象存储,同时记录filename、purpose、created_at元数据。超限返回 413content_size_limit_exceeded;GET /v1/files/{file_id}/content(retrieve_file_content):从存储读取原始内容并以附件形式返回,未找到返回 404;GET /v1/files/{file_id}(retrieve_file_metadata):仅读取元数据(大小、类型、创建时间等),不下载内容;HEAD /v1/files/{file_id}:以 HTTP 头形式返回元数据(X-File-ID、X-File-Name、ETag等);DELETE /v1/files/{file_id}:删除文件,先head_object探测,不存在则返回 404;GET /v1/files(list_files):支持purpose过滤、limit(1–100,默认 20)与基于file_id的游标分页。
注意files.py源码注释明确说明:"The implementation is for e2e test only for now, and can upload batch input file only"—— 即当前 Files API 以支撑批量测试为主要目标。
5.2 Batch API:/v1/batches系列端点
实现在 python/aibrix/aibrix/metadata/api/v1/batch.py:
POST /v1/batches(create_batch):接收 OpenAI 形状的BatchSpec(input_file_id、endpoint、completion_window等),通过BatchSpec.newBatchJobSpec()转换为内部BatchJobSpec,交给request.app.state.batch_driver.create_job()创建任务,并返回 OpenAI 格式的BatchResponse;GET /v1/batches/{batch_id}(get_batch):查询任务详情,未找到返回 404;POST /v1/batches/{batch_id}/cancel(cancel_batch):取消任务,已处于终态时返回 409;GET /v1/batches(list_batches):支持after游标与limit(1–100,默认 20)分页,返回object == "list"的结构,与测试断言一致。
5.3 任务状态机:测试轮询逻辑的底层依据
测试中轮询所依赖的状态串(scheduling → validating → in_progress → finalizing → completed/failed/...)在服务端有明确映射,见_batch_job_to_openai_response()(batch.py)中的注释说明:
- 内部
CREATED状态对外呈现为scheduling(任务已接受、等待资源准入); VALIDATING到IN_PROGRESS几乎是瞬时的(在admit()内完成);FINALIZED是终态伞形状态,实际结果由status.condition决定:completed/failed/expired/cancelled;output_file_id仅在任务完成且有成功结果(rc.completed > 0)时才对外暴露,error_file_id仅在存在失败请求时暴露——这正是 OpenAI 的契约行为,测试在completed分支中对output_file_id与request_counts的断言与此一一对应。
5.4 批处理系统架构
从 python/aibrix/aibrix/batch/README.md 的架构说明可以了解测试背后的完整链路:
- BatchDriver:任务生命周期的主编排器;
- JobManager:任务状态管理与追踪;
- JobDriver:连接任务管理与推理执行(如 Kubernetes Job 运行时);
- BatchWorker:以 sidecar 模式运行在 Kubernetes Job 中的 worker 脚本,等待同 Pod 内 vLLM 引擎在
localhost:8000就绪后执行任务,直到FINALIZING状态后以退出码 0 结束。
元数据服务需要相应 RBAC 权限才能创建/轮询/删除 Kubernetes Job,见 config/metadata/metadata.yaml 中aibrix-metadata-service-readerClusterRole 对batch/jobs资源的get/list/create/patch/delete授权。
六、配置说明
测试的服务地址是硬编码在测试函数中的:service_healthfixture 中base_url = "http://localhost:8888",这是 README 明确指出的当前行为("The service URL is hardcoded in the test functions")。如果你希望指向其他地址或端口,需要修改 test_batch_api.py 中的base_url后重新运行。
七、预期输出解读
7.1 服务可用时的成功输出
tests/e2e/test_batch_api.py::test_batch_api_service_availability PASSED tests/e2e/test_batch_api.py::test_batch_api_e2e_real_service PASSED tests/e2e/test_batch_api.py::test_batch_api_error_handling_real_service PASSED以上是 README 记录的历史用例输出形态;当前源码对应的实际输出为
test_batch_api_e2e_real_service与test_batch_api_openai_batch等用例的 PASSED 结果。使用-v -s运行时,你还会看到诸如Step 1: Uploading batch input file...、✅ File uploaded successfully with ID: ...、Attempt N: Status = ...、🎉 E2E test completed successfully!等逐步日志。
7.2 服务不可用时的跳过输出
当localhost:8888上没有可用服务时,service_healthfixture 会跳过全部依赖它的用例:
tests/e2e/test_batch_api.py::test_batch_api_service_availability SKIPPED tests/e2e/test_batch_api.py::test_batch_api_e2e_real_service SKIPPED tests/e2e/test_batch_api.py::test_batch_api_error_handling_real_service SKIPPED这种"服务不可用即跳过而非失败"的设计,保证了 CI 或本机环境未启动服务时测试套件仍能快速、无害地退出,也提醒使用者:E2E 测试的结果有效性完全依赖于前置条件是否就绪。
八、排障建议
结合源码可以给出几个常见的排查方向:
- 服务不通:确认
kubectl port-forward仍在运行,且目标 Service 名称与命名空间正确;可用curl http://localhost:8888/v1/batches手工验证; - 存储凭据缺失:运行
python -m scripts.generate_secrets list检查 Secret 是否已创建;S3 场景先确认aws configure已配置有效凭据; - 批量任务长时间停留在
scheduling:从源码注释看,scheduling表示任务在等待资源准入,应检查集群资源与 Job 运行时是否正常(参考 batch.py 的状态说明); - 任务
failed:GET /v1/batches/{batch_id}返回的errors字段会给出具体错误,可结合 python/aibrix/aibrix/batch/README.md 中的调试章节查看 worker 与 vLLM 的 Pod 日志。
通过本文介绍的环境准备、运行方式与源码对应关系,你可以快速上手并深入理解 AIBrix 的 OpenAI Batch API 端到端测试,进而将其复用到自己的部署验证流程中。
【免费下载链接】aibrixCost-efficient and pluggable Infrastructure components for GenAI inference项目地址: https://gitcode.com/GitHub_Trending/ai/aibrix
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考