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

时间:2013-09-09 22:26:51

标签: python performance memory-management multiprocessing string-matching

我有一个模糊的字符串匹配脚本,在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()

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

1 个答案:

答案 0 :(得分:4)

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

这意味着runMatch只需要一个needleID和needleCompany,它所做的就是调用findNeedle,然后执行# Then store results部分的任何操作。然后主程序变得更加简单:

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, needleCompanies.iteritems(), 
                                          chunkSize=NUMBER_TWEAKED_IN_TESTING):
            # Store result

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

if __name__ == '__main__':
    totalTargets = len(getTargets('all'))
    targetsPerBatch = totalTargets / numProcesses
    processes = [Process(target=runMatch, 
                         args=(targetsPerBatch,
                               xrange(0, 
                                      totalTargets,
                                      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, haystack[needleID], max(scores)]

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

但在这种情况下,你可以将整个事情简化得更远。你只是寻找得分最高的关键(resultID)和值(得分),对吧?所以:

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

这也消除了score上的所有重复搜索,这应该可以节省一些CPU。


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

要尝试的第一件事就是使用更小的批次 - 而不是input_size / cpu_count,尝试1.内存使用量是否下降?如果没有,我们已经排除了这一部分。

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

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

另一方面,如果您的总VM使用率仍然很低(请注意,shelve无法直接查看您的总VM大小 - getrusage通常是一个有用的近似值,尤其是如果ru_maxrss为0),则在第一个任务运行时,问题在于任务本身。

首先,ru_nswap任务函数的参数和返回的值。如果它们很大,特别是如果它们要么随着每个任务变得越来越大或者变化很大,那么它可能只是腌制和去除数据需要太多的内存,最终它们中的8个一起大到足以达到极限。 / p>

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