dask-optimization

Dask - 高度な最適化とクラスターのチューニング

並列コンピューティングは「無料」ではありません。分散環境では、データの移動（ネットワーク I/O）やタスクのスケジューリングにかかるコストが、計算時間を上回ることがよくあります。このガイドでは、オーバーヘッドを最小限に抑え、スループットを最大化する方法に焦点を当てます。

使用する状況

• Dask のジョブが「KilledWorker」または「OutOfMemory」エラーで失敗する場合。
• Dask Dashboard に「赤」（通信）や「灰色」（アイドル状態）の時間が多く表示される場合。
• クラスター全体の RAM の 10～100 倍のデータセットを処理する必要がある場合。
• dask.delayed または Futures を使って、カスタムの分散アルゴリズムを構築する場合。
• 特定のワークロードに合わせてリソース割り当て（CPU 対スレッド）を最適化する必要がある場合。

リファレンスドキュメント

• ベストプラクティス: https://docs.dask.org/en/latest/best-practices.html
• 分散診断: https://distributed.dask.org/en/latest/diagnosing-performance.html
• メモリ管理: https://distributed.dask.org/en/latest/worker.html#memory-management
• 検索パターン: client.scatter、dask.compute(optimize_graph=True)、repartition、client.restart

基本原則

1. 通信は最大のボトルネックです

最も高速な分散タスクとは、別のマシンからデータを必要としないタスクです。データの局所性を目指してください。

2. ちょうどよいチャンクサイズ

• 小さすぎる場合: スケジューラーが何百万もの小さなタスクに圧倒されます（タスクのオーバーヘッド）。
• 大きすぎる場合: タスクがメモリに収まらず、ディスクへのスピルやワーカーのクラッシュが発生します。
• 目安: ほとんどの数値データでは、チャンクあたり 100MB～300MB です。

3. 計算とシリアライズ

ワーカーに送信するオブジェクトはすべてシリアライズ（pickle 化）する必要があります。引数として渡す大きな Python オブジェクト（複雑な dict など）は、クラスターの処理速度を大幅に低下させることがあります。

クイックリファレンス：パフォーマンスのプロファイリング

from dask.distributed import Client, performance_report

client = Client("tcp://scheduler-address:8786")

# Generate a detailed HTML report of the computation
with performance_report(filename="dask-report.html"):
    result = big_computation.compute()

# Tip: Check the "Task Stream" for gaps. Gaps mean workers are idle 
# waiting for the scheduler or network.

重要なルール

✅ 推奨事項

• client.scatterを使用する - 複数のタスクが同じ大きなデータを必要とする場合は、一度だけワーカーに送信します。
• map_partitionsを優先する - DataFrame では、行単位の処理ではなく、チャンクごとに最適化された pandas 操作を 1 回実行できます。
• 分岐にはpersist()を使用する - 同じ中間結果を異なる 2 つの計算で使う場合は、メモリに persist() して、グラフ全体を 2 回再計算しないようにします。
• Dashboardでプロファイリングする - 「Memory」と「Task Stream」タブをリアルタイムで確認します。
• ワークロードに合わせてスレッド数を調整する - I/O バウンドのタスク（ウェブスクレイピング、ファイルの読み込み）には多くのスレッドを使用し、CPU バウンドのタスク（NumPy の演算、ML のトレーニング）では、Global Interpreter Lock（GIL）の問題を避けるため、ワーカーあたり 1 スレッドを使用します。

❌ 非推奨事項

• ループ内でcompute()を使用しない - データがローカルマシンに引き寄せられ、並列処理の効率が損なわれます。
• 大きなデータを引数として渡さない - delayed(func)(large_df) の代わりに、large_future = client.scatter(large_df) を使ってから delayed(func)(large_future) を使用します。
• データの偏りを無視しない - あるワーカーに 10GB のデータがあり、他のワーカーには 100MB しかない場合、クラスターの速度は最も遅いワーカーに左右されます。repartition または rechunk を使用します。
• list(dask_collection)を使用しない - すべての要素をローカルメモリに即座に計算して取り込みます。

データの局所性と通信

scatterとFuturesの使用

# ❌ BAD: Large object sent to every task (High Overhead)
large_lookup = load_heavy_dict()
results = [delayed(process)(x, large_lookup) for x in data]

# ✅ GOOD: Scatter once, use reference
[large_future] = client.scatter([large_lookup], broadcast=True)
results = [delayed(process)(x, large_future) for x in data]

メモリ管理のチューニング

「KilledWorker」への対処

ワーカーのメモリには、次のしきい値があります。

• Target（0.6）: Dask はこの値を下回るようにします。
• Spill（0.7）: Dask はデータをディスクに移動し始めます。
• Pause（0.8）: ワーカーは新しいタスクの受け付けを停止します。
• Terminate（0.95）: OS または Dask がワーカーを終了させます。

# Adjusting worker limits via config (distributed.yaml or code)
import dask
dask.config.set({"distributed.worker.memory.target": 0.45})
dask.config.set({"distributed.worker.memory.spill": 0.55})

タスクグラフの最適化

操作の融合

Dask はスケジューラーのオーバーヘッドを減らすため、多くの操作を自動的に 1 つのタスクに「融合」します。

# Multiple operations on a DataArray
da = da + 1
da = da * 2
da = da.sum()

# When computing, Dask optimizes the graph
# You can inspect it:
da.visualize(filename='graph.pdf', optimize_graph=True)

実践的なワークフロー

1. 大規模な結合の最適化（シャッフル）

結合ではワーカー間でデータを移動する（シャッフルする）必要があるため、処理コストが高くなります。

def optimized_join(left_ddf, right_ddf):
    # 1. Ensure right side is small enough to be broadcasted
    # or ensure both are partitioned by the join key.
    
    # If one DF is small (e.g. < 100MB)
    # result = left_ddf.map_partitions(lambda df: df.merge(small_df_local, on='key'))
    
    # 2. If both are large, set the index first (triggers a shuffle once)
    left_ddf = left_ddf.set_index('key') 
    right_ddf = right_ddf.set_index('key')
    
    # 3. Subsequent joins will be "locally aligned" (Zero communication)
    return left_ddf.merge(right_ddf, left_index=True, right_index=True)

2. Parquetによる高スループットのI/O

def fast_save(ddf, path):
    # 1. Categorical columns save massive space
    ddf = ddf.categorize() 
    
    # 2. Write with efficient compression
    # 'snappy' is usually the best balance for speed
    ddf.to_parquet(path, engine='pyarrow', compression='snappy', 
                   write_metadata_file=True)

3. 長時間稼働するクラスターの管理

# Prevent a "dirty" cluster from slowing down
def clean_cluster_state():
    client.cancel(list(client.futures)) # Clear all references
    client.restart() # Hard reset all workers
    import gc
    gc.collect() # Local cleanup

高度な設定

リソースのタグ付け

特定のタスクを特定のワーカー（GPU を搭載したワーカーなど）でのみ実行するよう、Dask に指示します。

# Start worker with: dask-worker ... --resources "GPU=1"

# Submit task requesting resource
future = client.submit(my_gpu_function, data, resources={'GPU': 1})

よくある落とし穴と解決策

「ゾンビワーカー」（シリアル化エラー）

タスクに必要なライブラリがワーカーにインストールされていない場合、タスクは繰り返し失敗します。

# ✅ Solution: Use pip_install or conda_install via client
from dask.distributed import PipInstall
client.register_plugin(PipInstall(packages=["scikit-learn"]))

管理対象外のメモリ

Python のガベージコレクターはすぐには動作しません。「管理対象外のメモリ」が解放されていないために、ワーカーのメモリがいっぱいに見えることがあります。

# ✅ Solution: Manually trigger GC on workers
def worker_gc():
    import gc
    return gc.collect()

client.run(worker_gc)

パーティションが多すぎる

1 GB のデータセットを 10,000 個のパーティションに分割すると、計算よりもスケジューリングに時間がかかってしまいます。

# ✅ Solution: Repartition to fewer pieces
ddf = ddf.repartition(npartitions=20)

Dask の最適化とは、リソースのバランスを取る技術です。ネットワークを介したデータの流れと、ワーカーのメモリの仕組みを理解すれば、産業レベルの信頼性で Python のロジックを地球規模のデータセットにスケールできます。