原文： https://pytorch.org/docs/stable/notes/rref.html

警告

RRef API 是實驗性的，隨時可能更改。

本說明描述了遠程引用協(xié)議的設(shè)計細節(jié)，并逐步介紹了不同情況下的消息流。 在繼續(xù)之前，請確保您熟悉分布式 RPC 框架。

背景

RRef 代表遠程參考。 它是位于本地或遠程工作人員上的對象的引用，并且透明地在內(nèi)部進行引用計數(shù)。 從概念上講，它可以視為分布式共享指針。 應(yīng)用程序可以通過調(diào)用 remote() 創(chuàng)建 RRef。 每個 RRef 都由 remote() 呼叫的被調(diào)用方工作者(即所有者）擁有，并且可以由多個用戶使用。 所有者存儲實際數(shù)據(jù)，并跟蹤全局參考計數(shù)。 每個 RRef 可以由全局RRefId唯一標識，該全局RRefId在創(chuàng)建時在 remote() 調(diào)用的調(diào)用方上分配。

在所有者工作程序中，只有一個OwnerRRef實例包含真實數(shù)據(jù)，而在用戶工作程序中，可以根據(jù)需要包含任意數(shù)量的UserRRefs，而UserRRef不保存數(shù)據(jù)。 所有者上的所有用法都將使用全局唯一的RRefId來檢索唯一的OwnerRRef實例。 在 rpc_sync() ， rpc_async() 或  remote() 調(diào)用中將UserRRef用作參數(shù)或返回值時，將創(chuàng)建該UserRRef 將會根據(jù)更新的參考計數(shù)通知所有者。 如果全局沒有UserRRef實例，并且所有者上也沒有對OwnerRRef的引用，則OwnerRRef及其數(shù)據(jù)將被刪除。

假設(shè)條件

RRef 協(xié)議的設(shè)計基于以下假設(shè)。

• 瞬態(tài)網(wǎng)絡(luò)故障：RRef 設(shè)計旨在通過重試消息來處理瞬態(tài)網(wǎng)絡(luò)故障。 節(jié)點崩潰或永久性網(wǎng)絡(luò)分區(qū)超出了范圍。 當(dāng)這些事件發(fā)生時，該應(yīng)用程序可能會關(guān)閉所有工作人員，還原到先前的檢查點，然后恢復(fù)訓(xùn)練。
• 非冪等 UDF ：我們假設(shè)提供給 rpc_sync() ， rpc_async() 或  remote() 的用戶功能(UDF） 不是冪等的，因此無法重試。 但是，內(nèi)部 RRef 控制消息將成為冪等且可重試。
• 消息傳遞無序：我們不假定任何一對節(jié)點之間的消息傳遞順序，因為發(fā)送者和接收者都使用多個線程。 無法保證首先處理哪個消息。

RRef 生命周期

該協(xié)議的目標是在適當(dāng)?shù)臅r候刪除OwnerRRef。 刪除OwnerRRef的正確時機是在沒有活動UserRRef實例且用戶代碼也沒有保存對OwnerRRef的引用的情況下。 棘手的部分是確定是否存在任何活動的UserRRef實例。

設(shè)計推理

用戶可以在以下三種情況下獲得UserRRef：

1. 從所有者那里收到了UserRRef。
2. 接收到另一個用戶的UserRRef。
3. 創(chuàng)建另一個工人擁有的新UserRRef。

情況 1 是最簡單的，所有者將其 RRef 傳遞給用戶，所有者調(diào)用 rpc_sync() ， rpc_async() 或  remote() 使用其 RRef 作為參數(shù)。 在這種情況下，將在用戶上創(chuàng)建一個新的UserRRef。 由于所有者是調(diào)用者，因此可以輕松地在OwnerRRef上更新其本地引用計數(shù)。

唯一的要求是任何UserRRef必須在銷毀時通知所有者。 因此，我們需要第一個保證：

G1。 刪除任何“ UserRRef”時，都會通知所有者。

由于郵件可能會延遲或出現(xiàn)亂序，因此我們還需要一項保證，以確保刪除郵件不會過早處理。 如果 A 向 B 發(fā)送涉及 RRef 的消息，我們將 A 上的 RRef 稱為父 RRef，將 B 上的 RRef 稱為子 RRef。

G2。 在所有者確認子 RRef 之前，不會刪除父 RRef。

在情況 2 和 3 中，所有者可能僅對 RRef 分支圖有部分了解或根本不了解。 例如，可以在用戶上構(gòu)建 RRef，并且在所有者收到任何 RPC 調(diào)用之前，創(chuàng)建者用戶可能已經(jīng)與其他用戶共享了 RRef，并且這些用戶可以進一步共享 RRef。 一個不變性是，任何 RRef 的派生圖始終都是一棵樹，因為派生 RRef 總是在被調(diào)用方上創(chuàng)建一個新的UserRRef實例(除非被調(diào)用方是所有者），因此每個 RRef 都有一個父級。

所有者對樹中任何UserRRef的視圖分為三個階段：

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

所有者對整棵樹的看法不斷變化。 擁有者認為沒有活動的UserRRef實例時，即刪除OwnerRRef實例時，所有UserRRef實例都可能確實被刪除或未知，因此所有者刪除了其OwnerRRef實例。 危險的情況是某些分叉未知，而另一些被刪除。

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 的父母。

如果在用戶上創(chuàng)建 RRef，事情會變得有些棘手：

OwnerRRef
    ^
    |
    A -> Y -> Z

如果 Z 在UserRRef上調(diào)用 to_here() ，則所有者至少知道刪除 Z 時的 A，否則， to_here() 不會結(jié)束。 如果 Z 沒有調(diào)用 to_here() ，則所有者可能在從 A 和 Y 發(fā)送任何消息之前就已從 Z 接收了所有消息。在這種情況下，由于尚未獲得OwnerRRef的真實數(shù)據(jù) 創(chuàng)建后，也沒有要刪除的內(nèi)容。 就像 Z 根本不存在一樣。 因此，仍然可以。

實作

G1 通過在UserRRef析構(gòu)函數(shù)中發(fā)送刪除消息來實現(xiàn)。 為了提供 G2 ，無論何時將父級UserRRef派生，都將其置于上下文中，并由新的ForkId對其進行索引。 僅當(dāng)父級UserRRef從子級收到確認消息(ACK）時，才會從上下文中刪除該父級UserRRef，并且只有當(dāng)所有者確認后，該子級才會發(fā)出 ACK。

協(xié)議方案

現(xiàn)在，讓我們討論以上設(shè)計如何在四種情況下轉(zhuǎn)換為協(xié)議。

用戶與所有者共享 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()

在這種情況下，在用戶工作程序 A 上創(chuàng)建UserRRef，然后將其與遠程消息一起傳遞給所有者工作程序 B，然后 B 創(chuàng)建OwnerRRef。 方法 remote() 立即返回，這意味著UserRRef可以在所有者了解之前被分叉/使用。

在所有者上，當(dāng)接收到 remote() 調(diào)用時，它將創(chuàng)建OwnerRRef，并返回一個 ACK 來確認{100, 1}(RRefId，ForkId）。 僅在收到此 ACK 后，A 才能刪除其UserRRef。 這涉及 G1 和 G2 。 G1 很明顯。 對于 G2 而言，OwnerRRef是UserRRef的子級，并且UserRRef直到收到所有者的 ACK 才被刪除。

上圖顯示了消息流，其中實心箭頭包含用戶功能，而虛線箭頭是內(nèi)置消息。 請注意，從 A 到 B 的前兩個消息 (remote() 和 to_here())可以按任何順序到達 B，但最終的刪除消息僅在以下情況下發(fā)送： ：

• B 確認UserRRef {100, 1}(G2），并且
• Python GC 同意刪除本地UserRRef實例。 當(dāng) RRef 不再在范圍內(nèi)并且可以進行垃圾回收時，就會發(fā)生這種情況。

用戶與所有者共享 RRef 作為參數(shù)

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 上創(chuàng)建UserRRef后，A 會將其用作對 B 的后續(xù) RPC 調(diào)用中的參數(shù)。A 將使UserRRef {100, 1}保持活動狀態(tài)，直到收到 B 的確認 (G2 ， 而不是 RPC 調(diào)用的返回值）。 這是必要的，因為在接收到所有先前的消息之前，A 不應(yīng)發(fā)出刪除消息，否則，OwnerRRef可以在使用前刪除，因為我們不能保證消息的傳遞順序。 這是通過創(chuàng)建 RRef 的子項ForkId并將其保存在地圖中，直到收到所有者確認該子項ForkId來完成的。 下圖顯示了消息流。

注意，UserRRef可以在功能完成甚至啟動之前在 B 上刪除。 但是，這是可以的，因為在 B 向子ForkId發(fā)送 ACK 時，它已經(jīng)獲取了OwnerRRef實例，這將防止它被過早刪除。

所有者與用戶共享 RRef

所有者對用戶是最簡單的情況，所有者可以在本地更新引用計數(shù)，并且不需要任何其他控制消息即可通知其他人。 關(guān)于 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, ))

上圖顯示了消息流。 請注意，當(dāng)在 rpc_async 調(diào)用之后OwnerRRef退出作用域時，將不會刪除它，因為如果存在任何已知的派生，則內(nèi)部會存在一個使它保持活動狀態(tài)的映射，在這種情況下為UserRRef {100, 1}。 (G2 )

用戶與用戶共享 RRef

這是最復(fù)雜的情??況，其中調(diào)用者用戶(父級UserRRef），被調(diào)用者用戶(子級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, ))

當(dāng) C 從 A 接收到子項UserRRef時，它向所有者 B 發(fā)送一個派生請求。稍后，當(dāng) B 確認 C 上的UserRRef時，C 將并行執(zhí)行兩個操作：1）發(fā)送子項 ACK 到 A，然后 2）運行用戶提供的功能。 在這段時間中，親本(A）將保持其UserRRef {100, 1}存活以實現(xiàn) G2 。