dask-optimization

用途

本技能是 Dask 的高级优化与集群调优子技能，面向分布式计算性能问题。并行计算并非“免费”：在分布式环境中，数据移动（网络 I/O）与任务调度的开销常常超过计算本身。本指南的目标是最小化开销、最大化吞吐量。

何时使用

Dask 作业因 “KilledWorker” 或 “OutOfMemory” 错误而失败。
Dask Dashboard 显示大量“红色”（通信）或“灰色”（空闲）时间。
需要处理比集群总内存大 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

核心原则

通信是杀手：最快的分布式任务是不需要从其他机器取数据的任务，应追求数据本地化。
分块大小的“金发姑娘”原则：太小会让调度器被海量小任务压垮（任务开销）；太大则任务放不进内存，导致磁盘溢写或 Worker 崩溃。多数数值数据的目标是每块 100MB - 300MB。
计算与序列化：发送给 Worker 的每个对象都必须序列化（pickle）。作为参数传递的大型 Python 对象（如复杂字典）会显著拖慢集群。

性能剖析快速参考

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：若多个任务需要同一份大数据，只向 Worker 发送一次。
优先使用 map_partitions：在 DataFrame 中，这允许每个分块运行一次优化过的 pandas 操作，而非逐行逻辑。
对分支计算使用 persist()：若同一中间结果用于两个不同计算，将其 persist() 到内存，避免整个图被计算两次。
用 Dashboard 剖析：实时观察 “Memory” 和 “Task Stream” 标签页。
让线程数匹配工作负载：I/O 密集型任务（网页抓取、读文件）用多线程；CPU 密集型任务（NumPy 数学、ML 训练）每个 Worker 用 1 个线程，以避免全局解释器锁（GIL）问题。

不应做：

不要在循环中使用 compute()：这会把数据拉回本地机器，破坏并行效率。
不要把大数据作为参数传递：不要用 delayed(func)(large_df)，而应使用 large_future = client.scatter(large_df)，再 delayed(func)(large_future)。
不要忽视数据倾斜：若一个 Worker 有 10GB 数据而其他只有 100MB，集群速度取决于最慢的 Worker。使用 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”：Worker 内存有若干阈值：

Target (0.6)：Dask 尽量保持在此之下。
Spill (0.7)：Dask 开始将数据移到磁盘。
Pause (0.8)：Worker 停止接受新任务。
Terminate (0.95)：操作系统或 Dask 杀死该 Worker。

# 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)

实用工作流

优化大型 Join（Shuffle）：Join 开销大，因为需要在 Worker 之间移动数据（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)

使用 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)

管理长时间运行的集群：

# 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 只在特定 Worker（例如带 GPU 的）上运行特定任务。

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

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

常见陷阱与解决方案

“僵尸 Worker”（序列化错误）：若任务需要 Worker 上未安装的库，任务会反复失败。

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

非托管内存：Python 的垃圾回收不是即时的。有时 Worker 看起来已满，是因为“非托管内存”尚未释放。

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

client.run(worker_gc)

分区过多：若 1GB 数据集有 10,000 个分区，调度时间会超过计算时间。

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

适用场景与限制

适用于需要调优 Dask 分布式集群性能、处理超大规模数据集、排查内存与调度瓶颈的场景。Dask 优化是平衡资源的艺术：理解数据在网络中的流动与 Worker 内存机制，才能以工业级可靠性将 Python 逻辑扩展到行星级数据集。具体版本兼容性、集群部署环境要求等细节暂未提供。