dask-optimization

# Dask - Advanced Optimization & Cluster Tuning Parallel computing is not "free". In a distributed environment, the cost of moving data (network I/O) and scheduling tasks can often exceed the computation time. This guide focuses on minimizing overhead and maximizing throughput. ## When to Use - Your Dask jobs are failing with "KilledWorker" or "OutOfMemory" errors. - The Dask Dashboard shows a lot of "red" (communication) or "gray" (idle) time. - You need to process datasets that are 10x-100x larger than the total RAM of your cluster. - You are building custom distributed algorithms using `dask.delayed` or Futures. - You need to optimize resource allocation (CPU vs. Threads) for specific workloads. ## Reference Documentation - **Best Practices**: https://docs.dask.org/en/latest/best-practices.html - **Distributed Diagnostics**: https://distributed.dask.org/en/latest/diagnosing-performance.html - **Memory Management**: https://distributed.dask.org/en/latest/worker.html#memory-management - **Search patterns**: `client.scatter`, `dask.compute(optimize_graph=True)`, `repartition`, `client.restart` ## Core Principles ### 1. Communication is the Killer The fastest distributed task is the one that doesn't need data from another machine. Aim for data locality. ### 2. The Goldilocks Chunk Size - **Too small**: Scheduler is overwhelmed by millions of tiny tasks (Task Overhead). - **Too large**: Tasks don't fit in memory, causing disk spilling or worker crashes. - **Target**: 100MB - 300MB per chunk for most numeric data. ### 3. Computation vs. Serialization Every object sent to a worker must be serialized (pickled). Large Python objects (like complex dicts) passed as arguments can slow down the cluster significantly. ## Quick Reference: Performance Profiling ```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. ``` ## Critical Rules ### ✅ DO - **Use client.scatter** - If multiple tasks need the same large piece of data, send it to workers once. - **Prefer map_partitions** - In DataFrames, this allows you to run a single optimized pandas operation per chunk instead of row-wise logic. - **Use persist() for branching** - If you use the same intermediate result in two different computations, `persist()` it in memory to avoid re-calculating the entire graph twice. - **Profile with the Dashboard** - Watch the "Memory" and "Task Stream" tabs in real-time. - **Match Threads to Workload** - Use many threads for I/O bound tasks (web scraping, reading files) and 1 thread per worker for CPU-bound tasks (NumPy math, ML training) to avoid Global Interpreter Lock (GIL) issues. ### ❌ DON'T - **Don't use compute() in loops** - This pulls data to the local machine and destroys parallel efficiency. - **Don't pass Large Data as Arguments** - Instead of `delayed(func)(large_df)`, use `large_future = client.scatter(large_df)` then `delayed(func)(large_future)`. - **Don't ignore Data Skew** - If one worker has 10GB of data and others have 100MB, the cluster is only as fast as the slowest worker. Use `repartition` or `rechunk`. - **Don't use list(dask_collection)** - This forces an immediate compute of all elements into local memory. ## Data Locality & Communication ### Using scatter and Futures ```python # ❌ 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] ``` ## Memory Management Tuning ### Fighting the "KilledWorker" Worker memory has thresholds: - **Target (0.6)**: Dask tries to stay below this. - **Spill (0.7)**: Dask starts moving data to disk. - **Pause (0.8)**: Worker stops accepting new tasks. - **Terminate (0.95)**: OS or Dask kills the worker. ```python # 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}) ``` ## Task Graph Optimization ### Fusing Operations Dask automatically "fuses" many operations into one task to reduce scheduler overhead. ```python # 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) ``` ## Practical Workflows ### 1. Optimizing Large Joins (Shuffling) Joins are expensive because they require moving data between workers (shuffling). ```python 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. High-Throughput I/O with Parquet ```python 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. Managing a Long-running Cluster ```python # 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 ``` ## Advanced Configuration ### Resource Tagging Tell Dask to run specific tasks only on specific workers (e.g., those with a GPU). ```python # Start worker with: dask-worker ... --resources "GPU=1" # Submit task requesting resource future = client.submit(my_gpu_function, data, resources={'GPU': 1}) ``` ## Common Pitfalls and Solutions ### The "Zombie Worker" (serialization error) If your task requires a library that isn't installed on the workers, the task will fail repeatedly. ```python # ✅ Solution: Use pip_install or conda_install via client from dask.distributed import PipInstall client.register_plugin(PipInstall(packages=["scikit-learn"])) ``` ### Unmanaged Memory Python's garbage collector isn't immediate. Sometimes workers appear full because "Unmanaged Memory" hasn't been freed. ```python # ✅ Solution: Manually trigger GC on workers def worker_gc(): import gc return gc.collect() client.run(worker_gc) ``` ### Too Many Partitions If you have 10,000 partitions for a 1GB dataset, you'll spend more time scheduling than calculating. ```python # ✅ Solution: Repartition to fewer pieces ddf = ddf.repartition(npartitions=20) ``` Dask Optimization is the art of balancing resources. By understanding the flow of data through the network and the mechanics of worker memory, you can scale Python logic to planetary-scale datasets with industrial reliability.