如何在CMake文件中设置CUDA_VISIBLE_DEVICES以利用多GPU执行测试

一、解决CUDA_VISIBLE_DEVICES在CMake中不生效的问题 #

你之前在py_test函数中设置环境变量无效，大概率是因为自定义的py_test没有将环境变量绑定到测试的属性上。正确的做法是通过CMake的set_tests_properties接口为测试单独设置环境变量，或者修改py_test函数使其支持GPU参数传递：

1. 直接为单个测试设置GPU #

如果是通过add_test定义测试，添加完测试后立即设置环境变量：

add_test(NAME test_demo COMMAND python test_demo.py)
# 指定该测试运行在GPU 1上
set_tests_properties(test_demo PROPERTIES ENVIRONMENT "CUDA_VISIBLE_DEVICES=1")

2. 修改自定义py_test函数支持GPU指定 #

如果你用的是项目自定义的py_test函数，可以修改它以接受GPU参数并自动设置环境变量：

function(py_test TEST_NAME)
    # 解析参数，提取GPU ID（默认0）
    set(GPU_ID 0)
    set(PY_ARGS "")
    foreach(arg ${ARGN})
        if(arg STREQUAL "GPU")
            list(GET ARGN ${CMAKE_CURRENT_LIST_INDEX}+1 GPU_ID)
        else()
            list(APPEND PY_ARGS ${arg})
        endif()
    endforeach()

    # 添加测试命令
    add_test(NAME ${TEST_NAME} COMMAND python ${PY_ARGS})
    # 绑定GPU环境变量
    set_tests_properties(${TEST_NAME} PROPERTIES ENVIRONMENT "CUDA_VISIBLE_DEVICES=${GPU_ID}")
endfunction()

调用时直接指定GPU：

py_test(test_model model_test.py GPU 2)
py_test(test_utils utils_test.py GPU 3)

二、实现多GPU充分利用的两种方案 #

1. 静态预分配GPU（简单直接） #

如果测试用例数量不多（≤GPU数量），可以给每个测试固定分配一块GPU，然后通过ctest -j N（N为GPU数量）并行运行测试：

# 给4个测试分别分配GPU 0-3
py_test(test_0 test_0.py GPU 0)
py_test(test_1 test_1.py GPU 1)
py_test(test_2 test_2.py GPU 2)
py_test(test_3 test_3.py GPU 3)

运行时执行：

cmake .. -DUSE_CUDA=ON && make && ctest -j4

这样4块GPU会同时工作，避免单GPU闲置。

2. 动态分配GPU（灵活适配多测试场景） #

如果测试用例远多于GPU数量，或者需要自动调度空闲GPU，可以写一个简单的Shell脚本（run_test_with_gpu.sh）来检测空闲GPU并设置环境变量：

#!/bin/bash
# 查找显存占用最低的GPU ID
GPU_ID=$(nvidia-smi --query-gpu=memory.free,index --format=csv,noheader,nounits | sort -nr | head -n1 | awk '{print $2}')
# 设置环境变量并运行测试
CUDA_VISIBLE_DEVICES=${GPU_ID} "$@"

然后在CMake中调用这个脚本运行测试：

add_test(NAME test_dynamic COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/run_test_with_gpu.sh python test_dynamic.py)
# 给所有测试都用这个脚本包装
function(py_test_dynamic TEST_NAME)
    add_test(NAME ${TEST_NAME} COMMAND ${CMAKE_CURRENT_SOURCE_DIR}/run_test_with_gpu.sh python ${ARGN})
endfunction()

这种方式会自动将测试分配到当前最空闲的GPU，最大化资源利用率。

三、预分配vs动态分配的选择 #

• 预分配：适合测试数量少、每个测试需要独占GPU资源的场景，配置简单，无资源竞争风险，便于调试单个测试的GPU问题。
• 动态分配：适合测试数量多、GPU资源需要灵活调度的场景，能自动适配不同负载，但需要额外的脚本逻辑，且调试时需要注意GPU分配的随机性。

内容的提问来源于stack exchange，提问作者Johnzy