Dask and the optimization workflow¶

Objectives

Build and inspect a lazy task graph.
Use Dask delayed and chunked collections without exceeding one node.
Select a thread or process scheduler from the work inside each task.
Measure task-size, chunk-size, and worker-count trade-offs.
Combine kernel acceleration and task parallelism in a reproducible recommendation.

Instructor note

40 min teaching/type-along
35 min exercises, knowledge check, and synthesis
Build the first lazy graph live and ask learners to predict scheduler behavior before each comparison.

Dask represents a computation as a graph whose nodes are tasks and whose edges are dependencies. A scheduler executes ready tasks while respecting those dependencies. On one node, this gives us a common interface for threads, processes, and chunked NumPy- or pandas-like collections.

Dask is not automatically faster than NumPy or a loop. It helps when the graph exposes useful parallelism, each task contains enough work, and the data movement fits the machine.

Lazy execution with `delayed`¶

Calling a delayed function constructs a task; it does not run the function:

from dask import compute, delayed

lazy_values = [delayed(abs)(value) for value in (-3, -2, -1)]
lazy_total = delayed(sum)(lazy_values)
result = lazy_total.compute()

The graph contains three independent abs tasks followed by one reduction. Calling .compute() or dask.compute(...) submits the graph and returns ordinary in-memory Python values.

Compute once

Construct the complete graph and call compute at the outer boundary. Calling compute inside a construction loop forces small serial executions and prevents Dask from seeing global parallelism.

The module example creates independent integrations for several phase values:

"""Run independent integration tasks with a configurable local scheduler."""

import argparse
from math import pi
from time import perf_counter

from dask import compute, delayed


def build_graph(function, n: int, phases: list[float]):
    """Return lazy integration tasks without executing them."""
    return [delayed(function)(n, phase) for phase in phases]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--implementation", choices=("python", "cython"), default="python")
    parser.add_argument("--scheduler", choices=("threads", "processes"), default="processes")
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--tasks", type=int, default=8)
    parser.add_argument("--steps", type=int, default=1_000_000)
    args = parser.parse_args()

    if args.implementation == "cython":
        from integrate_cython import integrate
    else:
        from integrate_python import integrate

    phases = [i / args.tasks for i in range(args.tasks)]
    graph = build_graph(integrate, args.steps, phases)

    start = perf_counter()
    results = compute(*graph, scheduler=args.scheduler, num_workers=args.workers)
    elapsed = perf_counter() - start

    assert all(abs(result - pi) < 1e-9 for result in results)
    print(
        f"{args.implementation=}, {args.scheduler=}, {args.workers=}, "
        f"{args.tasks=}, elapsed={elapsed:.6f} s"
    )


if __name__ == "__main__":
    main()

After building the Cython extension, compare configurations from content/episodes/code:

$ python dask_batch.py --implementation python --scheduler processes --workers 2
$ python dask_batch.py --implementation python --scheduler threads --workers 2
$ python dask_batch.py --implementation cython --scheduler threads --workers 2

The pure-Python loop normally favors processes because it holds the GIL. The Cython loop explicitly releases the GIL, so threads can run kernels concurrently while sharing memory. Treat that as a hypothesis and verify it on the allocated node.

Threads, processes, and `LocalCluster`¶

Dask’s local schedulers are convenient for short scripts:

results = compute(*tasks, scheduler="threads", num_workers=4)
results = compute(*tasks, scheduler="processes", num_workers=4)

For diagnostics and a dashboard, create a local distributed cluster:

from dask.distributed import Client, LocalCluster

cluster = LocalCluster(
    n_workers=4,
    threads_per_worker=1,
    processes=True,
    dashboard_address=":8787",
)
client = Client(cluster)

# Build and compute work here.

client.close()
cluster.close()

This is still single-node execution. processes=True gives each worker its own Python interpreter. For native kernels that release the GIL and share large read-only arrays, one process with several threads may reduce serialization and memory duplication.

Chunked collections¶

Dask Array divides an array into NumPy chunks. Operations remain lazy until computed:

import dask.array as da

x = da.random.default_rng(42).normal(size=(8_000, 8_000), chunks=(1_000, 1_000))
column_means = (x * x).mean(axis=0)
print(column_means)          # a lazy Dask Array
answer = column_means.compute()

Choose chunks by working backward from memory and useful task duration:

a chunk and its temporaries must fit comfortably in memory;
there should be enough chunks to keep workers occupied;
each chunk should contain much more work than scheduling it;
chunk boundaries should match later operations to avoid expensive rechunking.

Dask DataFrame applies the same idea to pandas partitions, and Dask Bag supports less-structured Python records. Use the narrowest collection that matches the data. For small in-memory data, plain NumPy or pandas is often simpler and faster.

Task-size experiment¶

Exercise

Use dask_batch.py with the pure-Python implementation and process scheduler.

Keep the total work approximately fixed while trying 2, 8, 32, and 128 tasks. Adjust --steps so tasks × steps remains constant.
Try 1, 2, and 4 workers, without exceeding the cores allocated to you.
Record the fastest elapsed time for each configuration and compute speedup against one worker.

Where do smaller tasks stop helping? Give two overheads that explain the result.

Solution

There is no universal optimum. Very few tasks can leave workers idle, while many short tasks increase graph construction, scheduler coordination, process communication, and result-handling costs. A sound conclusion includes the actual CPU allocation and configuration rather than claiming one task count is always best.

Scheduler-selection experiment¶

Exercise

Build the Cython extension and compare all four combinations of implementation (python, cython) and scheduler (threads, processes) with the same workers and total work.

Before running, predict which scheduler suits each implementation. Check results and explain any difference between prediction and measurement.

Solution

The Python kernel holds the GIL and should generally favor processes for CPU parallelism. The compiled kernel’s loop uses with nogil, so threads can execute it concurrently and avoid process serialization. On a memory-bandwidth-limited or oversubscribed system, neither may scale linearly; startup costs can also dominate small runs.

Array chunking exercise¶

Note

This is an optional extension for a longer workshop or independent study. The four-hour delivery uses the task-size experiment, scheduler-selection experiment, integrated challenge, and knowledge check.

Exercise

Run the Dask Array example with chunks of (250, 250), (1_000, 1_000), and (4_000, 4_000). Before computing, inspect column_means.npartitions and the graph size with len(column_means.dask). Record elapsed time and peak-memory observations from the dashboard or system monitor.

Which choice gives enough parallel work without creating a very large graph or oversized chunks?

Solution

The answer depends on the machine. Tiny chunks create many scheduled tasks; very large chunks reduce parallel slack and increase memory per task. A defensible choice reports task count, worker configuration, approximate bytes per chunk, and measured time.

Keypoints

Dask constructs graphs lazily and executes them when compute is called.
Threads suit I/O and native kernels that release the GIL; processes suit CPU-heavy Python code.
Chunked collections extend familiar array and table operations, but add scheduling overhead.
Task and chunk sizes must balance parallel slack, memory, and overhead.
Configure only the cores and memory allocated on the current node.

Putting Cython and Dask together¶

The two tools address different layers:

Cython reduces interpreter overhead inside the integration kernel.
Dask schedules multiple independent integrations across the node.

The combined design is useful only if the application contains both a costly kernel and enough independent work. A small workload may favor serial Cython because scheduling overhead exceeds the time saved.

Integrated optimization challenge

Produce a short performance report using the code in content/episodes/code.

Define a representative workload and correctness tolerance.
Measure a serial pure-Python baseline.
Measure serial Cython, parallel pure Python with Dask processes, and parallel Cython with Dask threads.
Keep total useful work fixed. Record CPU model/allocation, Python and package versions, worker count, task count, steps per task, scheduler, and native thread settings.
Report elapsed time, speedup against the baseline, and whether results agree.
Recommend one version and identify the next likely bottleneck.

Solution

A valid recommendation is evidence-based, not predetermined. On many machines, Cython removes the largest per-iteration cost and threaded Dask can then run independent nogil kernels without process copies. A smaller workload may favor serial Cython. Pure-Python Dask processes may be useful when changing the kernel is impractical. Every comparison must preserve total work and verify results against \(\pi\).

A reusable optimization workflow¶

Establish a correct, representative baseline.
Locate the expensive region and identify whether it is Python execution, native computation, I/O, memory, or independent tasks.
Improve algorithms and use existing optimized libraries before adding machinery.
Move a stable loop-heavy kernel to Cython when native types and operations can remove Python overhead.
Use Dask when work naturally forms tasks or chunks, selecting threads or processes from the code inside each task.
Tune task size, chunk size, worker count, and native threads on the allocated node.
Recheck correctness and report the complete measurement context.

Keypoints

Kernel acceleration and task parallelism solve different problems and can be composed.
The simplest implementation meeting the performance target is usually the easiest to maintain.
A speedup claim is incomplete without a baseline, fixed workload, configuration, and correctness check.

Dask and the optimization workflow¶

Lazy execution with delayed¶

Threads, processes, and LocalCluster¶

Chunked collections¶

Task-size experiment¶

Scheduler-selection experiment¶

Array chunking exercise¶

Putting Cython and Dask together¶

A reusable optimization workflow¶

Lazy execution with `delayed`¶

Threads, processes, and `LocalCluster`¶