magmalake/dista

dista (el tin dista-mojo, 0.2.0, Mojo 1.2.0) te permite marcar una función Mojo como remota, llamarla y hacer que se ejecute en otros procesos o pods de Kubernetes. Cada proceso ejecuta el mismo binario compilado, por lo que un mensaje de tarea es un ID de función más argumentos codificados: no se envía código, no se hace pickle, y los datos en columnas viajan como Arrow.

Lo que dista no es: un administrador de clústeres. No crea pods ni máquinas. Utiliza cada trabajador que se une, drena aquellos a los que se les indica detenerse y dice cuántos quiere. Kubernetes, un bucle de shell u orca los inician y detienen. Tampoco reconstruye resultados perdidos con un trabajador caído (sin reconstrucción de linaje): el job falla, nombrando el resultado, y un orquestador lo vuelve a ejecutar.

El entorno instalado contiene dista solo como un dista.mojoc compilado. Lee su fuente en GitHub, magmalake/dista.mojo, en la rev que fija tu pixi.toml, o en la caché git de pixi (ls -d ~/Library/Caches/rattler/cache/git-v0/checkouts/*/<primeros 7 de la rev>* en macOS, ~/.cache/rattler/... en Linux). Las rutas de archivo citadas aquí están en ese checkout.

pixi shelf add dista-mojo    # trae arrow-mlake, memory-region y objectstore

Escribe el Mojo en sí con la habilidad mojo-syntax cargada; para trabajo rápido fila por fila dentro de las tareas, consulta mojo-data-performance.

Declara tareas, luego llámalas

Una tarea es un def de nivel superior declarado una vez con sus tipos de argumento y resultado, porque Mojo 1.2 no puede inferir el tipo de resultado de una función. El registro es la lista de parámetros del Runtime. Cada muestra completa en este archivo fue construida y ejecutada contra dista 0.2.0.

from dista import ObjectRef, Runtime, Seq, Task, Task2

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

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

def scale(x: Int, factor: Float64) -> Float64:
    return Float64(x) * factor

comptime Square = Task[Int, Int, square]
comptime Total = Task[Seq[Int], Int, total]
comptime Scale = Task2[Int, Float64, Float64, scale]

def main() raises:
    var rt = Runtime[Square, Total, Scale](workers=4)
    var r = rt.remote[Square](7)              # un ObjectRef[Int], de inmediato
    print(rt.get(r))                          # 49: espera el resultado
    print(rt.get(rt.remote[Scale](5, 2.5)))   # 12.5
    var maps = List[ObjectRef[Int]]()
    for i in range(1000):
        maps.append(rt.remote[Square](i))
    print(rt.get(rt.remote[Total](maps)))     # 332833500
    print(rt.get(rt.remote[Square](r)))       # 2401: una ref como argumento
    rt.release(maps)

• Aridad: Task[A, R, f] toma un argumento, Task2[A, B, R, f] dos, Task3[A, B, C, R, f] tres. Pasa más como una sola estructura. Llamar a un Task2 con un argumento lanza takes 2 arguments en la llamada.
• La función puede lanzar o no. Una tarea que lanza vuelve a lanzar su mensaje en rt.get, y el trabajador continúa.
• Las referencias alimentan tareas sin que el controlador obtenga nada. Pasa un ObjectRef[A] donde la tarea recibe una A, o una List[ObjectRef[X]] donde recibe una Seq[X] (reunidas en orden). La tarea se ejecuta una vez que cada entrada existe, y falla con el error de la entrada si alguna falló.
• rt.release(ref) (o una lista) libera un resultado en la cabeza y en su trabajador una vez que ninguna tarea pendiente lo necesita; get sobre él luego lanza. Los resultados se conservan hasta que el runtime termina.
• También en Runtime: rt.ready(r) (¿ya terminó?, sin esperar), rt.location(r) (el ip:puerto del trabajador que lo tiene y el archivo del almacén de nodos, o "" si el resultado regresó en línea), rt.stats(), rt.bytes_fetched() y rt.timeline(path), que escribe un trace de Chrome para Perfetto.

Shuffles: varios resultados, o un valor junto a referencias

Un TaskMany[A, X, f], donde f devuelve una Seq[X], es una tarea de mapa: rt.remote_many[T](arg, n) devuelve n referencias, una por partición, cada una almacenada, obtenida y liberada por separado; f debe devolver exactamente n elementos. Un TaskWith[A, X, R, f], donde f toma una A y una Seq[X], es un reductor: rt.remote[T](j, refs). Una muestra que compila, y los detalles, están en references/shuffles.md.

Qué puede ser un argumento o un resultado

El códec elige una codificación en tiempo de compilación:

• Int, IntN, UIntN, Float32, Float64, Bool y String;
• Seq[T], una lista que pasa por el códec (List no: la reflexión no puede recuperar su tipo de elemento), también como campo de estructura;
• cualquier estructura simple de lo anterior, campo por campo mediante reflexión (los campos deben ser Movable; los tipos de argumento y resultado necesitan Movable & Deinitable);
• un tipo que implementa el trait Codec (su propio encode y decode);
• RecordBatch de arrow-mlake, enviado como flujo Arrow IPC. Úsalo para cualquier cosa grande; un lote por grupo de filas de Parquet es el grano previsto.

Un ObjectRef dentro de una estructura o un Seq es un error de compilación que nombra las tres formas que funcionan: la ref como argumento, una lista de refs, o un valor junto a refs con TaskWith. Los resultados de 64 KiB o menos regresan a través de la cabeza; los más grandes permanecen en el trabajador que los creó y se mapean desde memoria compartida en el mismo host o se obtienen entre hosts.

Actores

Un actor es estado en un trabajador propio. Sus llamadas se ejecutan una a la vez, en el orden en que se hicieron, y devuelven ObjectRefs como las tareas.

from dista import Actor, ActorHandle, Client, CtxTask, Method, ObjectRef, Runtime

@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

def bump(mut c: Client, h: ActorHandle[CounterActor]) raises -> Int:
    var r = c.call[Add](h, 1)    # una tarea que llama al actor que se le pasó
    var n = c.get(r)
    c.release(r)
    return n

comptime CounterActor = Actor[Counter, Int, new_counter]
comptime Add = Method[CounterActor, Int, Int, add]
comptime Bump = CtxTask[ActorHandle[CounterActor], Int, bump]

def main() raises:
    var rt = Runtime[CounterActor, Add, Bump](workers=3)
    var c = rt.create[CounterActor](0, checkpoint_every=100)
    var refs = List[ObjectRef[Int]]()
    for _ in range(10):
        refs.append(rt.call[Add](c, 5))
    print(rt.get(refs[9]))           # 50
    print(rt.get(rt.remote[Bump](c)))  # 51
    rt.kill(c)

• Un actor ocupa un trabajador entero, y las tareas nunca se ejecutan en él. Un actor creado cuando todos los trabajadores ya albergan uno espera a que se una un trabajador, así que en modo local da más trabajadores que actores.
• Los manejadores son valores: pasa un ActorHandle a un CtxTask (def f(mut c: Client, a: A) raises -> R) o guárdalo en el estado de otro actor, y llama a través del Client. Un CtxMethod (def f(mut s: S, mut c: Client, a: A) raises -> R) permite que un actor llame a otros; un ciclo de llamadas de actores causa un interbloqueo.
• Puntos de control: checkpoint_every=N en create, o rt.checkpoint(c), cuyo ObjectRef[Int] es el número de llamadas cubiertas. Pasan por el códec a DISTA_CHECKPOINT_URI, un directorio o s3://; sin él rt.checkpoint falla. Un tipo de estado que implementa Codec elige qué guardar. Cuando el trabajador de un actor muere, regresa en otro trabajador desde su último punto de control (o su constructor) y las llamadas posteriores se ejecutan de nuevo, por lo que los métodos deben ser seguros de repetir.

Modo local y modo clúster

Runtime(...) lee el entorno para elegir el rol del proceso:

Entorno	Rol
ninguna variable	modo local: bifurca una cabeza y workers trabajadores, luego los dirige
DISTA_ROLE=head	ejecuta la cabeza en DISTA_PORT (por defecto 6390; 0 elige uno y lo imprime), nunca retorna
DISTA_ROLE=worker, DISTA_HEAD=host:port	sirve a esa cabeza, nunca retorna
DISTA_HEAD=host:port	un controlador que envía a esa cabeza

Así que crea el Runtime primero en main: en una cabeza o trabajador el constructor nunca retorna, y cualquier cosa antes de él se ejecuta en todos los roles. En modo local, los trabajadores se bifurcan dentro del constructor y heredan el entorno establecido antes de él. Para ejecutar un clúster a mano, como lo hace tests/cluster.sh:

DISTA_ROLE=head DISTA_PORT=6390 ./app               # la cabeza
DISTA_ROLE=worker DISTA_HEAD=headhost:6390 ./app    # una vez por trabajador
DISTA_HEAD=headhost:6390 ./app                      # el controlador

Los trabajadores y controladores reintentan alcanzar la cabeza durante DISTA_CONNECT_TIMEOUT segundos (por defecto 60), por lo que el orden de inicio no importa. Una cabeza iniciada con dista head en su lugar es genérica: sirve varios programas a la vez y da las tareas de cada controlador solo a trabajadores construidos a partir del mismo binario.

El comando dista

Construido desde el repositorio dista con pixi run cli (en build/dista). Cada comando encuentra su cabeza mediante --head, o si no DISTA_HEAD, o si no localhost:6390.

Comando	Qué hace
dista head [--port 6390] [--pipeline 2]	una cabeza genérica para cualquier programa
dista status	trabajadores con su hash de registro, controladores, actores
dista status --json	la demanda, en una línea, para un autoescalador
dista drain <ip:port>	drena ese trabajador, como lo hace SIGTERM; rechaza el último mientras quede trabajo
dista scale --once	imprime la decisión de la política para la demanda actual
dista timeline -o trace.json	las tareas terminadas de cada controlador, como un trace de Chrome
dista stop	apaga la cabeza y sus trabajadores

Con DISTA_METRICS_PORT configurado, la cabeza y cada trabajador sirven métricas de Prometheus en /metrics. Cada variable, nombre de métrica y los manifiestos de Kubernetes están en references/operations.md.

Drenaje y autoescalado

Un trabajador que recibe SIGTERM (cómo Kubernetes detiene un pod) drena en lugar de morir: no acepta trabajo nuevo, termina el que tiene, sus resultados se mueven a otros trabajadores y su actor se mueve mediante un punto de control, luego sale con 0. Nada falla ni se ejecuta dos veces (tests/drain.sh). DISTA_DRAIN_SECONDS (por defecto 25) lo acota; un segundo SIGTERM lo termina de inmediato.

dista status --json da tasks_ready, tasks_waiting, tasks_running, actors_waiting, pipeline, checkpoints, y por trabajador address, node, running, idle_seconds, held_bytes, actor y draining. dista.scale (impórtalo desde dista.scale; no se reexporta desde dista) lo convierte en una decisión, sin reloj ni E/S propios:

from dista.scale import Demand, Policy, decide

def main() raises:
    var d = Demand.parse('{"tasks_ready":100,"pipeline":2,"workers":[]}')
    var got = decide(d, Policy(min_workers=0, max_workers=6), now=10)
    print(got.want, got.start, len(got.drain))  # 6 4 0

Campos y valores por defecto de Policy: min_workers=1 (0 está permitido: el último trabajador drena una vez que no queda ninguna tarea y no contiene nada), max_workers=4, tasks_per_worker=0 (la profundidad de pipeline de la cabeza), up_after=5 segundos de tareas listas sin ranura libre, idle_before_down=60 segundos, y max_step=4 trabajadores iniciados por decisión. decide(demand, policy, now, pending=0, short_since=0) devuelve start, drain (direcciones, las que tienen menos bytes retenidos primero), want y short_since, que el proveedor pasa de vuelta la próxima vez con pending, los trabajadores iniciados pero que aún no se han unido. dista scale --once --min 0 --max 6 ... lo imprime como líneas; tests/autoscale.sh dirige un job desde 0 trabajadores hasta 6 y de vuelta a 0 con él.

Bajo orca

Un pipeline de orca declara sus tareas y trabajadores, y orca inicia el runtime una vez por ejecución. Un paso nunca construye su propio Runtime; toma prestado el de la ejecución (un extracto):

comptime Trips = Task[Day, Int, trips]
comptime FanoutTasks = Tasks[Trips]                   # desde orca

def month_trips(mut ctx: Ctx, m: Month) raises -> Int:
    ref rt = ctx.runtime[FanoutTasks]()
    var r = rt.remote[Trips](Day(m.copy(), 1))   # luego rt.get, como siempre
    ...

comptime Fanout = Pipeline["fanout", MonthTrips, AllTrips, tasks=FanoutTasks, workers=3]

El programa completo es examples/fanout.mojo en orca.mojo.

orca ejecuta los pasos de una ejecución en orden y las particiones de un paso mapeado una tras otra, por lo que un paso que llama a una tarea y espera por ella es serial. Envía todo el trabajo de un paso con rt.remote primero, luego rt.get cada referencia.

Errores comunes

• registry mismatch: this peer was built with different tasks. El controlador y los trabajadores provienen de compilaciones diferentes, o enumeran tareas distintas en Runtime[...]. El hash del registro cubre los nombres de las tareas en orden, así que recompila una vez y ejecuta ese único binario en cada rol.
• ... is not registered with this runtime. Cada tarea, actor y método que llames debe estar en Runtime[...]. Las llamadas a través del Client de un CtxTask no se verifican en la llamada, por lo que una no registrada falla en el trabajador con unknown task ID o unknown method ID.
• Una función anidada o clausura como tarea. Usa un def de nivel superior.
• Suponer que una tarea se ejecuta una vez. Las tareas de un trabajador muerto se ejecutan de nuevo, hasta DISTA_TASK_RETRIES (por defecto 3): mantén las tareas y métodos seguros de repetir.
• Esperar que un resultado perdido regrese. Un resultado retenido solo por un trabajador muerto falla con result <id> lost: its worker went away, al igual que cada tarea que espera por él. Vuelve a ejecutar el job, o deja que orca reintente el paso.
• Referencias no liberadas mantienen resultados pequeños en la memoria de la cabeza. Y remote bloqueando más allá de DISTA_MAX_QUEUED tareas (por defecto 100.000) es contrapresión, no un cuelgue.