Chenfan Blog

Do cool things that matter.

0%

TensorFlow 拆包（五）：Distributed

Posted on Edited on In

接上篇：

单节点的运行流程基本上已经有个大体印象了，接着就要来拆我所关注的重点所在——分布式运行时了。

Architecture

进入代码部分之前，首先看一下官方文档中，对整个 TensorFlow 结构的介绍。以下是一个典型的分布式 TensorFlow 的架构图：

这里的 Master 和 Worker 服务只在分布式运行时中有，可以认为单节点的 Session 是包含了这两个服务的全部内容。

Client 指的是面向用户的前端编程接口，通常能用的就是 Python 和 C++ 了，client 完成运算图的构建，然后把图的定义通过 session 对象用tf.GraphDef这个 Protobuf 结构传给后面的 Master 服务来跑（即 Python 层定义好计算图，然后通过 Protobuf 的接口进入 C 部分的运行时）。

所以源码中会有/tensorflow/python/client这个目录，其中的内容做的就是架构图中 client 这个概念的任务。

分布式环境中的 Master 有以下几个任务：

  1. 精简并且优化计算图，根据当前次 client 提交运行的输入输出目标，提取出一个子图来
  2. 把子图划分到硬件上（graph -> Partition -> Device）
  3. 缓存前面那一步的结果，以便以后的 steps 能够重用而不用再把上面两步再执行一遍

这些事情看上去好眼熟啊…没错！这就是DirectSession 中 Executor 的任务啊！

在 tensorflow::Partition()中，划分子图到硬件上，对于不在同一个设备上的边需要补充一对 send/recv 的 op 接口，例如上面那张目标图：

Master 接下来再把任务分给具体的 Worker 服务来完成。每一个 Worker 服务都有自己对应的 tasks，ps 负责存储数据，worker 负责具体的计算。

注意 Worker Service 和 Worker task 的区别，Worker Service 可以有 ps 和 worker 这两种 tasks。

截两张实际运行中 Client、Master、Worker 的关系图。

上面一种做法是在 ps 和 worker1 中调用 Server.join()，建图等等的事情全部由 worker0 这边的一份代码完成，大概可以理解成某些代码中的 “Single Session” 的模式。

更常见的写法是下面这种，在 ps 上开 Server.join()，然后每个 worker 分别跑一遍完整的 Python 脚本，自己构建自己本地的计算图。

那我们看到 Client（Python 脚本）运行时是通过本地的 Master Service 来管理整个计算过程，Master 相当于单节点环境的 Session 的角色，由它去分配任务给本地的 worker 或者远程的 worker。

Worker 服务有以下几个任务：

  1. 处理 Master 交过来的请求
  2. 拿到自己的子图之后，调度其中的 op 完成具体的计算
  3. 处理与其他 task（即其他的 Worker 服务）之间的数据通信问题

第 2 步的详细处理也是前面分析过的，即ExecutorState 的 RunAsync() 和 ScheduleReady()部分处理的事情了。

关于 send/recv：

  • CPU 和 GPU 之间通过 cudaMemcpyAsync() 来 overlap 计算和数据传输
  • 两个本地 GPU 之间通过 DMA 直接传输
  • 在 task 之间（不同的 Worker 服务、不同的计算节点之间）通过 gRPC 或者后来增加的 RDMA 来传输

Master 和 Worker 简单地想可以认为是把 DirectSession 中的 Executor 相关的结构功能给拆了出来。

接下来看看具体的代码实现。

tf.train.Supervisor()

代码从单节点改造成分布式只需要替换掉几个固定的 API 即可，先从 Python 层 API 的 Supervisor 说起。

Supervisor 是一个对 Coordinator、Saver、SessionManager 等结构的封装类，用于管理运行的分布式 Session，在运行中建立检查点，并处理异常情况的恢复等等。

1.6.0 版用这个 API 的时候会有警告说将在未来移除，建议换成 tf.train.MonitoredTrainingSession，但是改用这个新 API 实测性能会下降一截，可能是配置方式需要做一下改变，暂时先放下不作研究。

Supervisor 的构造函数有一堆输入参数，挑几个比较重要的记一下：

  • graph：运算图，不指定则使用默认图（这个跟单机版一致）
  • is_chief：分布式环境中可能存在多个 worker 节点，但是其中需要有一个作为 chief worker 节点。chief 需要负责初始化整个运行图，其他 worker 节点将从 chief 节点获取计算图的信息
  • init_op：图中用于初始化所有变量的 op
  • summary_op：用于收集整个运算过程中的有关信息的 op
  • saver：chief 将把有关的信息写到 log 中去
  • global_step：在分布式环境中全局共享的一个变量，标识当前跑到了第几次迭代
  • session_manager：用于管理执行具体运行的 session，也负责异常恢复等等，如果不指定则会创建一个

创建结束时，Supervisor 所关联的计算图将会被锁定，不再允许修改，因为这个图可能会被多个线程共享。

Session

分布式环境下的 C 运行时中存在 3 种 Session 结构，分别是 WorkerSession、MasterSession 以及 GrpcSession，基本上跟前面的 Architecture 是能对应起来的。下面从它们在代码中的调用顺序开始分析：

WorkerSession

WorkerSession 在创建 tf.train.Server()时就被构造出来。

C 层面的 Server 是一个用于管理当前进程中的 Master 和 Worker 服务的结构，通过Start()、Stop()、Join()构成了下图的状态机：

1
2
3
4
5
6
7
//                 Join()            Join()
//                  ___               ___
//      Start()     \ /    Stop()     \ /
// NEW ---------> STARTED --------> STOPPED
//   \                          /
//    \________________________/
//            Stop(), Join()

GrpcServer 在被初始化时：

  • 检查当前可用的所有计算设备，构建 device 列表（与 DirectSession 中做的 AddDevices()一致）
  • 创建了 RpcRendezvousMgr
  • 检查传入的 cluster 信息中，其他 tasks 的端口等等的信息
  • 注册一个 Grpc 的通信 server
  • 创建 Master 以及 GrpcMasterService
  • 创建 GrpcWorker 以及 GrpcWorkerService
  • 启动 Grpc 的通信server
  • 创建 WorkerCache
  • 创建一个 SessionMgr，并随后在这个 SessionMgr 中创建 WorkerSession
  • 这里没有马上创建 MasterSession，而是保存好创建 MasterSession 所需要的信息（大概是因为 ps 中不需要 Master？而 Worker 是所有节点都要有的）
  • 创建 LocalMaster

Work 类用于管理 WorkerSession、处理子图、运行子图、接收 Tensor 数据。GrpcWorker 继承了 Worker 类之后重载了其中的数据传输部分，添加的是一个额外的传输方法，用于在传输大数据时不经过 Protobuf 序列化而直接传（调用 send/recv op 的接口的话，应该是默认要序列化之后再传吧）。

GrpcWorkerService 重载的是 AsyncServiceInterface 这个类，AsyncServiceInterface 抽象的是一个异步等待服务，即创建一个新的线程，用 polling 循环来等待传入的 RPC 请求。

GrpcWorkerService 底层关联的是 WorkerService 这个通过 Protobuf 定义用于 RPC 的结构。

WorkerSession 相对而言反而是个比较简单的结构：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
// WorkerSession encapsulates all of the state relating to a given session.
struct WorkerSession {
  // The name of the session.
  const string session_name;

  // The name of the worker. E.g., /job:mnist/replica:0/task:1.
  const string worker_name;

  // Object from which WorkerInterface instances can be obtained.
  const std::unique_ptr<WorkerCacheInterface> worker_cache;

  // Collection of local devices. These devices are typically RenamedDevices
  // in all except the SessionMgr.legacy_session_. legacy_session_.device_mgr
  // == worker_env_.device_mgr, which holds the true devices.
  const std::unique_ptr<DeviceMgr> device_mgr;

  // graph_mgr keeps track of the registered graphs of this session.
  //
  // Note: graph_mgr must be deleted before rendezvous_mgr!
  // Note: graph_mgr must be deleted before device_mgr!
  const std::unique_ptr<GraphMgr> graph_mgr;

  std::unique_ptr<ClusterFunctionLibraryRuntime> cluster_flr;

  WorkerSession(const string& session_name, const string& worker_name,
                std::unique_ptr<WorkerCacheInterface> worker_cache,
                std::unique_ptr<DeviceMgr> device_mgr,
                std::unique_ptr<GraphMgr> graph_mgr);
};

保存了名字啊、worker_cache啊、device_mgr啊、graph_mgr啊这样的内容。

GrpcSession

下一个断点首先是 GrpcSession 再被触发。

分布式环境下对应的 Session 结构为 Supervisor 中创建的managed_session()，对于 chief 节点，调用自己 SessionManager 中的 _restore_checkpoint() 来在 C 层面创建出 GrpcSession 结构，并且负责完成图的构建等等，之后检查本次运行是否有对应的检查点，有则把检查点的信息恢复出来。而非 chief 节点调用的是wait_for_session() ，创建 GrpcSession 之后等待 chief 节点完成图的构建。

GrpcSession 是从 Session 类继承出来的，其负责的任务跟单机版中的 DirectSession 很像，跟它是同一个层级的东西。

或者说 Session 类在整个 TensorFlow 架构中更确切的应该叫它 Client Session，它们与 Python 层的 sess = tf.Session() 这种结构是直接对应的，是用户编程界面与 TF 运行时的入口。

但 DirectSession 发挥功能的函数都是在本身中直接定义出来的，而这里的 GrpcSession 却可以说基本上是围绕 MasterService 的封装。通过 MasterInterface 来调用 MasterService 的功能来完成任务，可以说 GrpcSession 只是最上图中架构中 client 与 Master 服务之间的接口。

这里的 Master 接口有两种，LocalMaster 用于进程间的直接通信，GrpcMaster 用于 Grpc 通信，GrpcSession 在创建时会根据选项选择所需的 MasterInterface。通常情况下，由于 GrpcSession 都是是直接跟本地的 Master 进行交互，所以默认添加的是 LocalMaster。

MasterSession

上面managed_session()在创建完 C 层面的 GrpcSession 返回之后，会很快执行一次 sess.run()，有检查点的情况是恢复检查点时的变量数据，没有检查点时是执行 init_op 来完成变量初始化。

这里执行的 sess.run()与单节点版本的行为相同，需要首先执行_extend_graph()，不同的是这里执行的是tensorflow::GrpcSession::Extend()，最终到tensorflow::LocalMaster::CreateSession()、tensorflow::Master::CreateSession()。

话说 TensorFlow 中跟 Master 这个概念相关的结构有一堆，一层套一层，而且功能上跟 Worker 又有很多区别的地方。类比起来，大概 MasterSession 也就是跟 Executor 比较像，每一次 Client Session 要 Run 一个子图时（sess.run(...)），启动一个 MasterSession。

MasterSession 追溯到最后是由 GrpcSession.Extend()、GrpcSession.Create()在构建运行图或者修改运行图的时候创建。调用栈大概是这个样子，层次看起来还是比较乱：

1
2
3
4
5
tensorflow::GrpcSession::Create() ->
tensorflow::GrpcSession::CreateImpl(): master_->CreateSession() ->
tensorflow::LocalMaster::CreateSession(): master_impl_->CreateSession() ->
tensorflow::Master::CreateSession() -> （在一个闭包中运行）
tensorflow::MasterSession::Create()

注释中对 MasterSession 的介绍是：

  1. 负责分配 node 到 device
  2. 添加额外的边（例如 send/recv）
  3. 发射 commands 给 worker 来运行

具体来看，还是从sess.run()入手：

1
2
3
4
5
6
7
8
tensorflow::GrpcSession::Run() ->
tensorflow::GrpcSession::RunHelper() （开始准备 req 和 resp，用于异步请求和响应的结构）->
tensorflow::GrpcSession::RunProto(): master_->RunStep() ->
tensorflow::LocalMaster::RunStep(): master_impl_->RunStep() ->
tensorflow::Master::RunStep() -> （在一个闭包中运行）
tensorflow::MasterSession::Run() ->
tensorflow::MasterSession::DoRunWithLocalExecution() ->
tensorflow::MasterSession::ReffedClientGraph::RunPartitions()

最后的 ReffedClientGraph 是与计算图和 Worker 相关的内容了，具体的实现相当复杂，封装层次也是特别多，大致看了下RunPartitions()这里的注释：

  • 匹配 fed tensors 和它们在 req 中的 index
  • 给每个 partition 准备一个将发给 worker 的 call
  • 通过tensorflow::MasterSession::ReffedClientGraph::Part::worker（这是一个 WorkerInterface）的RunGraphAsync()方法，把运行的 call 提交给 worker 跑
  • 等待 RunGraph 的 calls 返回结果
  • 最后处理收到的运行结果

画张图稍微理一下上面这些结构的关系：

然后还有来自这里的一张图：

WorkerInterface

这两个类是作为 TensorFlow 运行时调用 gRPC 的接口基类。

从源码中可以看到，WorkerInterface 类定义了一堆诸如GetStatusAsync()、CreateWorkerSessionAsync()、DeleteWorkerSessionAsync()等等这样的虚函数接口，可以认为是跟 GrpcWorkerService 支持的 GrpcWorkerMethod 是一一对应的：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
// Names of worker methods.
enum class GrpcWorkerMethod {
  kGetStatus,
  kCreateWorkerSession,
  kDeleteWorkerSession,
  kRegisterGraph,
  kDeregisterGraph,
  kRunGraph,
  kCleanupGraph,
  kCleanupAll,
  kRecvTensor,
  kLogging,
  kTracing,
};

当然这个同时也是要跟 Protobuf 的配置要一一对应。

具体的实现在它的两个继承类 Worker 和 GrpcRemoteWorker 里面。

从代码上来看，GrpcRemoteWorker 类中的每一个函数都是调用 IssueRequest() 发起一个异步的 gRPC 调用，远程的 GrpcWorkerService 作为守护进程处理传入的 gRPC 请求。

Worker 类中的对应实现则都是直接在本地做。

Work Flow

最后回到前面的运行部分。

在tensorflow::MasterSession::ReffedClientGraph::RunPartitions()中，MasterSession 运行每一个已经划分好的 partitions 用的是 part.worker->RunGraphAsync() 调用。

part.worker 是每个 partitions 对应的 WorkerInterface 对象，很容易猜想到如果分配在远程对应的应该是 GrpcRemoteWorker 实例，否则对应的应该是 Worker 实例。

那再看数据收发部分的send/recv，之前已经知道了数据传输由recv部分发起，最终调的是RpcRemoteRendezvous::RecvFromRemoteAsync()：

继续往下看，检查各项参数，准备 RpcRecvTensorCall，之后启动 call->Start()，Start()里面调的是StartRTCall()：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
void StartRTCall(std::function<void()> recv_done) {
  resp_.InitAlloc(dst_device_, alloc_attrs_);
  using namespace std::placeholders;
  StatusCallback cb = std::bind(
      [this](std::function<void()> recv_done,
             // Begin unbound arguments.
             const Status& s) {
        if (!s.ok()) {
          mutex_lock l(mu_);
          status_.Update(s);
        }
        recv_done();
      },
      std::move(recv_done), _1);
  wi_->RecvTensorAsync(&opts_, &req_, &resp_, std::move(cb));
}

wi_ 同样是一个 WorkerInterface 的结构。

这样就很清晰了，无论是 Master、Worker 相互之间的控制还是send/recv的数据传输都是通过 WorkerInterface 的派生类作为接口完成的，接口的另一头是底层的 gRPC 通信库。

那么再看到响应 gRPC 调用的那一边，在 GrpcWorkerService 创建时，守护进程HandleRPCsLoop()就启动了：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
void HandleRPCsLoop() {
  // TODO(ncteisen): This may require performance engineering. We can
  // change the number of threads, the number of handlers per thread,
  // or even decide to specialize certain threads to certain methods.
  ENQUEUE_REQUEST(GetStatus, false);
  ENQUEUE_REQUEST(CreateWorkerSession, false);
  ENQUEUE_REQUEST(DeleteWorkerSession, false);
  ENQUEUE_REQUEST(CleanupAll, false);
  ENQUEUE_REQUEST(RegisterGraph, false);
  ENQUEUE_REQUEST(DeregisterGraph, false);

  // TODO(ncteisen): Determine a better policy for enqueuing the
  // appropriate number of each request type.
  for (int i = 0; i < 1000; ++i) {
    EnqueueRecvTensorRequestRaw();
  }
  for (int i = 0; i < 100; ++i) {
    ENQUEUE_REQUEST(RunGraph, true);
  }
  for (int i = 0; i < 100; ++i) {
    ENQUEUE_REQUEST(CleanupGraph, false);
  }

  ENQUEUE_REQUEST(Logging, false);
  ENQUEUE_REQUEST(Tracing, false);

  void* tag;
  bool ok;

  while (cq_->Next(&tag, &ok)) {
    UntypedCall<GrpcWorkerServiceThread>::Tag* callback_tag =
        static_cast<UntypedCall<GrpcWorkerServiceThread>::Tag*>(tag);
    CHECK(callback_tag);
    callback_tag->OnCompleted(this, ok);
  }
}

首先准备好一系列 gRPC 调用的等待队列，11 种调用请求与前面的 GrpcWorkerMethod 一一对应，插入完成之后就是 gRPC 部分的任务了。每个方法对应的处理过程的代码也都列在后面，随便挑一个举例：

1
2
3
4
5
6
7
8
void GetStatusHandler(
    WorkerCall<GetStatusRequest, GetStatusResponse>* call) {
  Schedule([this, call]() {
    Status s = worker_->GetStatus(&call->request, &call->response);
    call->SendResponse(ToGrpcStatus(s));
  });
  ENQUEUE_REQUEST(GetStatus, false);
}

响应 gRPC 请求时这里把要做的任务都封装到线程池里面去执行，然后向队列中重新补充一个相同的等待调用。具体执行的是 worker_（其实是一个 GrpcWorker），完成后向调用方返回一个 gRPC 的 Response。

最后的一个 while 循环是读取 gRPC 完成队列中的内容，处理 gRPC 调用完成之后的收尾工作，RequestReceived、ResponseSent、Cancelled这三种状态。

话说这种完成队列的方式跟 RDMA 的还是挺像的。

MasterInterface

MasterInterface 的结构跟 WorkerInterface 基本类似，不过话说从代码上能看出来不像是一拨人做的啊（命名风格等等），很奇怪。

支持的一些调用：

1
2
3
4
5
6
7
8
9
static const char* grpcMasterService_method_names[] = {
    "/tensorflow.MasterService/CreateSession",
    "/tensorflow.MasterService/ExtendSession",
    "/tensorflow.MasterService/PartialRunSetup",
    "/tensorflow.MasterService/RunStep",
    "/tensorflow.MasterService/CloseSession",
    "/tensorflow.MasterService/ListDevices",
    "/tensorflow.MasterService/Reset",
};

它所派生出来的两个类 GrpcRemoteMaster 和 LocalMaster 从名字上就能够看出来是分别针对远程和本地的调用接口。

乍一看 GrpcRemoteWorker 和 GrpcRemoteMaster 实现远程调用的写法居然完全不一样，很尴尬。仔细往下分析会发现 GrpcRemoteWorker 的 IssueRequest 里面封装的 RPCState 里面的内容跟 GrpcRemoteMaster 的 Call 中的内容很类似。所以为什么不用统一的写法呢。。。

然后 LocalMaster 这个类竟然只是个壳你敢信？。。。里面真正实现本地功能的是 Master 类。

话说前面 Worker 这个类实现的是本地功能，但是 Worker 类是直接继承的 WorkerInterface，到了这里 Master 类跟 MasterInterface 类没有关系，继承 MasterInterface 的是 LocalMaster 类，但是你又发现这个 LocalMaster 类居然是 Master 类的壳。。。相当于跟 Worker 差不多的结构，但是中间多包了一层。

再来看到 GrpcMasterService 的守护进程：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
void HandleRPCsLoop() override {
  ENQUEUE_REQUEST(CreateSession, true);
  ENQUEUE_REQUEST(ExtendSession, false);
  for (int i = 0; i < 100; ++i) {
    ENQUEUE_REQUEST(PartialRunSetup, false);
    ENQUEUE_REQUEST(RunStep, true);
  }
  ENQUEUE_REQUEST(CloseSession, false);
  ENQUEUE_REQUEST(ListDevices, false);
  ENQUEUE_REQUEST(Reset, false);

  void* tag;
  bool ok;
  while (cq_->Next(&tag, &ok)) {
    UntypedCall<GrpcMasterService>::Tag* callback_tag =
        static_cast<UntypedCall<GrpcMasterService>::Tag*>(tag);
    if (callback_tag) {
      callback_tag->OnCompleted(this, ok);
    } else {
      // NOTE(mrry): A null `callback_tag` indicates that this is
      // the shutdown alarm.
      cq_->Shutdown();
    }
  }
}

基本的结构跟前面 Worker 是一致的。

Worker 的远程调用实际发生在：

  • 本地 Master 处理好计算图的 partition 情况
  • 根据 partition 是在本地还是远端，分别请求本地 Worker 或者 GrpcRemoteWorker 来执行
  • 远程的 GrpcWorkerService 守护进程收到请求之后，调用自己本地的 Worker 进行处理，完成后将结果返回

话说 GrpcRemoteMaster 我还没找到到底是在什么情况下用到的。

后续：

Related Posts