magmalake Dista

Dista · API

dista.task

Tasks: remote functions, their IDs, and how a worker runs one.

A task wraps a top-level def with its argument and result types:

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

comptime Square = Task[Int, Int, square]

The types are spelled out because Mojo can't infer a function's result type from the function value. The function may raise or not; both bind to def(A) thin raises -> R.

Every node runs the same binary, so a task message names the task by ID: a 64-bit FNV-1a hash of the function's linkage name, which includes its module and argument types. The hash is fixed, unlike hash(), so a driver and a worker started separately agree on it.

An actor is a state type and the function that builds it, and each of its methods is a function taking the state mut:

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]

Actors and methods go in the runtime's registry with the tasks. The worker hosting an actor keeps its state behind a StatePtr, which only the actor's own create, destroy, save and restore and its methods' call cast back to its type. A checkpoint is the state through the codec, so a state type that implements Codec chooses what it saves.

Structs

Task

struct Task[A_: Deinitable & Movable, R_: Deinitable & Movable, f: def(A_) raises thin -> R_]

A remote function taking one A_ and returning an R_.

Implements Movable, TaskLike

comptime A = A_
comptime R = R_

name

def name() -> StaticString

id

def id() -> UInt64

run

def run(args: Span[UInt8], mut out: List[UInt8])

Task2

struct Task2[A_: Deinitable & Movable, B_: Deinitable & Movable, R_: Deinitable & Movable, f: def(A_, B_) raises thin -> R_]

A remote function taking two arguments: rt.remote[T](a, b).

One type per argument count, as Mojo's own library does it until its variadic type parameters settle, when Task can take any number (https://github.com/magmalake/dista.mojo/issues/71). On the wire the arguments are encoded one after the other.

Implements Movable, Task2Like, TaskLike

comptime A = A_
comptime B = B_
comptime R = R_

name

def name() -> StaticString

id

def id() -> UInt64

arity

def arity() -> Int

run

def run(args: Span[UInt8], mut out: List[UInt8])

Task3

struct Task3[A_: Deinitable & Movable, B_: Deinitable & Movable, C_: Deinitable & Movable, R_: Deinitable & Movable, f: def(A_, B_, C_) raises thin -> R_]

A remote function taking three arguments: rt.remote[T](a, b, c). See Task2.

Implements Movable, Task2Like, Task3Like, TaskLike

comptime A = A_
comptime B = B_
comptime C = C_
comptime R = R_

name

def name() -> StaticString

id

def id() -> UInt64

arity

def arity() -> Int

run

def run(args: Span[UInt8], mut out: List[UInt8])

TaskWith

struct TaskWith[A_: Deinitable & Movable, X: Copyable & Deinitable, R_: Deinitable & Movable, f: def(A_, Seq[X]) raises thin -> R_]

A task taking a value and other tasks' results, gathered in order: f(a, xs), submitted as rt.remote[T](a, refs). A shuffle's reducer is one: its partition number, and the map outputs holding that partition.

Implements Movable, TaskLike

comptime A = A_
comptime R = R_

name

def name() -> StaticString

id

def id() -> UInt64

run

def run(args: Span[UInt8], mut out: List[UInt8])

TaskMany

struct TaskMany[A_: Deinitable & Movable, X: Copyable & Deinitable, f: def(A_) raises thin -> Seq[X]]

A task with several results, each its own object: rt.remote_many[T](a, n) returns n ObjectRef[X]s, and f must return a Seq of exactly n. Each is stored, fetched and freed on its own, so a shuffle's map task can make one partition per reducer.

On the wire its result is a count, then each part as a length and its encoding, so the worker splits it without knowing X.

Implements ManyLike, Movable, TaskLike

comptime A = A_
comptime R = Seq[X]
comptime Part = X

name

def name() -> StaticString

id

def id() -> UInt64

is_many

def is_many() -> Bool

run

def run(args: Span[UInt8], mut out: List[UInt8])

Actor

struct Actor[S_: Deinitable & Movable, A_: Deinitable & Movable, new: def(A_) raises thin -> S_]

An actor whose state is an S_, built by new from an A_.

Implements ActorLike, Movable, TaskLike

comptime S = S_
comptime A = A_
comptime R = NoneType

name

def name() -> StaticString

id

def id() -> UInt64

is_actor

def is_actor() -> Bool

create

def create(args: Span[UInt8]) -> StatePtr

destroy

def destroy(state: Pointer[NoneType, MutUntrackedOrigin])

save

def save(state: Pointer[NoneType, MutUntrackedOrigin], mut out: List[UInt8])

restore

def restore(bytes: Span[UInt8]) -> StatePtr

Method

struct Method[Act_: ActorLike, A_: Deinitable & Movable, R_: Deinitable & Movable, f: def(mut Act_.S, A_) raises thin -> R_]

A method of actor Act_: f on its state and an A_, returning R_.

Implements MethodLike, Movable, TaskLike

comptime A = A_
comptime R = R_
comptime Act = Act_

name

def name() -> StaticString

id

def id() -> UInt64

owner

def owner() -> UInt64

call

def call(state: Pointer[NoneType, MutUntrackedOrigin], args: Span[UInt8], mut out: List[UInt8])

Traits

TaskLike

What the runtime needs from a task, an actor or a method, without knowing its function. Each implements the part that applies to it.

comptime A

The argument type.

comptime R

The result type.

name

def name() -> StaticString

The function's linkage name.

id

def id() -> UInt64

The task's wire ID.

run

def run(args: Span[UInt8], mut out: List[UInt8])

Decodes the argument, calls the function, appends the encoded result.

arity

def arity() -> Int

How many arguments it takes: 2 for a Task2, 3 for a Task3.

is_actor

def is_actor() -> Bool

is_many

def is_many() -> Bool

Whether its result is several objects (TaskMany).

owner

def owner() -> UInt64

For a method, its actor's ID; 0 for anything else.

create

def create(args: Span[UInt8]) -> StatePtr

An actor: decodes the argument and builds the state.

destroy

def destroy(state: Pointer[NoneType, MutUntrackedOrigin])

An actor: destroys the state create built.

save

def save(state: Pointer[NoneType, MutUntrackedOrigin], mut out: List[UInt8])

An actor: appends its state's encoding, a checkpoint.

restore

def restore(bytes: Span[UInt8]) -> StatePtr

An actor: builds the state again from a checkpoint save wrote.

call

def call(state: Pointer[NoneType, MutUntrackedOrigin], args: Span[UInt8], mut out: List[UInt8])

A method: runs on the actor's state, appends the encoded result.

Task2Like

Implements TaskLike

comptime B

The second argument's type; A is the first's.

Task3Like

Implements Task2Like, TaskLike

comptime C

The third argument's type.

ManyLike

Implements TaskLike

comptime Part

The type of each of its results.

ActorLike

Implements TaskLike

comptime S

The actor's state type.

MethodLike

Implements TaskLike

comptime Act

The actor whose method this is.

Functions

fnv1a64

def fnv1a64(s: StringSpan) -> UInt64

The 64-bit FNV-1a hash of s's bytes.

split_parts

def split_parts(result: Span[UInt8]) -> List[List[UInt8]]

A TaskMany result's parts, each as its own encoded object.

Aliases

comptime StatePtr = Pointer[NoneType, MutUntrackedOrigin]

An actor's state on its worker, with its type erased.