Docs

Shuffles

A shuffle moves data between two sets of tasks: map tasks each split their output by key, and each reduce task takes its share from every map task. The partitions stay on the workers that made them, and the reducers fetch them from there.

Map and reduce

Two kinds of task make a shuffle work. A TaskMany returns several results, each an object of its own, stored, fetched and freed separately, like Ray's num_returns. A TaskWith takes a value beside a list of other tasks' results, gathered in order into a Seq:

from dista import ObjectRef, Seq, TaskMany, TaskWith

def map_rows(m: Int) -> Seq[Seq[Pair]]:          # one partition per reducer
    ...
def reduce_part(j: Int, parts: Seq[Seq[Pair]]) -> Seq[Pair]:
    ...

comptime MapRows = TaskMany[Int, Seq[Pair], map_rows]
comptime ReducePart = TaskWith[Int, Seq[Pair], Seq[Pair], reduce_part]

var outputs = List[List[ObjectRef[Seq[Pair]]]]()
for m in range(8):
    outputs.append(rt.remote_many[MapRows](m, 4))   # 4 results each
var reduced = List[ObjectRef[Seq[Pair]]]()
for j in range(4):
    var mine = List[ObjectRef[Seq[Pair]]]()
    for m in range(8):
        mine.append(outputs[m][j])
    reduced.append(rt.remote[ReducePart](j, mine))  # partition j of every map

A TaskMany must return exactly as many values as asked for, or every one of its results fails. The repository's tests/test_shuffle.mojo runs a group-by this way, 8 map tasks for 4 reducers, and checks it against one process.

Partitioning again

A TaskMany can wait on other tasks' results too: rt.remote_many[T](ref, n) on one result, or rt.remote_many[T](refs, n) when it takes a Seq. A later stage can then partition again what earlier tasks made, without the driver fetching any of it.

Placement by data location

The head learns each worker's node from the results it stores: a node is a host name, or DISTA_NODE where it's set. When one node holds more than half the bytes of a task's stored inputs, the task prefers it, and a free worker on that node takes it first. A task waits only while a worker on its own node is free to take it, so a busy node's tasks still go elsewhere. A step after a map saves the whole fetch; an all-to-all shuffle gains little, since each reducer's inputs are spread evenly.

Spilling

A shuffle's intermediate data can be larger than memory. A worker keeps up to DISTA_STORE_BYTES of results in its node store (no limit when unset) and writes the rest to DISTA_SPILL_DIR on disk, by default $TMPDIR, else /tmp. A store write that fails spills too. Readers don't change: one on the same host maps the file, one elsewhere fetches it. Tuning covers the store's settings.

The shuffle benchmark, on Dista and on Ray, is in bench/shuffle, and JOB=shuffle bench/cloud/bench matrix runs it on AWS.