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.
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.
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.
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.
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.
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.
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.
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 partitions | Too 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.
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.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.getall 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.putshared 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
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.
Interview prompts
- Walk through the three MapReduce stages and which one dominates cost at scale. (§1, §5 — map creates parallelism cheaply over partitions; shuffle is the all-to-all group-by-key, network-bound and O(P²), and dominates; reduce concentrates pressure per key.)
- Why is
ray.get(all_partials)then reducing on the driver an anti-pattern? (§3 — every byte funnels through one process/network link and the merge is an O(N) serial loop on one core — the inherently-serial section Amdahl caps; adding workers doesn't help.) - What does a tree reduce buy you, and what property of the operation does it require? (§4 — cuts critical-path depth from O(N) to O(log N) and removes the driver as a funnel; requires an associative, commutative reduce so order/grouping can be rearranged. 1024 partials → ~10 levels vs 1023 merges.)
- What is a combiner and why is it the highest-leverage optimization? (§5 — map-side pre-aggregation that emits partial aggregates instead of raw records, cutting shuffle volume (the dominant cost) by orders of magnitude; valid only for associative/commutative ops.)
- Why does shuffle cost scale with O(P²) and bandwidth rather than with compute? (§5 — every mapper sends a slice to every reducer, an all-to-all of P×P logical transfers, and the bytes cross the network, the scarcest cluster resource.)
- How do you choose partition count? (§5 — small multiple (2–4×) of core count: enough for balancing/retries and to avoid stragglers, coarse enough that each task does ≥ tens of ms and P² shuffle doesn't explode.)
- Where does this exact pattern reappear later in Ray? (§6 — Ray Data (blocks/groupby), Train (allreduce is a tree reduce of gradients), Tune (parallel trials reduced into a scheduling decision).)