在 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 資源。