Prhub

#44761 [ROCm][CI] Stabilizing teardown and timeout of flaky tests to prevent rare OOMs

原始 PR 作者 AndreasKaratzas 合并时间 2026-06-08 14:11 文件变更 4 提交数 1 评论 1 代码增减 +130 / -65

执行摘要

稳定 ROCm 测试清理流程,修复单位不匹配和超时问题

PR body 中指出:'several hybrid/APC and pooling tests can leave enough residual VRAM on gfx942 to trip later test startup or teardown checks, even though the model allocations have already been released',因此需要稳定测试的清理和 teardown 过程以防止罕见的 OOM。

值得精读,尤其是 tests/utils.py 中关于 ROCm 内存清理的阈值的注释,以及对跨平台 GPU 内存查询单位不一致的修复模式。此 PR 展示了如何系统性解决测试环境中的 GPU 资源泄漏问题,是 CI 稳定性的良好实践。

讨论亮点

无实质性讨论;仅 claude[bot] 自动评论(因 fork 无法运行),DarkLight1337 直接批准。

实现拆解

  1. 统一内存等待函数:在 tests/utils.py 中新增 wait_for_rocm_memory_to_settle 封装,内部调用 wait_for_gpu_memory_to_clear 并内置平台检查与合理的默认参数(4 GiB 最小阈值),避免各测试文件重复相同的等待逻辑。
  2. 修复 MiB/bytes 单位不一致_get_gpu_memory_used 中 ROCm 分支原直接返回 amdsmi 给出的 MiB 值,而 CUDA 分支返回 bytes,导致 _wait_for_gpu_memory_release 的字节比较永远提前满足。现在将 MiB 显式乘以 1024*1024 转为 bytes,使等待逻辑生效。
  3. 增强 HfRunner/VllmRunner 的 teardown:在 tests/conftest.py 中,HfRunner.__exit__VllmRunner.__exit__ 于清理后调用 wait_for_rocm_memory_to_settleVllmRunner.__exit__ 中为 engine_core.shutdown 设置了 60 秒超时(仅 ROCm),避免硬杀死导致的 VRAM 残留。同时将 VllmRunner._wait_for_rocm_memory_release 中的等待逻辑替换为对共享函数的调用。
  4. 简化 hybrid 测试清理:删除 test_hybrid.py 中的本地 _wait_for_rocm_memory_to_settle 函数,将其集成到共享工具中;移除 _owned_vLLM_runner 的 finally 等待块,因为清理责任已上移至高层的 runner teardown。
  5. ColBERT 测试增加上下文管理器:在 test_colbert.py 中新增 _hf_colbert_model 上下文管理器,确保 HF 模型和权重在测试退出时被显式删除并通过 cleanup_dist_env_and_memorywait_for_rocm_memory_to_settle 释放 GPU 内存,避免模块级 fixture 跨测试残留 VRAM。
文件 模块 状态 重要度
tests/utils.py 测试工具 modified 6.74
tests/conftest.py 测试配置 modified 5.94
tests/models/language/generation/test_hybrid.py 混合测试 modified 5.91
tests/models/language/pooling/test_colbert.py ColBERT 测试 modified 6.18

关键符号

wait_for_rocm_memory_to_settle wait_for_gpu_memory_to_clear RemoteOpenAIServer._get_gpu_memory_used _hf_colbert_model HfRunner.__exit__ VllmRunner.__exit__ VllmRunner._wait_for_rocm_memory_release

关键源码片段

tests/utils.py test-coverage

新增核心工具函数 `wait_for_rocm_memory_to_settle` 并修复单位转换 bug,是整个变更的基础。

def wait_for_rocm_memory_to_settle(threshold_ratio: float = 0.01) -> None:
    """Wait until ROCm GPU memory usage drops below a threshold ratio.    On ROCm, VRAM is freed asynchronously by the driver.  After a model is
    destroyed, the driver may keep the allocation resident for a short time.
    Calling this before starting a new test avoids an OOM on the startup
    guard.
    """
    from vllm.platforms import current_platform
    if not current_platform.is_rocm():
        return
    num_gpus = current_platform.device_count()
    if num_gpus == 0:
        return
    # Use a conservative 4 GiB minimum threshold because the ROCm driver
    # itself keeps ~2.5 GiB resident even when no model is loaded. A strict
    # 1% ratio would never be satisfied and would hang the test.
    wait_for_gpu_memory_to_clear(
        devices=list(range(num_gpus)),
        threshold_ratio=threshold_ratio,
        timeout_s=120,
    )
tests/conftest.py test-coverage

修改了 `HfRunner` 和 `VllmRunner` 的 teardown 逻辑,增加统一的清理等待和更长的 shutdown 超时。

def __exit__(self, exc_type, exc_value, traceback):
    from tests.utils import wait_for_rocm_memory_to_settle
    # 显式删除模型,然后等待 ROCm 异步释放 VRAM,
    # 避免下一个 runner 启动时因残留内存而导致 OOM。
    del self.model
    cleanup_dist_env_and_memory()
    wait_for_rocm_memory_to_settle()

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

所有变更仅限于测试基础设施,不影响产品代码。主要风险包括:

  • 新增的 60 秒 shutdown 超时可能让测试总耗时增加,但仅对 ROCm 生效;
  • 统一的内存等待函数如果阈值选择不当,可能导致测试超时或等待不足,但已有充分注释并基于经验值(4 GiB)设置;
  • 单位转换修复可能暴露之前被掩盖的其他问题,但这是正确的行为。

影响范围:仅 ROCm CI 上的 hybrid/APC/pooling 测试。影响程度:预期显著降低随机 OOM 导致的失败率,提升 CI 可靠性。对其他平台无影响。对用户无影响,团队可获得更稳定的测试反馈。

测试超时增加 阈值可能不适用所有 GPU

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论