如何通过KeyboardInterrupt正常关闭Python的ThreadPoolExecutor？

我使用Python的ThreadPoolExecutor实现了一个异步客户端，脚本会定期向线程池提交同步客户端可调用对象，期望通过KeyboardInterrupt（Ctrl+C）停止循环。但第一次按下Ctrl+C时程序会挂起，需要再按两次才会最终退出并出现错误。

原代码 #

#!/usr/bin/env python

import numpy as np
import tritonclient.http as tritonclient

import argparse
import itertools
import logging
import random
import sys
import time
from concurrent.futures import ThreadPoolExecutor

distributions = {
    'poisson': lambda w: random.expovariate(1/w),
    'uniform': lambda w: random.uniform(0, 2*w),
}

class Client:
    def __init__(self, url, model):
        self.client = tritonclient.InferenceServerClient(url)
        config = self.client.get_model_config(model)
        self.inputs = config['input']
        self.outputs = [output['name'] for output in config['output']]
        self.model = model

    def __call__(self):
        inputs = []
        for config in self.inputs:
            assert config['data_type'] == 'TYPE_FP32'
            shape = [1] + config['dims']
            datatype = config['data_type'].removeprefix('TYPE_')
            input = tritonclient.InferInput(config['name'], shape, datatype)
            array = np.random.default_rng().random(shape, dtype=np.float32)
            input.set_data_from_numpy(array)
            inputs.append(input)
        result = self.client.infer(self.model, inputs)
        for output in self.outputs:
            result.get_output(output)

def benchmark(fn):
    t_i = time.time()
    fn()
    t_f = time.time()
    print(t_i, t_f - t_i)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--distribution', choices=distributions.values(),
                        type=distributions.get, default=lambda w: w)
    parser.add_argument('-n', '--nrequests', default=-1, type=int)
    parser.add_argument('-o', '--open', '--open-loop', action='store_true')
    parser.add_argument('-u', '--url', default='localhost:8000')
    parser.add_argument('-v', '--verbose', action='count', default=0)
    parser.add_argument('model')
    rate = parser.add_mutually_exclusive_group()
    rate.add_argument('-w', '--wait', '--delay', '-l', '--lambda',
                      default=0, type=float)
    rate.add_argument('-r', '--rate', '-f', '--frequency', type=float)
    args = parser.parse_args()

    level = (logging.DEBUG if args.verbose > 1
            else logging.INFO if args.verbose
            else logging.WARNING)
    logging.basicConfig(level=level)

    if args.rate:
        args.wait = 1/args.rate
    logging.debug(args)

    client = Client(args.url, args.model)

    with ThreadPoolExecutor() as executor:
        try:
            for _ in (itertools.count() if args.nrequests < 0
                        else range(args.nrequests)):
                if args.open:
                    executor.submit(benchmark, client)
                else:
                    benchmark(client)
                time.sleep(args.distribution(args.wait))
        except KeyboardInterrupt:
            pass
        except BrokenPipeError:
            pass

错误日志 #

1717617460.23863 0.003475189208984375
1717617460.250774 0.0033867359161376953
1717617460.2690861 0.0033500194549560547
^C^CTraceback (most recent call last):
  File "/data/pcoppock/mlos/apps/tritonclient", line 73, in <module>
    with ThreadPoolExecutor() as executor:
  File "/usr/lib/python3.10/concurrent/futures/_base.py", line 649, in __exit__
    self.shutdown(wait=True)
  File "/usr/lib/python3.10/concurrent/futures/thread.py", line 235, in shutdown
    t.join()
  File "/usr/lib/python3.10/threading.py", line 1096, in join
    self._wait_for_tstate_lock()
  File "/usr/lib/python3.10/threading.py", line 1116, in _wait_for_tstate_lock
    if lock.acquire(block, timeout):
KeyboardInterrupt
^CException ignored in: <module 'threading' from '/usr/lib/python3.10/threading.py'>
Traceback (most recent call last):
  File "/usr/lib/python3.10/threading.py", line 1537, in _shutdown
    atexit_call()
  File "/usr/lib/python3.10/concurrent/futures/thread.py", line 31, in _python_exit
    t.join()
  File "/usr/lib/python3.10/threading.py", line 1096, in join
    self._wait_for_tstate_lock()
  File "/usr/lib/python3.10/threading.py", line 1116, in _wait_for_tstate_lock
    if lock.acquire(block, timeout):
KeyboardInterrupt: 

linux$ 

问题原因 #

ThreadPoolExecutor的with语句在退出时会自动调用shutdown(wait=True)，该参数会强制主线程等待线程池中所有已提交的任务执行完成后才会退出。当按下Ctrl+C触发KeyboardInterrupt后，主线程进入except块执行pass，随后进入with语句的退出逻辑，此时如果线程池中还有任务在运行（比如Triton的infer调用未完成），主线程会挂起等待，直到任务结束或再次触发中断。

修复方案 #

手动管理ThreadPoolExecutor的生命周期，捕获中断信号后主动调用shutdown(wait=False)，让线程池停止接受新任务且不等待现有任务完成，直接终止程序。

修改后的核心代码如下：

if __name__ == '__main__':
    # ... 其余代码保持不变 ...

    client = Client(args.url, args.model)

    executor = ThreadPoolExecutor()
    try:
        for _ in (itertools.count() if args.nrequests < 0
                    else range(args.nrequests)):
            if args.open:
                executor.submit(benchmark, client)
            else:
                benchmark(client)
            time.sleep(args.distribution(args.wait))
    except KeyboardInterrupt:
        print("\n收到中断信号，正在退出...")
    except BrokenPipeError:
        print("\n管道断开，正在退出...")
    finally:
        # 不等待未完成任务，立即关闭线程池
        executor.shutdown(wait=False)

说明 #

1. 不再使用with语句自动管理线程池，改为手动创建和销毁，避免with退出时强制等待任务完成。
2. 在finally块中调用shutdown(wait=False)，确保无论是否触发异常，线程池都会被正确关闭。
3. wait=False参数会让线程池直接停止所有工作线程，主线程无需等待任务结束，从而实现按下Ctrl+C后快速退出。

内容的提问来源于stack exchange，提问作者Patrick Coppock