极简状态化进程池

最近在课题中遇到一个典型场景：需要在异构 GPU 服务器上部署高频推理服务。简单来说，就是一台服务器上插着多张算力不同的显卡，我们要在上面运行图像分类推理任务，并对外提供高吞吐的 API 服务。面对海量请求，如何让这些“能力参差不齐”的 GPU 协同工作，最大化整体吞吐量，成了关键挑战。

这其实是一个数据分发的负载均衡问题，最朴素的思路是随机分发请求到不同 GPU。但问题很明显：由于显卡算力不均，整体吞吐会被最慢的那张卡拖累；而且随机策略本身存在“热点”风险——连续多个请求可能撞到同一张卡，导致其他卡空闲，资源利用率难以提升。

熟悉并发编程的朋友一定对“生产者-消费者”模型不陌生：通过维护一个共享任务队列，让各工作单元以自主竞争任务，既能自然实现负载均衡，又能利用队列削峰填谷，显著提升资源利用率。然而，这种模式的简单实现存在一个问题：由于任务由空闲设备“谁先抢到谁执行”，结果天然乱序，对于需要保持请求顺序的服务场景并不友好。

另一方面，Python 标准库中的 concurrent.futures.ProcessPoolExecutor 本可简化任务分发与结果收集，但它面向的是无状态函数调用——每个 worker 仅执行纯函数，无法维护设备上下文（例如将模型加载到指定 GPU）。而这恰恰是异构推理场景的核心需求：每个进程必须绑定特定设备并持有模型状态。

为解决上述问题，我实现了一个极简的工具库：stateful-pool。整个库仅 200 余行代码，实现了两大关键能力：

有状态工作进程的池化管理
有序的任务提交与结果回收

text

+-----------------------+              +----------------------+
|     Main Process      |              |     Process Pool     |
| +--------+ +--------+ |              | +------------------+ |
| | Thread | | Thread | | --- Task --> | | Worker (Process) | |
| |   1    | |   2    | | <-- Resp --- | +------------------+ |
| +--------+ +--------+ |              | +------------------+ |
|      ...     ...      |              | | Worker (Process) | |
| +--------+ +--------+ |              | +------------------+ |
| | Thread | | Thread | | --- Task --> |         ...          |
| |  N-1   | |   N    | | <-- Resp --- | +------------------+ |
| +--------+ +--------+ |              | | Worker (Process) | |
|                       |              | +------------------+ |
+-----------------------+              +----------------------+

(主进程维护一个任务队列，多个工作进程从队列中抢占任务并执行，结果通过事件通知主进程。)

通过这种基于Actor的设计，stateful-pool提升吞吐效率的同时可以降低代码复杂度。使用时工作进程只需实现两个接口：

spawn：负责进程初始化，例如将模型加载至指定 GPU，建立设备上下文；
execute：处理具体推理请求。

不到80行就可以实现一个多进程负载均衡的图像分类服务： https://github.com/MenxLi/stateful_pool/blob/main/benchmark/exp/server_spool.py

实测相较多进程随机分发策略，这个小工具在双卡服务器上运行100次推理任务 提升了1/3的吞吐量的同时降低了1/3的代码量，延迟也更低更稳定。

有状态进程池是很多应用场景的共性需求，该库已经开源在 GitHub 上，也发布了 PyPI 包，欢迎大家试用和反馈！链接如下：

GitHub: https://github.com/MenxLi/stateful_pool

使用时直接安装：

bash

pip install stateful-pool

p.s. 上述比较实际上是对比随机分发与队列满载运行的情况下程序运行速度，随机分发的平均吞吐量会被最慢的设备拖累，而队列满载则能充分利用所有设备的算力，从而提升整体效率。在数学上可以证明：如果有N个相同设备，有K个任务（ $K ≫ N$ ），随机分发情况下单个设备最多的任务数期望为：

E [m a x] \approx \frac{K}{N} + \sqrt{\frac{2 K \log N}{N}}

所以使用队列满载的方式可以将单个设备的任务数控制在 $λ = \frac{K}{N}$ ，避免了随机分发中可能出现的负载不均问题，从而提升整体吞吐量，提升幅度大约为：

\frac{E [m a x] - λ}{E [m a x]} \approx \frac{\sqrt{2 λ \log N}}{λ + \sqrt{2 λ \log N}}

实际上设备之间性能不同、限流策略等因素也会影响实际提升幅度，尤其是这种理论分析假设了任务不间断且设备性能相同的理想情况，实际环境中可能会有更多复杂因素，例如任务可能是间歇性一波接一波的。但总体来说队列满载的方式能显著提升吞吐效率，尤其在设备数量较多、突发密集任务量较小的情况下效果更明显。

附：预期提升计算器

任务数量 (K):
设备数量 (N):

预期吞吐提升: 15.43%

极简状态化进程池 ​

极简状态化进程池