TensorFlow中的通訊機制——Rendezvous（一）本地傳輸

時間  2019-11-17
標籤 tensorflow 通訊 機制 rendezvous 本地 傳輸 简体版
原文   原文鏈接

背景

[做者：DeepLearningStack，阿里巴巴算法工程師，開源TensorFlow Contributor]node

在TensorFlow源碼中咱們常常能看到一個奇怪的詞——Rendezvous。若是從仔細統計該單詞出現的頻率和模塊，你會發現不管在單機仍是分佈式，不管在core目錄仍是contrib目錄都存在它的身影，所涉及的模塊很是多。Rendezvous是一個法語單詞，發音也比較特殊，通常直譯爲「約會、相會、會和」，而在TensorFlow中，Rendezvous是用來完成消息傳輸的通訊組件。大部分源碼讀者在起初閱讀時通訊部分的代碼時可能會以爲有點懵圈，爲何不使用Communicator這樣簡單明瞭的單詞來代表通訊過程，反而使用這樣一個晦澀的法語詞做爲抽象呢？其實在瞭解TensorFlow消息通訊的原理後就會發現，使用Rendezvous做爲這一過程的抽象是很是貼切的。算法

由於Rendezvous所涉及的模塊組件較多，爲了讓讀者按部就班地理解TensorFlow中的通訊機制，決定將Rendezvous分紅多個系列，由淺入深分開梳理。這樣作的目的不但能讓讀者閱讀時對總體層次結構有較好的把握，並且簡短的篇幅也便於閱讀，因此建議讀者按順序閱讀本系列。 本文是TensorFlow通訊機制系列的第一篇文章，側重總體結構和本地傳輸通訊的梳理。緩存

消息傳輸的惟一標識符——ParsedKey

在TensorFlow中不管是單機仍是分佈式都涉及到消息傳輸，而且消息傳輸老是從發送端Send，接收端Recv。那麼這裏就存在一個消息的對應問題：在多組消息同時發送接收時，須要對每一對Send和Recv梳理一個對應關係，即Send端發送的消息與Recv端接收的消息不能有錯位。若是Recv端本打算接收的消息是A，但因爲消息對應錯誤致使接收到了B，那麼整個訓練過程就會出現錯誤。其實解決這個問題也很是簡單，由於每一對Send和Recv所處理消息都是同一個，因此只要讓某個消息在被Send前加上一個惟一標識符，而Recv在接收消息前也可以按照某種規則拼出同樣的惟一標識符，這個對應關係就完美解決了。在TensorFlow中確實定義了這樣一種標識符，它就是結構體ParsedKey。安全

ParsedKey結構體

在tensorflow/core/framework/rendezvous.h的Rendezvous類內定義告終構體ParsedKey，它內容很是簡短卻又十分全面，不但包含了消息傳輸的全部必須的內容，還具有惟一性，在咱們直接分析其源代碼結構。併發

 1 // Parses the key constructed by CreateKey and parse src/dst device
 2 // names into structures respectively.
 3 struct ParsedKey {
 4   StringPiece src_device;
 5   DeviceNameUtils::ParsedName src;
 6   uint64 src_incarnation = 0;
 7   StringPiece dst_device;
 8   DeviceNameUtils::ParsedName dst;
 9   StringPiece edge_name;
10 
11   ParsedKey() {}
12   ParsedKey(const ParsedKey& b) { *this = b; }
13 
14   ParsedKey& operator=(const ParsedKey& b);
15   StringPiece FullKey() const { return buf_; }
16 
17  private:
18   friend class Rendezvous;
19   friend class SendOp;
20   friend class RecvOp;
21   string buf_;
22 };

能夠看到其結構很是簡單，一個完備的ParsedKey要包括六個部分。app

src_device：消息發送源的字符串信息，形如/job:localhost/replica:0/task_id:0/device:GPU:0框架

src：和src_device的信息量相同，只不過是結構體的表示方法異步

src_incarnation：通常來講這個字段沒有什麼做用，可是當某個worker重啓後，該值會發生變化，用來和以前掛掉的worker作區分，這便於debugasync

dst_device：消息發送的接收方字符串信息，格式和src_device相同分佈式

dst：和dst_device的信息量相同，只不過是結構體的表示方法

edge_name：這個字段是該Key最特殊的地方，它能夠靈活指定爲任何字符串，實現不一樣Key的區分。好比它能夠是Tensor的名字，也能夠是具備某種特殊意義的固定字符串

CreateKey過程與ParseKey過程

通常狀況下，在TensorFlow中應該優先使用CreateKey函數來構造能夠解析的Key字符串，而後通過ParseKey過程將該字符串的每一個信息解析到ParsedKey結構體中，之因此使用CreateKey函數構造Key字符串是由於這是最安全保險的方式，下面是CreateKey函數構造Key字符串的過程展示。

CreateKey只要接受五個參數便可安全構造字符串形式的Key，這裏面特殊之處有兩個，a. 參數中frame_and_iter通常直接取自OpKernelContext中的FrameAndIter對象；b. src_incarnation要作一個十六進制的字符串轉換。CreateKey函數的輸出是以分號（";"）爲分隔符的字符串，該字符串一樣包含五個域。CreateKey是一個static函數，代碼比較簡單，就不在這裏列出。隨後咱們這個字符串傳入ParseKey函數便可完成結構體ParsedKey的解析，解析過程以下。

ParseKey對輸入字符串的前四個域作了映射，拋棄了第五個域，可是在提供Key字符串時須要提供完整的五個域，不然會檢查報錯。和CreateKey相同，ParseKey過程也是一個static函數，代碼以下所示。

 1 /* static */
 2 Status Rendezvous::ParseKey(StringPiece key, ParsedKey* out) {
 3   if (key.data() == out->buf_.data()) {
 4     // Caller used our buf_ string directly, so we don't need to copy.  (The
 5     // SendOp and RecvOp implementations do this, for example).
 6     DCHECK_EQ(key.size(), out->buf_.size());
 7   } else {
 8     // Make a copy that our StringPieces can point at a copy that will persist
 9     // for the lifetime of the ParsedKey object.
10     out->buf_.assign(key.data(), key.size());
11   }
12   StringPiece s(out->buf_);
13   StringPiece parts[5];
14   for (int i = 0; i < 5; i++) {
15     parts[i] = ConsumeNextPart(&s, ';');
16   }
17   if (s.empty() &&          // Consumed the whole string
18       !parts[4].empty() &&  // Exactly five parts
19       DeviceNameUtils::ParseFullName(parts[0], &out->src) &&
20       strings::HexStringToUint64(parts[1], &out->src_incarnation) &&
21       DeviceNameUtils::ParseFullName(parts[2], &out->dst) &&
22       !parts[3].empty()) {
23     out->src_device = StringPiece(parts[0].data(), parts[0].size());
24     out->dst_device = StringPiece(parts[2].data(), parts[2].size());
25     out->edge_name = StringPiece(parts[3].data(), parts[3].size());
26     return Status::OK();
27   }
28   return errors::InvalidArgument("Invalid  rendezvous key: ", key);
29 }

Rendezvous

在瞭解ParsedKey以後，咱們就能夠窺探Rendezvous這個類的內部結構和實現了。最基本的Rendezvous類被定義在了tensorflow/core/framework/rendezvous.h文件中，它對外提供了最基本的Send、Recv和RecvAsync接口和實現。整體來講這個類仍是比較抽象的，在不一樣的通訊場景下須要提供不一樣的實現。好比對於本地傳輸來講，TensorFlow提供了LocalRendezvous和IntraProcessRendezvous實現類，對於使用跨進程通訊場景來講，TensorFlow提供了RemouteRendezvous實現系列。不一樣通訊場景的實現細節差異至關大，因此本系列將對這些作逐個梳理，本文只關注本地傳輸部分。若是對跨進程傳輸感興趣，那麼請關注該系列的下一篇文章。Rendezvous類中最重要的函數是Send和Recv系列，它們的簽名和註釋以下代碼所示。

 1 // The caller is a tensor producer and it sends a message (a tensor
 2 // "val" and a bool "is_dead") under the given "key".
 3 //
 4 // {val, is_dead} is bundled as a message sent and received.
 5 // Typically, is_dead is set by some control flow nodes
 6 // (e.g., a not-taken branch).  args is passed by Send to the
 7 // Recv function to communicate any information that the Recv
 8 // function might need.  This is typically only necessary for
 9 // Send/Recv on the same worker.
10 //
11 // Send() never blocks.
12 virtual Status Send(const ParsedKey& key, const Args& args, const Tensor& val, const bool is_dead) = 0;
13 
14 virtual void RecvAsync(const ParsedKey& key, const Args& args, DoneCallback done) = 0;
15 
16 // Synchronous wrapper for RecvAsync.
17 Status Recv(const ParsedKey& key, const Args& args, Tensor* val, bool* is_dead, int64 timeout_ms);
18 Status Recv(const ParsedKey& key, const Args& args, Tensor* val, bool* is_dead);

TensorFlow中的Recv有兩種，一種是同步版本，換一種是異步版本。一般狀況下爲了計算和通訊的overlap，TensorFlow普遍使用了RecvAsync函數。而且在後面一節中咱們能夠知道，Send過程並非真的參與數據通訊，全部的通訊過程均由RecvAsync完成。

Rendezvous相關類結構

在瞭解通訊過程以前，應該先熟悉下Rendezvous相關的類結構。下面的類圖展現了當期TensorFlow系統中全部的Rendezvous相關類圖結構。

全部的Rendezvous相關類都以Rendezvous基類爲核心，LocalRendezvous和IntraProcessRendezvous是咱們本文分析的重點，SimpleRendezvous實現很是簡單，讀者能夠在熟悉前兩個實現以後自行分析該類。而BaseRemoteRendezvous類以及相關類是跨進程通訊相關的組件，這部份內容將在下一篇文章中分析。

Rendezvous基類中的Recv函數

由於Recv函數只是RecvAsync函數的同步版本封裝，所以在每一個實現類繼承從新函數時，只須要提供Send函數的實現和RecvAsync函數實現便可，下面的代碼是Rendezvous基類中同步版本實現。

 1 Status Rendezvous::Recv(const ParsedKey& key, const Args& recv_args,
 2                         Tensor* val, bool* is_dead, int64 timeout_ms) {
 3   Status ret;
 4   Notification n;
 5   RecvAsync(key, recv_args,
 6             [&ret, &n, val, is_dead](const Status& s, const Args& send_args,
 7                                      const Args& recv_args, const Tensor& v,
 8                                      const bool dead) {
 9               ret = s;
10               *val = v;
11               *is_dead = dead;
12               n.Notify();
13             });
14   if (timeout_ms > 0) {
15     int64 timeout_us = timeout_ms * 1000;
16     bool notified = WaitForNotificationWithTimeout(&n, timeout_us);
17     if (!notified) {
18       return Status(error::DEADLINE_EXCEEDED,
19                     "Timed out waiting for notification");
20     }
21   } else {
22     n.WaitForNotification();
23   }
24   return ret;
25 }

能夠看出，不管RecvAsync的實現內容是什麼，Recv函數均可以將RecvAsync視爲黑盒，在其上層封裝成爲與RecvAsync相同實現的同步函數版本。

本地傳輸過程

使用本地傳輸過程包括LocalRendezous和IntraProcessRendezvous兩個實現類，可是後者是前者的封裝，所以本文分析的重點在於LocalRendezvous實現類。

消息隊列的緩存——Table

在TensorFlow中，幾乎每一個Rendezvous實現類都有本身的消息隊列緩存，而幾乎每種消息隊列緩存都是依靠Table實現的。Rendezvous的發送(Send)和接收(Recv)都將經過Table完成，這完美地闡釋了「約會、相會、會和」的釋義，這也是爲何TensorFlow使用這樣一個法語詞來抽象通訊過程。下圖形象化的表示了Table以及Table中的每一個Item。

在LocalRendezvous實現類中，Send端和Recv端使用的是同一個Rendezvous對象，因此他們共享同一個Table，因此Table屬於臨界資源，應該加鎖造成互斥訪問。Item這個結構中其實有不少內容，在上圖中只解釋兩個比較重要的部分。

Value：這就是參與通訊Tensor本體

Waitor：這是在確認Tensor被接收端完成接收後的處理函數，也就是consumer處理該Tensor的函數過程

傳輸過程分析

不管是Send過程仍是Recv過程，它們都將藉助Table完成Tensor的轉發。Send過程做爲Tensor的生產者，它負責將待發送的Tensor送入Table中，並將ParsedKey做爲該Item的鍵。而Recv過程做爲消費者，它也會根據本身所需拼出相同的ParsedKey，而後從Table中查看是否已經存在該項。

應該注意的是，Tensor雖然由Send端生產，可是Table中的Item卻不必定是由Send端插入。由於在TensorFlow中，Send和RecvAsync兩者的相對順序是不能保證前後的，常常出現需求比供給在時間片上先到的狀況，那麼這時就會出現RecvAsync先拼出了ParsedKey而後當即查表的狀況。應對這種狀況的一種方案是，RecvAsync放棄這次查詢，開啓另外一個線程輪詢該表直到Send端產生爲止，而後執行consumer的waiter函數，但這是一個很是消耗資源的實現方式。TensorFlow爲了保證異步性，使用另外一種無需CPU輪詢消耗資源的實現方式。

咱們知道，在Send和RecvAsync順序相對異步的狀況下，waitor函數的執行時機只有兩種狀況，它取決於Send的供給和RecvAsync的需求哪個先到達。若生產者先到達，那麼waiter函數的調用由RecvAsync執行。若消費者的需求先到達，那麼waiter函數的調用由Send執行。簡而言之，老是遲到的一方執行waiter函數。那麼能夠這樣設計：和Send端相同，容許RecvAsync將所需的Item插入到Table中，並連同waiter函數一塊兒發送到該表裏。若是Send端後到達，那麼Send函數將從表中取出該Item，並執行waiter函數，反之，則由RecvAsync函數取出本身所須要的Item，而後執行waiter函數，下面的圖展現了這個過程。

Send過程源碼

瞭解上述的過程後，咱們能夠直接看Send函數的源碼了。下面是LocalRendezvous的Send函數源碼展現。

 1 Status Send(const ParsedKey& key, const Args& send_args, const Tensor& val,
 2               const bool is_dead) override {
 3     uint64 key_hash = KeyHash(key.FullKey());
 4     VLOG(2) << "Send " << this << " " << key_hash << " " << key.FullKey();
 5 
 6     mu_.lock();
 7     if (!status_.ok()) {
 8       // Rendezvous has been aborted.
 9       Status s = status_;
10       mu_.unlock();
11       return s;
12     }
13 
14     ItemQueue* queue = &table_[key_hash];
15     if (queue->empty() || queue->front()->IsSendValue()) {
16       // There is no waiter for this message. Append the message
17       // into the queue. The waiter will pick it up when arrives.
18       // Only send-related fields need to be filled.
19       Item* item = new Item;
20       item->value = val;
21       item->is_dead = is_dead;
22       item->send_args = send_args;
23       if (item->send_args.device_context) {
24         item->send_args.device_context->Ref();
25       }
26       queue->push_back(item);
27       mu_.unlock();
28       return Status::OK();
29     }
30 
31     // There is an earliest waiter to consume this message.
32     Item* item = queue->front();
33     queue->pop_front();
34     mu_.unlock();
35 
36     // Notify the waiter by invoking its done closure, outside the
37     // lock.
38     DCHECK(!item->IsSendValue());
39     item->waiter(Status::OK(), send_args, item->recv_args, val, is_dead);
40     delete item;
41     return Status::OK();
42   }

RecvAsync過程源碼

下面是LocalRendezvous的RecvAsync函數源碼展現。

 1   void RecvAsync(const ParsedKey& key, const Args& recv_args,
 2                  DoneCallback done) override {
 3     uint64 key_hash = KeyHash(key.FullKey());
 4     VLOG(2) << "Recv " << this << " " << key_hash << " " << key.FullKey();
 5 
 6     mu_.lock();
 7     if (!status_.ok()) {
 8       // Rendezvous has been aborted.
 9       Status s = status_;
10       mu_.unlock();
11       done(s, Args(), recv_args, Tensor(), false);
12       return;
13     }
14 
15     ItemQueue* queue = &table_[key_hash];
16     if (queue->empty() || !queue->front()->IsSendValue()) {
17       // There is no message to pick up.
18       // Only recv-related fields need to be filled.
19       Item* item = new Item;
20       item->waiter = std::move(done);
21       item->recv_args = recv_args;
22       if (item->recv_args.device_context) {
23         item->recv_args.device_context->Ref();
24       }
25       queue->push_back(item);
26       mu_.unlock();
27       return;
28     }
29 
30     // A message has already arrived and is queued in the table under
31     // this key.  Consumes the message and invokes the done closure.
32     Item* item = queue->front();
33     queue->pop_front();
34     mu_.unlock();
35 
36     // Invokes the done() by invoking its done closure, outside scope
37     // of the table lock.
38     DCHECK(item->IsSendValue());
39     done(Status::OK(), item->send_args, recv_args, item->value, item->is_dead);
40     delete item;
41   }

關於IntraProcessRendezvous的Send和RecvAsync函數

其實本質上IntraProcessRendezvous和LocalRendezvous是同一個函數實現，只是前者對後者作了一層封裝。咱們從源碼中看到，LocalRendezvous是IntraProcessRendezvous的成員之一，只是在回調函數中多了一些簡單的處理而已，好比它會仔細考量Tensor的生產方和消費方是存在於CPU仍是GPU，是否能夠經過P2P直接拷貝，仍是須要經過Host作中轉，關於拷貝過程使用的是下面的函數，其餘地方大同小異，所以再也不贅述。有興趣的讀者能夠到tensorflow/core/common_runtime/目錄下參考rendezvous_mgr.h、rendezvous_mgr.cc和copy_tensor.h與copy_tensor.cc這幾個文件。

 1 // Copies "input" to "output" between devices accessible to the
 2   // local process via some DMA-like method.  "edge_name" is the name
 3   // of the tensor being copied, for debugging purposes. Depending on
 4   // the type of devices and memory in use, the copy may be performed
 5   // synchronously or asynchronously.  'done' will be invoked only
 6   // after the copy is actually complete.
 7   static void ViaDMA(StringPiece edge_name, DeviceContext* send_dev_context,
 8                      DeviceContext* recv_dev_context, Device* src, Device* dst,
 9                      const AllocatorAttributes src_alloc_attr,
10                      const AllocatorAttributes dst_alloc_attr,
11                      const Tensor* input, Tensor* output,
12                      int dev_to_dev_stream_index, StatusCallback done);

總結

本文是TensorFlow通訊機制系列的第一篇文章，先經過拋出高併發狀況下消息通訊兩端的對應問題引出TensorFlow中的ParsedKey結構設計的必要性，而後給出了Rendezvous全局類圖，最後詳細的分析了LocalRendezvous的消息傳輸實現過程。TensorFlow的通訊機制的完美的闡釋了Rendezvous一詞的含義——不管是Send端仍是Recv端都須要在臨界資源Table中「約會」，進行消息的傳輸。隨後還着重分析了異步狀況下，本屬於consumer的waiter函數調用時機設計問題——爲了保證waiter函數的執行不被阻塞，從設計上採起Late invoke的方案。IntraProcessRendezous本質是LocalRendezvous的一層封裝，它在數據拷貝上面作了更多的工做，藉助LocalRendezvous實現了Send和Recv處於不一樣或相同種類Device狀況下，對上層徹底透明的拷貝過程。因爲篇幅緣由，特地將TensorFlow通訊機制分爲多個系列分析，做爲第一篇文章，本篇介紹了Rendezvous的基本框架。在該系列以後的文章中，還會對跨進程的通訊進行詳細地分析。

相關文章
相關標籤/搜索
每日一句
    每一个你不满意的现在，都有一个你没有努力的曾经。
本站公眾號
   歡迎關注本站公眾號,獲取更多信息
聯繫我們 最近搜索 最新文章 沪ICP备13005482号-10 MyBatis教程 SQL 教程 MySQL教程 Java 教程 Thymeleaf 教程 Hibernate教程 Spring教程 Redis教程