magmalake/dista

dista（dista-mojo tin，0.2.0，Mojo 1.2.0）讓你將 Mojo 函式標記為遠端、呼叫它，並讓它在其他程序或 Kubernetes Pod 上執行。每個程序都執行同一份編譯後的二進位檔，因此任務訊息只包含函式 ID 加上編碼後的參數：不傳輸程式碼、不進行 pickle、且欄位式資料以 Arrow 格式傳輸。

dista 不是什麼：它不是叢集管理員。它不會建立 Pod 或機器。它使用所有加入的工作節點，對那些被要求停止的工作節點進行 draining，並告知需要多少個工作節點。Kubernetes、Shell 迴圈或 orca 負責啟動和停止它們。它也不會重建因工作節點死亡而遺失的結果（沒有 lineage 重建）：任務會失敗，並指明是哪個結果，然後由編排器重新執行。

已安裝的環境僅將 dista 作為編譯後的 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/...）。此處引用的檔案路徑均位於該 checkout 目錄中。

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：將 ref 作為參數傳遞
    rt.release(maps)

• 參數數量： Task[A, R, f] 接受一個參數，Task2[A, B, R, f] 接受兩個，Task3[A, B, C, R, f] 接受三個。如果需要更多參數，請將它們包裝在一個 struct 中傳遞。使用一個參數呼叫 Task2 會在呼叫點引發 takes 2 arguments 錯誤。
• 函式可以引發異常或不引發。 一個會引發異常的任務會在 rt.get 中重新引發其錯誤訊息，而工作節點會繼續執行。
• 將 Ref 傳遞給任務時，驅動程式無需自行擷取任何資料。 在任務接受 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 追蹤檔案）。

Shuffles：多個結果，或者一個包含 refs 的值

一個 TaskMany[A, X, f]，其中 f 返回 Seq[X]，是一個映射任務：rt.remote_many[T](arg, n) 返回 n 個 refs，每個對應一個分區，各自獨立地被儲存、擷取和釋放；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 不會：反射無法恢復其元素型別），也可以作為 struct 的欄位；
• 任何由上述型別組成的普通 struct，透過反射逐欄位進行（欄位必須是 Movable 的；參數和結果型別需要是 Movable & Deinitable）；
• 實作 Codec trait 的型別（具有自己的 encode 和 decode 方法）；
• arrow-mlake 的 RecordBatch，以 Arrow IPC 串流的形式傳送。對於任何大型資料都應使用此型別；預期的粒度是每個 Parquet row group 一個 batch。

在 struct 或 Seq 內部的 ObjectRef 是一個編譯錯誤，錯誤訊息會說明可行的三種形式：將 ref 作為參數、ref 的列表，或者使用 TaskWith 將一個值與 refs 並列。大小為 64 KiB 或更小的結果會通過主節點回傳；較大的結果會停留在產生它們的工作節點上，並在同一主機上透過共享記憶體映射，或在不同主機之間進行擷取。

Actors

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 數量更多的工作節點。
• Handle 是值： 你可以將 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:// URI；沒有設定此變數時，rt.checkpoint 會失敗。實作 Codec trait 的狀態型別可以選擇要儲存的內容。當 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>	對該工作節點進行 draining，如同 SIGTERM 的做法；當仍有工作未完成時，會拒絕 draining 最後一個工作節點
dista scale --once	根據當前需求印出策略的決定
dista timeline -o trace.json	每個驅動程式已完成的任務，以 Chrome trace 格式輸出
dista stop	關閉主節點及其所有工作節點

設定 DISTA_METRICS_PORT 後，主節點和每個工作節點都會在 /metrics 端點提供 Prometheus 指標。所有變數、指標名稱及 Kubernetes 部署清單檔案均在 references/operations.md 中。

Draining 與自動擴展

收到 SIGTERM 訊號的工作節點（這是 Kubernetes 停止 Pod 的方式）會進行 draining 而非直接終止：它不接受新工作、完成現有工作、將其結果移至其他工作節點、其 actor 透過檢查點遷移，然後以退出碼 0 結束。不會有任何任務失敗或重複執行（參見 tests/drain.sh）。DISTA_DRAIN_SECONDS（預設 25 秒）限制了 draining 的時間；收到第二個 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：最後一個工作節點會在沒有剩餘任務且不持有任何資料時進行 draining）、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（已啟動但尚未加入的工作節點數量）傳回。dista scale --once --min 0 --max 6 ... 會將其印製為多行文字；tests/autoscale.sh 使用它將一個任務從 0 個工作節點驅動到 6 個，再回到 0 個。

在 orca 中使用

orca 管線會宣告其任務與工作節點，orca 會在每次執行時啟動執行時期。步驟不會自行建立 Runtime；它會借用該次執行的 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 送出所有工作，然後再對每個 ref 使用 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 重試該步驟。
• 未釋放的 refs 會將小結果保留在主節點的記憶體中。此外，remote 在超過 DISTA_MAX_QUEUED 個任務（預設 100,000）時會阻塞，這是背壓機制，並非掛起。