Python 多进程 进程池

Python 多进程 进程池

在前文Python 多进程编程介绍了Python的多进程库multiprocessing中的进程池类Pool，但是后来使用的时候还是发现一些不好用的地方。这里我自己做了一个自定义的进程池 MultiProcessPool类来方便处理。

系统的进程池类 Pool

如下代码，所有的子进程返回结果都是保存在一个名为tasks的list中。

处理子进程返回结果的时候，是按照子进程在list中的顺序来处理的，因为get()是个阻塞函数。

如果子进程运行的时间不确定，会导致有些进程已经运行结束，但是等不到处理。

#!/usr/bin/python

from __future__ import print_function  # at top of module
from __future__ import division, unicode_literals, with_statement

import time
import os
from multiprocessing import Pool


def task_func(timeout):
    name = 'task[%d]' % os.getpid()
    print('%s start' % name)
    for i in range(timeout):
        time.sleep(1)
        print('%s waiting(%d/%d)' % (name, i, timeout))
    print('%s end' % name)
    # 通过return返回数据 
    return {'task': name, 'timeout': timeout}


def main():
    print('main start')

    # 创建进程池, 允许同时运行4个进程
    pool = Pool(4)

    # 创建6个任务， 并全部非阻塞启动（实际只有4个开始运行， 另外2个在等待）
    tasks = [pool.apply_async(task_func, args=(i,)) for i in range(6)]

    # 读取每个任务的返回结果
    for task in tasks:
        print(task.get())

    print('main end')


if __name__ == '__main__':
    main()

自定义的进程池 MultiProcessPool类

下面我自己做了在Pool类的基础上重新包装了一个进程池MultiProcessPool类，主要方法为：

• assign_task() 创建任务（不是立即创建子进程，仅仅保存入口函数和参数）
• apply_async() 运行所有任务（此时才批量创建子进程）
• join() 阻塞等待所有子进程运行结束（可以指定回调函数来处理子进程返回的结果）

MultiProcessPool类的代码如下

import multiprocessing


class MultiProcessPool():

    def __init__(self, ncpu=None):
        """create a pool with same size of CPUs, if ncpu is None"""
        self._pool = multiprocessing.Pool(processes=ncpu) # 系统的进程池
        self._tasks = {}  # 当前子进程的结果对象
        self._cmds = []  # 保存所有task的入口函数和参数

    def __len__(self):
        """how many task is runing"""
        return len(self._tasks)

    def __bool__(self):
        return True

    @property
    def processes(self):
        """how many task can be runing as same time"""
        return self._pool._processes

    def assign_task(self, task_func, args=None):
        """create a task (not run immediately, run by apply_async() later)"""
        item = task_func, args
        self._cmds.append(item)  # 已tuple的形式保存任务的函数入口和参数

    def apply_async(self):
        """start all task by async, use assign_task() to create task"""
        while self._cmds:
            cmd = self._cmds.pop()
            func, args = cmd
            if args and not isinstance(args, tuple):
                raise TypeError('args should be a tuple, not %s' % type(args))
            # 创建子进程， 存在dict中，通过ID索引方便操作
            self._tasks[id(cmd)] = self._pool.apply_async(func, args=args)
        return len(self._tasks)

    def join(self, task_handle=None, args=None, pre_handle=None, post_handle=None,
             wait_handle=None, wait_internal=0.5, timeout=None):
        """wait for all task done, if timeout is None"""
        t0 = time.time()
        if pre_handle:
            pre_handle(args)
        while self._tasks:
            # 遍历所有子进程， 找到已经运行完成子进程
            done = [tid for tid, task in self._tasks.items() if task.ready()]
            for tid in done:
                # 获取每一个完成的子进程运行结果
                ret = self._tasks[tid].get()
                del self._tasks[tid]  # 从子进程dict里删除
                # 调用用户的回调函数
                if task_handle:
                    task_handle(args, ret)
            time.sleep(wait_internal)
            if wait_handle:
                wait_handle(args)
            if timeout and time.time() - t0 > timeout:
                break  # 超时退出
        if post_handle:
            post_handle(args)
        return len(self._tasks)

MultiProcessPool 进程池会以周期性的（默认0.5秒）从所有任务中找到已经完成的任务，并进行处理。

使用的时候，如下

def main():
    # 回调函数， 保存任务运行结果
    def __handle_result(args, result):
        ret = args
        ret.update(result)
        return

    print('main start')
    # 创建进程池
    pool = MultiProcessPool()
    # 创建任务
    for i in range(6):
        pool.assign_task(task_func, args=(i,))
    # 启动任务
    pool.apply_async()
    # 等待所有任务结束们，并处理结果
    ret = {}
    pool.join(__handle_result, args=(ret,))

    return

if __name__ == '__main__':
    main()

看起来比之前的用法简单了很多。