Docs

Using Dista

Dista turns ordinary Mojo functions into tasks that run across a cluster. This page covers what a program needs: declaring tasks and actors, passing results around, and running the same binary locally and on Kubernetes. The design itself is in the repository's DESIGN.md.

Install

Dista isn't published as a package yet. A pixi project depends on it from source, at a pinned commit, since a git dependency without a revision never updates:

[dependencies]
dista-mojo = { git = "https://github.com/magmalake/dista.mojo", rev = "<commit>" }

It targets Mojo 1.2 and tracks the nightly toolchain. Reading Parquet and Iceberg tables uses magmalake's parquet-mojo and iceberg-mojo alongside it.

Tasks

A task is a top-level def taking one argument, declared once with its argument and result types, and registered in the runtime's parameter list:

from dista import Runtime, Task

def square(x: Int) -> Int:
    return x * x

def describe(x: Int) raises -> String:
    if x < 0:
        raise Error("negative")
    return String("x = ", x)

comptime Square = Task[Int, Int, square]
comptime Describe = Task[Int, String, describe]

def main() raises:
    var rt = Runtime[Square, Describe](workers=4)
    var r = rt.remote[Square](7)
    print(rt.get(r))  # 49

rt.remote returns at once with an ObjectRef; rt.get waits for the result, or raises the task's error. rt.ready(r) asks without waiting. Every process in a cluster runs the same binary, so tasks are named on the wire by a hash of the function, and a process built with different tasks is refused when it connects.

Results

A small result (up to 64 KB) comes back to the driver through the head. A larger one stays on the worker that computed it: in a shared-memory file on that host, which a reader on the same host maps without copying, or in the worker's memory, which a reader elsewhere fetches from that worker directly. rt.location(r) says where a result lives.

The head counts references to every result: the driver's own, until rt.release(r), and one for every pending task that takes it as input. At zero the result is deleted wherever it is. rt.stats() reports what the head holds.

Tasks on results

A task can take another task's result instead of a value. The head runs it once the input exists, and the worker running it fetches the input from wherever it lives. A list of refs is gathered, in order, into a Seq:

from dista import Seq

def total(xs: Seq[Int]) -> Int:
    var t = 0
    for i in range(len(xs)):
        t += xs[i]
    return t

comptime Total = Task[Seq[Int], Int, total]

var squares = List[ObjectRef[Int]]()
for i in range(100):
    squares.append(rt.remote[Square](i))
var sum = rt.remote[Total](squares)   # runs once all 100 exist
var next = rt.remote[Square](sum)     # a task on a task's result

If an input fails, every task that depends on it fails with that error.

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. Two kinds of task make that 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:

from dista import 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
for j in range(4):
    var mine = List[ObjectRef[Seq[Pair]]]()
    for m in range(8):
        mine.append(outputs[m][j])
    var r = 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. Each large partition stays on the worker that made it, and the reducer fetches it from there.

What crosses the wire

Arguments and results are encoded by type, at compile time. Scalars, Bool and String have fixed encodings. Any struct is encoded field by field. Arrow RecordBatches travel as Arrow IPC, so columnar data keeps its format between tasks. A type that needs its own encoding implements the Codec trait. A bare List doesn't encode: send Seq, a struct, or Arrow instead.

Actors

An actor is state that lives on one worker between calls: a type, the function that builds it, and methods that take it mut. Each actor gets a worker of its own, and its calls run one at a time, in the order they were made. A method's result is an ObjectRef like a task's, so a task can take it as input.

from dista import Actor, Method

@fieldwise_init
struct Counter(Movable):
    var total: Int

def new_counter(start: Int) -> Counter:
    return Counter(start)

def add(mut c: Counter, n: Int) -> Int:
    c.total += n
    return c.total

comptime CounterActor = Actor[Counter, Int, new_counter]
comptime Add = Method[CounterActor, Int, Int, add]

var rt = Runtime[CounterActor, Add, Square](workers=2)
var c = rt.create[CounterActor](0)
var n = rt.call[Add](c, 5)            # ObjectRef[Int]
print(rt.get(rt.remote[Square](n)))   # 25
rt.kill(c)                             # after its earlier calls; frees the worker

Calling a method on the wrong actor's handle is a compile error. A handle is a value, so a task or another actor can call the actor too: declare it with CtxTask or CtxMethod, and its function gets a Client to call through:

from dista import Client, CtxTask

def bump(mut c: Client, h: ActorHandle[CounterActor]) raises -> Int:
    return c.get(c.call[Add](h, 1))

comptime Bump = CtxTask[ActorHandle[CounterActor], Int, bump]
var n = rt.remote[Bump](c)   # a task calling the actor

When an actor's worker goes away, the head places the actor on another worker, builds it again from its last checkpoint, or its constructor if it has none, and runs every call since then again, in order. Each call is answered once. A call may run twice on the actor's state, so a method should be safe to repeat; that's the price of not losing calls.

Checkpoints

A checkpoint is the actor's state through the codec, written to DISTA_CHECKPOINT_URI: a directory every node sees, or s3://bucket/prefix. Ask for one every so many calls, or when it suits:

var c = rt.create[CounterActor](0, checkpoint_every=100)
...
var calls = rt.get(rt.checkpoint(c))   # written; covers this many calls

A restart then replays only the calls since the checkpoint, and the head stops keeping the ones it covers. A state type that implements Codec decides what it saves: a cache can save nothing, a loaded model its path. Checkpoints also carry an actor across a head restart: the new head restores it from its newest checkpoint, and its callers send again the calls no checkpoint includes. Killing an actor deletes its checkpoints.

Running it

Locally, Runtime(workers=n) forks a head and n workers and drives them from the same process, which is how tests and laptops run. On a cluster, the same binary takes its role from the environment: one head, any number of workers, and drivers that connect to the head.

The head doesn't run your code, so it can be the generic one: dista head, from the dista command (pixi run cli). It serves drivers and workers built from any program, several at once, and gives each driver's tasks and actors only to workers built from the same program. A head built into your program, DISTA_ROLE=head, serves only that program.

The repository's deploy/ folder builds one image for every role, with your program at /opt/dista/bin/app and the command at /opt/dista/bin/dista; the manifests run the head as dista head. tests/kind.sh runs a cluster on kind, and bench/cloud/ runs one on EC2 with k3s. Give the workers' /dev/shm room for the large results they keep. A deleted pod's process exits on its SIGTERM.

Runtime(pipeline=n) sets how many tasks the head keeps in flight per worker (default 2).

Watching a cluster

dista status prints the head's view: workers, with the program each was built from and what it's running or hosting; drivers; and actors, with their worker, calls and last checkpoint. dista timeline -o trace.json writes every driver's finished tasks, and dista stop shuts the head and its workers down. Each takes --head host:port, or reads DISTA_HEAD.

With DISTA_METRICS_PORT set, the head and every worker serve Prometheus metrics at /metrics: tasks submitted and finished by outcome, task durations as a histogram, queue depth, objects and bytes held, workers, drivers, actors, restarts and checkpoints. The kind manifests turn it on at port 9090, for Kubernetes monitoring to scrape.

The timeline

The head stamps each task's submission, dispatch and completion, and the worker its start and end. rt.timeline(path) writes every finished task, in Chrome's trace format, which Perfetto and the browser's trace viewer open. Each worker is a track, with a second track for time spent waiting for inputs. The driver can add its own work, so a gap between tasks says what it was:

var t = rt.clock()
var tasks = scan.plan_files()
rt.span("plan", t)          # from t to now, on the driver's track
...
_ = rt.timeline("trace.json")

When things fail

A task that raises fails its ObjectRef with the error, and the tasks that depend on it. A worker that dies has the tasks it was running sent to another worker, up to DISTA_TASK_RETRIES times (default 3). The results it held are lost: Dista doesn't rebuild them by re-running the tasks that made them, so a task that needs one fails, naming it, and the job is run again, by Orca or by hand. In local mode, where nothing replaces a dead worker, tasks fail with no workers left once every worker is gone, rather than waiting. A head that dies is rebuilt from its peers when it comes back: workers rejoin and report what they hold, and drivers reconnect under their old number and resubmit what they still need. A task running at that moment may run twice, so tasks should be safe to repeat, as a pure function is. Actors come back too, from their checkpoints (see above). A chaos test kills the head at random throughout a run and checks every answer, and on kind the actor's pod is deleted mid-run and its count must come out exact.