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)个任务而阻塞时,这是背压机制,而不是挂起。