如何在Python/Numpy中自动化一个依赖于环境变量的BLAS基准测试？

Question

我需要帮助解决如何在Python中自动化基准测试的问题，涉及到IT技术。我正在测试通过numpy调用BLAS库时线程数量对性能的影响。在Linux环境中，OpenBLAS中的线程是通过环境变量OMP_NUM_THREADS控制的。我想进行一项测试，将OMP_NUM_THREADS从1逐步增加到最大值，在每个线程计数下测量程序运行时间，并最终处理所有线程计数的总时间。但遇到的问题是：虽然在Python中可以设置环境变量，但这只影响子进程或子shell。因此，我可以使用以下驱动代码正确地运行我的基准测试：

#!/usr/bin/env python                                                                                                     # driver script for thread test
import os

thread_set =[1,2,4,8,16]
for thread in thread_set:

    os.environ['OMP_NUM_THREADS']='{:d}'.format(thread)
    os.system("echo $OMP_NUM_THREADS")
    os.system("numpy_test")

以及numpy_test脚本：

#!/usr/bin/env python
#timing test for numpy dot product (using OpenBLAS)                                                      
#based on https://dev59.com/Hmgu5IYBdhLWcg3wRlL0
import sys
import timeit

setup = "import numpy; x = numpy.random.random((1000,1000))"
count = 5

t = timeit.Timer("numpy.dot(x, x.T)", setup=setup)
dot_time = t.timeit(count)/count
print("dot: {:7.3g} sec".format(dot_time))

然而，分析这个过程是一个非常手动的过程。

特别地，我无法将numpy_test中的dot_time值传回到我的外部包装函数，因此我无法以任何自动化方式分析测试结果。例如，我想绘制dot_time与线程数之间的关系图，或者评估dot_time/ 线程数是否恒定。

如果我尝试在Python实例内完全通过定义Python测试函数（避免上面的os.system()方法），然后在thread in thread_set循环内运行测试函数，则所有测试函数实例都会继承相同的OMP_NUM_THREADS值（即父Python shell的值）。因此，这个测试失败了：

#!/usr/bin/env python
#attempt at testing threads that doesn't work
#(always uses inherited value of OMP_NUM_THREADS)
import os

import sys
import timeit

def test_numpy():
    setup = "import numpy; x = numpy.random.random((1000,1000))"
    count = 5

    t = timeit.Timer("numpy.dot(x, x.T)", setup=setup)
    dot_time = t.timeit(count)/count
    print("dot: {:7.3g} sec".format(dot_time))
    return dot_time

thread_set =[1,2,4,8,16]
for thread in thread_set:
    os.environ['OMP_NUM_THREADS']='{:d}'.format(thread)
    os.system("echo $OMP_NUM_THREADS")
    time_to_run = test_numpy()
    print(time_to_run)

这种方法会失败，因为每个thread实例执行的时间都相同，test_numpy()总是继承父环境中OMP_NUM_THREADS的值，而不是通过os.environ()设置的值。然而，如果类似于这样的解决方案有效，那么我需要进行的分析将变得非常简单。

在真正的测试中，我将运行几千个排列，所以自动化至关重要。基于此，我希望能得到以下任一问题的答案:

1. 如何从子进程返回值(dot_time)？是否有更优雅的解决方案，而不是读/写文件？

2. 是否有更好的结构化方式来进行这种(依赖于环境变量)的测试?

提前感谢您的回答。

- bpbrown

Answer 1

您可以像这样做：

import subprocess

os.environ['OMP_NUM_THREADS'] = '{:d}'.format(thread)
proc = subprocess.Popen(["numpy_test"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout, stderr = proc.communicate()

然后您将在stdout中获得 numpy_test 脚本的输出。一般来说，我认为 subprocess.call 和 subprocess.Popen 优于 os.system 。

- Bi Rico

---

这个可以通过一些字符串解析来实现。在实现时需要注意：proc.comunicate() 中有拼写错误，应该是 proc.communicate()。 - bpbrown

Answer 2

如果你想从子进程得到输出，使用subprocess.check_output，例如，替换

。

os.system("numpy_test")

使用

dot_output = subprocess.check_output(["numpy_test"])
dot_time = ... # extract time from dot_output

- Ray