Slurm分配16核运行Python脚本却仅占用1CPU的问题咨询

核心原因说明：CPython存在全局解释器锁（GIL），默认单进程代码仅能使用1个CPU核心，Slurm分配的多核资源需要代码做多进程适配才能利用。

改造方案1：并行化字典处理阶段（改造成本最低） #

你的代码分为两个阶段：读文件建字典（IO密集型，并行改造成本高）、字典条目分析输出（CPU密集型，易并行），优先改造第二阶段即可用上全部16核资源。
改造后的代码示例：

import csv
from multiprocessing import Pool

# 把原来的分析逻辑抽成单独的进程函数
def process_entry(entry):
    k, v = entry
    # 这里保留你原来的分析逻辑即可
    nl = [不同变量的计算逻辑]
    return nl

if __name__ == "__main__":
    # 第一阶段：单进程读文件建字典，和原有逻辑一致
    with open(huge_file,'r') as hugefile:
        reader = csv.reader(hugefile, delimiter="\t")
        d = {} 
        for r in reader:
            v = r[0]+r[1]
            if v not in d:
                d[v] = [[r[5], r[4]]]
            else:
                d[v].append([r[5], r[4]])
    
    # 第二阶段：16进程并行处理所有字典条目
    with Pool(processes=16) as pool, open(final_file, 'w') as final:
        writer = csv.writer(final, delimiter="\t")
        # imap_unordered会按进程完成顺序返回结果，比map更快
        for result in pool.imap_unordered(process_entry, d.items()):
            writer.writerow(result)

注意：如果需要输出结果的顺序和字典key的顺序一致，把imap_unordered换成map即可，速度会稍慢。

改造方案2：全流程并行（极致提速） #

如果第一阶段读文件建字典耗时占比很高，可以对大文件做分块并行读取，进一步压缩耗时：

1. 先通过bash命令把大文件拆成16个等行的分块文件：

split -l $(( $(wc -l < huge_file) / 16 )) huge_file huge_chunk_

2. 改造Python代码，多进程读分块文件建局部字典，主进程合并后再并行处理：

import csv
import glob
from multiprocessing import Pool

def process_chunk(chunk_path):
    # 每个进程单独处理一个分块文件，返回局部字典
    d = {}
    with open(chunk_path, 'r') as f:
        reader = csv.reader(f, delimiter="\t")
        for r in reader:
            v = r[0]+r[1]
            if v not in d:
                d[v] = [[r[5], r[4]]]
            else:
                d[v].append([r[5], r[4]])
    return d

def process_entry(entry):
    k, v = entry
    nl = [不同变量的计算逻辑]
    return nl

if __name__ == "__main__":
    # 16进程并行处理所有分块文件，得到16个局部字典
    chunk_paths = glob.glob("huge_chunk_*")
    with Pool(processes=16) as pool:
        local_dicts = pool.map(process_chunk, chunk_paths)
    
    # 主进程合并所有局部字典
    d = {}
    for local_d in local_dicts:
        for k, v in local_d.items():
            if k not in d:
                d[k] = v
            else:
                d[k].extend(v)
    
    # 并行处理字典条目，和方案1逻辑一致
    with Pool(processes=16) as pool, open(final_file, 'w') as final:
        writer = csv.writer(final, delimiter="\t")
        for result in pool.imap_unordered(process_entry, d.items()):
            writer.writerow(result)

额外优化建议 #

• 原有代码中if v not in d.keys()可以改成if v not in d，省去调用keys()的开销，运行速度更快
• 无需修改代码，直接用PyPy解释器运行脚本，针对大循环场景可比CPython提速3~10倍
• 如果字典过大内存不足，可以用diskcache库替代原生字典，把数据落地到磁盘避免OOM

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