使用joblib多线程比单线程慢的问题排查与优化咨询

首先贴出原代码：

from tqdm import tqdm
from joblib import Parallel, delayed
import numpy as np
import time

def func(x,a,b,c):
    """x³+ax²+bx+c"""
    return x**3+a*x**2+b*x+c

def func_prime(x,a,b):
    """3x²+2ax+b"""
    return 3*x**2+2*a*x+b

def check(i):
    a,b,c = i
    if func_prime(-1/4, a,b) != -1/4:
        return False
    if func_prime(1/4,a,b) >=0:
        return False
    for i in range(-300,300):
        if func(i-1, a,b,c)*func(i+1,a,b,c)<0:
            return False
    return a,b,c

def gen():
    a=np.arange(-1, 1, 1/128)
    b=np.arange(-1, 1, 1/128)
    c=np.arange(-1, 1, 1/128)
    for i_a in a:
        for i_b in b:
            for i_c in c:
                yield i_a, i_b, i_c

# single thread. it spend 11 seconds. 
result = set()
for i in tqdm(gen()):
    result.update([check(i)])

# using joblib. it spend 313 seconds. 
start_time = time.time()
_ = Parallel(n_jobs=10)(delayed(check)(i) for i in gen())
end_time = time.time()
print(end_time - start_time)

1. 为何多线程执行速度更慢？ #

• GIL限制：Python的全局解释器锁（GIL）导致同一时间只有一个线程能执行Python字节码，CPU密集型任务的多线程本质是串行执行，还额外增加了线程切换的开销。
• 任务粒度太小：每个check任务的计算量很小，但joblib的delayed需要对每个任务进行序列化、分发、结果回收，这些额外开销远大于任务本身的计算时间，叠加后总耗时反而飙升。
• 数据传输开销：生成器gen产生的大量小数据在线程间传递，进一步拖慢了整体速度。

2. 如何优化才能让并行执行快于单线程？ #

• 改用多进程而非多线程：多进程可以绕过GIL，适合CPU密集型任务。修改joblib调用方式，指定多进程后端并设置批量处理，减少分发开销：

  _ = Parallel(n_jobs=10, backend='multiprocessing', batch_size=1000)(delayed(check)(i) for i in gen())
  

• 提前数学过滤减少任务量：通过前两个条件推导a和b的约束关系，直接砍掉大量无效组合：
  从func_prime(-1/4, a,b) = -1/4化简可得：b = a/2 - 7/16，只需生成满足该公式的a、b组合，不用遍历全部取值范围。
• 优化check函数：
  • 循环中一旦发现func(i-1)*func(i+1)<0就立即返回，不用跑完整个循环；
  • 用numpy向量化运算替代Python循环，加快计算速度；
  • 提前计算重复用到的变量，避免冗余计算。

3. 如何实现找到解后立即终止程序的提前停止功能？ #

单线程版本 #

直接在遍历过程中检查结果，找到有效解就立即退出：

result = None
for i in tqdm(gen()):
    res = check(i)
    if res:
        result = res
        break
print("找到解：", result)

多进程版本（支持提前终止） #

使用concurrent.futures.ProcessPoolExecutor配合as_completed，一旦找到第一个有效结果，就取消所有未完成的任务：

from concurrent.futures import ProcessPoolExecutor, as_completed

def main():
    executor = ProcessPoolExecutor(max_workers=10)
    futures = [executor.submit(check, i) for i in gen()]
    for future in as_completed(futures):
        res = future.result()
        if res:
            print("找到解：", res)
            # 取消所有未完成任务并终止进程池
            for f in futures:
                if not f.done():
                    f.cancel()
            executor.shutdown(wait=False)
            return

if __name__ == "__main__":
    main()

注意：多进程下需将执行代码放在if __name__ == "__main__"块中，避免子进程重复初始化。

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