在 16 个 CPU 而不是 1 个 CPU 上运行 python 脚本

Question

我有一个激活 python 脚本的 bash 脚本：

#!/bin/bash
#SBATCH -J XXXXX
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=16

python my_python_script.py

python 脚本正在扫描一个非常大的文件（约 480,000,000 行）并创建一个字典，该字典稍后将被写入 output 文件：

with open (huge_file,'r') as hugefile, open (final_file, 'w') as final:
  reader= csv.reader (hugefile, delimiter="\t")
  writer= csv.writer (final, delimiter="\t")

  d={} 

  for r in reader:
      v=r[0]+r[1]
      if v not in d.keys():
        d[v]=[r[5],r[4]]
      else:
        d[v].append([r[5],r[4]])

  for k,v in d.items():
    #analyses    
    nl = [different variables]
    writer.writerow(nl)

由于文件的大小，我想使用 16 个 CPU 来运行，但即使我在 bash 脚本中定义了 16 个 CPU，它也只使用 1 个 CPU。

我读了很多关于subprocess的内容，但它似乎不适用于这种情况。 我很想听听任何建议。

Answer 1

我建议使用多处理池来填充字典。

from multiprocessing import Pool

d = dict()

def func(r):
    v = r[0]+r[1]
    if v not in d:
        d[v] = [r[5], r[4]]
    else:
        d[v].append([r[5], r[4]])

with Pool(16) as p:
    p.map(func, reader)

类似地，可以将 Pool 应用于 dict d 和分析 function 来完成分析。

Answer 2

核心在这里帮不了你，因为字典操作很简单而且速度非常快。

这里有一个 I/O 问题，读取和写入文件是瓶颈。

如果您使用多处理模块，您可能会遇到其他问题。 构建的字典将彼此独立，因此您将拥有与其他数据重复的键。 如果必须保留 CSV 数据的顺序，可能是因为它是时间序列数据，则必须将字典中的 arrays 作为附加步骤进行合并然后排序，除非在合并 16 个字典时考虑到这个问题。 这也意味着您将 CSV 分成 16 个块并在每个内核上单独处理它们，以便您可以跟踪排序。

您是否考虑过将巨大的 CSV 文件读入 SQLite 数据库？ 这至少可以让您更好地控制数据的访问方式，因为 16 个进程可以在指定排序时同时访问数据。

我真的怀疑这里有什么可以并行化的。 即使您使用多处理模块，您也需要编写整个文件同时考虑整个字典，这限制了您并行化此任务的方式。

Answer 3

多处理很难应用，因为所有内容都需要分类到一个中心 dict d 中。 几个进程必须始终知道字典中已经有哪些键，这使得它变得非常复杂。 因此，更简单的解决方案是尝试加快处理速度，同时保持在一个进程内。 dict 和列表理解似乎是一个很好的前进方向：

# prepare dict keys and empty list entries:
d = {r[0]+r[1]: [] for r in reader}

# fill dict
[d[r[0]+r[1]].append([r[5], r[4]]) for r in reader]

# d is ready for analysis

Answer 4

您可以使用线程来执行此操作。 首先，将 function 中的代码分开，执行如下操作：

import threading

def your_func_name(result, reader, index):
    sliced_d = {}
    for r in reader:
        v=r[0]+r[1]
        if v not in sliced_d.keys():
            sliced_d[v]=[r[5],r[4]]
        else:
            sliced_d[v].append([r[5],r[4]])

    for k,v in sliced_d.items():
    #analyses    
        nl = [different variables]
        writer.writerow(nl)
    result[index] = sliced_d

现在，定义您想要使用的 CPU 数量并相应地对阅读器进行切片。 然后，将切片发送到线程中。

d = {}
cpus = 16
slice = len(reader)//cpus
results = [None] * slice
for i in range(cpus):
    if i < spawns-1:
        threads.append(Thread(your_func_name, (result, reader[i*slice, (i+1)*slice], i)))
    else:
        threads.append(Thread(your_func_name (result, reader[i*tot_threads:len(reader)], i))


for i in range(cpus):
    threads[i].join()

for i in range(cpus):
    d.update(result[i])

注意：代码可能有一些错误，因为我写了这样一个例子。

Answer 5

这是一个如何使用多个进程的想法（尚未使用大文件进行测试，但确信它会在调试时工作）。

第 1 步是使用 Linux split function 将大文件拆分为段：

 bash> split -l 10000000 hugefile segment

这将创建每个有 10,000,000 行的文件，名称将是segmentaa, segmentab, .... （参见split的man页）

现在 Python 程序读取这些文件段，每个文件段启动一个进程，然后将结果合并到一个字典中：

import multiprocessing as mp
import csv


# define process function working on a file segment
def proc_target(q, filename):
    with open(filename, 'r') as file_segment:
        reader = csv.reader(file_segment, delimiter="\t")

        dd = dict()

        def func(r):
            key = r[0] + r[1]
            if key in dd:
                dd[key].append([r[5], r[4]])
            else:
                dd[key] = [r[5], r[4]]

        [func(r) for r in reader]

        # send result via queue to main process
        q.put(dd)


if __name__ == '__main__':
    segment_names = ['segmentaa', 'segmentab', 'segmentac']:  # maybe there are more file segments ...
    processes = dict()  # all objects needed are stored in this dict
    mp.set_start_method('spawn')

    # launch processes
    for fn in segment_names :
        
        processes[fn] = dict()

        q = mp.Queue()
        p = mp.Process(target=proc_target, args=(q, fn))
        p.start()

        processes[fn]["process"] = p
        processes[fn]["queue"] = q

    # read results
    for fn in segment_names:
        processes[fn]["result"] = processes[fn]["queue"].get()
        processes[fn]["process"].join()

    # consolidate all results
    # start with first segment result and merge the others into it
    d = processes[segment_names[0]]["result"]

    # helper function for fast execution using list comprehension
    def consolidate(key, value):
        if key in d:
            d[key].append(value)
        else:
            d[key] = value

    # merge other results into d
    for fn in segment_names[1:]:
        [consolidate(key, value) for key, value in processes[fn]["result"]]

    # d is ready

为了避免 I/O 瓶颈，明智的做法是将段分布在多个磁盘上，并让并行进程并行访问不同的 I/O 资源。