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 / 指定 rank | flatten("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 是权威来源。
