执行摘要
- 一句话:NIXL EP 导入金丝雀测试,验证 CUDA 运行时链接一致性
- 推荐动作:建议阅读 test_nixl_imports.py 的实现,了解如何通过
ldd 命令在测试中验证动态链接库的依赖是否符合预期。这种金丝雀测试模式值得借鉴,可以在类似的多版本 CUDA 包冲突场景中复用。
功能与动机
Issue #42525 报告 nightly 镜像中 nixl_ep 导入失败,原因是 libcudart.so.12 缺失。根因是 nixl-cu12 和 nixl-cu13 包同时存在导致链接不一致。PR #42542 修复了 Docker 安装顺序,此 PR 添加测试覆盖以避免回归。
实现拆解
- 新建测试文件
tests/v1/kv_connector/nixl_integration/test_nixl_imports.py,提供 _print_distribution_version 辅助函数和 test_nixl_and_nixl_ep_imports 主测试函数。
- 测试函数中依次:打印
torch.version.cuda 和 nixl 包版本;使用 importlib.import_module 导入 nixl、nixl._api、nixl._bindings 和 nixl_ep;通过 ldd 检查 nixl_ep_cpp*.so 链接的 libcudart.so 是否与 torch.version.cuda 主版本一致。
- 修改
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh,在配置扫荡循环前加入 python3 -m pytest -s -x 执行该金丝雀测试,若失败则提前终止 CI 任务。
- 测试跳过条件:若
torch.version.cuda is None(非 CUDA 环境),则跳过整个测试。
关键文件:
tests/v1/kv_connector/nixl_integration/test_nixl_imports.py(模块 导入测试;类别 test;类型 test-coverage;符号 _print_distribution_version, test_nixl_and_nixl_ep_imports): 新增的金丝雀测试文件,实现了 NIXL 导入和 CUDA 运行时链接检查的核心逻辑。
tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh(模块 精度扫荡;类别 test;类型 test-coverage): 在配置精度扫荡测试入口处添加了金丝雀测试的调用,确保扫荡前先执行导入检查。
关键符号:test_nixl_and_nixl_ep_imports, _print_distribution_version
关键源码片段
tests/v1/kv_connector/nixl_integration/test_nixl_imports.py
新增的金丝雀测试文件,实现了 NIXL 导入和 CUDA 运行时链接检查的核心逻辑。
# SPDX-License-Identifier: Apache-2.0
"""NIXL 导入金丝雀测试,用于验证 CUDA wheel 选择一致性。"""
import importlib
import importlib.metadata as metadata
import pathlib
import subprocess
import sys
import pytest
import torch
def _print_distribution_version(package_name: str) -> None:
"""打印包版本,用于调试输出。"""
try:
version = metadata.version(package_name)
except metadata.PackageNotFoundError:
version = "not installed"
print(f"{package_name}: {version}")
@pytest.mark.skipif(torch.version.cuda is None, reason="CUDA NIXL EP canary")
def test_nixl_and_nixl_ep_imports() -> None:
"""验证核心 NIXL 和 NIXL EP 扩展都能正常导入。"""
print(f"torch cuda: {torch.version.cuda}")
# 打印 nixl 包的版本信息,帮助定位版本冲突
for package_name in ("nixl", "nixl-cu12", "nixl-cu13"):
_print_distribution_version(package_name)
# 导入 nixl 核心包,由 NixlConnector 使用
nixl = importlib.import_module("nixl")
print(f"nixl: {nixl.__file__}")
importlib.import_module("nixl._api")
importlib.import_module("nixl._bindings")
# 导入 nixl_ep 扩展,由 fused MoE 专家并行模块使用
nixl_ep = importlib.import_module("nixl_ep")
print(f"nixl_ep: {nixl_ep.__file__}")
# 找到 nixl_ep 的 C++ 扩展文件
assert nixl_ep.__file__ is not None
extension_dir = pathlib.Path(nixl_ep.__file__).parent
extension_files = sorted(extension_dir.glob("nixl_ep_cpp*.so"))
assert extension_files, f"No nixl_ep_cpp extension found in {extension_dir}"
# 使用 ldd 检查扩展链接的 CUDA 运行时版本
extension_file = extension_files[0]
completed = subprocess.run(
["ldd", str(extension_file)],
capture_output=True,
check=False,
text=True,
)
print(completed.stdout)
if completed.stderr:
print(completed.stderr, file=sys.stderr)
assert completed.returncode == 0
# 确保链接的 libcudart 与 PyTorch 的 CUDA 版本主号一致
if torch.version.cuda is not None:
cuda_major = torch.version.cuda.split(".", maxsplit=1)[0]
expected_cudart = f"libcudart.so.{cuda_major}"
assert expected_cudart in completed.stdout
assert f"{expected_cudart} => not found" not in completed.stdout
评论区精华
主要讨论来自 reviewer ovidiusm 对 pytestmark = pytest.mark.distributed 行的疑问(Is this needed?)。但最终 reviewer 并未要求删除该行,并给出了 Overall LGTM 的批准。其他 review 无实质反馈。
- pytestmark 装饰器是否需要 (question): 评论者最终批准 PR,未要求删除该行,因此保留。
风险与影响
- 风险:风险极低:变更仅涉及新增测试文件和 CI 脚本,不修改任何业务逻辑。需关注的是:测试依赖
ldd 命令,在非 Linux 环境无法运行,但已有 CUDA skip 条件;测试需导入 nixl 和 nixl_ep 包,若这些包未安装会直接失败,这正是意图所在。
- 影响:仅影响 CI 流程:在 NIXL 集成扫荡配置精度测试前增加一道快速检查,确保 nixl_ep 扩展链接的 CUDA 运行时与 PyTorch 一致。对用户运行无影响,也不会引入额外的部署依赖。
- 风险标记:ldd 外部依赖, CUDA 环境假设
关联脉络
- PR #42525 [Bug]: Regression: vllm/vllm-openai:nightly fails to import nixl_ep due to missing libcudart.so.12: 本 PR 旨在为此 bug 添加回归测试覆盖。
- PR #42542 [PD] Fix broken NIXL EP installation: 本 PR 为此修复的配套测试,确保修复后不会再次被破坏。
参与讨论