快捷方式

远程引用协议

本说明介绍了远程引用协议的设计细节,并逐步介绍了不同场景下的消息流。在继续之前,请确保您熟悉分布式 RPC 框架

背景

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

在所有者工作器上,只有一个 OwnerRRef 实例,其中包含实际数据,而在用户工作器上,可以有尽可能多的 UserRRefs,并且 UserRRef 不保存数据。所有者上的所有使用都将使用全局唯一的 RRefId 检索唯一的 OwnerRRef 实例。当 rpc_sync()rpc_async()remote() 调用中使用 RRef 作为参数或返回值时,将创建一个 UserRRef,并且所有者将根据需要更新引用计数而收到通知。当全局没有 UserRRef 实例并且所有者也没有对 OwnerRRef 的引用时,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(父 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

如果 Z 的所有消息(包括删除消息)都在 Y 的消息之前由所有者处理,那么所有者将在知道 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 进行索引。只有在收到来自子节点的确认消息(ACK)时,父 UserRRef 才会从上下文中删除,而子节点只有在被所有者确认后才会发送 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 创建 OwnerRRefremote() 立即返回,这意味着 UserRRef 可以分叉/使用,而所有者对此一无所知。

在所有者上,当收到 remote() 调用时,它将创建 OwnerRRef 并返回一个 ACK 来确认 {100, 1} (RRefId, ForkId)。只有在收到此 ACK 后,A 才能删除它的 UserRRef。这涉及 G1G2G1 是显而易见的。对于 G2OwnerRRefUserRRef 的子节点,并且 UserRRef 在收到来自所有者的 ACK 之前不会被删除。

user_to_owner_ret.png

上面的图表显示了消息流,其中实线箭头包含用户函数,虚线箭头是内置消息。请注意,来自 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 在后续的 RPC 调用到 B 时使用它作为参数。A 将保持 UserRRef {100, 1} 存活,直到它收到来自 B 的确认(G2,而不是 RPC 调用的返回值)。这是必要的,因为 A 不应该在收到所有先前消息之前发送删除消息,否则,OwnerRRef 可能在使用之前被删除,因为我们不保证消息传递顺序。这是通过创建 RRef 的子 ForkId、将它们保存在一个映射中,直到收到所有者确认子 ForkId 来完成的。下图显示了消息流。

user_to_owner_arg.png

请注意,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, ))
owner_to_user.png

上图显示了消息流。请注意,当 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, ))
user_to_user.png

当 C 从 A 接收子 UserRRef 时,它向所有者 B 发送一个分叉请求。稍后,当 B 确认 C 上的 UserRRef 时,C 将并行执行两个操作:1) 向 A 发送子 ACK,以及 2) 运行用户提供的函数。在此期间,父节点 (A) 将保持其 UserRRef {100, 1} 存活以实现 G2

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源