Windows11下Python多进程共享大型数据集的解决方案咨询

我在Windows 11上加速一个CPU密集型Python脚本，由于Python线程无法利用多核CPU，只能采用multiprocessing方案。我有一个从文件加载后内存占用11GB的大型字典，需要检查计算结果是否存在于该字典中；输入计算的数据来自100GB的文件，可分批通过pool.map分配给进程，但无法将字典复制到每个进程（内存不足），因此需要找到进程共享字典的方法。

伪代码流程 #

--主线程--
- 从文件加载字典结构   # 内存占用11GB
- ...
- 未加载完所有数据块时循环
-    从文件加载计算数据块   # 每块10000行
-    将计算数据块分配（map）给进程
-    等待进程完成所有数据块处理

--子进程--
- 遍历子数据块中的每个元素
-    执行计算
-    检查计算结果是否在字典中  # 此处存在问题！
-    将结果存入文件

已实现代码及问题 #

采纳建议后我写了如下代码，但仍存在问题：

def ReadDictFromFile()
    cnt=0
    print("Reading dictionary from " + dictfilename)
    with open(dictfilename, encoding=("utf-8"), errors=("replace")) as f:
        next(f) #跳过首行（表头）
        for line in f:
            s = line.rstrip("\n")
            (key,keyvalue) = s.split()
            shared_dict[str(key)]=keyvalue
            cnt = cnt + 1
            if ((cnt % 1000000) == 0): #每加载1000000条记录打印进度
                print(cnt)
                return #临时返回以加速测试，暂不加载完整字典
    print("Done loading dictionary")        


def checkqlist(qlist)
    print(str(os.getpid()) + "-" + str(len(qlist)))
    
    for li in qlist:
        try:
            checkvalue = calculations(li)
        
            (found, keyval) = InMem(checkvalue)
                
            if (found):
                print("FOUND!!! " + checkvalue + ' ' + keyvalue)            
        except Exception as e:
            print("(" + str(os.getpid()) + ")Error log: %s" % repr(e))
            time.sleep(15)


def InMem(checkvalue):
    if(checkvalue in shared_dict):
        return True, shared_dict[checkvalue]
    else:
        return False, ""


if __name__ == "__main__":
    start_time = time.time()

    global shared_dict 
    manager = Manager()
    shared_dict = manager.dict()

    ReadDictFromFile()

    chunksize=5
    nr_of_processes = 10
    with open(filetocheck, encoding=("utf-8"), errors=("replace")) as f:
        qlist = []
        for line in f:
            s = line.rstrip("\n")
            qlist.append(s)
            if (len(qlist) >= (chunksize * nr_of_processes)):
                chunked_list = [qlist[i:i+chunk_size] for i in range(0, len(qlist), chunk_size)]
                try:
                    with multiprocessing.Pool() as pool:
                        pool.map(checkqlist, chunked_list, nr_of_processes)          #问题：qlist是单个字符串，而非约416个字符串的列表。  
                except Exception as e:
                    print("error log: %s" % repr(e))
                    time.sleep(15)
    logit("Completed! " + datetime.datetime.now().strftime("%I:%M%p on %B %d, %Y"))
    print("--- %s seconds ---" % (time.time() - start_time))

核心问题 #

1. 拼写错误：chunk_size应为定义好的chunksize，导致分块逻辑彻底失效
2. 参数传递错误：pool.map的第三个参数是单任务分块大小，不是进程数，传入nr_of_processes会打乱任务分配
3. 内存泄漏风险：处理完批次数据后未清空qlist，导致数据持续累积
4. 未定义变量：checkqlist中打印的keyvalue未定义，应使用InMem返回的keyval
5. 缺失依赖导入：代码未导入time、os等必要模块，运行会报错
6. 性能瓶颈：Manager.dict()依赖进程间通信（IPC），11GB大字典的频繁访问会导致严重延迟

修复方案及优化代码 #

修复后可运行代码 #

import time
import multiprocessing
from multiprocessing import Manager
import os
import datetime

def logit(msg):
    print(msg)

def calculations(val):
    # 替换为你的实际计算逻辑
    return val

def ReadDictFromFile(shared_dict, dictfilename):
    cnt = 0
    print(f"Reading dictionary from {dictfilename}")
    with open(dictfilename, encoding="utf-8", errors="replace") as f:
        next(f)  # 跳过首行表头
        for line in f:
            s = line.rstrip("\n")
            # 处理值含空格的情况，仅按第一个空格分割键值
            key, keyvalue = s.split(maxsplit=1)
            shared_dict[str(key)] = keyvalue
            cnt += 1
            if cnt % 1000000 == 0:
                print(f"Loaded {cnt} entries")
                # return  # 测试时临时返回，正式运行请注释
    print("Done loading dictionary")        

def checkqlist(args):
    qlist, shared_dict = args
    pid = os.getpid()
    print(f"{pid} - Processing {len(qlist)} items")
    
    for li in qlist:
        try:
            checkvalue = calculations(li)
            found, keyval = InMem(checkvalue, shared_dict)
                
            if found:
                print(f"FOUND!!! {checkvalue} {keyval}")            
        except Exception as e:
            print(f"({pid})Error log: {repr(e)}")
            time.sleep(15)

def InMem(checkvalue, shared_dict):
    if checkvalue in shared_dict:
        return True, shared_dict[checkvalue]
    else:
        return False, ""

if __name__ == "__main__":
    start_time = time.time()

    # 配置文件路径
    dictfilename = "your_dict_file.txt"
    filetocheck = "your_data_file.txt"

    # 创建进程共享字典
    manager = Manager()
    shared_dict = manager.dict()

    # 加载字典到共享内存
    ReadDictFromFile(shared_dict, dictfilename)

    chunksize = 5
    nr_of_processes = 10
    batch_size = chunksize * nr_of_processes

    with open(filetocheck, encoding="utf-8", errors="replace") as f:
        qlist = []
        for line in f:
            s = line.rstrip("\n")
            qlist.append(s)
            if len(qlist) >= batch_size:
                # 按chunksize拆分任务
                chunked_list = [qlist[i:i+chunksize] for i in range(0, len(qlist), chunksize)]
                # 打包任务与共享字典（pool.map仅支持单参数）
                tasks = [(chunk, shared_dict) for chunk in chunked_list]
                try:
                    with multiprocessing.Pool(processes=nr_of_processes) as pool:
                        pool.map(checkqlist, tasks)
                except Exception as e:
                    print(f"error log: {repr(e)}")
                    time.sleep(15)
                # 清空当前批次，避免累积
                qlist = []
        # 处理剩余不足一个批次的数据
        if qlist:
            chunked_list = [qlist[i:i+chunksize] for i in range(0, len(qlist), chunksize)]
            tasks = [(chunk, shared_dict) for chunk in chunked_list]
            with multiprocessing.Pool(processes=nr_of_processes) as pool:
                pool.map(checkqlist, tasks)

    logit(f"Completed! {datetime.datetime.now().strftime('%I:%M%p on %B %d, %Y')}")
    print(f"--- {time.time() - start_time} seconds ---")

性能优化建议 #

针对11GB大字典的性能瓶颈，推荐以下方案替代Manager.dict()：

1. 内存映射文件：将字典键排序后存储为索引文件，值存储为二进制文件，每个进程通过mmap映射文件，用二分查找快速定位键，实现接近本地内存的访问速度，且无需复制字典。
2. 共享内存对象：使用pyarrow等库将字典序列化为共享内存对象，所有进程直接访问物理内存中的数据，避免IPC开销。
3. 查询服务进程：单独启动一个进程负责字典查询，其他进程通过队列发送请求，适合查询频率较低的场景。

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