magmalake/dista

dista (o tin dista-mojo, 0.2.0, Mojo 1.2.0) permite marcar uma função Mojo como remota, chamá-la e fazê-la executar em outros processos ou pods Kubernetes. Cada processo executa o mesmo binário compilado, então uma mensagem de tarefa é um ID de função mais argumentos codificados: nenhum código é enviado, nada é empacotado (pickled), e dados colunares viajam como Arrow.

O que dista não é: um gerenciador de cluster. Ele não cria pods ou máquinas. Ele utiliza todos os workers que entram, drena aqueles que recebe ordem de parar, e informa quantos deseja. Kubernetes, um loop de shell ou orca inicia e para eles. Ele também não reconstrói resultados perdidos com um worker morto (sem reconstrução de linhagem): o job falha, nomeando o resultado, e um orquestrador o executa novamente.

O ambiente instalado mantém dista apenas como um dista.mojoc compilado. Leia seu código-fonte no GitHub, magmalake/dista.mojo, na rev que seu pixi.toml fixa, ou no cache git do pixi (ls -d ~/Library/Caches/rattler/cache/git-v0/checkouts/*/<primeiros 7 da rev>* no macOS, ~/.cache/rattler/... no Linux). Os caminhos de arquivos citados aqui estão nesse checkout.

pixi shelf add dista-mojo    # traz arrow-mlake, memory-region e objectstore

Escreva o Mojo em si com a skill mojo-syntax carregada; para trabalho rápido linha a linha dentro de tarefas, veja mojo-data-performance.

Declare tarefas, depois chame-as

Uma tarefa é uma def de nível superior declarada uma vez com seus tipos de argumento e resultado, porque o Mojo 1.2 não pode inferir o tipo de resultado de uma função. O registro é a lista de parâmetros do Runtime. Toda amostra completa neste arquivo foi construída e executada 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)              # um ObjectRef[Int], imediatamente
    print(rt.get(r))                          # 49: espera pelo 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: uma ref como argumento
    rt.release(maps)

• Aridade: Task[A, R, f] recebe um argumento, Task2[A, B, R, f] dois, Task3[A, B, C, R, f] três. Passe mais como uma única struct. Chamar uma Task2 com um argumento levanta takes 2 arguments na chamada.
• A função pode levantar exceção ou não. Uma tarefa que levanta exceção relança sua mensagem em rt.get, e o worker continua.
• Refs alimentam tarefas sem que o driver busque nada. Passe um ObjectRef[A] onde a tarefa recebe um A, ou um List[ObjectRef[X]] onde ela recebe um Seq[X] (coletados em ordem). A tarefa executa assim que cada entrada existe, e falha com o erro da entrada se alguma falhou.
• rt.release(ref) (ou uma lista) libera um resultado no head e em seu worker assim que nenhuma tarefa pendente precisar dele; chamar get nele depois levanta exceção. Resultados são mantidos até o runtime terminar.
• Também em Runtime: rt.ready(r) (já terminou, sem esperar), rt.location(r) (ip:porta do worker que o mantém e arquivo do node-store, ou "" se o resultado veio inline), rt.stats(), rt.bytes_fetched() e rt.timeline(caminho), que escreve um trace Chrome para Perfetto.

Shuffles: vários resultados, ou um valor junto com refs

Um TaskMany[A, X, f], onde f retorna um Seq[X], é uma tarefa de mapa: rt.remote_many[T](arg, n) retorna n refs, uma por partição, cada uma armazenada, buscada e liberada individualmente; f deve retornar exatamente n itens. Um TaskWith[A, X, R, f], onde f recebe um A e um Seq[X], é um redutor: rt.remote[T](j, refs). Uma amostra que compila, e os detalhes, estão em references/shuffles.md.

O que pode ser um argumento ou um resultado

O codec escolhe uma codificação em tempo de compilação:

• Int, IntN, UIntN, Float32, Float64, Bool e String;
• Seq[T], uma lista que passa pelo codec (List não: reflection não pode recuperar seu tipo de elemento), também como campo de struct;
• qualquer struct simples dos itens acima, campo por campo por reflection (os campos devem ser Movable; tipos de argumento e resultado precisam de Movable & Deinitable);
• um tipo que implementa a trait Codec (seus próprios encode e decode);
• RecordBatch do arrow-mlake, enviado como stream IPC Arrow. Use-o para qualquer coisa grande; um batch por grupo de linhas Parquet é o grão pretendido.

Um ObjectRef dentro de uma struct ou de um Seq é um erro de compilação que nomeia as três formas que funcionam: a ref como argumento, uma lista de refs, ou um valor junto com refs usando TaskWith. Resultados de 64 KiB ou menos voltam pelo head; os maiores ficam no worker que os criou e são mapeados da memória compartilhada no mesmo host ou buscados entre hosts.

Atores

Um ator é um estado em um worker próprio. Suas chamadas executam uma de cada vez, na ordem em que foram feitas, e retornam ObjectRefs como as tarefas fazem.

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)    # uma tarefa chamando o ator que recebeu
    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)

• Um ator ocupa um worker inteiro, e tarefas nunca executam nele. Um ator criado quando todo worker já hospeda um espera um worker entrar, então dê ao modo local mais workers do que atores.
• Handles são valores: passe um ActorHandle para uma CtxTask (def f(mut c: Client, a: A) raises -> R) ou mantenha-o no estado de outro ator, e chame através do Client. Um CtxMethod (def f(mut s: S, mut c: Client, a: A) raises -> R) permite que um ator chame outros; um ciclo de chamadas de ator causa deadlock.
• Checkpoints: checkpoint_every=N no create, ou rt.checkpoint(c), cujo ObjectRef[Int] é o número de chamadas cobertas. Eles passam pelo codec para DISTA_CHECKPOINT_URI, um diretório ou s3://; sem ele rt.checkpoint falha. Um tipo de estado que implementa Codec escolhe o que salvar. Quando o worker de um ator morre, ele volta em outro worker a partir de seu último checkpoint (ou seu construtor) e as chamadas desde então são executadas novamente, então os métodos devem ser seguros para repetir.

Modo local e modo cluster

Runtime(...) lê o ambiente para escolher o papel do processo:

Ambiente	Papel
nenhuma variável	modo local: bifurca um head e workers workers, depois os dirige
DISTA_ROLE=head	executa o head em DISTA_PORT (padrão 6390; 0 escolhe um e imprime), nunca retorna
DISTA_ROLE=worker, DISTA_HEAD=host:porta	serve aquele head, nunca retorna
DISTA_HEAD=host:porta	um driver submetendo àquele head

Portanto crie o Runtime primeiro no main: em um head ou worker o construtor nunca retorna, e qualquer coisa antes dele executa em todos os papéis. Em modo local os workers são bifurcados dentro do construtor e herdam o ambiente definido antes dele. Para executar um cluster manualmente, como tests/cluster.sh faz:

DISTA_ROLE=head DISTA_PORT=6390 ./app               # o head
DISTA_ROLE=worker DISTA_HEAD=headhost:6390 ./app    # uma vez por worker
DISTA_HEAD=headhost:6390 ./app                      # o driver

Workers e drivers tentam novamente alcançar o head por DISTA_CONNECT_TIMEOUT segundos (padrão 60), então a ordem de início não importa. Um head iniciado com dista head em vez disso é genérico: ele serve vários programas ao mesmo tempo e fornece as tarefas de cada driver apenas a workers construídos a partir do mesmo binário.

O comando dista

Construído a partir do repositório dista com pixi run cli (para build/dista). Cada comando encontra seu head por --head, senão DISTA_HEAD, senão localhost:6390.

Comando	O que faz
dista head [--port 6390] [--pipeline 2]	um head genérico para qualquer programa
dista status	workers com seu hash de registro, drivers, atores
dista status --json	a demanda, em uma linha, para um autoscaler
dista drain <ip:porta>	drena aquele worker, como SIGTERM faz; recusa o último enquanto houver trabalho
dista scale --once	imprime a decisão da política para a demanda atual
dista timeline -o trace.json	tarefas finalizadas de cada driver, como um trace Chrome
dista stop	desliga o head e seus workers

Com DISTA_METRICS_PORT definido, o head e cada worker servem métricas Prometheus em /metrics. Cada variável, nome de métrica e os manifiestos Kubernetes estão em references/operations.md.

Drenagem e autoescalonamento

Um worker que recebe SIGTERM (como o Kubernetes para um pod) drena em vez de morrer: não aceita novos trabalhos, termina o que tem, seus resultados se movem para outros workers e seu ator se move através de um checkpoint, então termina com código 0. Nada falha ou executa duas vezes (tests/drain.sh). DISTA_DRAIN_SECONDS (padrão 25) o limita; um segundo SIGTERM o encerra imediatamente.

dista status --json fornece tasks_ready, tasks_waiting, tasks_running, actors_waiting, pipeline, checkpoints, e por worker address, node, running, idle_seconds, held_bytes, actor e draining. dista.scale (importe-o de dista.scale; não é re-exportado de dista) transforma isso em uma decisão, sem relógio ou E/S próprios:

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 e padrões de Policy: min_workers=1 (0 é permitido: o último worker drena quando não há tarefa restante e ele não detém nada), max_workers=4, tasks_per_worker=0 (a profundidade do pipeline do head), up_after=5 segundos de tarefas prontas sem slot livre, idle_before_down=60 segundos, e max_step=4 workers iniciados por decisão. decide(demand, policy, now, pending=0, short_since=0) retorna start, drain (endereços, menos bytes retidos primeiro), want e short_since, que o provedor passa de volta na próxima vez com pending, os workers iniciados mas que ainda não entraram. dista scale --once --min 0 --max 6 ... imprime como linhas; tests/autoscale.sh aciona um job de 0 workers a 6 e de volta a 0 com ele.

Sob orca

Um pipeline orca declara suas tarefas e workers, e o orca inicia o runtime uma vez por execução. Uma etapa nunca constrói seu próprio Runtime; ela pega o runtime da execução emprestado (um trecho):

comptime Trips = Task[Day, Int, trips]
comptime FanoutTasks = Tasks[Trips]                   # do 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))   # depois rt.get, como de costume
    ...

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

O programa completo é examples/fanout.mojo em orca.mojo.

Orca executa as etapas de uma execução em ordem e as partições de uma etapa mapeada uma após a outra, então uma etapa que chama uma tarefa e espera por ela é serial. Envie todo o trabalho de uma etapa com rt.remote primeiro, depois rt.get cada ref.

Erros comuns

• registry mismatch: this peer was built with different tasks. O driver e os workers vieram de compilações diferentes, ou listam tarefas diferentes em Runtime[...]. O hash do registro cobre os nomes das tarefas em ordem, então recompile uma vez e execute esse único binário em todos os papéis.
• ... is not registered with this runtime. Toda tarefa, ator e método que você chama deve estar em Runtime[...]. Chamadas através do Client de uma CtxTask não são verificadas no momento da chamada, então uma não registrada falha no worker em vez disso, com unknown task ID ou unknown method ID.
• Uma função aninhada ou closure como tarefa. Use uma def de nível superior.
• Assumir que uma tarefa executa uma vez. As tarefas de um worker morto executam novamente, até DISTA_TASK_RETRIES (padrão 3): mantenha tarefas e métodos seguros para repetir.
• Esperar que um resultado perdido volte. Um resultado mantido apenas por um worker morto falha com result <id> lost: its worker went away, assim como toda tarefa esperando por ele. Reexecute o job, ou deixe o orca tentar a etapa novamente.
• Refs não liberados mantêm resultados pequenos na memória do head. E remote bloqueando além de DISTA_MAX_QUEUED tarefas (padrão 100.000) é contrapressão, não um travamento.