Python 多进程管理工具类

前言

项目当中刚好用到了，放出来，方便下次使用。

功能

1. 创建进程，支持自定义进程id,创建之后自动运行。具备轮询等待，当当前队列长度大于进程长度时，进入堵塞状态，直到当前任务加入线程池进入执行状态
2. 支持根据进程id终止某一进程
3. 支持更具进程id重启进程
4. 支持不定参数，只需要传入待运行函数以及对应参数即可开启任务

优点：

1. 实现简单，开箱即用
2. 适用计算密集型任务

实现

import queue
import signal
import multiprocessing as mp

class ProcessPool:
    def __init__(self, max_processes=mp.cpu_count()):
        self.max_processes = max_processes
        self.tasks = queue.Queue()
        self.processes = {
    
    }
        self._stop_event = mp.Event()

        # 注册 Ctrl+C 信号处理函数，确保在程序被终止时能够正常清理资源
        signal.signal(signal.SIGINT, self._handle_signals)
        signal.signal(signal.SIGTERM, self._handle_signals)

    def add_process(self, process_id, target, args=()):

        if len(target.__code__.co_varnames) > len(args):
            raise ValueError("Too few arguments")
        p = mp.Process(target=self._wrapper, args=(process_id, target, args), daemon=True)
        p.start()
        self.processes[process_id] = p
        while len(self.processes) >= self.max_processes:
            for process_id, p in list(self.processes.items()):
                if not p.is_alive():
                    del self.processes[process_id]
                    break
            else:
                time.sleep(0.1)

        if not p.is_alive():
            return

        while not self.tasks.empty():
            process_id, target, args = self.tasks.get()
            p = mp.Process(target=self._wrapper, args=(process_id, target, args), daemon=True)
            p.start()
            self.processes[process_id] = p

    def stop(self):
        """
        停止进程池。
        """
        self._stop_event.set()
        for p in self.processes.values():
            p.terminate()
        self.processes.clear()

    def restart(self, process_id):
        """
        重启指定 id 的进程。
        :param process_id: 进程id。
        """
        if process_id not in self.processes:
            raise ValueError("Process not found")
        p = self.processes[process_id]
        p.terminate()
        p.join()
        target, args = self.tasks.queue[0][1:]
        p = mp.Process(target=self._wrapper, args=(process_id, target, args), daemon=True)
        p.start()
        self.processes[process_id] = p

    def kill(self, process_id):

        """
        结束指定 id 的进程。
        :param process_id: 进程id。
        """
        if process_id not in self.processes:
            raise ValueError("Process not found")
        p = self.processes[process_id]
        p.terminate()

    def _handle_signals(self, signum, frame):
        """
        信号处理函数，用于在进程被终止时正常清理资源。
        """
        self.stop()

    @staticmethod
    def _wrapper(process_id, target, args):
        """
        包装目标函数，以便能够支持动态参数。
        :param process_id: 进程id。
        :param target: 进程要执行的目标函数。
        :param args: 目标函数的参数列表。
        """
        target(*args)

测试

#进程测试函数
def func1():
    print("func1 start")
    time.sleep(1)
    print("func1 end")


def func2(a, b):
    print("func2 start")
    time.sleep(1)
    print(f"func2 end, a = {
      
      a}, b = {
      
      b}")

if __name__ == '__main__':

    pool = ProcessPool(max_processes=2)
    pool.add_process(2, func2, args=(1, 2))
    pool.add_process(3, func2, args=(3, 4))

总结

以后这种类型的博文，我将写的非常简短，不会再介绍任何代码相关的实现，原理，因为会的自然就会，不会的一时半会也不见得可以理解，开箱即用即可，莫问东西，博文长了，我懒得写，你懒得看，我下次也懒得翻。