远程引用协议 ¶

本笔记描述了远程引用协议的设计细节，并遍历了不同场景下的消息流程。在继续之前，请确保您熟悉分布式 RPC 框架。

背景 ¶

RRef 代表远程引用。它是指向位于本地或远程工作者的对象的引用，在底层透明地处理引用计数。从概念上讲，它可以被视为一个分布式共享指针。应用程序可以通过调用 remote() 来创建 RRef。每个 RRef 都由 remote() 调用的被调用者工作者（即所有者）拥有，并可被多个用户使用。所有者存储真实数据并跟踪全局引用计数。每个 RRef 都可以通过全局 RRefId 唯一标识，该标识在 remote() 调用的调用者处创建时分配。

在所有者工作者上，只有一个 OwnerRRef 实例，其中包含真实数据，而在用户工作者上，可以有必要的 UserRRefs 个， UserRRef 不持有数据。所有对所有者的使用都将通过全局唯一的 OwnerRRef 检索唯一的 RRefId 实例。当它作为 UserRRef 、 rpc_sync() 或 rpc_async() 调用中的参数或返回值时，将创建一个 remote() ，所有者将根据更新引用计数进行通知。当没有全局 OwnerRRef 实例，且所有者上也没有对 UserRRef 的引用时，将删除 OwnerRRef 及其数据。

假设 §

RRef 协议设计基于以下假设。

• 临时网络故障：RRef 设计通过重试消息来处理临时网络故障。它无法处理节点崩溃或永久网络分区。当这些事件发生时，应用程序应关闭所有工作进程，回滚到上一个检查点，并继续训练。

• 非幂等 UDF：我们假设提供给 rpc_sync() 、 rpc_async() 或 remote() 的用户函数（UDF）不是幂等的，因此不能重试。然而，内部 RRef 控制消息是幂等的，在消息失败时会被重试。

• 消息顺序错误交付：我们不假设任何一对节点之间的消息交付顺序，因为发送方和接收方都在使用多个线程。无法保证哪个消息将被首先处理。

RRef 生命周期 ¶

协议的目标是在适当的时候删除一个 OwnerRRef 。删除一个 OwnerRRef 的正确时机是在没有活跃的 UserRRef 实例，并且用户代码也没有持有 OwnerRRef 的引用时。棘手的部分是确定是否存在任何活跃的 UserRRef 实例。

设计理由 ¶

用户可以在三种情况下获得一个 UserRRef ：

1. 接收来自所有者的 UserRRef 。

2. 接收来自另一个用户的 UserRRef 。

3. 创建一个由另一个工人拥有的新 UserRRef 。

情况 1 是最简单的，其中所有者将其 RRef 传递给用户，所有者调用 rpc_sync() 、 rpc_async() 或 remote() 并使用其 RRef 作为参数。在这种情况下，将在用户处创建一个新的 UserRRef 。由于所有者是调用者，它可以轻松更新其本地引用计数在 OwnerRRef 上。

唯一的要求是任何 UserRRef 在销毁时必须通知所有者。因此，我们需要第一个保证：

G1. 当任何 UserRRef 被删除时，所有者将被通知。

由于消息可能会延迟或顺序混乱，我们需要另一个保证以确保删除消息不会被过早处理。如果 A 向 B 发送涉及 RRef 的消息，我们将在 A（父 RRef）和 B（子 RRef）上调用 RRef。

G2. 父 RRef 将在子 RRef 被所有者确认之前不会被删除。

在案例 2 和 3 中，所有者可能只有部分或完全没有关于 RRef 分叉图的了解。例如，一个 RRef 可能是在用户上构建的，在所有者收到任何 RPC 调用之前，创建用户可能已经将 RRef 与其他用户共享，而这些用户可能进一步将 RRef 共享出去。一个不变量是任何 RRef 的分叉图始终是一棵树，因为分叉 RRef 总是在被调用者上创建一个新的 UserRRef 实例（除非被调用者是所有者），因此每个 RRef 都有一个父节点。

所有者对树中任何 UserRRef 的看法有三个阶段：

1) unknown -> 2) known -> 3) deleted.

所有者对整个树的观点不断变化。当所有者认为没有活跃的 UserRRef 实例时，它会删除其 OwnerRRef 实例，即当 OwnerRRef 被删除时，所有 UserRRef 实例可能确实被删除或未知。危险的情况是当一些分叉未知而其他分叉被删除时。

G2 可以简单地保证在所有者知道其所有子 UserRRef 实例之前，不能删除任何父 UserRRef 。然而，可能发生的情况是，在所有者知道其父 UserRRef 之前，子 UserRRef 可能已经被删除。

考虑以下示例，其中 OwnerRRef 分支到 A，然后 A 分支到 Y，Y 再分支到 Z：

OwnerRRef -> A -> Y -> Z

如果在 Y 的消息之前，包括删除消息，Z 的所有消息都由所有者处理，那么所有者将在知道 Y 存在之前了解到 Z 的删除。然而，这并不会引起任何问题。因为至少 Y 的一个祖先（A）是活着的，并且它会阻止所有者删除 OwnerRRef 。更具体地说，如果所有者不知道 Y，由于 G2，A 不能被删除，而所有者知道 A，因为它是 A 的父节点。

如果 RRef 是在用户上创建的，事情会变得有些复杂：

OwnerRRef
    ^
    |
    A -> Y -> Z

如果 Z 在 UserRRef 上调用 to_here() ，那么当 Z 被删除时，所有者至少知道 A，因为否则 to_here() 就无法完成。如果 Z 没有调用 to_here() ，所有者可能会在收到任何来自 A 和 Y 的消息之前收到 Z 的所有消息。在这种情况下，由于 OwnerRRef 的实际数据尚未创建，也没有什么可以删除的。这就像 Z 根本不存在一样。因此，仍然没问题。

实现 ¶

G1 通过在 UserRRef 析构函数中发送删除消息来实现。为了提供 G2，每当父进程 UserRRef 被克隆时，都会将其放入上下文中，并按新的 ForkId 进行索引。父进程 UserRRef 只有在从子进程接收到确认消息（ACK）时才会从上下文中移除，而子进程只有在被所有者确认后才会发送 ACK。

协议场景 ¶

现在我们来讨论上述设计如何在四个场景中转化为协议。

用户将 RRef 与所有者共享作为返回值 ¶

import torch
import torch.distributed.rpc as rpc

# on worker A
rref = rpc.remote('B', torch.add, args=(torch.ones(2), 1))
# say the rref has RRefId 100 and ForkId 1
rref.to_here()

在这种情况下， UserRRef 是在用户工作者 A 上创建的，然后它与远程消息一起传递给所有者工作者 B，然后 B 创建 OwnerRRef 。方法 remote() 立即返回，这意味着在所有者知道之前， UserRRef 可以被分叉/使用。

在所有者这边，当接收到 remote() 调用，它会创建 OwnerRRef ，并返回一个 ACK 来确认 {100, 1} （ RRefId ， ForkId ）。只有在接收到这个 ACK 之后，A 才能删除它的 UserRRef 。这涉及到 G1 和 G2。G1 很明显。对于 G2， OwnerRRef 是 UserRRef 的子节点， UserRRef 不会删除，直到它从所有者那里收到 ACK。

上面的图显示了消息流，其中实线箭头包含用户函数，虚线箭头是内置消息。请注意，A 发送到 B 的前两条消息（ remote() 和 to_here() ）可以以任何顺序到达 B，但最终的删除消息只有在以下情况下才会发送：

• B 确认 UserRRef {100, 1} （G2），并且

• Python 垃圾回收器同意删除局部 UserRRef 实例。这发生在 RRef 不再在作用域内且符合垃圾回收条件时。

用户将 RRef 与所有者作为参数共享

import torch
import torch.distributed.rpc as rpc

# on worker A and worker B
def func(rref):
  pass

# on worker A
rref = rpc.remote('B', torch.add, args=(torch.ones(2), 1))
# say the rref has RRefId 100 and ForkId 1
rpc.rpc_async('B', func, args=(rref, ))

在这种情况下，在 A 上创建 UserRRef 之后，A 将其用作后续 RPC 调用 B 的参数。A 将保持 UserRRef {100, 1} 存活，直到从 B 收到确认（G2，不是 RPC 调用的返回值）。这是必要的，因为 A 应该在收到所有之前的消息之后再发送删除消息，否则， OwnerRRef 可能在使用前被删除，因为我们不能保证消息的传输顺序。这是通过创建 RRef 的子 ForkId ，将它们保存在映射中，直到收到所有者的确认来实现的。下面的图示显示了消息流。

注意， UserRRef 可能在 func 完成或开始之前在 B 上被删除。然而这没关系，因为当 B 发送子 ForkId 的 ACK 时，它已经获取了 OwnerRRef 实例，这将防止它过早地被删除。

与用户共享所有者 RRef ¶

所有者到用户是最简单的情况，所有者可以在本地更新引用计数，无需任何额外的控制消息来通知他人。至于 G2，与父节点立即从所有者那里收到 ACK 相同，因为父节点就是所有者。

import torch
import torch.distributed.rpc as RRef, rpc

# on worker B and worker C
def func(rref):
  pass

# on worker B, creating a local RRef
rref = RRef("data")
# say the rref has RRefId 100
dist.rpc_async('C', func, args=(rref, ))

上图显示了消息流。请注意，当 OwnerRRef 在 rpc_async 调用后退出作用域时，它不会被删除，因为内部有一个映射来保持其活跃状态，如果存在已知的分支，则情况为 UserRRef {100, 1} 。（G2）

用户与用户共享 RRef ¶

这是最复杂的情况，调用者用户（父用户 UserRRef ）、被调用者用户（子用户 UserRRef ）和所有者都需要参与。

import torch
import torch.distributed.rpc as rpc

# on worker A and worker C
def func(rref):
  pass

# on worker A
rref = rpc.remote('B', torch.add, args=(torch.ones(2), 1))
# say the rref has RRefId 100 and ForkId 1
rpc.rpc_async('C', func, args=(rref, ))

当 C 从 A 接收子用户 UserRRef 后，它会向所有者 B 发送一个分叉请求。稍后，当 B 确认 C 的 UserRRef 后，C 将并行执行两个动作：1）向 A 发送子用户确认，2）运行用户提供的函数。在此期间，父用户（A）将保持其 UserRRef {100, 1} 活跃以实现 G2。