all_lessons/ray/05lesson 6 / 17

Part I - Core mental model

Distributed Patterns: Map, Shuffle, Reduce

Lesson 04 gave you the runtime under the API: a scheduler that places tasks, an object store that carries their results, and the ownership model that lets a value live on a node rather than route through your driver. You now have all the parts but no pattern to assemble them into. This lesson installs the one composite shape that recurs at every scale in distributed computing — map, then shuffle, then reduce. We build it by hand as a word count, then expose the two facts that decide whether it scales: the shuffle, not the map, is where the cost lives; and reducing on the driver is the bottleneck that silently caps every naive attempt. Get this shape right and Ray Data, Train, and Tune later read as variations, not new systems.

Book source
Chapter 2, 'A Simple MapReduce Example with Ray' - mapping, shuffling document data, and reducing word counts. We keep the book's word-count spine and push past it to the engineering that the toy hides: shuffle cost as O(P²) network connections, tree-reduce depth, combiners, and partition tuning.
Linear position
Prerequisite: Lesson 04 (Runtime, scheduling, the object store) — you know that f.remote() returns an ObjectRef immediately, that results live in the per-node object store, and that the driver is the first thing you turn into a bottleneck.
New capability: you can take any "aggregate over a large dataset" problem, decompose it into map / shuffle / reduce, predict which stage will dominate, and structure the reduce so the driver and the network do not become the ceiling.
The plan
Five moves. (1) Frame map → shuffle → reduce as a three-stage graph and define each stage by the problem it solves. (2) Build a real Ray word count, mapping over document partitions with a shared read-only payload via ray.put. (3) Expose the driver-bottleneck anti-pattern — ray.get everything, reduce on one process — and show why it serializes the program. (4) Fix it with a tree reduce and work the depth and throughput numbers (1024 partials: 1023 sequential merges vs ~10 parallel levels). (5) Cost the shuffle (O(P²) connections, network-bound), add combiners to shrink it, tune partition count, and connect the shape forward.

1 · Three stages, each named by its job

Almost every "compute one answer from a dataset too big for one machine" problem decomposes into the same three stages. They are not arbitrary; each exists to solve a distinct problem, and each leaves a distinct cost.

Map
Run the same function independently over each partition (a slice of the input). No partition needs any other, so all of them run in parallel as Ray tasks. This is where parallelism is created. Cheap to scale: more partitions → more concurrent tasks.
Shuffle
Shuffle = the all-to-all redistribution that moves every record from the mapper that produced it to the reducer responsible for its key. It is the only stage where data crosses the network between workers in bulk. This is where the cost lives.
Reduce
Aggregate all values that now share a key into one result — sum the counts, average the gradients, concatenate the rows. Parallel across keys, but the per-key work and the final merge are where pressure concentrates.

The canonical example, and the book's, is counting words across a large corpus. Map: each task reads a partition of documents and emits a partial count, a dict like {"ray": 12, "task": 7}. Shuffle: all the partial counts for the word "ray" must end up together so they can be summed. Reduce: add them up into one global count per word. Hold this concrete picture; everything below is about making the shuffle and reduce stages not destroy the parallelism the map stage created.

One term to pin down now, because the whole reduce strategy turns on it. A reduce operation is associative if grouping does not matter — (a + b) + c = a + (b + c) — and commutative if order does not matter — a + b = b + a. Sum, count, max, min, and set-union are all both. This is not pedantry: only an associative, commutative reduce can be split across a tree of workers and combined in any order, which is exactly the freedom §4 spends. A reduce that is neither (say, "concatenate strings in original document order") forbids that rearrangement and forces a more careful, often serial, merge.

2 · The map stage, written in Ray

Map is the easy stage and Ray makes it nearly free to express: one remote function, one call per partition. The only subtlety is shared read-only data. Suppose every map task needs the same 200 MB stop-word list and vocabulary table. If you pass it as a normal argument to 256 tasks, Ray serializes and ships 200 MB × 256 ≈ 51 GB across the cluster — for a table that never changes. The fix from lesson 04 is ray.put: place the value in the object store once, get back an ObjectRef, and pass the ref. Now each task reads it zero-copy if it is on the same node, and Ray fetches it at most once per node.

import ray
ray.init()

# Shared, read-only payload: put it in the object store ONCE.
# Without this, the 200 MB table is re-serialized into every map task.
vocab = load_vocab()                      # ~200 MB dict/set
vocab_ref = ray.put(vocab)                # one copy in the object store

@ray.remote
def map_count(doc_partition, vocab_ref):
    # Ray auto-dereferences a top-level ObjectRef arg before the task runs,
    # so `vocab` here is the real object, read zero-copy on this node.
    vocab = vocab_ref
    counts = {}
    for doc in doc_partition:             # a partition = a slice of documents
        for word in doc.split():
            if word in vocab:
                counts[word] = counts.get(word, 0) + 1
    return counts                         # a PARTIAL count: dict per partition

# Fan out: one map task per partition. Returns instantly — refs, not results.
partition_refs = [map_count.remote(p, vocab_ref) for p in partitions]

Each map_count.remote(...) returns an ObjectRef immediately, so the loop fires all the map tasks without blocking. With 256 partitions and, say, 64 CPU cores in the cluster, ~64 run at once and the map stage finishes in roughly 256 / 64 = 4 waves. The map stage scales by adding partitions and cores; it is not where scaling breaks. The next two stages are.

3 · The driver-bottleneck anti-pattern

The obvious way to finish the word count is to pull every partial dict back to the driver and merge them in a loop. It is also the single most common way to throw away everything Ray just bought you.

# ANTI-PATTERN: shuffle + reduce both happen on the driver.
partials = ray.get(partition_refs)        # pulls ALL 256 dicts to the driver
total = {}
for counts in partials:                   # serial merge, one process, one core
    for word, n in counts.items():
        total[word] = total.get(word, 0) + n

Two distinct failures hide in those five lines. First, every byte flows through one process. The 256 partial dicts are not free — say each is 8 MB of distinct vocabulary counts, so ray.get drags ~2 GB into the driver's heap and deserializes all of it on one core. The driver's inbound network link and its single Python thread become the funnel the entire cluster waits behind. Second, the merge loop is serial. Merging 256 dicts is 255 sequential merge steps on one core; merging 1,000,000 partitions would be 999,999 sequential steps. This is precisely the inherently-serial section that Amdahl's law (lesson 00) says caps your speedup — you parallelized the map and then handed the result to a single-threaded reduce. Add 1,000 more workers and this code does not get faster, because the ceiling is the driver, not the worker count.

The real contract: a reduce is itself a distributed problem
The instinct "fan out the expensive part, collect the answers, combine them here" feels safe because the combine looks cheap. At scale it is not. If the map stage produced N partial results, a driver-side reduce is O(N) bytes through one network link and O(N) sequential merge steps on one core. The reduce must be parallelized too — and because our merge is associative and commutative (§1), we are free to do exactly that.

4 · Tree reduce — turning O(N) into O(log N)

A tree reduce merges partial results in pairs (or small groups) using reduce tasks, level by level, instead of all at once on the driver. Level 0 merges 256 partials into 128; level 1 merges 128 into 64; and so on until one result remains. Each level runs its merges in parallel across the cluster, and the driver only ever touches the single final result. Because the merge is associative and commutative, the order and grouping are free to rearrange this way without changing the answer.

@ray.remote
def merge_counts(*partials):              # combine a small group of dicts
    out = {}
    for counts in partials:
        for word, n in counts.items():
            out[word] = out.get(word, 0) + n
    return out                            # itself a partial — feeds the next level

def tree_reduce(refs, fanin=2):
    # Each level launches parallel merge tasks; depth is log_fanin(N).
    while len(refs) > 1:
        next_level = []
        for i in range(0, len(refs), fanin):
            group = refs[i:i + fanin]     # ObjectRefs; Ray dereferences them
            next_level.append(merge_counts.remote(*group))
        refs = next_level
    return refs[0]

total_ref = tree_reduce(partition_refs)   # build the tree, don't block yet
total = ray.get(total_ref)                # ONE small object reaches the driver

The driver now receives exactly one object — the final merged dict — instead of all 256, and it never runs a merge itself. The intermediate dicts move worker-to-worker through the object store, and the scheduler can place each merge near the bytes it consumes.

Worked number — linear vs tree reduce on 1024 partials
Take N = 1024 partial counts. Linear (driver) reduce: 1,024 − 1 = 1023 sequential merges, all on one core, all bytes through one process — the critical path is 1023 merge-times long. Tree reduce, fan-in 2: depth is log₂(1024) = 10 levels. Level 0 runs 512 merges in parallel, level 1 runs 256, … level 9 runs 1. The critical path is 10 merges, not 1023. If one pairwise merge takes 20 ms, the linear reduce's critical path is 1023 × 20 ms ≈ 20.5 s of serial work; the tree's critical path is 10 × 20 ms = 200 ms — a ~100× shorter critical path, and the driver's inbound bytes drop from ~all-of-them to one final dict. The total work is similar (~1023 merges either way), but the tree spreads it across the cluster and cuts the depth from O(N) to O(log N).

Two practical notes. A larger fan-in (say 8) makes a shallower tree — log₈(1024) ≈ 3.3 levels — at the cost of bigger, slower individual merges; fan-in is a knob trading depth against per-task size. And you do not always have to wait for the whole tree: ray.wait lets you consume completed sub-results as they finish, which matters when the reduce feeds a streaming consumer rather than a single final answer.

5 · The shuffle is the expensive stage — and why

The word count above quietly folded the shuffle into the merge, because summing dicts does the grouping-by-key implicitly. But name the shuffle explicitly, because in real pipelines it is a separate, dominating stage. A true shuffle with P map partitions and P reduce partitions has every mapper send a slice of its output to every reducer — an all-to-all exchange. That is P × P = P² logical connections, and crucially, the data moves across the network between workers rather than staying local.

This is why shuffle, not map, dominates at scale. Map cost grows linearly in the data and parallelizes cleanly. Shuffle cost grows with the all-to-all fan-out and is bounded by network bandwidth, the scarcest resource in a cluster. A straggler — one task far slower than its peers, here usually a reducer that drew an unlucky share of bytes or a hot key — stalls the whole stage, because the next stage cannot start until the slowest shuffle target finishes.

Worked number — shuffle data volume and the network ceiling
Suppose the map stage emits 500 GB of intermediate key-value pairs that must be shuffled across P = 200 partitions on a cluster wired at 10 Gigabit Ethernet (≈ 1.25 GB/s usable per node, ~16 nodes). In a full shuffle, essentially all 500 GB crosses the network once. Aggregate cluster bandwidth ≈ 16 × 1.25 = 20 GB/s, so the shuffle alone takes at least 500 / 20 = 25 s of pure data movement — and that is the floor, before any skew or straggler. The connection count is 200 × 200 = 40,000 logical transfers. Now compare the map stage: if mapping that 500 GB is 500 GB ÷ (16 nodes × ~2 GB/s compute-read) ≈ 16 s but fully parallel and local, the shuffle's 25 s of network-bound, all-to-all movement is the stage that sets the wall-clock. Double P to 400 and the connections quadruple to 160,000 while the bytes stay ~500 GB — more, smaller transfers and more scheduling overhead, for no bandwidth relief.

Combiners — shrink the shuffle before it happens

A combiner is map-side pre-aggregation: instead of emitting one record per occurrence, each map task locally aggregates its own partition first and emits the partial result. Our map_count already does this — it returns {"ray": 12}, not twelve copies of "ray". That is the highest-leverage optimization in the whole pattern, because it attacks the dominant cost directly: it cuts the bytes that enter the shuffle.

Worked number — what the combiner saves
A 1 TB corpus over a 50,000-word vocabulary, 200 partitions. Without a combiner: each map task emits one (word, 1) pair per word occurrence — roughly the full 1 TB of pairs flows into the shuffle. At 20 GB/s aggregate that is 1000 / 20 = 50 s of shuffle just to move raw pairs. With a combiner: each of the 200 map tasks emits at most one entry per word it saw, so each partial dict is bounded by the 50,000-word vocabulary — say ~8 MB. Total shuffle volume ≈ 200 × 8 MB = 1.6 GB, which crosses 20 GB/s in 1.6 / 20 = 0.08 s. The combiner cut shuffle volume from ~1 TB to ~1.6 GB — a ~600× reduction — and only works because count is associative and commutative (§1), so partial sums are valid.

Partition count is a real knob

How many partitions P should you choose? It is a genuine trade-off with a wrong answer on both ends.

Too few partitionsToo many partitions
Fewer tasks than cores → idle workers, no parallelism. A 64-core cluster with 8 partitions runs 56 cores idle.Tasks finer than ~tens of ms → per-task scheduling overhead (lesson 04) dominates the real work.
Coarse partitions → big, uneven tasks → one slow partition is a straggler the whole stage waits on.P² shuffle connections explode → many tiny transfers, metadata and connection-setup overhead swamp bandwidth gains.
Less room to retry: losing one coarse task re-does a lot of work.Each shuffle slice is tiny → poor network utilization, lots of bookkeeping per byte moved.

The usual sweet spot: enough partitions that every core has several tasks to chew through (so a straggler can be rebalanced and a lost task is cheap to retry), but coarse enough that each task does at least tens of milliseconds of work and the P² shuffle does not blow up. A common rule of thumb is a small multiple — 2× to 4× — of the cluster's core count, then tuned by watching for stragglers (raise P) or scheduling overhead (lower P).

6 · The same shape, everywhere later

This is the payoff for building it by hand. The map → shuffle → reduce graph is not a one-off recipe; it is the skeleton under most of the libraries in Part II.

Ray Data (lesson 09)Datasets are lists of blocks; map_batches is the map stage, groupby/repartition is the shuffle, and aggregations are the reduce — pipelined so data bigger than cluster RAM still flows. Same graph, hardened and streamed.
Ray Train (lesson 10)Data-parallel training is a reduce in disguise: each worker computes a local gradient (map), then allreduce sums them across workers (shuffle + tree reduce) so every replica gets the same averaged gradient. Allreduce is literally a ring/tree reduce optimized for the network.
Ray Tune (lesson 08)Many trials run in parallel (map), and their metrics are aggregated and compared by the scheduler (reduce) to decide which trials to keep. The fan-out/fan-in is the same shape, with the reduce making a control decision rather than a sum.

When lesson 10 says "allreduce," you will recognize a tree reduce that someone tuned for GPU interconnects; the cross-track note on collectives in system_ml makes that mechanism explicit. The shape you just built is the thing those systems are specializing.

7 · Failure modes & checklist

Failure modes

  • Driver-side reduce. ray.get all partials then merge in a loop. All bytes through one process, an O(N) serial merge — the classic Amdahl ceiling. Use a tree reduce.
  • No combiner. Emitting raw per-record output into the shuffle when a partial aggregate would do. Shuffle volume can be orders of magnitude larger than necessary; pre-aggregate map-side first.
  • Hot-key skew. One key (a stop-word like "the", or one popular product) lands all its bytes on a single reducer, which becomes a straggler stalling the stage. Salt the key or aggregate it hierarchically.
  • Task count ≠ parallelism. 10,000 tasks with 9,900 holding trivial bytes is not 10,000× parallel; check bytes-per-task, not just task count.
  • Reducing a non-associative op as if it were one. Tree reduce silently reorders work; if the operation depends on order or grouping, the answer is wrong, not just slow.

Implementation checklist

  • Is my reduce associative and commutative? (If not, can I make it so, or must the merge stay ordered?)
  • Am I reducing on the driver, or in a tree of reduce tasks?
  • Did I add a combiner so map tasks emit partial aggregates, not raw records?
  • Did I ray.put shared read-only data once instead of passing it into every map task?
  • Is partition count a small multiple of core count — enough for balancing, coarse enough to avoid overhead?
  • Have I measured input, intermediate (shuffle), and output bytes separately to find the dominant stage?
  • Is any single key carrying a disproportionate share of the shuffle bytes?

Checkpoint exercise

Try it
You are counting words over a 1 TB corpus on a 64-core, 16-node cluster (10 GbE). (a) Pick a partition count and justify it against the trade-table. (b) Estimate the shuffle volume with and without a combiner, and the resulting shuffle time at ~20 GB/s aggregate bandwidth. (c) For the reduce, compute the critical-path depth of a fan-in-2 tree versus a driver-side merge. Now flip one assumption: the corpus is dominated by a single token that appears in 40% of all documents. Which of your answers must change, and what do you do about the hot key? (The right answer should change once skew enters — the tree depth is fine, but one reducer is now a straggler, so you salt or hierarchically aggregate that key.)

Where this points next

You now have the full Core toolkit: tasks and refs (02), actors (03), the runtime and object store (04), and the map-shuffle-reduce pattern that composes them (05). Lesson 06 spends it on a real application — an RL maze built by hand — where the four roles (environment simulation, policy, rollout collection, and the learner) map onto exactly the fan-out/object-store/fan-in shape you just learned, and where you will watch the bottleneck migrate from compute to coordination as you add rollout workers. That hand-built application is the thing RLlib (07) will then hide behind an API.

Takeaway
Map → shuffle → reduce is the composite shape behind most distributed aggregation. Map creates parallelism cheaply by running independent tasks over partitions; the shuffle — the all-to-all redistribution that groups records by key — is the expensive, network-bound stage (O(P²) connections, bounded by cluster bandwidth) and is what dominates at scale; the reduce concentrates pressure. The two mistakes that cap a naive attempt are reducing on the driver (all bytes through one process, an O(N) serial merge — pure Amdahl ceiling) and skipping the combiner (oversized shuffle). Fix the first with a tree reduce that cuts critical-path depth from O(N) to O(log N) (1024 partials: ~10 levels vs 1023 sequential merges), valid because the merge is associative and commutative; fix the second with map-side pre-aggregation that can shrink shuffle volume by orders of magnitude. Tune partition count to a small multiple of core count. This exact shape reappears as Ray Data's blocks, Train's allreduce, and Tune's trial aggregation.

Interview prompts