Dista

Dista

Mojo functions, across a cluster

Dista is for developers who have written fast Mojo and now need more than one machine. A task is an ordinary Mojo function. The program that calls it runs unchanged on a laptop, where Dista starts its own workers, and on a Kubernetes cluster, where every node runs the same small image. There is no Python in the path and no second language to learn. 

from dista import Runtime, Task

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

comptime Square = Task[Int, Int, square]

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

Results that are large stay on the worker that computed them. A task on the same host maps them from shared memory, and one elsewhere fetches them from that worker directly, without going through the head. A task can take other tasks' results as its input, so a pipeline is a graph of calls rather than a sequence of round trips to the driver.

Performance so far

12 s a pass, against 28.7 s for Rust and 38.6 s for Ray

The same eight queries over 1.4 billion taxi trips in an Iceberg table in S3, on three worker nodes, with the answers checked against each other. A day earlier Dista took 56.7 s on the same run; rewriting how it reads from object storage made it the fastest of the three. The performance page has every query, the before-and-after timeline, the earlier runs where Dista lost, and how to reproduce each number.

Median pass of the eight queries, 3 × AWS c7g.4xlarge (16 Graviton3 vCPUs, 1 November 2026. Lower is better. One run on one day: the comparison frameworks varied about 10% between runs.

Status

Early, and measured as it goes

Dista has passed its fourth phase's gate: the taxi queries on a Kubernetes cluster stay exact with a worker killed mid-run, and a worker stopped the way Kubernetes stops one drains without anything running twice. Before that came tasks, results, a head that restarts, and actors that survive their worker or the head restarting. What it doesn't do yet is listed here rather than left to be found out.

Works today

  • Tasks: a top-level Mojo function, declared once, called with rt.remote; Task2 and Task3 for two or three arguments
  • Results that stay on the worker that computed them, mapped from shared memory on the same host or fetched from that worker directly
  • Tasks on other tasks' results: one ObjectRef, or a list gathered into a Seq
  • Placement by data location: a task goes to the node holding most of its input, when one does
  • Tasks whose worker dies run again on another, up to DISTA_TASK_RETRIES times
  • Draining: a worker that Kubernetes stops finishes its tasks, and its results and its actor move to other workers, so nothing fails or runs twice
  • Spilling: past DISTA_STORE_BYTES, a worker's results go to disk, and readers find them there
  • Backpressure: with too much queued, rt.remote waits for the workers to catch up
  • Arrow record batches as task arguments and results
  • Freeing: reference counts at the head, rt.release
  • A head that restarts: workers rejoin with what they hold, drivers resume where they were
  • Actors: stateful structs on a worker of their own, calls in order, handles passed to tasks and other actors
  • Actors that survive their worker or the head restarting, restored from checkpoints in a directory or S3
  • Shuffles: several results from one task, and a task taking a value beside other tasks' results; a shuffle benchmark against Ray
  • A task timeline in Chrome's trace format, the driver's own work included
  • Prometheus metrics on the head and every worker
  • The dista command: one head for several programs; dista status, timeline and stop
  • One container image for every role, 32 MB compressed, on Kubernetes
  • A published package: pixi shelf add dista-mojo

Not yet

  • Refs nested inside a struct or list in a task's argument, which is a compile error for now (deferred, issue 68)
  • Rebuilding results lost with a worker, as Ray's lineage does: the job fails, naming them, and orca runs it again (deferred, issue 69)
  • A running job asking for more workers: orca sizes each run's cluster (deferred, issue 76)

A longer write-up, with the design and the measurements behind it, follows with the second benchmark, which spends its time in user code and a shuffle.