记一次Python多进程调优实践（自己实现并行KNN算法）

记一次Python多进程调优实践（自己实现并行KNN算法）

环境

• 服务器硬件：16核CPU，192GB内存
• 操作系统：Ubuntu 16.04
• Python版本：3.6.5
• numpy版本：1.14.3

为什么使用多进程和共享内存

工作中需要写一个knn算法，用numpy实现，数据量N过大（N ~= 3,400,000）。knn算法输入的特征矩阵有8G，knn的一个中间步骤是计算特征相似度矩阵（用于knn算法中选取相似度最大的k个topk），空间复杂度为O(N * N)，用float32存储的相似度矩阵内存占用达46T，这在计算上是不可行的，所以需要分batch操作，即，每个batch只计算一部分特征与所有特征的相似度，并求topk.
显然，这是一个计算密集型而非IO密集型的任务，Python的CPython实现由于全局解释器锁（GIL）的存在，只能使用多进程而非多线程来实现将任务并行地执行在多核上。所以，必须采用多进程。
随之而来的问题是，多进程中，每个进程都有一份数据的拷贝，如果将算法的输入：8G的特征矩阵在每个进程中拷贝一份，是非常消耗内存的。实际上，输入的特征在算法中是只读的，完全可以将特征矩阵作为共享内存，而且无需加锁进行同步，这样可以大大减少多进程的内存消耗，使得服务器192G内存可以运行更多的进程。

batch_size调优

batch_size参数是一个非常重要的待调优参数。在串行情况下，batch_size越大，由于numpy本身对于矩阵运算的优化，程序执行时间越短；但是在多进程情况下，过大的batch_size会导致每个进程的内存占用过大，从而导致需要使用Swp区作为虚拟内存，进程频繁产生缺页等待和磁盘IO，从而降低执行速度。因此，调优目的是：找到一个合适的batch_size参数，使得CPU与物理内存同时达到满载，此时可以近似认为达到最大执行速度
经过测试，单进程情况下batch_size与一个batch程序执行时间的关系如下：

batch_size	一个batch时间(秒)
10	1.6
100	12.2
1000	75.1

batch_size=1000

此时物理内存耗尽，频繁产生缺页等待，大量进程处于"D状态"（Uninterruptible sleep）,执行速度极慢：需要减小batch_size

batch_size=100

此时CPU几乎满载，物理内存仍有较多空闲：需要增大batch_size

batch_size=300

此时所有16个CPU满载，物理内存几乎完全占满，达到调优目的

代码

以下是knn算法多进程并行的一个实现（省略了输入输出），主要应用了进程池技术和共享内存。由于multiprocessing库实现上的一些问题，共享内存变量只能在进程池初始化（_pool_init）的时候作为全局变量（global _shm_feat