magmalake/dista

dista (la boîte dista-mojo, 0.2.0, Mojo 1.2.0) vous permet de marquer une fonction Mojo comme distante, de l'appeler, et de l'exécuter sur d'autres processus ou pods Kubernetes. Chaque processus exécute le même binaire compilé, donc un message de tâche est un ID de fonction plus des arguments encodés : aucun code n'est expédié, rien n'est picklé, et les données columnaires voyagent comme Arrow.

Ce que dista n'est pas : un gestionnaire de cluster. Il ne crée pas de pods ou de machines. Il utilise chaque worker qui rejoint, vide ceux qu'on lui dit d'arrêter, et indique combien il en veut. Kubernetes, une boucle shell ou orca les démarre et les arrête. Il ne reconstruit pas non plus les résultats perdus avec un worker mort (pas de reconstruction de lignage) : le travail échoue, nomme le résultat, et un orchestrateur l'exécute à nouveau.

L'environnement installé ne contient dista que comme un dista.mojoc compilé. Lisez son code source sur GitHub, magmalake/dista.mojo, à la rev que votre pixi.toml épingle, ou dans le cache git de pixi (ls -d ~/Library/Caches/rattler/cache/git-v0/checkouts/*/<7 premiers de rev>* sur macOS, ~/.cache/rattler/... sur Linux). Les chemins de fichiers cités ici sont dans ce checkout.

pixi shelf add dista-mojo    # apporte arrow-mlake, memory-region et objectstore

Écrivez le Mojo lui-même avec la compétence mojo-syntax chargée ; pour un travail rapide ligne par ligne dans les tâches, voir mojo-data-performance.

Déclarez les tâches, puis appelez-les

Une tâche est un def de niveau supérieur déclaré une fois avec ses types d'argument et de résultat, car Mojo 1.2 ne peut pas inférer le type de résultat d'une fonction. Le registre est la liste de paramètres du Runtime. Chaque exemple complet de ce fichier a été compilé et exécuté avec 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], immédiatement
    print(rt.get(r))                          # 49 : attend le résultat
    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 : une ref comme argument
    rt.release(maps)

• Arité : Task[A, R, f] prend un argument, Task2[A, B, R, f] deux, Task3[A, B, C, R, f] trois. Passez-en plus comme une seule structure. Appeler une Task2 avec un argument lève takes 2 arguments à l'appel.
• La fonction peut lever ou non. Une tâche qui lève relève son message dans rt.get, et le worker continue.
• Les refs alimentent les tâches sans que le pilote ne récupère quoi que ce soit. Passez un ObjectRef[A] là où la tâche prend un A, ou une List[ObjectRef[X]] là où elle prend un Seq[X] (rassemblés dans l'ordre). La tâche s'exécute une fois que chaque entrée existe, et échoue avec l'erreur de l'entrée si une a échoué.
• rt.release(ref) (ou une liste) libère un résultat à la tête et sur son worker dès qu'aucune tâche en attente n'en a besoin ; get dessus lève ensuite. Les résultats sont autrement conservés jusqu'à la fin de l'exécution.
• Aussi sur Runtime : rt.ready(r) (fini ou non, sans attendre), rt.location(r) (le ip:port du worker qui le détient et le fichier du magasin de nœuds, ou "" si le résultat est revenu en ligne), rt.stats(), rt.bytes_fetched() et rt.timeline(path), qui écrit une trace Chrome pour Perfetto.

Shuffles : plusieurs résultats, ou une valeur à côté de refs

Un TaskMany[A, X, f], où f retourne un Seq[X], est une tâche de map : rt.remote_many[T](arg, n) retourne n refs, une par partition, chacune stockée, récupérée et libérée individuellement ; f doit retourner exactement n éléments. Un TaskWith[A, X, R, f], où f prend un A et un Seq[X], est un réducteur : rt.remote[T](j, refs). Un exemple qui compile, et les détails, sont dans references/shuffles.md.

Ce qui peut être un argument ou un résultat

Le codec choisit un encodage à la compilation :

• Int, IntN, UIntN, Float32, Float64, Bool et String ;
• Seq[T], une liste qui passe par le codec (List ne passe pas : la réflexion ne peut pas retrouver son type d'élément), également comme champ de structure ;
• toute structure simple des types ci-dessus, champ par champ par réflexion (les champs doivent être Movable ; les types d'argument et de résultat ont besoin de Movable & Deinitable) ;
• un type implémentant le trait Codec (ses propres encode et decode) ;
• le RecordBatch d'arrow-mlake, envoyé comme un flux IPC Arrow. Utilisez-le pour tout ce qui est volumineux ; un lot par groupe de lignes Parquet est la granularité prévue.

Un ObjectRef à l'intérieur d'une structure ou d'un Seq est une erreur de compilation qui nomme les trois formes qui fonctionnent : la ref comme argument, une liste de refs, ou une valeur à côté de refs avec TaskWith. Les résultats de 64 Kio ou moins reviennent via la tête ; les plus gros restent sur le worker qui les a créés et sont mappés depuis la mémoire partagée sur le même hôte ou récupérés entre hôtes.

Acteurs

Un acteur est un état sur un worker qui lui est propre. Ses appels s'exécutent un par un, dans l'ordre où ils ont été faits, et retournent des ObjectRef comme les tâches.

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)    # une tâche appelant l'acteur qu'on lui a passé
    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 acteur prend un worker entier, et les tâches ne s'exécutent jamais dessus. Un acteur créé alors que chaque worker en héberge un attend qu'un worker rejoigne, donc donnez au mode local plus de workers que d'acteurs.
• Les handles sont des valeurs : passez un ActorHandle à un CtxTask (def f(mut c: Client, a: A) raises -> R) ou gardez-le dans l'état d'un autre acteur, et appelez via le Client. Un CtxMethod (def f(mut s: S, mut c: Client, a: A) raises -> R) permet à un acteur d'en appeler d'autres ; un cycle d'appels d'acteurs se bloque.
• Points de contrôle : checkpoint_every=N sur create, ou rt.checkpoint(c), dont l'ObjectRef[Int] est le nombre d'appels couverts. Ils passent par le codec vers DISTA_CHECKPOINT_URI, un répertoire ou s3:// ; sans cela rt.checkpoint échoue. Un type d'état implémentant Codec choisit quoi sauvegarder. Quand le worker d'un acteur meurt, il revient sur un autre worker depuis son dernier point de contrôle (ou son constructeur) et les appels depuis sont réexécutés, donc les méthodes doivent être sûres à répéter.

Mode local et mode cluster

Runtime(...) lit l'environnement pour choisir le rôle du processus :

Environnement	Rôle
aucune variable	mode local : fork une tête et workers workers, puis les pilote
DISTA_ROLE=head	exécute la tête sur DISTA_PORT (défaut 6390 ; 0 en choisit un et l'imprime), ne retourne jamais
DISTA_ROLE=worker, DISTA_HEAD=host:port	sert cette tête, ne retourne jamais
DISTA_HEAD=host:port	un pilote soumettant à cette tête

Donc créez le Runtime en premier dans main : dans une tête ou un worker le constructeur ne retourne jamais, et tout ce qui précède s'exécute dans chaque rôle. En mode local les workers sont forkés à l'intérieur du constructeur et héritent de l'environnement défini avant lui. Pour exécuter un cluster à la main, comme le fait tests/cluster.sh :

DISTA_ROLE=head DISTA_PORT=6390 ./app               # la tête
DISTA_ROLE=worker DISTA_HEAD=headhost:6390 ./app    # une fois par worker
DISTA_HEAD=headhost:6390 ./app                      # le pilote

Les workers et pilotes réessayent de joindre la tête pendant DISTA_CONNECT_TIMEOUT secondes (défaut 60), donc l'ordre de démarrage n'a pas d'importance. Une tête démarrée avec dista head est générique : elle sert plusieurs programmes à la fois et donne les tâches de chaque pilote uniquement aux workers construits à partir du même binaire.

La commande dista

Construite à partir du dépôt dista avec pixi run cli (vers build/dista). Chaque commande trouve sa tête par --head, sinon DISTA_HEAD, sinon localhost:6390.

Commande	Ce qu'elle fait
dista head [--port 6390] [--pipeline 2]	une tête générique pour tout programme
dista status	workers avec leur hash de registre, pilotes, acteurs
dista status --json	la demande, sur une ligne, pour un autoscaler
dista drain <ip:port>	vide ce worker, comme le fait SIGTERM ; refuse le dernier s'il reste du travail
dista scale --once	imprime la décision de la politique pour la demande actuelle
dista timeline -o trace.json	les tâches terminées de chaque pilote, comme une trace Chrome
dista stop	arrête la tête et ses workers

Avec DISTA_METRICS_PORT défini, la tête et chaque worker servent des métriques Prometheus sur /metrics. Toutes les variables, noms de métriques et les manifests Kubernetes sont dans references/operations.md.

Drainage et autoscaling

Un worker recevant SIGTERM (comment Kubernetes arrête un pod) se vide au lieu de mourir : il ne prend pas de nouveau travail, finit ce qu'il a, ses résultats passent à d'autres workers et son acteur se déplace via un point de contrôle, puis il sort avec le code 0. Rien n'échoue ou ne s'exécute deux fois (tests/drain.sh). DISTA_DRAIN_SECONDS (défaut 25) le limite ; un deuxième SIGTERM l'arrête immédiatement.

dista status --json donne tasks_ready, tasks_waiting, tasks_running, actors_waiting, pipeline, checkpoints, et par worker address, node, running, idle_seconds, held_bytes, actor et draining. dista.scale (importez-le depuis dista.scale ; il n'est pas réexporté depuis dista) transforme cela en une décision, sans horloge ni E/S propres :

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

Champs de Policy et valeurs par défaut : min_workers=1 (0 est autorisé : le dernier worker se vide quand il n'y a plus de tâche et qu'il ne détient rien), max_workers=4, tasks_per_worker=0 (la profondeur du pipeline de la tête), up_after=5 secondes de tâches prêtes sans emplacement libre, idle_before_down=60 secondes, et max_step=4 workers démarrés par décision. decide(demand, policy, now, pending=0, short_since=0) retourne start, drain (adresses, octets détenus les plus faibles en premier), want et short_since, que le fournisseur repasse la prochaine fois avec pending, les workers démarrés mais pas encore rejoints. dista scale --once --min 0 --max 6 ... l'imprime sous forme de lignes ; tests/autoscale.sh pilote un travail de 0 workers à 6 et retour à 0 avec.

Sous orca

Un pipeline orca déclare ses tâches et workers, et orca démarre l'exécution une fois par exécution. Une étape ne construit jamais son propre Runtime ; elle emprunte celui de l'exécution (un extrait) :

comptime Trips = Task[Day, Int, trips]
comptime FanoutTasks = Tasks[Trips]                   # depuis 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))   # puis rt.get, comme d'habitude
    ...

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

Le programme complet est examples/fanout.mojo dans orca.mojo.

orca exécute les étapes d'une exécution dans l'ordre et les partitions d'une étape mappée l'une après l'autre, donc une étape qui appelle une tâche et l'attend est sérialisée. Envoyez tout le travail d'une étape avec rt.remote d'abord, puis rt.get chaque ref.

Erreurs courantes

• registry mismatch: this peer was built with different tasks. Le pilote et les workers viennent de compilations différentes, ou listent des tâches différentes dans Runtime[...]. Le hash du registre couvre les noms de tâches dans l'ordre, donc recompilez une fois et exécutez ce seul binaire dans chaque rôle.
• ... is not registered with this runtime. Chaque tâche, acteur et méthode que vous appelez doit être dans Runtime[...]. Les appels via le Client d'un CtxTask ne sont pas vérifiés à l'appel, donc un non enregistré échoue sur le worker, avec unknown task ID ou unknown method ID.
• Une fonction imbriquée ou une fermeture comme tâche. Utilisez un def de niveau supérieur.
• Supposer qu'une tâche s'exécute une fois. Les tâches d'un worker mort s'exécutent à nouveau, jusqu'à DISTA_TASK_RETRIES (défaut 3) : gardez les tâches et méthodes sûres à répéter.
• S'attendre à ce qu'un résultat perdu revienne. Un résultat détenu uniquement par un worker mort échoue avec result <id> lost: its worker went away, comme chaque tâche l'attendant. Relancez le travail, ou laissez orca réessayer l'étape.
• Refs non libérées gardent les petits résultats dans la mémoire de la tête. Et remote bloquant au-delà de DISTA_MAX_QUEUED tâches (défaut 100 000) est une contre-pression, pas un blocage.