背景
[作者:DeepLearningStack,阿里巴巴算法工程師,開源TensorFlow Contributor]
歡迎大家關注我的公眾號,“互聯網西門二少”,我將繼續輸出我的技術干貨~
在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要包括六個部分。
src_device:消息發送源的字符串信息,形如/job:localhost/replica:0/task_id:0/device:GPU:0
src:和src_device的信息量相同,只不過是結構體的表示方法
src_incarnation:一般來說這個字段沒有什么作用,但是當某個worker重啟后,該值會發生變化,用來和之前掛掉的worker做區分,這便於debug
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的基本框架。在該系列之后的文章中,還會對跨進程的通信進行詳細地分析。