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教程