怎样实现分配进程池任务执行脚本

本文目录导读：

1. 基础方法：multiprocessing.Pool
2. 使用 concurrent.futures.ProcessPoolExecutor
3. 实际脚本执行场景
4. 高级用法：动态任务分配
5. 最佳实践建议

我来介绍几种在Python中实现进程池分配任务执行脚本的常用方法：

基础方法：multiprocessing.Pool

示例1：简单任务分配

import multiprocessing
import time
import random
def worker_task(task_id):
    """工作函数"""
    print(f"进程 {multiprocessing.current_process().name} 开始执行任务 {task_id}")
    time.sleep(random.random() * 2)  # 模拟耗时操作
    result = task_id * 2
    print(f"任务 {task_id} 完成，结果: {result}")
    return result
if __name__ == "__main__":
    # 创建进程池，默认使用CPU核心数
    pool = multiprocessing.Pool(processes=4)
    # 准备任务列表
    tasks = list(range(10))
    # 方法1: 同步执行所有任务
    results = pool.map(worker_task, tasks)
    print(f"所有任务结果: {results}")
    # 方法2: 异步执行
    async_results = [pool.apply_async(worker_task, args=(i,)) for i in range(5)]
    for r in async_results:
        result = r.get()  # 获取结果
        print(f"异步结果: {result}")
    pool.close()
    pool.join()

使用 concurrent.futures.ProcessPoolExecutor

这是更现代的API,使用起来更方便：

from concurrent.futures import ProcessPoolExecutor, as_completed
import time
def process_script(script_name, params):
    """执行脚本函数"""
    print(f"执行脚本: {script_name} 参数: {params}")
    time.sleep(1)
    return f"{script_name} 执行完成"
def main():
    # 定义要执行的脚本
    scripts = [
        ("script1.py", {"param1": "value1"}),
        ("script2.py", {"param2": "value2"}),
        ("script3.py", {"param3": "value3"}),
        ("script4.py", {"param4": "value4"}),
    ]
    # 使用进程池执行
    with ProcessPoolExecutor(max_workers=3) as executor:
        # 提交所有任务
        future_to_script = {
            executor.submit(process_script, name, params): name 
            for name, params in scripts
        }
        # 获取完成的任务结果
        for future in as_completed(future_to_script):
            script_name = future_to_script[future]
            try:
                result = future.result()
                print(f"脚本 {script_name} 结果: {result}")
            except Exception as e:
                print(f"脚本 {script_name} 执行失败: {e}")
if __name__ == "__main__":
    main()

实际脚本执行场景

更接近生产环境的示例：

import subprocess
from concurrent.futures import ProcessPoolExecutor, as_completed
import os
import logging
logging.basicConfig(level=logging.INFO)
class ScriptExecutor:
    def __init__(self, max_workers=4):
        self.max_workers = max_workers
        self.executor = ProcessPoolExecutor(max_workers=max_workers)
    def execute_script(self, script_path, args=None):
        """执行外部脚本"""
        try:
            cmd = ['python', script_path]
            if args:
                cmd.extend(args)
            # 使用subprocess执行外部脚本
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=300  # 5分钟超时
            )
            if result.returncode == 0:
                return {
                    "script": script_path,
                    "status": "success",
                    "output": result.stdout,
                    "error": result.stderr
                }
            else:
                return {
                    "script": script_path,
                    "status": "failed",
                    "output": result.stdout,
                    "error": result.stderr
                }
        except subprocess.TimeoutExpired:
            return {
                "script": script_path,
                "status": "timeout",
                "error": "Script execution timed out"
            }
        except Exception as e:
            return {
                "script": script_path,
                "status": "error",
                "error": str(e)
            }
    def run_batch(self, scripts_list):
        """批量执行脚本"""
        futures = []
        for script_info in scripts_list:
            script_path = script_info.get('path')
            args = script_info.get('args', [])
            future = self.executor.submit(self.execute_script, script_path, args)
            futures.append((future, script_path))
        results = []
        for future, script_path in futures:
            try:
                result = future.result()
                logging.info(f"Script {script_path} completed: {result['status']}")
                results.append(result)
            except Exception as e:
                logging.error(f"Script {script_path} failed: {e}")
                results.append({
                    "script": script_path,
                    "status": "error",
                    "error": str(e)
                })
        return results
    def cleanup(self):
        """清理资源"""
        self.executor.shutdown(wait=True)
# 使用示例
if __name__ == "__main__":
    # 创建执行器
    executor = ScriptExecutor(max_workers=3)
    # 准备要执行的脚本列表
    scripts = [
        {"path": "scripts/task1.py", "args": ["--param1", "value1"]},
        {"path": "scripts/task2.py", "args": ["--param2", "value2"]},
        {"path": "scripts/task3.py", "args": ["--param3", "value3"]},
    ]
    # 执行所有脚本
    results = executor.run_batch(scripts)
    # 处理结果
    for result in results:
        print(f"Script: {result['script']}")
        print(f"Status: {result['status']}")
        if result['output']:
            print(f"Output: {result['output'][:100]}...")  # 只显示前100个字符
    # 清理
    executor.cleanup()

高级用法：动态任务分配

import queue
from multiprocessing import Process, Queue, cpu_count
from queue import Empty
import time
import random
class DynamicTaskPool:
    def __init__(self, num_workers=None):
        self.num_workers = num_workers or cpu_count()
        self.task_queue = Queue()
        self.result_queue = Queue()
        self.processes = []
    def worker(self):
        """工作进程"""
        while True:
            try:
                # 获取任务，超时时间1秒
                task = self.task_queue.get(timeout=1)
                task_id, func, args, kwargs = task
                try:
                    result = func(*args, **kwargs)
                    self.result_queue.put({
                        'task_id': task_id,
                        'status': 'success',
                        'result': result
                    })
                except Exception as e:
                    self.result_queue.put({
                        'task_id': task_id,
                        'status': 'failed',
                        'error': str(e)
                    })
            except Empty:
                # 队列为空，停止工作
                break
            except Exception as e:
                print(f"Worker error: {e}")
                break
    def start_workers(self):
        """启动工作进程"""
        for i in range(self.num_workers):
            p = Process(target=self.worker, name=f"Worker-{i}")
            p.start()
            self.processes.append(p)
    def add_task(self, task_id, func, args=(), kwargs=None):
        """添加任务到队列"""
        if kwargs is None:
            kwargs = {}
        self.task_queue.put((task_id, func, args, kwargs))
    def get_results(self, timeout=None):
        """获取所有结果"""
        results = []
        start_time = time.time()
        while len(results) < self.task_queue.qsize() + len(results):
            try:
                remaining = None
                if timeout:
                    elapsed = time.time() - start_time
                    remaining = max(0, timeout - elapsed)
                result = self.result_queue.get(timeout=remaining)
                results.append(result)
            except Empty:
                break
        return results
    def shutdown(self):
        """关闭所有进程"""
        for p in self.processes:
            if p.is_alive():
                p.terminate()
            p.join()
# 使用示例
def example_task(task_id, data):
    """示例任务函数"""
    time.sleep(random.random() * 2)
    return f"Task {task_id}: processed {data}"
if __name__ == "__main__":
    # 创建动态任务池
    pool = DynamicTaskPool(num_workers=3)
    # 添加任务
    tasks = [
        (1, "data1"),
        (2, "data2"),
        (3, "data3"),
        (4, "data4"),
        (5, "data5"),
    ]
    for task_id, data in tasks:
        pool.add_task(task_id, example_task, (task_id, data))
    # 启动工作进程
    pool.start_workers()
    # 等待并获取结果
    pool.wait_for_completion()
    results = pool.get_results()
    # 输出结果
    for result in results:
        if result['status'] == 'success':
            print(f"Task {result['task_id']}: {result['result']}")
        else:
            print(f"Task {result['task_id']} failed: {result['error']}")
    # 关闭池
    pool.shutdown()

最佳实践建议

1. 任务粒度控制：确保每个任务执行时间适中，避免过短任务导致过多调度开销
2. 资源限制：使用 imap_unordered 等流式处理方法处理大量任务
3. 错误处理：始终捕获异常，避免一个任务失败影响其他任务
4. 性能监控：添加日志和监控，跟踪任务执行情况

这些方法提供了从简单到复杂的进程池实现,可以根据你的具体需求选择合适的方式。