如何在Numba+CUDA的Python并行程序中实现进度条？

问题背景 #

在Windows系统上使用Numba+CUDA运行耗时并行程序时，需要控制台实时显示进度条。尝试过tqdm、numba-progress均无法兼容CUDA环境，自行编写的基于类的方案因无法将类传入核函数失败。

原示例代码：

from __future__ import print_function, absolute_import

from numba import cuda
from numba.cuda.random import create_xoroshiro128p_states, xoroshiro128p_uniform_float32
import numpy as np
from math import gamma, exp, ceil


# 目标是给这个函数加进度条
@cuda.jit
def generate_samples(rng_states, out, rate):
    thread_id = cuda.grid(1)

    def poission_sample(rate, random_number): 
        probability_sum = 0
        index = -1
        while probability_sum < random_number:
            index += 1
            probability_sum += ((rate**index)/gamma(index+1)) * exp(-rate)
            
        return index
    
    # 希望在这里能更新全局进度计数器
    out[thread_id] = poission_sample(rate, xoroshiro128p_uniform_float32(rng_states, thread_id))

number_of_samples = 10000000

threads_per_block = 512
blocks = ceil(number_of_samples/threads_per_block)
rng_states = create_xoroshiro128p_states(threads_per_block * blocks, seed=1)
out = np.zeros(threads_per_block * blocks, dtype=np.float32)

generate_samples[blocks, threads_per_block](rng_states, out, 5)
    
print('Average Sample:', out.mean())

解决方案 #

核心思路是设备端原子计数器+主机端轮询更新：用CUDA原子操作让每个线程完成任务后递增设备上的计数器，主机端定期读取计数器值并刷新控制台进度条。

修改后的完整代码 #

from __future__ import print_function, absolute_import

from numba import cuda
from numba.cuda.random import create_xoroshiro128p_states, xoroshiro128p_uniform_float32
import numpy as np
from math import gamma, exp, ceil
import time


@cuda.jit
def generate_samples(rng_states, out, rate, progress_counter):
    thread_id = cuda.grid(1)
    # 跳过超出样本数的线程，避免无效计数
    if thread_id >= out.shape[0]:
        return

    def poission_sample(rate, random_number): 
        probability_sum = 0
        index = -1
        while probability_sum < random_number:
            index += 1
            probability_sum += ((rate**index)/gamma(index+1)) * exp(-rate)
            
        return index
    
    out[thread_id] = poission_sample(rate, xoroshiro128p_uniform_float32(rng_states, thread_id))
    
    # 原子递增计数器，保证多线程操作安全
    cuda.atomic.add(progress_counter, 0, 1)


def update_progress(current, total, bar_length=50):
    # 自定义控制台进度条，用回车覆盖实现实时更新
    percent = float(current) / total
    filled_length = int(round(bar_length * percent))
    bar = '=' * filled_length + '-' * (bar_length - filled_length)
    print(f'\rProgress: [{bar}] {percent:.1%} ({current}/{total})', end='', flush=True)


number_of_samples = 10000000

threads_per_block = 512
blocks = ceil(number_of_samples/threads_per_block)
total_threads = threads_per_block * blocks

# 创建设备端计数器（数组形式，因为原子操作需要内存地址）
progress_counter = cuda.to_device(np.array([0], dtype=np.int32))
rng_states = create_xoroshiro128p_states(total_threads, seed=1)
out = np.zeros(total_threads, dtype=np.float32)

# 异步启动CUDA核函数
generate_samples[blocks, threads_per_block](rng_states, out, 5, progress_counter)

# 主机端轮询更新进度条
current_progress = 0
while current_progress < number_of_samples:
    # 从设备端拷贝当前计数到主机
    host_counter = progress_counter.copy_to_host()
    current_progress = host_counter[0]
    update_progress(current_progress, number_of_samples)
    time.sleep(0.1)  # 控制刷新频率，避免CPU占用过高

# 确保核函数完全执行完毕
cuda.synchronize()
print('\nAverage Sample:', out.mean())

关键细节说明 #

1. 设备端计数器：必须用数组形式创建，因为Numba CUDA的原子操作需要操作内存地址，单个变量无法直接使用cuda.atomic.add。
2. 原子操作：cuda.atomic.add保证多个线程同时更新计数器时不会出现数据竞争，计数结果准确。
3. 异步执行+轮询：CUDA核函数启动后是异步运行的，主机端可以立即进入循环读取计数器，通过copy_to_host获取实时进度。
4. 线程边界检查：添加if thread_id >= out.shape[0]: return，防止超出样本数的线程执行无效操作并错误更新计数器。

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