Python多进程Pool执行Cassandra批量加载时Worker闲置问题求助

我之前也碰到过类似的大体积数据下多进程池闲置的问题，结合你的场景——中等数据量（1GB）正常运行、超2.5GB数据时Worker闲置，再加上4核16GB的环境配置，大概率是主进程数据分发阻塞、进程间序列化开销过大或者内存资源瓶颈导致的，下面给你拆解具体原因和对应的解决办法：

1. 避免一次性加载全量数据，改用迭代式分发任务 #

问题根源 #

multiprocessing.Pool.map()方法会先把所有输入数据全部序列化，再一次性分发给Worker进程。当数据量超过2.5GB时，主进程会卡在数据序列化/内存加载阶段，根本没机会把任务传递给Worker，导致Worker全程闲置。而1GB的数据刚好卡在序列化和内存的阈值内，所以之前能正常运行。

解决办法 #

改用imap()或imap_unordered()代替map()，这两个方法是迭代式分发任务，会分批把数据传给Worker，不需要先把全量数据加载到内存。同时配合生成器函数分批读取数据，彻底避免主进程内存过载：

import multiprocessing
from cassandra.cluster import Cluster

def data_generator(file_path):
    # 分批读取大文件/数据源，比如每次读1000条
    with open(file_path, 'r') as f:
        batch = []
        for line in f:
            batch.append(line.strip())
            if len(batch) == 1000:
                yield batch
                batch = []
        if batch:
            yield batch

def mp_worker(batch_data):
    # 每个任务处理一批数据，而不是单条
    cluster = Cluster(['<ip_address>'])
    session = cluster.connect('<your_keyspace>')
    # 批量插入逻辑，比如用Cassandra的BatchStatement
    batch_stmt = session.prepare("INSERT INTO your_table (...) VALUES (...)")
    for data in batch_data:
        session.execute(batch_stmt, (data,))
    cluster.shutdown()

if __name__ == '__main__':
    with multiprocessing.Pool(processes=4) as pool:
        # 用imap迭代处理生成器的输出
        for _ in pool.imap(mp_worker, data_generator('large_data.txt')):
            pass

2. 优化进程间数据传递，让Worker自行读取数据 #

问题根源 #

即使使用迭代式分发，大体积数据的序列化/反序列化开销依然很大，主进程可能还是会卡在数据传递环节。而且多进程间传递大对象本身就不是高效的做法。

解决办法 #

让主进程只传递数据的存储路径/标识，Worker进程自行读取数据。比如把大文件拆分成多个小文件，主进程把文件名列表传给Worker，Worker自己打开文件处理：

def mp_worker(file_name):
    cluster = Cluster(['<ip_address>'])
    session = cluster.connect('<your_keyspace>')
    with open(file_name, 'r') as f:
        # 读取当前文件的数据并插入Cassandra
        for line in f:
            session.execute(...)
    cluster.shutdown()

if __name__ == '__main__':
    # 假设大文件已经拆分成多个small_data_*.txt
    file_list = ['small_data_1.txt', 'small_data_2.txt', ...]
    with multiprocessing.Pool(processes=4) as pool:
        pool.map(mp_worker, file_list)

这种方式下，进程间传递的只是字符串文件名，开销几乎可以忽略，Worker能快速拿到任务并执行。

3. 复用Cassandra连接，避免重复初始化开销 #

问题根源 #

如果你的mp_worker函数每次都初始化Cluster和Session，大体积数据下重复创建/销毁连接会导致大量耗时，甚至可能触发Cassandra的连接数限制，导致Worker卡在连接阶段，看起来像是闲置。

解决办法 #

使用Pool的initializer和initargs参数，让每个Worker进程只初始化一次Cassandra连接，全程复用：

def init_worker():
    # 每个Worker进程初始化一次连接，存在全局变量中
    global session, cluster
    cluster = Cluster(['<ip_address>'])
    session = cluster.connect('<your_keyspace>')

def mp_worker(data):
    # 直接使用全局的session处理数据
    session.execute(...)

if __name__ == '__main__':
    with multiprocessing.Pool(processes=4, initializer=init_worker) as pool:
        pool.imap(mp_worker, data_generator('large_data.txt'))
    # 所有Worker结束后关闭连接
    cluster.shutdown()

4. 排查内存瓶颈，避免主进程内存过载 #

问题根源 #

16GB内存看似能容纳2.5GB数据，但如果数据在加载后转换成字典、列表等结构，内存占用会大幅膨胀（比如CSV数据转换成字典后，内存占用可能是原文件的3-5倍），导致主进程触发系统OOM或者严重的内存交换（swap），进程被卡住无法分发任务。

解决办法 #

• 用psutil库监控主进程的内存占用，确认是否出现内存过载：

  import psutil
  process = psutil.Process()
  print(f"当前内存占用: {process.memory_info().rss / 1024 / 1024:.2f} MB")
  

• 进一步缩小每批处理的数据量，比如从1000条改成500条，降低单批数据的内存占用。
• 避免在主进程中存储不必要的中间数据，比如读取一条就传递一条（生成器的方式），不要把全量数据存在列表中。

调试小技巧 #

• 给主进程和Worker进程加日志，打印任务分发和执行的节点，比如主进程打印“开始分发第X批任务”，Worker打印“开始处理第X批任务”，这样能快速定位是主进程没发任务还是Worker没接收。
• 用htop或top查看进程状态：如果Worker进程一直处于S（睡眠）状态，说明主进程没发任务；如果处于R（运行）状态但CPU占用低，可能是Worker卡在Cassandra连接或数据读取环节。

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