代码之家 › 专栏 › 技术社区 › Stefan

带有池工作程序的Python多进程-内存使用优化

string-matching multiprocessing memory-management performance python

Stefan · 技术社区 · 12 年前

我有一个模糊的字符串匹配脚本,在400万个公司名称的大海捞针中寻找大约3万个针头。虽然脚本运行良好,但由于内存不足,我尝试在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 回复 | 直到 9 年前

abarnert 12 年前

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

这意味着 runMatch 只需要一个针线ID和针线公司,它所做的就是打电话 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()

或者,如果结果很小,那么不要让所有的进程(大概)为一些共享的存储结果而争吵,只需返回它们。那你就不需要了 运行匹配 总之,只是:

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() 可直接迭代。(事实上,在这种情况下,看起来您使用的是Python 2.x keys 和 values 是已经列表,所以你只是在没有充分理由的情况下额外复制一份。)

但在这种情况下,您可以进一步简化整个过程。你只是在寻找得分最高的键(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,所以这就是攻击ite的方法。例如,使用 shelve 数据库而不是普通数据库 dict .

也可以尝试转储内存使用情况(使用 resource 单元 getrusage(RUSAGE_SELF) )在初始值设定项函数的开始和结束处。如果最后的干草堆只有,比如说,0.3GB,但你又分配了1.3GB来构建它,那就是攻击的问题。例如,您可以派生一个子进程来构建和pickle dict,然后让池初始值设定项打开并取消pickle 搁置 db,并在初始值设定项中以只读方式打开它。无论哪种方式,这也意味着你只做一次CSV解析/dict构建工作,而不是8次。

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

第一 getsizeof 任务函数的参数和返回的值。如果它们很大,特别是如果它们随着每项任务的进行而不断变大,或者变化很大,那么可能只是因为数据占用了太多内存,最终其中8个数据的大小足以达到极限。

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