優化 dask.distributed 調度以減少數據

Question

我有一個關於 dask.distributed 中任務的調度/執行順序的問題，以應對大型原始數據集的強數據縮減。

我們將 dask.distributed 用於從電影幀中提取信息的代碼。 它的具體應用是結晶學，但通常的步驟是：

1. 將 HDF5 文件（或其中一些連接的）中存儲為 3D 數組的電影幀讀入 dask 數組。 這顯然是相當重的 I/O
2. 將這些幀分組到通常包含 10 個移動靜止圖像的連續子堆棧中，其中的幀被聚合（求和或平均），從而產生單個 2D 圖像。
3. 在 2D 圖像上運行幾個計算量大的分析函數（例如某些特征的位置），返回一個結果字典，與電影本身相比，它可以忽略不計。

我們通過在步驟 1 和 2 中使用 dask.array API 來實現這一點（后者使用map_blocks ，塊/塊大小為一個或幾個聚合子堆棧），然后將數組塊轉換為 dask.delayed 對象（使用to_delayed ）將其傳遞給 function 執行實際數據縮減 function（步驟 3）。 我們注意正確對齊 HDF5 arrays 的塊、dask 計算和步驟 2 中的聚合范圍，以便每個最終延遲的 object（任務元素）的tasks圖非常干凈。 這是示例代碼：

def sum_sub_stacks(mov):
    # aggregation function
    sub_stk = []
    for k in range(mov.shape[0]//10):
        sub_stk.append(mov[k*10:k*10+10,...].sum(axis=0, keepdims=True))
    return np.concatenate(sub_stk)

def get_info(mov):
    # reduction function
    results = []
    for frame in mov:
        results.append({
            'sum': frame.sum(),
            'variance': frame.var()
            # ...actually much more complex/expensive stuff
        })
    return results

# connect to dask.distributed scheduler
client = Client(address='127.0.0.1:8786')

# 1: get the movie
fh = h5py.File('movie_stack.h5')
movie = da.from_array(fh['/entry/data/raw_counts'], chunks=(100,-1,-1))

# 2: sum sub-stacks within movie
movie_aggregated = movie.map_blocks(sum_sub_stacks, 
                                    chunks=(10,) + movie.chunks[1:],
                                    dtype=movie.dtype)

# 3: create and run reduction tasks
tasks = [delayed(get_info)(chk) 
         for chk in movie_aggregated.to_delayed().ravel()]

info = client.compute(tasks, sync=True)

理想的操作調度顯然是讓每個工作人員在單個塊上執行 1-2-3 序列，然后繼續下一個，這將保持 I/O 負載恆定，CPU 最大化和 memory 低。

相反，首先所有工作人員都試圖從文件中讀取盡可能多的塊（步驟 1），這會造成 I/O 瓶頸並迅速耗盡工作人員 memory 導致本地驅動器崩潰。 通常，在某些時候，worker 最終會移動到第 2/3 步，這會迅速釋放 memory 並正確使用所有 CPU，但在其他情況下，worker 會以不協調的方式被殺死或整個計算停止。 中間情況也會發生，幸存的工人僅在一段時間內表現得合理。

是否有任何方法可以提示調度程序以如上所述的首選順序處理任務，或者是否有其他方法可以改善調度行為？ 還是這種代碼/做事方式本身就很愚蠢？

Answer 1

首先，你所做的事情本身並沒有什么愚蠢的！

一般來說，Dask 試圖減少它所持有的臨時對象的數量，並且它還通過可並行性（圖的寬度和工作人員的數量）來平衡這一點。 調度很復雜，Dask 還使用了另一種優化，它將任務融合在一起，使它們更加優化。 有很多小塊，您可能會遇到問題： https://docs.dask.org/en/latest/array-best-practices.html?highlight=chunk%20size#select-a-good-chunk-size

Dask 確實有許多優化配置，我建議在考慮其他塊大小后使用它們。 我還鼓勵您通讀以下問題，因為圍繞調度配置進行了健康的討論。

最后，您可能會考慮額外的memory 配置您的工作人員，因為您可能希望更嚴格地控制每個工作人員應該使用多少 memory