Python 中用 subprocess 管理子进程

Python 中用 subprocess 管理子进程

由 Python 所启动的子进程可以平行地运行，这能够充分利用计算机的每一个 CPU 核心，来尽量提升程序的处理效率。虽然 Python 解释器本身有可能会局限在一个 CPU 上面，但还是很容易就能通过 Python 来驱动并协调 CPU 密集型（CPU-intensive）的任务。

Python 里面有许多方式都可以运行子进程（例如 os.popen 函数以及 os.exec* 系列的函数），其中最好的办法是通过内置的 subprocess 模块来管理。

下面就通过 run 函数启动一条进程，然后确认该进程已经正常终止，最后打印它的输出值。

import subprocess

result = subprocess.run(
    ['echo', 'Hello from the child!'],
    capture_output=True,
    encoding='utf-8')

result.check_returncode()  # No exception means clean exit
print(result.stdout)
# >>>
# Hello from the child!

子进程可以独立于父进程而运行，这里的父进程指 Python 解释器所在的那条进程。假如刚才那条子进程不是通过 run 函数启动，而是由 Popen 类启动的，那么就可以在它启动之后，让 Python 程序去做别的任务，每做一段时间就来查询一次子进程的状态以决定要不要继续执行任务。

proc = subprocess.Popen(['sleep', '1'])
while proc.poll() is None:
    print('Working...')
    # Some time-consuming work here
    ...
print('Exit status', proc.poll())
# >>>
# Working...
# Working...
# Working...
# Working...
# Exit status 0

把子进程从父进程中剥离，可以让程序平行地运行多条子进程。例如，可以像下面这样，先把需要运行的这些子进程用 Popen 启动起来。

import time

start = time.time()
sleep_procs = []
for _ in range(10):
    proc = subprocess.Popen(['sleep', '1'])
    sleep_procs.append(proc)

然后，在主进程里调用每条子进程的 communicate 方法，等待这条子进程把它的 I/O 工作处理完毕。

for proc in sleep_procs:
    proc.communicate()

end = time.time()
delta = end - start
print(f'Finished in {delta:.3} seconds')

# >>>
# Finished in 1.05 seconds

从统计结果可以看出，这 10 条子进程确实表现出了平行的效果。假如它们是按顺序执行的，那么程序耗费的总时长至少应该是 10 秒，而不是现在看到的 1 秒左右。

还可以在 Python 程序里面把数据通过管道发送给子进程所运行的外部命令，然后将那条命令的输出结果获取到 Python 程序之中。而且，在执行外部命令的这个环节中，可以平行地运行多条命令。例如，要用 oepnssl 这样的命令行工具来加密数据。首先以适当的命令行参数构建一批子进程，并配置好相应的 I/O 管道，这在 Python 里很容易就能做到。

import os

def run_encrypt(data):
    env = os.environ.copy()

    env['password'] = 'zf7ShyBhZOraQDdE/FiZpm/m/8f9X+M1'
    proc = subprocess.Popen(
        ['openssl', 'enc', '-des3', '-pass', 'env:password'],
        env=env,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE)
    proc.stdin.write(data)
    proc.stdin.flush()  # Ensure that the child gets input
    return proc

为了演示这个程序，随机生成一些字节数据，并传给刚才的加密函数去处理。在实际工作中，输入给管道的那些数据可能并不是随机生成的，而是由用户所输入，或是从文件句柄（file handle）与网络套接字（network socket）等渠道获取的。

procs = []
for _ in range(3):
    data = os.urandom(10)
    proc = run_encrypt(data)
    procs.append(proc)

这些子进程能够平行地运行，并各自处理它所接到的那份数据。现在逐个等待每一条子进程完工，然后获取并打印这条子进程的加密结果。可以看到，这样写确实能够对随机生成的数据进行加密。

for proc in procs:
    out, _ = proc.communicate()
    print(out[-10:])

# >>>
# b'\x8c(\xd\xc7m1\xf0F4\xe6'
# b'\x0eD\x97\xe9\x10h{\xbd\xf0'
# b'g\x93)\x14U\xa9\xdc\xdd\x04\xd2'

这些平行运行的子进程还可以分别与另一套平行的子进程对接，形成许多条平行的管道（pipe）。这种管道与 UNIX 管道类似，能够把一条子进程的输出端同另一条子进程的输入端连接起来。下面，写这样一个函数，让它开启一条子进程来运行 openssl 命令，这条命令会根据输入端所发来的数据在输出端生成 Whirlpool 哈希。

def run_hash(input_stdin):
    return subprocess.Popen(
        ['openssl', 'dgst', '-whirlpool', '-binary'],
        stdin=input_stdin,
        stdout=subprocess.PIPE)

现在，可以先启动一批进程来加密数据，然后启动另一批进程根据前面那些进程的加密结果生成哈希码。请注意，前面那批进程（也就是上游进程）的 stdout 实例必须谨慎地处理，用它把相应的哈希进程（也就是下游进程）启动之后，就应该及时关闭（close）并将其设为 None。

encrypt_procs = []
hash_procs = []
for _ in range(3):
    data = os.urandom(100)

    encrypt_proc = run_encrypt(data)
    encrypt_procs.append(encrypt_proc)

    hash_proc = run_hash(encrypt_proc.stdout)
    hash_procs.append(hash_proc)

    # Ensure that the child consumes the input stream and
    # the communicate() method doesn't inadvertently steal
    # input from the child. Also lets SIGPIPE propagate to
    # the upstream process if the downstream process dies.
    encrypt_proc.stdout.close()
    encrypt_proc.stdout = None

只要上、下游的子进程都启动起来，两者之间的 I/O 管道就会自动打通。所要做的仅仅是等待这两批子进程完工并把最终结果打印出来。

for proc in encrypt_procs:
    proc.communicate()
    assert proc.returncode == 0

for proc in hash_procs:
    out, _ = proc.communicate()
    print(out[-10:])
    assert proc.returncode == 0

# >>>
# b'\xe2j\x98h\xfd\xec\xe7T\xd84'
# b'\xf3.i\x01\xd74\xf2\x94E'
# b'5_n\xc3-\xe6j\xe[i'

如果子进程有可能一直不结束，或者由于某种原因卡在输入端或输出端，那么可以在调用 communicate 方法时指定 timeout 参数。这样的话，子进程若是没能在指定时间内结束，程序就会抛出异常，有机会把这条不正常的子进程停掉。

proc = subprocess.Popen(['sleep', '10'])
try:
    proc.communicate(timeout=0.1)
except subprocess.TimeoutExpired:
    proc.terminate()
    proc.wait()

print('Exit status', proc.poll())
# >>>
# Exit status -15