magmalake Dista

Docs

Using Dista

Dista turns ordinary Mojo functions into tasks that run across a cluster. One binary runs as the driver, the head and every worker, on a laptop or on Kubernetes. This page installs it and runs a first program; the pages below cover the rest. How it works is in the repository's DESIGN.md.

The docs

  • Tasks and results: Declaring tasks, where results live and how they're freed, tasks on other tasks' results, cancelling and timeouts, and what crosses the wire.
  • Shuffles: Tasks with several results, tasks taking a value beside results, placement by data location, and spilling.
  • Actors and checkpoints: State on a worker of its own, handles passed to tasks, restarts, checkpoints, and actors moving off a draining worker.
  • Running a cluster: Roles from the environment, the generic head, Kubernetes and EC2, draining a worker, and autoscaling.
  • The dista command: dista head, status, drain, scale, timeline and stop; Prometheus metrics; the task timeline.
  • Failure and recovery: What happens when a task raises, a worker dies or goes silent, a task hangs, a driver leaves, or the head restarts.
  • Tuning: Tasks in flight per worker, backpressure, the node store and spilling, memory held for drivers, and large results.

Install it

Dista is the mojoshelf tin dista-mojo, version 0.2.2. With pixi and mojoshelf's pixi shelf installed, a new project adds it, and the tins it needs (arrow-mlake, memory-region and objectstore) come with it: 

pixi init -c https://conda.modular.com/max-nightly -c conda-forge app && cd app
sed -i.bak 's/^\[workspace\]/[workspace]\npreview = ["pixi-build"]/' pixi.toml   # tins build with pixi-build
pixi add "mojo-compiler==1.2.0"
pixi shelf add dista-mojo

It targets Mojo 1.2.0, and its CI also runs on the 1.3 nightly. Reading Parquet and Iceberg tables in tasks uses magmalake's parquet-mojo and iceberg-mojo beside it.

A first program

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

from dista import ObjectRef, 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 refs = List[ObjectRef[Int]]()
    for i in range(1000):
        refs.append(rt.remote[Square](i))
    var total = 0
    for r in refs:
        total += rt.take(r)  # get and release: read once, then freed
    print(total)  # 332833500

Build it with mojo build and run it. With nothing set in its environment, it forks a head and four workers on this machine and drives them. Point the same binary at a head and it joins a cluster instead: 

DISTA_ROLE=head DISTA_PORT=6390 ./app                  # the head
DISTA_ROLE=worker DISTA_HEAD=headhost:6390 ./app       # once per worker
DISTA_HEAD=headhost:6390 ./app                         # the driver

Tasks and results says what a task can take and return, and Running a cluster how the roles fit together.

From source

The repository builds and tests with pixi. pixi run cli builds the dista command to build/dista. 

git clone https://github.com/magmalake/dista.mojo && cd dista.mojo
pixi run test               # codec, frame, runtime and cluster tests, Mojo 1.2.0
pixi run -e nightly test    # the same on the 1.3 nightly
pixi run -e data test-data  # tasks over Parquet and Iceberg; needs those tins checked out next door
pixi run bench              # task throughput
pixi run test-kind          # the smoke test on a kind cluster (needs docker, kind, kubectl)

What each release changed is in CHANGELOG.md, and what Dista is for, with its phases, in VISION.md.