Протокол удалённых ссылок

Обоснование проектирования

Пользователь может получить UserRRef в трёх случаях:

1. Получение UserRRef от владельца.
2. Получение UserRRef от другого пользователя.
3. Создание нового UserRRef, принадлежащего другому рабочему узлу.

Случай 1 — самый простой: владелец передаёт свой RRef пользователю, вызывая rpc_sync(), rpc_async() или remote() и используя свой RRef в качестве аргумента. В этом случае на стороне пользователя создаётся новый UserRRef. Поскольку владелец является вызывающей стороной, он может легко обновить локальный счётчик ссылок на OwnerRRef.

Единственное требование — любой UserRRef должен уведомить владельца при удалении. Поэтому необходимо следующее первое условие:

G1. Владелец будет уведомлён об удалении любого UserRRef.

Поскольку сообщения могут задерживаться или приходить не по порядку, необходимо ещё одно условие, гарантирующее, что сообщение об удалении не будет обработано преждевременно. Если A отправляет B сообщение, связанное с RRef, RRef на стороне A (родительский RRef) и RRef на стороне B (дочерний RRef) называются соответственно родительским и дочерним.

G2. Родительский RRef НЕ будет удалён, пока дочерний RRef не будет подтверждён владельцем.

В случаях 2 и 3 владелец может иметь лишь частичную информацию о графе ветвления RRef или не иметь её вовсе. Например, RRef может быть создан на пользовательском узле, и до того, как владелец получит какой-либо RPC-вызов, создавший его пользователь может уже поделиться RRef с другими пользователями, а те, в свою очередь, могут передать его дальше. Один из инвариантов состоит в том, что граф ветвления любого RRef всегда является деревом: при ветвлении RRef на вызываемой стороне всегда создаётся новый экземпляр UserRRef (за исключением случая, когда вызываемая сторона — владелец), поэтому у каждого RRef только один родитель.

Представление владельца о любом UserRRef в дереве проходит три этапа:

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

Представление владельца обо всём дереве постоянно меняется. Владелец удаляет свой экземпляр OwnerRRef, когда считает, что действующих экземпляров UserRRef больше нет; то есть после удаления OwnerRRef все экземпляры UserRRef либо действительно удалены, либо неизвестны владельцу. Опасный случай возникает, когда некоторые ветвления неизвестны, а другие уже удалены.

G2 очевидным образом гарантирует, что родительский UserRRef не может быть удалён до того, как владелец узнает обо всех его дочерних экземплярах UserRRef. Однако дочерний UserRRef может быть удалён до того, как владелец узнает о его родительском UserRRef.

Рассмотрим следующий пример: OwnerRRef ветвится на A, затем A ветвится на Y, а Y — на Z:

OwnerRRef -> A -> Y -> Z

Если все сообщения от Z, включая сообщение об удалении, будут обработаны владельцем раньше сообщений от Y, владелец узнает об удалении Z, прежде чем узнает о существовании Y. Тем не менее это не вызовет проблем. По крайней мере один из предков Y (A) останется активным и не позволит владельцу удалить OwnerRRef. Точнее, если владельцу неизвестен Y, A не может быть удалён согласно G2, а владелец знает об A, поскольку является его родителем.

Ситуация становится немного сложнее, если RRef создаётся на пользовательском узле:

OwnerRRef
    ^
    |
    A -> Y -> Z

Если Z вызывает to_here() для UserRRef, владелец как минимум знает об A к моменту удаления Z, поскольку иначе to_here() не завершился бы. Если Z не вызывает to_here(), владелец может получить все сообщения от Z раньше любого сообщения от A и Y. В этом случае реальные данные 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, ))

В этом случае после создания UserRRef на A узел A использует его в качестве аргумента последующего RPC-вызова узлу B. Узел A будет поддерживать UserRRef {100, 1} активным до получения подтверждения от B (G2, а не возвращаемого значения RPC-вызова). Это необходимо, поскольку A не должен отправлять сообщение об удалении до получения всех предыдущих сообщений; иначе OwnerRRef может быть удалён до использования, поскольку порядок доставки сообщений не гарантируется. Для этого создаётся дочерний ForkId от RRef, который хранится в map до подтверждения владельцем дочернего ForkId. На рисунке ниже показан поток сообщений.

Обратите внимание, что UserRRef может быть удалён на B до завершения или даже начала func. Это допустимо: к моменту отправки B ACK для дочернего ForkId он уже получил экземпляр OwnerRRef, который не позволит удалить его преждевременно.

Владелец передаёт RRef пользователю

Передача от владельца пользователю — самый простой случай: владелец может обновить счётчик ссылок локально, и ему не нужны дополнительные управляющие сообщения для уведомления других. Что касается G2, это равносильно немедленному получению ACK от владельца родительским RRef, поскольку родителем является сам владелец.

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, он не удаляется: внутри существует map, которая поддерживает его активность, если известны какие-либо ветвления, в данном случае — 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, ))

Получив от A дочерний UserRRef, C отправляет владельцу B запрос на ветвление. Позже, когда B подтверждает UserRRef на C, C параллельно выполняет два действия: 1) отправляет A дочерний ACK и 2) запускает предоставленную пользователем функцию. В это время родительский RRef (A) будет поддерживать свой UserRRef {100, 1} активным, обеспечивая G2.