Skip to content

Cookbook:常用配方速查 ​

面对"我该怎么写 XX"的问题,这一章给出经过官方 CI 验证的标准答案。每个配方都来自官方 Cookbook(docs/source/cookbook.md),对应 python/tests/test_cookbook.py 中真实运行的测试代码——文档即测试,代码永不过时。

配方 1:获取网格的 rank-0 Actor ​

问题:经常需要拿到网格的 rank-0,比如作为 SPMD 作业的协调者。

python
# counters 是一个多 Actor 网格
rank_0 = counters.flatten("rank").slice(rank=0)
await rank_0.increment.call_one()

为什么这样写:

  • flatten("rank") 把所有维度折叠成一个名为 rank 的维度——无论网格原始形状如何都能工作;
  • 直接 slice(hosts=0, gpus=0) 的问题:维度一旦改名就报错,flatten 写法更健壮;
  • 切片结果是单例网格,所以可以用 call_one。

配方 2:调用自己网格中的指定 rank ​

问题:Actor 常需要联系自己 mesh 里的另一个 rank(比如 rank-0 协调者)。

python
class Worker(Actor):
    mesh: "Worker"

    @endpoint
    async def set_mesh(self, mesh: "Worker") -> None:
        # Actor 构造时无法引用自己的 mesh,
        # 所以 mesh 拉起后用端点把引用广播给每个成员
        self.mesh = mesh

    @endpoint
    async def call_coordinator(self) -> str:
        coordinator = self.mesh.flatten("rank").slice(rank=0)
        return await coordinator.ping.call_one()

要点:

  • Actor 在 __init__ 里拿不到自己的 mesh,标准做法是拉起后调用 set_mesh 把句柄发进去;
  • 对自己 mesh 的句柄切片,得到目标 rank 的引用,call_one 就是点对点直呼。

配方 3:用 Port 做异步扇入 ​

问题:call 是同步收集——所有响应凑齐才返回。想让各 rank 各自就绪、各自上报呢?

python
from monarch.actor import Channel, Port

class Aggregator(Actor):
    mesh: "Aggregator"

    @endpoint
    async def set_mesh(self, mesh: "Aggregator") -> None:
        self.mesh = mesh

    @endpoint
    async def collect(self) -> list[int]:
        # 打开一个通道:port 用于发送,receiver 用于接收
        port, receiver = Channel[int].open()
        # 把发送端广播给除自己外的所有 rank
        others = self.mesh.flatten("rank").slice(rank=slice(1, None))
        others.produce.broadcast(port)
        expected = others.size() * STREAM_LEN
        # 谁先完成谁先到,逐个收齐
        return sorted([await receiver.recv() for _ in range(expected)])

    @endpoint
    async def produce(self, port: Port[int]) -> None:
        rank = context().actor_instance.rank["gpus"]

        async def stream() -> None:
            for step in range(STREAM_LEN):
                port.send(rank * 10 + step)   # 就绪即发,不等别人

        self.task = asyncio.create_task(stream())

要点:

  • Channel[int].open() 返回一对 (port, receiver),port 可以作为参数发给任何 endpoint;
  • 各 rank 的 produce 立即返回,后台任务在自己准备好时才 port.send;
  • 协调者按预期数量 recv,天然形成扇入聚合。

配方 4:Actor 内获取自己的身份 ​

python
from monarch.actor import current_rank

class MyActor(Actor):
    def __init__(self):
        self.rank = current_rank().rank    # Point 对象,类 dict

    @endpoint
    def whoami(self):
        # context() 的 message_rank 是"这条消息"的坐标,
        # 与自身 rank 可能不同(消息可能发给切片)
        return {
            "my_rank": current_rank().rank,
            "msg_rank": context().message_rank,
        }

配方 5:初始化时互相配对的两组 Actor ​

python
class Trainer(Actor):
    def __init__(self, dataloader: DataLoader):
        # 每个 Trainer 用自己的 rank 坐标选中对应的 DataLoader
        my_rank = context().actor_instance.rank
        self.dataloader = dataloader.slice(**my_rank)

dataloaders = trainer_procs.spawn("dataloader", DataLoader)
dtrainers = trainer_procs.spawn("dtrainers", Trainer, dataloaders)

slice(**my_rank) 把类 dict 的坐标展开成关键字参数,形状相同的两个网格逐一配对。

配方 6:显式响应端口(乱序响应) ​

问题:消息按序到达,但有时想让后到的请求先响应——比如推理引擎永远用最新权重。

python
class InferenceEngine(Actor):
    def __init__(self):
        self.q = Queue()
        self.new_weights = None
        self.worker = threading.Thread(target=self.run)
        self.worker.start()

    @endpoint
    def update_weights(self, new_weights):
        self.new_weights = new_weights

    @endpoint(explicit_response_port=True)
    def infer(self, port: Port[int]):
        # 声明 explicit_response_port 后,端点拿到 port 自行决定何时回复
        self.q.put(port)

    def run(self):
        while True:
            request = self.q.get()
            result = f"inferred with: {self.new_weights}"  # 永远用最新权重
            time.sleep(1)  # 假装推理耗时
            request.send(result)   # 就绪时才发送响应

调用方无感知:engine.infer.call_one() 照常拿到 Future,只是响应顺序由 Actor 内部掌控。

小结 ​

场景配方
取 rank-0 / 指定 rankflatten("rank").slice(rank=N)
Actor 找同网格伙伴set_mesh 注入句柄 + 切片
异步扇入Channel.open() + port.send / receiver.recv
身份查询current_rank() / context()
网格配对slice(**my_rank)
乱序响应@endpoint(explicit_response_port=True)

更多配方持续增加中,官方仓库的 cookbook.md 是权威来源。

Monarch中文教程