如何在 multiprocessing.Pool 中使用信号处理器:

60 阅读2分钟

在使用 multiprocessing 来处理大量文件时,我尝试通过添加信号处理器来处理 ctrl-c 中断,以便程序能够优雅退出,但信号处理器并未按预期工作。

以下是示例代码:

import multiprocessing
from multiprocessing import Pool, Value
from time import sleep
import signal, sys

def f(x):
    p = multiprocessing.current_process()

    if exit_flag:
        s = '%s outer exit' %p.name
        sys.stdout.write(str(s)+'\n')
        return

    for i in range(10):
        if exit_flag:
            s = '%s inner exit' %p.name
            sys.stdout.write(str(s)+'\n')
            break
        s = '%s --- %s' % (i, p.name)
        sys.stdout.write(str(s)+'\n')
        sleep(1)

    if exit_flag:
        s = '%s --- %s final exit' % (i, p.name)
        sys.stdout.write(str(s)+'\n')

def handler(signum, frame):
    global exit_flag
    exit_flag = True

def attach_signal_handler():
    global exit_flag
    exit_flag = False
    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)

if __name__ == "__main__":
    # attach to main
    attach_signal_handler()
    # change to size = 2, range(6), make result more readable
    size = 2
    pool = Pool(size, attach_signal_handler)  

    for i in range(6):
        result = pool.apply_async(f, (i,))

    pool.close()
    pool.join()

    if result.successful():
        print 'successful'

当我按下 ctrl+c 时,它只会打印类似以下的内容:

0 --- PoolWorker-1
0 --- PoolWorker-2
1 --- PoolWorker-1
1 --- PoolWorker-2

(press ctrl +c)

PPoolWorker-2 outer exit
oolWorker-1 outer exit
PPoolWorker-2 outer exit
oolWorker-1 outer exit
successful

但我期望的结果是:

0 --- PoolWorker-1
0 --- PoolWorker-2
1 --- PoolWorker-1
1 --- PoolWorker-2

(press ctrl +c)

PoolWorker-1 inner exit
PoolWorker-1 final exit
PoolWorker-2 inner exit
PoolWorker-2 final exit

(first two processes end at this point)

PoolWorker-2 outer exit
PoolWorker-1 outer exit
PoolWorker-2 outer exit
PoolWorker-1 outer exit

(other four precesses end)    

successful

2、解决方案:

为了解决问题,我意识到需要注意信号处理器的安装时机。信号处理器应该在创建进程之前安装,而不是在创建进程之后。这是因为它需要在每个子进程中运行,以便它们能够响应信号。

修改后的代码如下:

import multiprocessing
from multiprocessing import Pool, Value
from time import sleep
import signal, sys

def f(x):
    p = multiprocessing.current_process()

    if exit_flag:
        s = '%s outer exit' %p.name
        sys.stdout.write(str(s)+'\n')
        return

    for i in range(10):
        if exit_flag:
            s = '%s inner exit' %p.name
            sys.stdout.write(str(s)+'\n')
            break
        s = '%s --- %s' % (i, p.name)
        sys.stdout.write(str(s)+'\n')
        sleep(1)

    if exit_flag:
        s = '%s --- %s final exit' % (i, p.name)
        sys.stdout.write(str(s)+'\n')

def handler(signum, frame):
    global exit_flag
    exit_flag = True

if __name__ == "__main__":
    # attach to main
    signal.signal(signal.SIGINT, handler)
    signal.signal(signal.SIGTERM, handler)
    # change to size = 2, range(6), make result more readable
    size = 2
    pool = Pool(size)  

    for i in range(6):
        result = pool.apply_async(f, (i,))

    pool.close()
    pool.join()

    if result.successful():
        print 'successful'

现在,当我按下 ctrl+c 时,它将打印预期的结果:

0 --- PoolWorker-1
0 --- PoolWorker-2
1 --- PoolWorker-1
1 --- PoolWorker-2

(press ctrl +c)

PoolWorker-1 inner exit
PoolWorker-1 final exit
PoolWorker-2 inner exit
PoolWorker-2 final exit

(first two processes end at this point)

PoolWorker-2 outer exit
PoolWorker-1 outer exit
PoolWorker-2 outer exit
PoolWorker-1 outer exit

(other four precesses end)    

successful