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 操作，而不是逐行处理。
• 使用 persist() 处理分支 - 如果在两种不同的计算中使用同一个中间结果，可将其 persist() 到内存中，避免重复计算整个图。
• 使用 Dashboard 进行性能分析 - 实时查看“Memory”和“Task Stream”选项卡。
• 根据工作负载匹配线程数 - 对 I/O 密集型任务（网页抓取、读取文件）使用多个线程；对 CPU 密集型任务（NumPy 数学运算、ML 训练）使用每个工作进程 1 个线程，以避免全局解释器锁（GIL）问题。

❌ 不应该做

• 不要在循环中使用 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”

工作进程内存有以下阈值：

• 目标（0.6）：Dask 会尽量保持在此阈值以下。
• 溢写（0.7）：Dask 开始将数据移至磁盘。
• 暂停（0.8）：工作进程停止接受新任务。
• 终止（0.95）：操作系统或 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 会自动将多个操作“融合”为一个任务，以减少调度开销。

# 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. 优化大型联接（Shuffle）

联接操作开销很大，因为需要在工作进程之间移动数据（Shuffle）。

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

高级配置

资源标记

让 Dask 只在特定工作节点（例如配备 GPU 的工作节点）上运行指定任务。

# 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 逻辑扩展到行星级数据集。