Community编程与开发github.com

magmalake/dista

使用 dista(magmalake 的 Ray 风格任务框架)编写和运行分布式 Mojo 程序,其中一份编译后的二进制文件同时充当驱动、主节点和每个工作节点。当您编写或审查导入 dista 或提及分布式 Mojo、远程任务、`Task`/`Task2`/`TaskMany`/`TaskWith`、`ObjectRef`、`rt.remote`、`rt.get`、`Seq`、带有方法或检查点的 actor、集群、主节点、工作节点、`DISTA_ROLE`/`DISTA_HEAD`、`dista` 命令(status、drain、scale、timeline、stop)、排空、自动扩缩或 orca 管道步骤中的 dista 运行时;以及当 dista 作业挂起、因注册表不匹配被拒绝或因结果丢失而失败时,请使用此技能。与 Modular 的 mojo-syntax 技能和 magmalake 的 mojo-data-performance 技能搭配使用。

dista 是什么?

dista 是 magmalake 开发的一个用于 Mojo 语言的分布式任务框架(类似于 Ray),它使得在集群或本地进程中并行执行 Mojo 函数变得简单。核心概念是所有进程运行相同的编译后二进制文件,通过函数 ID 和序列化参数进行通信,无需发送代码或 pickle 对象,并且列式数据使用 Arrow 格式传输。dista 支持本地模式(自动 fork 工作进程)和集群模式(通过环境变量 `DISTA_ROLE` 和 `DISTA_HEAD` 指定角色),并提供了一个 `dista` 命令行工具用于管理集群状态、排空工作节点、自动扩缩和生成时间线。该技能详细涵盖了任务和 Actor 的声明与调用、远程引用(`ObjectRef`)的使用、数据序列化规则、检查点机制、与 orca 管道框架的集成,以及常见的错误排查和最佳实践。

兼容平台~Claude Code~Codex CLI~Cursor
npx skills add https://github.com/magmalake/skills/tree/HEAD/dista

Installed? Explore more 编程与开发 skills: steipete/bluebubbles, steipete/eightctl, steipete/blucli · View all 6 →

在你喜欢的 AI 中提问

打开一个已预加载此 Agent Skill 的新对话。

文档

dista(dista-mojo tin,0.2.0,Mojo 1.2.0)允许您将一个 Mojo 函数标记为远程,调用它,并让它在其他进程或 Kubernetes pod 上运行。每个进程运行 相同的编译后二进制文件,因此任务消息是一个函数 ID 加上编码后的参数:没有代码发送,没有 pickle 操作,列式数据以 Arrow 格式传输。

dista 不是什么:它不是集群管理器。它不会创建 pod 或机器。它使用每一个加入的工作节点,排空那些被告知停止的节点,并声明它需要多少个。Kubernetes、shell 循环或 orca 负责启动和停止它们。它也不会重建因工作节点死亡而丢失的结果(没有血统重建):作业失败,会指出哪个结果丢失,然后由编排器重新运行。

安装的环境中只包含一个编译后的 dista.mojoc 文件。在 GitHub magmalake/dista.mojo 上读取其源码,版本为你 pixi.toml 中固定的 rev,或者在 pixi 的 git 缓存中(macOS 上为 ls -d ~/Library/Caches/rattler/cache/git-v0/checkouts/*/<rev 的前 7 个字符>*,Linux 上为 ~/.cache/rattler/...)。此处引用的文件路径均位于该检出目录中。

pixi shelf add dista-mojo    # 引入 arrow-mlake、memory-region 和 objectstore

使用加载了 mojo-syntax 技能来编写 Mojo 代码;关于在任务内部进行快速的逐行操作,请参阅 mojo-data-performance。

声明任务,然后调用它们

任务是一个顶级的 def,一次性声明其参数和结果类型,因为 Mojo 1.2 无法推断函数的结果类型。注册表是 Runtime 的参数列表。本文件中的所有完整示例均已针对 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)              # 一个 ObjectRef[Int],立即返回
    print(rt.get(r))                          # 49:等待结果
    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:将引用作为参数传入
    rt.release(maps)
  • 参数个数: Task[A, R, f] 接收一个参数,Task2[A, B, R, f] 接收两个,Task3[A, B, C, R, f] 接收三个。更多参数可作为单个结构体传递。若使用一个参数调用 Task2,会在调用点引发 takes 2 arguments 错误。
  • 函数可以引发异常,也可以不引发异常。 引发异常的任务会在 rt.get 中重新抛出其异常消息,工作节点继续运行。
  • 引用可以直接输入任务,无需驱动获取任何内容。 如果任务需要 A 类型的参数,您可以传入 ObjectRef[A];如果任务需要 Seq[X](按顺序收集),则可以传入 List[ObjectRef[X]]。该任务将在每个输入都存在后运行一次,如果某个输入出错,则该任务会失败并返回该输入的错误。
  • rt.release(ref)(或一个引用列表)释放结果,释放时机为结果在主节点上且其工作节点上没有任何待处理任务需要它时;之后对其执行 get 将引发异常。否则,结果将一直保留,直到运行时结束。
  • 此外,Runtime 上还有:rt.ready(r)(检查是否已完成,不等待)、rt.location(r)(返回持有结果的工作节点的 ip:port 和节点存储文件,如果结果以内联方式传回则返回 "")、rt.stats()、rt.bytes_fetched() 和 rt.timeline(path)(为 Perfetto 写入 Chrome 跟踪信息)。

数据混洗:多个结果,或一个带有引用的值

TaskMany[A, X, f],其中 f 返回 Seq[X],是一个映射任务:rt.remote_many[T](arg, n) 返回 n 个引用,每个分区一个,每个引用独立存储、获取和释放;f 必须恰好返回 n 个元素。TaskWith[A, X, R, f],其中 f 接收一个 A 和一个 Seq[X],是一个归约器:rt.remote[T](j, refs)。一个可编译的示例及详细信息在 references/shuffles.md 中。

什么可以作为参数或结果

编解码器在编译时选择编码方式:

  • Int、IntN、UIntN、Float32、Float64、Bool 和 String;
  • Seq[T],一个通过编解码器的列表(List 不行:反射无法恢复其元素类型),也可以作为结构体字段;
  • 任何由上述类型组成的普通结构体,通过反射逐个字段处理(字段必须是 Movable;参数和结果类型需要满足 Movable & Deinitable);
  • 实现 Codec 特质的类型(拥有自己的 encode 和 decode 方法);
  • arrow-mlake 的 RecordBatch,以 Arrow IPC 流的形式发送。对于任何大型数据请使用此方式;每个 Parquet 行组对应一个批次是预期的粒度。

在结构体或 Seq 内部使用 ObjectRef 会导致编译错误,并指出可行的三种形式:将引用作为参数、引用列表、或使用 TaskWith 将值与引用并列。大小不超过 64 KiB 的结果会通过主节点返回;更大的结果会留在产生它们的工作节点上,并在同一主机上通过共享内存映射,或跨主机获取。

Actor

Actor 是在其工作节点上拥有独立状态的对象。它的调用按调用顺序一次运行一个,并像任务一样返回 ObjectRef。

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)    # 一个任务调用它收到的 actor
    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)
  • 一个 Actor 独占一个完整的工作节点, 并且任务永远不会在其上运行。在一个 Actor 创建时,如果每个工作节点都托管着一个 Actor,那么该 Actor 会等待一个工作节点加入,因此在本地模式下要提供比 Actor 数量更多的工作节点。
  • 句柄是值: 将 ActorHandle 传递给 CtxTask(def f(mut c: Client, a: A) raises -> R)或将其存储在另一个 Actor 的状态中,并通过 Client 进行调用。CtxMethod(def f(mut s: S, mut c: Client, a: A) raises -> R)允许一个 Actor 调用其他 Actor;Actor 调用循环会导致死锁。
  • 检查点: 在 create 时设置 checkpoint_every=N,或调用 rt.checkpoint(c),其返回的 ObjectRef[Int] 是已覆盖的调用次数。检查点通过编解码器保存到 DISTA_CHECKPOINT_URI(一个目录或 s3://);如果未设置该变量,rt.checkpoint 会失败。实现了 Codec 特质的状态类型可以自行选择要保存的内容。当 Actor 的工作节点死亡时,它会从最近的检查点(或其构造函数)在另一个工作节点上恢复,并且 自检查点以来的调用将重新运行,因此方法应该是可安全重复执行的。

本地模式与集群模式

Runtime(...) 读取环境变量以选择进程的角色:

环境变量角色
两者均未设置本地模式:fork 一个主节点和 workers 个工作节点,然后驱动它们
DISTA_ROLE=head在 DISTA_PORT(默认 6390;设置为 0 则选择一个可用端口并打印它)上运行主节点,永不返回
DISTA_ROLE=worker,DISTA_HEAD=host:port服务于该主节点,永不返回
DISTA_HEAD=host:port一个向该主节点提交作业的驱动

因此,在 main 中首先创建 Runtime:在主节点或工作节点中,构造函数永不返回,而之前的任何代码都会在每个角色中运行。在本地模式中,工作节点在构造函数内部被 fork,并继承在其之前设置的环境变量。要手动运行一个集群,可以像 tests/cluster.sh 那样:

DISTA_ROLE=head DISTA_PORT=6390 ./app               # 主节点
DISTA_ROLE=worker DISTA_HEAD=headhost:6390 ./app    # 每个工作节点一次
DISTA_HEAD=headhost:6390 ./app                      # 驱动

工作节点和驱动会在 DISTA_CONNECT_TIMEOUT 秒(默认 60)内重试连接主节点,因此启动顺序无关紧要。以 dista head 命令启动的主节点是通用的:它可以同时服务于多个程序,并且只将每个驱动的任务分配给由相同二进制构建的工作节点。

dista 命令

从 dista 仓库使用 pixi run cli 构建(输出到 build/dista)。每个命令通过 --head 查找其主节点,否则使用 DISTA_HEAD,再否则使用 localhost:6390。

命令作用
dista head [--port 6390] [--pipeline 2]一个适用于任何程序的通用主节点
dista status显示工作节点及其注册表哈希、驱动、actor 信息
dista status --json在一行中显示需求,供自动扩缩器使用
dista drain <ip:port>排空该工作节点(类似于 SIGTERM 信号的处理);如仍有未完成的工作,则拒绝排空最后一个工作节点
dista scale --once根据当前需求打印策略的决策
dista timeline -o trace.json每个驱动已完成的任务的 Chrome 跟踪信息
dista stop关闭主节点及其所有工作节点

如果设置了 DISTA_METRICS_PORT,主节点和每个工作节点都会在 /metrics 提供 Prometheus 指标。所有变量、指标名称和 Kubernetes 清单都在 references/operations.md 中。

排空与自动扩缩

收到 SIGTERM(Kubernetes 停止 pod 的方式)的工作节点会进行排空而不是死亡:它不再接收新工作,完成现有工作,将其结果迁移到其他工作节点,其 Actor 通过检查点迁移,然后以退出码 0 退出。没有工作会失败或重复运行(参见 tests/drain.sh)。DISTA_DRAIN_SECONDS(默认 25)限制了排空时间;第二次收到 SIGTERM 会立即终止它。

dista status --json 输出包含 tasks_ready、tasks_waiting、tasks_running、actors_waiting、pipeline、checkpoints,以及每个工作节点的 address、node、running、idle_seconds、held_bytes、actor 和 draining。dista.scale(从 dista.scale 导入,它不会从 dista 重新导出)将这些数据转化为一个决策,该函数本身不涉及时钟或 I/O:

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

Policy 字段及其默认值:min_workers=1(允许为 0:当没有剩余任务且不持有任何内容时,最后一个工作节点会排空)、max_workers=4、tasks_per_worker=0(使用主节点的管道深度)、up_after=5 秒(在有就绪任务但无空闲槽位时)、idle_before_down=60 秒,以及 max_step=4(每次决策最多启动的工作节点数)。decide(demand, policy, now, pending=0, short_since=0) 返回 start、drain(地址列表,按持有字节数从小到大排序)、want 和 short_since。提供者将在下次调用时通过 pending(已启动但尚未加入的工作节点数)将 short_since 传递回来。dista scale --once --min 0 --max 6 ... 会将其打印为行;tests/autoscale.sh 演示了如何驱动一个作业从 0 个工作节点扩展到 6 个,再缩减回 0 个。

在 orca 下使用

orca 管道声明其任务和工作节点,orca 每次运行启动一个运行时。步骤从不自行构建 Runtime;它借用运行时的实例(以下为摘录):

comptime Trips = Task[Day, Int, trips]
comptime FanoutTasks = Tasks[Trips]                   # 来自 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))   # 然后照常使用 rt.get
    ...

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

完整的程序位于 orca.mojo 仓库的 examples/fanout.mojo 中。

orca 按顺序运行一个运行的相关步骤,并且映射步骤的分区是一个接一个运行的,因此如果一个步骤只调用一个任务并等待它,那么该步骤是串行的。应该先使用 rt.remote 发送步骤的所有工作,然后再逐一 rt.get 每个引用。

常见错误

  • registry mismatch: this peer was built with different tasks。 驱动和工作节点来自不同的构建,或者在 Runtime[...] 中列出了不同的任务。注册表哈希按顺序覆盖任务名称,因此请一次性构建,并在所有角色中运行同一个二进制文件。
  • ... is not registered with this runtime。 您调用的每个任务、actor 和方法都必须出现在 Runtime[...] 中。通过 CtxTask 的 Client 进行的调用在调用时不会被检查,因此未注册的调用会在工作节点上失败,并显示 unknown task ID 或 unknown method ID。
  • 将嵌套函数或闭包用作任务。 请使用顶级的 def。
  • 假设任务只运行一次。 死亡工作节点的任务会重新运行,最多重试 DISTA_TASK_RETRIES 次(默认 3):请确保任务和方法是安全可重复执行的。
  • 期望丢失的结果能够恢复。 仅由死亡工作节点持有的结果会失败并显示 result <id> lost: its worker went away,所有等待该结果的任务也会失败。请重新运行作业,或让 orca 重试该步骤。
  • 未释放的引用会将小型结果保留在主节点的内存中。此外,当 remote 调用因 DISTA_MAX_QUEUED(默认 100,000)个任务而阻塞时,这是背压机制,而不是挂起。

Individual skills in this repo

This repo contains 2 individual skills — each has its own dedicated page.

相关技能