池工作者的Python多进程 – 内存使用优化

我有一个模糊的字符串匹配脚本,在400万公司名称的大海捞针中寻找大约30K针.虽然脚本工作正常,但我在AWS h1.xlarge上通过并行处理加速处理的尝试失败了,因为我的内存不足.

我不想试图获得更多的内存,如回应my previous question所述,我想找出如何优化工作流程 – 我对此很新,所以应该有足够的空间.顺便说一句,我已经尝试过queues(也工作过,但遇到了同样的MemoryError,再看了一堆非常有用的SO贡献,但还没到那里.

这是与代码最相关的内容.我希望它足以澄清逻辑 – 很高兴根据需要提供更多信息：

def getHayStack():
    ## loads a few million company names into id: name dict
    return hayCompanies

def getNeedles(*args):
    ## loads subset of 30K companies into id: name dict (for allocation to workers)
    return needleCompanies

def findNeedle(needle,haystack):
    """ Identify best match and return results with score """
    results = {}
    for hayID,hayCompany in haystack.iteritems():
        if not isnull(haystack[hayID]):
            results[hayID] = levi.setratio(needle.split(' '),hayCompany.split(' '))
    scores = list(results.values())
    resultIDs = list(results.keys())
    needleID = resultIDs[scores.index(max(scores))]
    return [needleID,haystack[needleID],max(scores)]

def runMatch(args):
    """ Execute findNeedle and process results for poolWorker batch"""
    batch,first = args
    last = first + batch
    hayCompanies = getHayStack()
    needleCompanies = getTargets(first,last)
    needles = defaultdict(list)
    current = first
    for needleID,needleCompany in needleCompanies.iteritems():
        current += 1
        needles[targetID] = findNeedle(needleCompany,hayCompanies)
    ## Then store results

if __name__ == '__main__':
    pool = Pool(processes = numProcesses)
    totalTargets = len(getTargets('all'))
    targetsPerBatch = totalTargets / numProcesses
    pool.map_async(runMatch,itertools.izip(itertools.repeat(targetsPerBatch),xrange(0,totalTargets,targetsPerBatch))).get(99999999)
    pool.close()
    pool.join()

所以我想问题是：我怎样才能避免为所有工人加载大海捞针 – 例如通过分享数据或采取不同的方法,例如将更大的干草堆划分为工人而不是针头？如何通过避免或消除混乱来改善内存使用？

最佳答案

你的设计有点令人困惑.您正在使用N个工作池,然后将M个工作分解为N个大小为M / N的任务.换句话说,如果你完全正确,那么你就是在工作流程之上构建的池之上模拟工作进程.为什么要这么麻烦？如果要使用进程,只需直接使用它们即可.或者,将池用作池,将每个作业作为自己的任务发送,并使用批处理功能以适当(和可调整)的方式批处理它们.

这意味着runMatch只需要一个needleID和needleCompany,它所做的就是调用findNeedle然后执行#Then store结果部分.然后主程序变得更简单：

if __name__ == '__main__':
    with Pool(processes=numProcesses) as pool:
        results = pool.map_async(runMatch,needleCompanies.iteritems(),chunkSize=NUMBER_TWEAKED_IN_TESTING).get()

或者,如果结果很小,而不是让所有进程(可能)都在争夺一些共享的结果存储事物,而只是返回它们.那么你根本就不需要runMatch,只需：

if __name__ == '__main__':
    with Pool(processes=numProcesses) as pool:
        for result in pool.imap_unordered(findNeedle,chunkSize=NUMBER_TWEAKED_IN_TESTING):
            # Store result

或者,如果您确实想要完成N个批次,只需为每个批次创建一个流程：

if __name__ == '__main__':
    totalTargets = len(getTargets('all'))
    targetsPerBatch = totalTargets / numProcesses
    processes = [Process(target=runMatch,args=(targetsPerBatch,targetsPerBatch))) 
                 for _ in range(numProcesses)]
    for p in processes:
        p.start()
    for p in processes:
        p.join()

此外,您似乎每次为每个任务调用getHayStack()(以及getNeedles).我不确定在同一时间最终获得这个实时的多个副本是多么容易,但考虑到它是迄今为止最大的数据结构,这将是我试图排除的第一件事.实际上,即使它不是内存使用问题,getHayStack也很容易成为一个重大的性能损失,除非你已经在进行某种缓存(例如,第一次将它显式存储在全局或可变的默认参数值中),然后只是使用它),所以它可能值得修复.

一次解决两个潜在问题的一种方法是在Pool构造函数中使用初始化程序：

def initPool():
    global _haystack
    _haystack = getHayStack()

def runMatch(args):
    global _haystack
    # ...
    hayCompanies = _haystack
    # ...

if __name__ == '__main__':
    pool = Pool(processes=numProcesses,initializer=initPool)
    # ...

接下来,我注意到您在多个地方显式生成列表,而实际上并不需要它们.例如：

scores = list(results.values())
resultIDs = list(results.keys())
needleID = resultIDs[scores.index(max(scores))]
return [needleID,max(scores)]

如果有一些结果,这是浪费;只需直接使用results.values()迭代. (事实上,看起来你正在使用Python 2.x,在这种情况下,键和值已经是列表,所以你只是在没有充分理由的情况下制作一个额外的副本.)

但在这种情况下,你可以进一步简化整个事情.你只是寻找得分最高的关键(resultID)和值(得分),对吧？所以：

needleID,score = max(results.items(),key=operator.itemgetter(1))
return [needleID,score]

这也消除了对得分的所有重复搜索,这应该节省一些CPU.

这可能无法直接解决内存问题,但应该可以更容易地进行调试和/或调整.

首先要尝试的是使用更小的批次 – 而不是input_size / cpu_count,尝试1.内存使用量是否下降？如果没有,我们已经排除了这一部分.

接下来,尝试sys.getsizeof(_haystack)并查看它的内容.如果它是1.6GB,那么你正在削减一些东西,试图将其他所有东西都压缩到0.4GB,这就是攻击它的方式 – 例如,使用shelve数据库而不是简单的dict.

还尝试在初始化函数的开始和结束时转储内存使用(使用resource模块,getrusage(RUSAGE_SELF)).如果最终的干草堆只有0.3GB,但你又分配了1.3GB,这就是攻击的问题.例如,你可以分离一个子进程来构建和pickle dict,然后让池初始化器打开它并取消它.或者在第一个子节点中组合两个构建一个搁置数据库,并在初始化程序中以只读方式打开它.无论哪种方式,这也意味着您只进行一次CSV解析/字典构建工作而不是8次.

另一方面,如果您的总VM使用率仍然很低(请注意,当第一个任务运行时,getrusage没有任何方式直接看到您的总VM大小-ru_maxrss通常是一个有用的近似值,特别是如果ru_nswap为0),问题在于任务本身.

首先,获取任务函数的参数和返回的值.如果它们很大,特别是如果它们要么随着每个任务变得越来越大或者变化很大,那么它可能只是腌制和解开数据会占用太多内存,最终它们中的8个一起大到足以达到极限.

否则,问题很可能出在任务函数本身.要么你有内存泄漏(你只能通过使用有缺陷的C扩展模块或ctypes进行真正的泄漏,但是如果你在调用之间保留任何引用,例如,在全局中,你可能只是永远保持事物不必要地),或者某些任务本身会占用太多内存.无论哪种方式,这应该是你可以通过拉出多处理并直接运行任务来更容易测试的东西,这样更容易调试.

池工作者的Python多进程 – 内存使用优化

相关推荐