Docs

Tasks and results

A task is a Mojo function that runs on a worker. Calling it returns a reference to a result that may not exist yet, which you read, pass to other tasks, or give up on.

Tasks

A task is a top-level def, declared once with its argument and result types and registered in the runtime's parameter list. Task takes one argument, Task2 two and Task3 three; more travel as a struct. A function that raises declares the same way:

from dista import Runtime, Task, Task2

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

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

def add(a: Int, b: Int) -> Int:
    return a + b

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

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

rt.remote returns at once with an ObjectRef. rt.take and rt.get wait for the result, or raise the task's error; rt.ready(r) asks without waiting. Every process in a cluster runs the same binary, so a task is 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, 64 KiB or less, 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 worker's ip:port, then its file if it has one, or nothing if it came back with the reply.

A result larger than 64 MiB crosses hosts in parts, so results past 1 GiB fetch too. A task's argument has to fit one frame: past 1 GiB, remote raises. Large data belongs in results, which stay on the workers.

Freeing results

A result is kept in the driver, at the head and on its worker until the driver releases its ref. rt.take(r) is get then release, the way to read a result used once, which is most of them; rt.take(refs) takes a list in order. rt.get(r) keeps the result, for reading again, until rt.release(r). A loop that only calls get grows all three by every result.

The head counts references to every result: the driver's own, until it's released, and one for every pending task that takes it as input. At zero the result is deleted wherever it is, so releasing a ref that a submitted task still needs is fine. rt.stats() reports what the head holds, and dista status shows how many objects each driver 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 ObjectRef, Runtime, Seq, Task

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 maps = List[ObjectRef[Int]]()
for i in range(1000):
    maps.append(rt.remote[Square](i))
var sum = rt.remote[Total](maps)   # runs once all 1,000 exist
rt.release(maps)                   # only Total needs them, and it still gets them
var next = rt.remote[Square](sum)  # a task on a task's result
print(rt.take(next))

If an input fails, every task that depends on it fails with that error, and never runs. A task given a ref takes it whole: a ref nested inside a struct is a compile error. The ways to pass refs are one ref, a list gathered into a Seq, and a value beside a list, which shuffles use.

Cancelling and timeouts

rt.cancel(r) gives up on a result. A task that hasn't run never does; one that's running finishes on its worker, which can't be interrupted, and its result is dropped. get then raises cancelled, and so does every task waiting on r. A result already in the driver is kept. An actor call is answered cancelled but still runs, since the calls after it expect its effect. rt.cancel(refs) cancels a list; a cancelled ref is released like any other.

var r = rt.remote[Slow](x)
try:
    print(rt.get(r, timeout=30.0))
except e:
    # "timed out after 30.0 s waiting for ...": r is still valid
    rt.cancel(r)
    rt.release(r)

rt.get(r, timeout=s) raises an error starting with timed out once s seconds pass, leaving the ref valid, to wait on again, cancel or release. get(r) without a timeout waits as long as it takes. A time limit on the task itself, however its worker stalls, is DISTA_TASK_TIMEOUT.

What crosses the wire

Arguments and results are encoded by type, at compile time. Scalars, Bool and String have fixed encodings, and any other struct is encoded field by field. Arrow RecordBatches from arrow-mlake travel as Arrow IPC, so columnar data keeps its format between tasks. A type that needs its own encoding implements the Codec trait, an encode(self, mut out) and a static decode(buf, mut pos). A bare List doesn't encode: send a Seq, a struct, or Arrow instead. Both ends are the same binary, so there's no schema evolution to manage.