远程引用协议¶
本文介绍了远程引用协议（Remote Reference protocol）的设计细节，并详细讲解了不同场景下的消息流。在继续之前，请确保您熟悉分布式 RPC 框架。
背景¶
RRef 是 Remote REFerence 的缩写。它是位于本地或远程工作器上的对象的引用，并在底层透明地处理引用计数。概念上，它可以被视为一个分布式共享指针。应用程序可以通过调用 remote() 创建 RRef。每个 RRef 都由 remote() 调用的被调用方工作器（即所有者，owner）拥有，并可供多个用户使用。所有者存储实际数据并跟踪全局引用计数。每个 RRef 都可以通过一个全局唯一的 RRefId 来唯一标识，该 ID 在 remote() 调用的调用方创建时分配。
在所有者工作器上，只有一个 OwnerRRef 实例，其中包含实际数据，而在用户工作器上，可以根据需要有任意数量的 UserRRef 实例，且 UserRRef 不持有数据。所有在所有者上的使用都将通过全局唯一的 RRefId 检索唯一的 OwnerRRef 实例。当 UserRRef 在 rpc_sync()、rpc_async() 或 remote() 调用中用作参数或返回值时，就会创建一个 UserRRef，并且所有者会收到通知以更新引用计数。当全局没有 UserRRef 实例且所有者上也没有对 OwnerRRef 的引用时，OwnerRRef 及其数据将被删除。
假设¶
RRef 协议基于以下假设设计。
瞬态网络故障：RRef 设计通过重试消息来处理瞬态网络故障。它无法处理节点崩溃或永久性网络分区。当这些事件发生时，应用程序应关闭所有工作器，回滚到先前的检查点，然后恢复训练。
非幂等用户自定义函数 (UDF)：我们假设提供给
rpc_sync()、rpc_async()或remote()的用户自定义函数 (UDF) 是非幂等的，因此不能重试。然而，内部 RRef 控制消息是幂等的，并在消息失败时重试。消息乱序交付：我们不假定任意一对节点之间的消息交付顺序，因为发送方和接收方都使用多线程。无法保证哪条消息将首先被处理。
RRef 生命周期¶
该协议的目标是在适当的时间删除 OwnerRRef。删除 OwnerRRef 的正确时机是当没有存活的 UserRRef 实例，并且用户代码也没有持有对 OwnerRRef 的引用时。棘手之处在于确定是否存在任何存活的 UserRRef 实例。
设计原理¶
用户可以通过三种情况获得 UserRRef
从所有者那里接收
UserRRef。从另一个用户那里接收
UserRRef。创建一个由另一个工作器拥有的新的
UserRRef。
情况 1 最简单，所有者将其 RRef 传递给用户，所有者调用 rpc_sync()、rpc_async() 或 remote() 并将其 RRef 用作参数。在这种情况下，将在用户工作器上创建一个新的 UserRRef。由于所有者是调用方，它可以轻松更新其在 OwnerRRef 上的本地引用计数。
唯一的要求是任何 UserRRef 在销毁时必须通知所有者。因此，我们需要第一个保证
G1. 当任何 UserRRef 被删除时，所有者将收到通知。
由于消息可能延迟或乱序到达，我们需要另一个保证来确保删除消息不会处理得太早。如果 A 发送一条涉及 RRef 的消息给 B，我们将 A 上的 RRef 称为父 RRef，将 B 上的 RRef 称为子 RRef。
G2. 父 RRef 在子 RRef 被所有者确认之前不会被删除。
在情况 2 和 3 中，所有者可能对 RRef 的分叉图（fork graph）仅有部分了解或完全不了解。例如，可以在用户工作器上构建一个 RRef，并且在所有者收到任何 RPC 调用之前，创建该 RRef 的用户可能已经将 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
如果 Z 的所有消息（包括删除消息）在所有者处理 Y 的消息之前被所有者处理。所有者会在知道 Y 存在之前得知 Z 已被删除。然而，这并不会导致任何问题。因为 Y 的至少一个祖先（A）将会存活，并且它会阻止所有者删除 OwnerRRef。更具体地说，如果所有者不知道 Y，A 由于 **G2** 不会被删除，而且所有者知道 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 GC (垃圾回收器) 同意删除本地
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 在后续对 B 的 RPC 调用中将其用作参数。A 会保持 UserRRef {100, 1} 存活，直到收到 B 的确认 (**G2**，而非 RPC 调用的返回值)。这是必要的，因为 A 不应在收到所有先前的消息之前发送删除消息，否则，由于我们不保证消息交付顺序，OwnerRRef 可能会在使用前被删除。这是通过创建 RRef 的子 ForkId，并将它们保存在一个映射中直到收到所有者确认该子 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 发送一个分叉请求（fork request）。稍后，当 B 确认 C 上的 UserRRef 后，C 将并行执行两个操作：1) 向 A 发送子 ACK，以及 2) 运行用户提供的函数。在此期间，父节点（A）将保持其 UserRRef {100, 1} 存活以实现 **G2**。
