AWS Batch中Multiprocessing.Pool迭代IMapIterator遇PicklingError问题

因公司框架限制，采用multiprocessing实现批量API请求，编写了Batcher类管理并发进程池。本地环境运行正常，但作为AWS Batch作业执行时，迭代pool.imap返回的IMapIterator对象时抛出序列化错误：

_pickle.PicklingError: Can't pickle <class 'boto3.resources.factory.s3.ServiceResource'>: attribute lookup s3.ServiceResource on boto3.resources.factory failed

调用栈：

for res in results
File "/usr/local/lib/python3.9/multiprocessing/pool.py", line 870, in next
 raise value
File "/usr/local/lib/python3.9/multiprocessing/pool.py", line 537, in _handle_tasks
 put(task)
File "/usr/local/lib/python3.9/multiprocessing/connection.py", line 211, in send
self._send_bytes(_ForkingPickler.dumps(obj))
File "/usr/local/lib/python3.9/multiprocessing/reduction.py", line 51, in dumps
cls(buf, protocol).dump(obj)
_pickle.PicklingError: Can't pickle <class 'boto3.resources.factory.s3.ServiceResource'>: attribute lookup s3.ServiceResource on boto3.resources.factory failed

相关代码 #

Batcher类实现 #

class Batcher:
    def __init__(self, concurrency: int = 8):
        self.concurrency = concurrency

    def _interprete_response_to_succ_or_err(self, resp: requests.Response) -> str:
        if isinstance(resp, str):
            if "Error:" in resp:
                return "dlq"
            else:
                return "err"
        if isinstance(resp, requests.Response):
            if resp.status_code == 200:
                return "succ"
            else:
                return "err"

    def _fetch_dat_data(self, id: str) -> requests.Response:
        
        try:
            resp = requests.get(API_ENDPOINT)
            return resp
        except Exception as e:
            return f"ID {id} -> Error: {str(e)}"

    def _dispatch_batch(self, batch: list) -> dict:
        pool = MPool(self.concurrency)
        results = pool.imap(self._fetch_dat_data, batch)
        pool.close()
        pool.join()
        return results

    def _run_batch(self, id):
        return self._dispatch_batch(id)

    def start(self, id_list: list):
        """ In real class, this function will create smaller
        batches from bigger chunks of data """
        results = self._run_batch(id_list)
        print(
            [
                res.text
                for res in results
                if self._interprete_response_to_succ_or_err(res) == "succ"
            ]
        )

调用代码 #

if __name__ == "__main__":
    """
    the source of ids is a csv file with single column in s3 that contains list
    of columns with single id per line
    """
    id_list = boto3_get_object_body(my_file_name).decode().split("\n") # custom function, works
    batcher = Batcher()
    batcher.start(id_list)

1. multiprocessing启动方式差异

   • 本地默认用fork模式：直接复制父进程地址空间，不需要序列化整个类实例，即使存在不可pickle的对象也不会触发序列化。
   • AWS Batch环境可能用spawn模式：重新启动Python解释器，需要序列化所有传递给子进程的对象（包括Batcher实例及关联的上下文），而boto3的S3 ServiceResource属于不可pickle的对象，因此报错。

2. 隐式的boto3资源引用泄漏
   boto3_get_object_body函数中创建的S3客户端/资源可能被绑定到全局变量或Batcher实例的隐式上下文（比如通过模块级别的变量），当子进程启动时，pickle尝试序列化这些对象导致失败。

3. AWS Batch环境的Python配置
   部分AWS Batch镜像可能默认设置了spawn作为multiprocessing启动方法，或者Python版本/环境变量配置与本地不同，导致序列化逻辑触发条件变化。

• 隔离boto3操作与多进程逻辑：在__main__中完成S3数据读取后，确保销毁所有boto3客户端/资源，只将纯数据（id_list字符串列表）传递给Batcher类，避免不可pickle对象进入子进程上下文。
• 显式设置启动方法：如果AWS Batch环境基于Unix/Linux，在代码开头添加：

  import multiprocessing as mp
  mp.set_start_method('fork')
  

  强制使用fork模式，避免序列化整个类实例。
• 将API请求函数改为独立函数：把_fetch_dat_data从类方法改为模块级别的独立函数，减少子进程需要序列化的对象范围。
• 检查boto3_get_object_body实现：确保函数内部创建的S3客户端/资源是局部变量，不会泄漏到全局或类上下文。

内容的提问来源于stack exchange，提问作者Bijay Regmi