Press "Enter" to skip to content

[源码解析] TensorFlow 分布式环境(2)—Master 静态逻辑

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

[源码解析] TensorFlow 分布式环境(2)—Master 静态逻辑

 

目录

[源码解析] TensorFlow 分布式环境(2)—Master 静态逻辑

在具体介绍 TensorFlow 分布式的各种 Strategy 之前,我们首先需要看看分布式的基础:分布式环境。只有把基础打扎实了,才能在以后的分析工作之中最大程度的扫清障碍,事半功倍。本文梳理下 Master 的静态逻辑。

 

本系列其他文章是:

 

[ 翻译] TensorFlow 分布式之论文篇 “TensorFlow : Large-Scale Machine Learning on Heterogeneous Distributed Systems”

 

[ 翻译] TensorFlow 分布式之论文篇 “Implementation of Control Flow in TensorFlow”

 

[ 源码解析] TensorFlow 分布式环境(1) — 总体架构

 

1. 总述

 

Server 上运行了两个 RPC 服务,分别是MasterService 和 WorkerService。如果 Client 接入到Server,那幺Server 就是 Master 角色,Client 访问的就是 MasterService 服务(MasterService 同时负责协调和控制多个 WorkerService 的执行过程)。

 

Master 这个角色的具体实现是 Master Service。Master Service是一个GRPC service,用于与一系列远端的分布式设备进行交互来协调多个worker service。

Master Service 对应了 “//tensorflow/core/protobuf/master_service.proto”,其内部有 CreateSession,RunStep 等接口,所有的 TensorFlow Server 都实现了 Master Service。
客户端可以与 Master Service 交互以执行分布式 TensorFlow 计算。
一个 Master Service 会跟踪多个 “主会话(master sessions)”。每个 master sessions 封装了一个计算图及其相关状态。
Master session 运行在 Master 之上,在会话建立后,master 返回一个句柄给客户端,该句柄可用于关联客户端和主会话。
每个 Master session 通常对应一个 “客户会话(client session)”。客户端可以通过调用 CreateSession 向 master 发送一个初始图,通过调用 ExtendSession 向图添加节点。
这里需要说明下,Master 即是一个概念角色,比如 Master 节点,也有一个具体 Master 类。

2. 接口

 

2.1 接口规范

 

Client 通过 GrpcSession 调用 Master Service,既然是 RPC 服务,那幺 Client 和 MasterService 之间就需要有一个接口规范。这个规范定义在 master_service.proto 文件中,其定义了各个接口的消息体。

 

service MasterService {
  // Creates a session.
  rpc CreateSession(CreateSessionRequest) returns (CreateSessionResponse);
  // Extends a session.
  rpc ExtendSession(ExtendSessionRequest) returns (ExtendSessionResponse);
  // Prepares future partial run calls.
  rpc PartialRunSetup(PartialRunSetupRequest) returns (PartialRunSetupResponse);
  // Drives the graph computation.
  rpc RunStep(RunStepRequest) returns (RunStepResponse);
  // Closes a session.
  rpc CloseSession(CloseSessionRequest) returns (CloseSessionResponse);
  // List the devices usable by the master.
  rpc ListDevices(ListDevicesRequest) returns (ListDevicesResponse);
  // Close and abandon all existing sessions.  Ongoing computations
  // will no longer affect fresh ones via the resources in containers listed in
  // the ResetRequest.  See ResetRequest for more details.
  rpc Reset(ResetRequest) returns (ResetResponse);
  // Registers a callable for execution with RunCallable.
  rpc MakeCallable(MakeCallableRequest) returns (MakeCallableResponse);
  // Executes a callable registered with MakeCallable.
  rpc RunCallable(RunCallableRequest) returns (RunCallableResponse);
  // Frees resources associated with a callable registered with MakeCallable.
  rpc ReleaseCallable(ReleaseCallableRequest) returns (ReleaseCallableResponse);
}

 

2.2 MasterInterface

 

Client 使用接口 MasterInterface 获取远端 MasterService 的服务。MasterInterface 是接口类,是 Client 与 TensorFlow Master service 进行通信的抽象接口。这个接口既支持基于 RPC 的 master 实现,也支持不需要 RPC 往返的进程内部的 master 实现。MasterInterface 所有接口都是同步接口,这样 Client 就像调用本地函数一样调用远端 MasterService 提供的服务。

 

MasterInterface有两种实现,都是用来和 Master service 进行通信,

LocalMaster 用于进程间的直接通信,此时 Client 和 Master 在同一个进程。
GrpcRemoteMaster 则使用 Grpc 来和 Master service 进行通信,此时 Client 和 Master 分别部署在两个不同进程。

可以调用工厂方法 NewGrpcMaster 生成 GrpcRemoteMaster 实例。
GrpcRemoteMaster 其实就实现了 gRPC 客户端,它通过 Stub 访问远端 Master 上的 MasterService 服务,具体服务是 GrpcMasterService。
因为 MasterInterface 都是同步接口,所以 Client 就好像访问本地函数一样访问 MasterService。

class MasterInterface {
 public:
  virtual ~MasterInterface() {}
  virtual Status CreateSession(CallOptions* call_options,
                               const CreateSessionRequest* request,
                               CreateSessionResponse* response) = 0;
  virtual Status ExtendSession(CallOptions* call_options,
                               const ExtendSessionRequest* request,
                               ExtendSessionResponse* response) = 0;
  virtual Status PartialRunSetup(CallOptions* call_options,
                                 const PartialRunSetupRequest* request,
                                 PartialRunSetupResponse* response) {
    return errors::Unimplemented("Partial run not implemented for this master");
  }
  virtual Status RunStep(CallOptions* call_options,
                         RunStepRequestWrapper* request,
                         MutableRunStepResponseWrapper* response) = 0;
  virtual Status RunStep(CallOptions* call_options,
                         const RunStepRequest* request,
                         RunStepResponse* response) {
    std::unique_ptr<RunStepRequestWrapper> wrapped_request(
        new ProtoRunStepRequest(request));
    std::unique_ptr<MutableRunStepResponseWrapper> wrapped_response(
        new NonOwnedProtoRunStepResponse(response));
    return RunStep(call_options, wrapped_request.get(), wrapped_response.get());
  }
  virtual MutableRunStepRequestWrapper* CreateRunStepRequest() {
    MutableProtoRunStepRequest* ret = new MutableProtoRunStepRequest;
    ret->request_.set_request_id(GetUniqueRequestId());
    return ret;
  }
  virtual MutableRunStepResponseWrapper* CreateRunStepResponse() {
    return new OwnedProtoRunStepResponse;
  }
  virtual Status CloseSession(CallOptions* call_options,
                              const CloseSessionRequest* request,
                              CloseSessionResponse* response) = 0;
  virtual Status ListDevices(CallOptions* call_options,
                             const ListDevicesRequest* request,
                             ListDevicesResponse* response) = 0;
  virtual Status Reset(CallOptions* call_options, const ResetRequest* request,
                       ResetResponse* response) = 0;
  virtual Status MakeCallable(CallOptions* call_options,
                              const MakeCallableRequest* request,
                              MakeCallableResponse* response) = 0;
  virtual Status RunCallable(CallOptions* call_options,
                             const RunCallableRequest* request,
                             RunCallableResponse* response) = 0;
  virtual Status ReleaseCallable(CallOptions* call_options,
                                 const ReleaseCallableRequest* request,
                                 ReleaseCallableResponse* response) = 0;
 protected:
  // NOTE: This should only be called by implementations of this
  // interface whose CreateRunStepResponse() method returns a
  // proto-based wrappers for the RunStepResponse message.
  RunStepResponse* get_proto_from_wrapper(
      MutableRunStepResponseWrapper* wrapper) {
    return wrapper->get_proto();
  }
};

 

具体使用如下,如果 Client 和 Master 在同一个进程,则直接使用 LocalMaster,否则使用 GrpcRemoteMaster 来利用 gRPC 访问远程 GrpcMasterService。图上两个矩形封装的 Master 代表实际的 Master 类,此类实现了具体 Master 功能。

 

 

图 1 Master 逻辑结构

 

2.3 调用

 

下面的伪代码说明了客户端如何与 master 交互,这其实就是分布式模式之中,使用 GrpcRemoteMaster 来通过 gRPC 与远端 MasterSerivce 服务交互的过程。

 

stub = NewStub("/job:mnist/replica:0/task:0")
{handle} = stub->CreateSession({graph_def})
do {
   stub->RunStep({handle, {feeds}, {fetches}})
   // The client can evaluate a predicate locally, based on the
   // result of fetches, to determine whether to terminate. For
   // example, it might fetch the loss and evaluate whether it is less
   // than some threshold.
} while (!should_stop({fetches}));
stub->CloseSession({handle})

 

3. LocalMaster

 

当 Client 调用时候,GrpcSession 使用 LocalMaster 获取本地master,如果没有得到,则才使用 GrpcRemoteMaster。此时 Client 和 master 没有跨节点,LocalMaster 使客户端和master之间能够直接进行进程内通信,这样就可以给同进程内部的Client提供更高效的Master服务。

 

3.1 定义

 

LocalMaster 定义如下,主要成员变量就是 master_impl_。LocalMaster 其实就是一个壳而已,直接转发给master_impl_。master_impl_ 是当 Client 和 master 没有跨节点时候,本地直接调用的类。

 

class LocalMaster : public MasterInterface {
 private:
  Master* master_impl_;  // Not owned.
  const int64 default_timeout_in_ms_;
  // See LocalMaster::Lookup for the factory function that creates
  // objects of this type.
  LocalMaster(Master* master_impl, const int64 default_timeout_in_ms);
  TF_DISALLOW_COPY_AND_ASSIGN(LocalMaster);
};

 

3.2 注册

 

LocalMaster 有一个静态变量 local_master_registry_ 用来注册。

 

typedef std::unordered_map<string, MasterInfo> LocalMasterRegistry;
LocalMasterRegistry* local_master_registry() {
  static LocalMasterRegistry* local_master_registry_ = new LocalMasterRegistry;
  return local_master_registry_;
}

 

在 GrpcServer 初始化时候,调用如下代码把 target=”grpc://” 生成的 Master 注册到本地 LocalMaster。

 

LocalMaster::Register(target(), master_impl_.get(), config.operation_timeout_in_ms());

 

就是把 master 注册到这个static变量 local_master_registry_ 之中。

 

/* static */void LocalMaster::Register(const string& target, Master* master,
                           int64 default_timeout_in_ms) {
  mutex_lock l(*get_local_master_registry_lock());
  local_master_registry()->insert(
      {target, MasterInfo(master, default_timeout_in_ms)});
}

 

3.3 查找

 

当调用 GrpcSession::Create 方法时候,如果 Client 和 Master 在同一个进程,Lookup 在本地能够找到注册的 Master,则会生成一个 LocalMaster 返回,同时 LocalMaster 的 master_impl_ 就配置成找到的 Master。如果找不到,就返回空,则 GrpcSession::Create 方法会创建一个 GrpcRemoterMaster,这样就同远端 Master 进行交互。

 

/* static */std::unique_ptr<LocalMaster> LocalMaster::Lookup(const string& target) {
  std::unique_ptr<LocalMaster> ret;
  mutex_lock l(*get_local_master_registry_lock());
  auto iter = local_master_registry()->find(target);
  if (iter != local_master_registry()->end()) {
    ret.reset(new LocalMaster(iter->second.master,
                              iter->second.default_timeout_in_ms));
  }
  return ret;
}

 

以下是同一个进程,Lookup 可以找到的情况,生成 LocalMaster 进行本地操作。

 

 

图 2 同进程 master 操作

 

我们看看不同进程的情况。此时进程 1 之中的 LocalMaster 没有指向任何 Master,因为本地没有启动 Server,所以 GrpcSession::Create 方法第一步 Lookup 调用失败,返回 Null,GrpcSession::Create 方法执行第二步骤,创建 GrpcRemoteMaster,进行远程交互。进程 2 之中,LocalMaster 因为没有客户端调用 GrpcSession::Create 方法,所以也没有指向任何 Master。

 

 

图 3 跨进程 master 操作

 

3.4 功能

 

LocalMaster 调用到其内部成员变量 master_impl_ 来完成业务功能。

 

Status LocalMaster::CreateSession(CallOptions* call_options,
                                  const CreateSessionRequest* request,
                                  CreateSessionResponse* response) {
  Notification n;
  Status ret;
  master_impl_->CreateSession(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}
Status LocalMaster::ExtendSession(CallOptions* call_options,
                                  const ExtendSessionRequest* request,
                                  ExtendSessionResponse* response) {
  Notification n;
  Status ret;
  master_impl_->ExtendSession(request, response, [&n, &ret](const Status& s) {
    ret.Update(s);
    n.Notify();
  });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}
Status LocalMaster::RunStep(CallOptions* call_options,
                            RunStepRequestWrapper* request,
                            MutableRunStepResponseWrapper* response) {
  Notification n;
  Status ret;
  master_impl_->RunStep(call_options, request, response,
                        [&n, &ret](const Status& s) {
                          ret.Update(s);
                          n.Notify();
                        });
  TF_RETURN_IF_ERROR(
      WaitForNotification(call_options, default_timeout_in_ms_, &n));
  return ret;
}

 

4. GrpcRemoteMaster

 

GrpcRemoteMaster 是 gRPC 客户端的一种实现, 其终通过 Stub 调用远端 Master 上的 GrpcMasterService 服务,这样调用行为就犹如本地函数调用一样。远端 GrpcMasterService 实现了 MasterService 服务定义的所有接口,是 MasterService 服务的真正实体。当创建 GrpcRemoteMaster 实例时候,需要通过 target 来指定 Master 服务的地址和端口,并且创建对应的 RPC 通道。GrpcSession 和 GrpcRemoteMaster 从严格意义上讲都是 Client 实现的一部分。

 

4.1 定义

 

GrpcRemoteMaster 具体定义如下,主要是使用了MasterServiceStub。

 

// GrpcRemoteMaster is an implementation of the MasterInterface
// that uses gRPC to talk to the Master service.
class GrpcRemoteMaster : public MasterInterface {
  using MasterServiceStub = grpc::MasterService::Stub;
 public:
  explicit GrpcRemoteMaster(const SharedGrpcChannelPtr& client_channel)
      : stub_(grpc::MasterService::NewStub(client_channel)) {}
  ~GrpcRemoteMaster() override {}
  std::unique_ptr<MasterServiceStub> stub_;
};

 

4.2 功能

 

GrpcRemoteMaster 的功能很简单,就是通过 gRPC 的一 个 stub 调用远端 Master 服务的相应接口。

 

4.2.1 CreateSession

 

我们使用 CreateSession 为例看看,是使用 CallWithRetry 完成功能。

 

Status CreateSession(CallOptions* call_options,
                     const CreateSessionRequest* request,
                     CreateSessionResponse* response) override {
  return CallWithRetry(call_options, request, response,
                       &MasterServiceStub::CreateSession);
}

 

CallWithRetry 代码如下,其又是调用 s = FromGrpcStatus((stub_.get()->*pfunc)(&ctx, *request, response)) 获取 Stub 来完成功能。

 

template <typename Request, typename Response>
Status CallWithRetry(CallOptions* call_options, const Request* request,
                     Response* response,
                     ::grpc::Status (MasterServiceStub::*pfunc)(
                         ::grpc::ClientContext*, const Request&, Response*),
                     string trace_string = {}) {
  absl::Duration timeout = absl::Milliseconds(call_options->GetTimeout());
  absl::Time expired_time = absl::FromUnixMicros(Env::Default()->NowMicros());
  if (timeout > absl::ZeroDuration()) {
    expired_time += timeout;
  }
  Status s;
  for (int num_retries = 0;; ++num_retries) {
    ::grpc::ClientContext ctx;
    std::unique_ptr<profiler::TraceMe> trace;
    if (!trace_string.empty()) {
      trace.reset(NewTraceRpc(trace_string, &ctx));
    }
    ctx.set_fail_fast(false);
    if (timeout > absl::ZeroDuration()) {
      // We do not modify the timeout here to match legacy behavior. However,
      // this could violate the contract of tensorflow::Session. If we retry
      // an RPC just before the deadline is exceeded, we will still set the
      // timeout to the original value. This leads to the overall timeout
      // being double what was expected.
      ctx.set_deadline(absl::ToChronoTime(absl::Now() + timeout));
    }
    s = FromGrpcStatus((stub_.get()->*pfunc)(&ctx, *request, response));
    if (!errors::IsUnavailable(s)) {
      return s;
    }
    // TODO(b/117162170): we may want to make this configurable.
    constexpr int kMaxRetries = 10;
    if (num_retries >= kMaxRetries) {
      return s;
    }
    absl::Time now = absl::FromUnixMicros(Env::Default()->NowMicros());
    const absl::Time deadline_with_backoff =
        now + absl::Microseconds(ComputeBackoffMicroseconds(num_retries));
    // Wait for a short period of time before retrying the RPC.  If our
    // backoff would put us past the RPC deadline, we truncate it to ensure
    // our RPC starts before the deadline.
    const auto backoff_until = (timeout <= absl::ZeroDuration() ||
                                expired_time > deadline_with_backoff)
                                   ? deadline_with_backoff
                                   : expired_time;
    Env::Default()->SleepForMicroseconds(
        absl::ToInt64Microseconds(backoff_until - now));
    now = absl::FromUnixMicros(Env::Default()->NowMicros());
    if (now > expired_time && timeout > absl::ZeroDuration()) {
      // If timeout_in_ms is set, exit the retry loop on timeout.
      return errors::DeadlineExceeded(ctx.debug_error_string());
    }
  }
}

 

4.2.2 Master Service Stub

 

接下来我们看看 Stub,这是依据 “//tensorflow/core/protobuf/master_service.proto” 来使用 grpc 实现的。

 

class Stub final : public StubInterface {
 public:
  Stub(const std::shared_ptr< ::grpc::ChannelInterface>& channel);
  ::grpc::Status CreateSession(::grpc::ClientContext* context,
                               const CreateSessionRequest& request,
                               CreateSessionResponse* response) override;
  ::grpc::Status ExtendSession(::grpc::ClientContext* context,
                               const ExtendSessionRequest& request,
                               ExtendSessionResponse* response) override;
  ::grpc::Status PartialRunSetup(::grpc::ClientContext* context,
                                 const PartialRunSetupRequest& request,
                                 PartialRunSetupResponse* response) override;
  ::grpc::Status RunStep(::grpc::ClientContext* context,
                         const RunStepRequest& request,
                         RunStepResponse* response) override;
  ::grpc::Status CloseSession(::grpc::ClientContext* context,
                              const CloseSessionRequest& request,
                              CloseSessionResponse* response) override;
  ::grpc::Status ListDevices(::grpc::ClientContext* context,
                             const ListDevicesRequest& request,
                             ListDevicesResponse* response) override;
  ::grpc::Status Reset(::grpc::ClientContext* context,
                       const ResetRequest& request,
                       ResetResponse* response) override;
  ::grpc::Status MakeCallable(::grpc::ClientContext* context,
                              const MakeCallableRequest& request,
                              MakeCallableResponse* response) override;
  ::grpc::Status RunCallable(::grpc::ClientContext* context,
                             const RunCallableRequest& request,
                             RunCallableResponse* response) override;
  ::grpc::Status ReleaseCallable(::grpc::ClientContext* context,
                                 const ReleaseCallableRequest& request,
                                 ReleaseCallableResponse* response) override;
 private:
  std::shared_ptr< ::grpc::ChannelInterface> channel_;
  const ::grpc::internal::RpcMethod rpcmethod_CreateSession_;
  const ::grpc::internal::RpcMethod rpcmethod_ExtendSession_;
  const ::grpc::internal::RpcMethod rpcmethod_PartialRunSetup_;
  const ::grpc::internal::RpcMethod rpcmethod_RunStep_;
  const ::grpc::internal::RpcMethod rpcmethod_CloseSession_;
  const ::grpc::internal::RpcMethod rpcmethod_ListDevices_;
  const ::grpc::internal::RpcMethod rpcmethod_Reset_;
  const ::grpc::internal::RpcMethod rpcmethod_MakeCallable_;
  const ::grpc::internal::RpcMethod rpcmethod_RunCallable_;
  const ::grpc::internal::RpcMethod rpcmethod_ReleaseCallable_;
};

 

具体远端的对应方法是:

 

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",
    "/tensorflow.MasterService/MakeCallable",
    "/tensorflow.MasterService/RunCallable",
    "/tensorflow.MasterService/ReleaseCallable",
};
std::unique_ptr<MasterService::Stub> MasterService::NewStub(
    const std::shared_ptr< ::grpc::ChannelInterface>& channel,
    const ::grpc::StubOptions& options) {
  std::unique_ptr<MasterService::Stub> stub(new MasterService::Stub(channel));
  return stub;
}

 

Stub 内部调用 grpc 完成发送功能。

 

::grpc::Status MasterService::Stub::CreateSession(
    ::grpc::ClientContext* context, const CreateSessionRequest& request,
    CreateSessionResponse* response) {
  return ::grpc::internal::BlockingUnaryCall(
      channel_.get(), rpcmethod_CreateSession_, context, request, response);
}

 

所以,如果是 GrpcRemoteMaster,则调用流程应该是:GrpcRemoteMaster 接收到 grpc session 的请求,转交给 grpc master service,这期间经历了 GrpcSession -> GrpcRemoteMaster -> GrpcMasterService -> Master -> MasterSession 一系列流程。

 

4.3 创建

 

当建立 GrpcSession 时候,create 方法之中会先查找有没有 Master。如果找到了就直接返回 LocalMaster,这部分我们前面介绍过。如果 Lookup 找不到。所以会调用 NewGrpcMaster 生成一个 GrpcRemoteMaster。

 

/* static */Status GrpcSession::Create(const SessionOptions& options,
                           std::unique_ptr<GrpcSession>* out_session) {
  std::unique_ptr<GrpcSession> session(new GrpcSession(options));
  std::unique_ptr<MasterInterface> master;
  // For testing, we enable the client to disable the use of the local
  // master registry, so that the RPC stack is exercised.
  if (!options.config.rpc_options().use_rpc_for_inprocess_master()) {
    master = LocalMaster::Lookup(options.target); 
  }
  if (!master) {
    SharedGrpcChannelPtr master_channel;
    TF_RETURN_IF_ERROR(
        NewHostPortGrpcChannel(options.target.substr(kSchemePrefixLength),
                               &options.config.rpc_options(), &master_channel));
    // 建立 GrpcRemoteMaster,与远端 Master 交互
    master.reset(NewGrpcMaster(master_channel));
  } else {
    session->is_local_ = true;
  }
  session->SetRemoteMaster(std::move(master));
  *out_session = std::move(session);
  return Status::OK();
}

 

NewGrpcMaster 方法具体如下:

 

MasterInterface* NewGrpcMaster(const SharedGrpcChannelPtr& channel) {
  return new GrpcRemoteMaster(channel);
}

 

5. GrpcMasterService

 

GrpcMasterService 实现了 RPC 对应的 MasterService。GrpcMasterService 会:

预先了解有哪些本地设备可以给客户使用,也会发现远端设备并且跟踪其统计数据。
维护/管理实时计算图会话(MasterSession),这些会话将调用本地或者远端设备来对收到的计算图进行计算。
会话功能是:对收到的计算图进行分析,剪枝,把节点放到可用设备上,通过调用 RunGraph 在工作者上进行图计算。

5.1 创建

 

GrpcServer 之中,master_service_ 是 GrpcMasterService 类型的变量。

 

// 创建 Master 以及对应的 GrpcMasterService
  master_impl_ = CreateMaster(&master_env_);
  master_service_ = NewGrpcMasterService(master_impl_.get(), config, &builder);

 

GrpcServer 使用 master_thread_ 线程来执行 GrpcMasterService 的 HandleRPCsLoop方法。

 

master_thread_.reset(
    env_->StartThread(ThreadOptions(), "TF_master_service",
                      [this] { master_service_->HandleRPCsLoop(); }));

 

5.2 定义

 

GrpcMasterService 定义如下,master_impl_ 是 Server 传入的 master 指针,是一个 Master 类的实例:

 

class GrpcMasterService : public AsyncServiceInterface {
  Master* master_impl_ = nullptr;  // Not owned.
  std::unique_ptr<::grpc::ServerCompletionQueue> cq_;
  grpc::MasterService::AsyncService master_service_;
  mutex mu_;
  bool is_shutdown_ TF_GUARDED_BY(mu_);
  const ConfigProto default_session_config_;
  ::grpc::Alarm* shutdown_alarm_ = nullptr;
  template <class RequestMessage, class ResponseMessage>
  using MasterCall = Call<GrpcMasterService, grpc::MasterService::AsyncService,
                          RequestMessage, ResponseMessage>;
}

 

GrpcMasterService 初始化时候,会得到 grpc 的消息队列 cq_。

 

GrpcMasterService(Master* master, const ConfigProto& default_session_config,
                  ::grpc::ServerBuilder* builder)
    : master_impl_(master),
      is_shutdown_(false),
      default_session_config_(default_session_config) {
  builder->RegisterService(&master_service_);
  cq_ = builder->AddCompletionQueue();
}

 

5.3 主循环

 

前面提到了,master_thread_ 线程来执行 GrpcMasterService 的 HandleRPCsLoop 方法。HandleRPCsLoop 会调用 GrpcMasterService 内部函数来进行处理RPC消息。主循环 HandleRPCsLoop 代码如下:

 

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);
  ENQUEUE_REQUEST(MakeCallable, false);
  for (int i = 0; i < 100; ++i) {
    ENQUEUE_REQUEST(RunCallable, true);
  }
  ENQUEUE_REQUEST(ReleaseCallable, 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();
    }
  }
}

 

上面代码之中有一些最佳实践,具体就是围绕 ENQUEUE_REQUEST 做了一些处理:

this->cq_ 是 grpc 队列。
ENQUEUE_REQUEST 宏会为给定的 RPC 方法名称创建一个新请求(比如 ENQUEUE_REQUEST(GetStatus, false) 就会生成一个 GetStatus 请求),这些请求将在 this->cq_ 之上进行排队。
预先把一定数量的要处理的任务放入 cq_,如果任务被任务响应 handler 调用,则 handler 会调用ENQUEUE_REQUEST() 往队列之中补充一个同样的调用,这样可以确保完成队列 cq_ 有足够的任务来处理传入的请求,这样处理将不会阻塞,整体处理速度会提高。
代码最后的 while 循环将读取 gRPC 队列中的内容,就是 gRPC 调用之后的收尾工作。

#define ENQUEUE_REQUEST(method, supports_cancel)                              \
  do {                                                                        \
    mutex_lock l(mu_);                                                        \
    if (!is_shutdown_) {                                                      \
      Call<GrpcMasterService, grpc::MasterService::AsyncService,              \
           method##Request, method##Response>::                               \
          EnqueueRequest(&master_service_, cq_.get(),                         \
                         &grpc::MasterService::AsyncService::Request##method, \
                         &GrpcMasterService::method##Handler,                 \
                         (supports_cancel));                                  \
    }                                                                         \
  } while (0)

 

5.4 消息处理

 

在具体消息响应之中,会调用 master_impl_ 进行处理,当 Master 处理完成之后,处理函数将回调一个 lambda 表达式,向 Client 返回的响应消息。可以看到,代码在最后会使用 ENQUEUE_REQUEST 再插入一个同样类型的请求,比如下面最后会返回给 Client 一个 CreateSessionResponse。

 

// RPC handler for creating a session.
void CreateSessionHandler(
    MasterCall<CreateSessionRequest, CreateSessionResponse>* call) {
  CreateSessionRequest* rewritten_req = new CreateSessionRequest;
  rewritten_req->mutable_config()->MergeFrom(default_session_config_);
  rewritten_req->MergeFrom(call->request);
  master_impl_->CreateSession(rewritten_req, &call->response,
                              [call, rewritten_req](const Status& status) {
                                call->SendResponse(ToGrpcStatus(status));
                                delete rewritten_req;
                              });
  ENQUEUE_REQUEST(CreateSession, true);
}

 

5.5 功能

 

GrpcMasterService 提供的 API 如下:

 

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",
    "/tensorflow.MasterService/MakeCallable",
    "/tensorflow.MasterService/RunCallable",
    "/tensorflow.MasterService/ReleaseCallable",
};

 

我们举出三个具体功能分析一下:

 

5.5.1 CreateSession

 

CreateSessionRequest 消息之中会带有 Client 设定的计算图和配置信息。Master 接收到请求之后,为这个 Client 建立一个 MasterSession 实例,并建立一个唯一地标识该 MasterSession 实例的 session_handle。这是通过 Master 类成员变量 std::unordered_map<string, MasterSession*> sessions_ 来完成的,session_handle 就是 string 类型。

 

Master 返回消息 CreateSessionResponse 给 Client。CreateSessionResponse 消息中携带:

session_handle。Client 的 GrpcSession 据此和 Master 端的 MasterSession 建立关联,后续交互之中,Client 在消息内均会携带此 session_handle,随后,Client 与 Master 的所有交互中,在请求 消息中通过携带 session_handle,Master 通过它在 std::unordered_map<string, MasterSession*> sessions_ 会找到相对应的 MasterSession 实例。
初始 graph_version。用于后续发起 ExtendSession 操作,往原始的计算图中追加新的节点。

 

图 4 CreateSession

 

具体响应代码如下:

 

// RPC handler for creating a session.
void CreateSessionHandler(
    MasterCall<CreateSessionRequest, CreateSessionResponse>* call) {
  CreateSessionRequest* rewritten_req = new CreateSessionRequest;
  rewritten_req->mutable_config()->MergeFrom(default_session_config_);
  rewritten_req->MergeFrom(call->request);
  master_impl_->CreateSession(rewritten_req, &call->response,
                              [call, rewritten_req](const Status& status) {
                                call->SendResponse(ToGrpcStatus(status));
                                delete rewritten_req;
                              });
  ENQUEUE_REQUEST(CreateSession, true);
}

 

5.5.2 ExtendSession

 

当建立 Session 之后,Client 可以通过 ExtendSession 告诉 Master 我需要拓展原有计算图的规模 (只能追加子图,不能修改或删除)。

 

在请求消息 ExtendSessionRequest 中有:

session_handle :用来查找哪一个 MasterSession 实例;
graph_def :需要加到计算图上的节点;
current_graph_version :需要拓展的计算图版本号;

在在响应消息 ExtendSessionResponse 中返回 new_graph_version,其用于下一此 ExtendSession 操作。

 

 

图 5 ExtendSession

 

具体代码如下:

 

// RPC handler for extending a session.
void ExtendSessionHandler(
    MasterCall<ExtendSessionRequest, ExtendSessionResponse>* call) {
  master_impl_->ExtendSession(&call->request, &call->response,
                              [call](const Status& status) {
                                call->SendResponse(ToGrpcStatus(status));
                              });
  ENQUEUE_REQUEST(ExtendSession, false);
}

 

5.5.3 RunStep

 

客户端会迭代执行 RunStep,请求消息 RunStepRequest 的变量较多,比如:

session_handle :用来查找哪一个 MasterSession 实例;
feed :输入的 NamedTensor 列表;
fetch :待输出 Tensor 的名称列表;
target :执行节点列表;

响应消息 RunStepResponse 主要携带:

tensor :输出的 Tensor 列表;

 

图 6 RunStep

 

消息定义具体如下:

 

message RunStepRequest {
  // REQUIRED: session_handle must be returned by a CreateSession call
  // to the same master service.
  string session_handle = 1;
  // Tensors to be fed in the step. Each feed is a named tensor.
  repeated NamedTensorProto feed = 2;
  // Fetches. A list of tensor names. The caller expects a tensor to
  // be returned for each fetch[i] (see RunStepResponse.tensor). The
  // order of specified fetches does not change the execution order.
  repeated string fetch = 3;
  // Target Nodes. A list of node names. The named nodes will be run
  // to but their outputs will not be fetched.
  repeated string target = 4;
  // Options for the run call.
  RunOptions options = 5;
  // Partial run handle (optional). If specified, this will be a partial run
  // execution, run up to the specified fetches.
  string partial_run_handle = 6;
  // If true then some errors, e.g., execution errors that have long
  // error messages, may return an OK RunStepResponse with the actual
  // error saved in the status_code/status_error_message fields of the
  // response body. This is a workaround since the RPC subsystem may
  // truncate long metadata messages.
  bool store_errors_in_response_body = 7;
  // Unique identifier for this request. Every RunStepRequest must
  // have a unique request_id, and retried RunStepRequest must have
  // the same request_id. If request_id is zero, retry detection is disabled.
  int64 request_id = 8;
}
message RunStepResponse {
  // NOTE: The order of the returned tensors may or may not match
  // the fetch order specified in RunStepRequest.
  repeated NamedTensorProto tensor = 1;
  // Returned metadata if requested in the options.
  RunMetadata metadata = 2;
  // If store_errors_in_response_body is true in the request, then
  // optionally the server may return an OK status for the RPC and
  // fill the true status into the fields below, to allow for messages
  // that are too long to fit in metadata.
  error.Code status_code = 3;
  string status_error_message = 4;
}

 

具体代码如下:

 

// RPC handler for running one step in a session.
void RunStepHandler(MasterCall<RunStepRequest, RunStepResponse>* call) {
  auto* trace = TraceRpc("RunStep/Server", call->client_metadata());
  CallOptions* call_opts = new CallOptions;
  if (call->request.options().timeout_in_ms() > 0) {
    call_opts->SetTimeout(call->request.options().timeout_in_ms());
  } else {
    call_opts->SetTimeout(default_session_config_.operation_timeout_in_ms());
  }
  RunStepRequestWrapper* wrapped_request =
      new ProtoRunStepRequest(&call->request);
  MutableRunStepResponseWrapper* wrapped_response =
      new NonOwnedProtoRunStepResponse(&call->response);
  call->SetCancelCallback([call_opts]() { call_opts->StartCancel(); });
  master_impl_->RunStep(
      call_opts, wrapped_request, wrapped_response,
      [call, call_opts, wrapped_request, trace](const Status& status) {
        call->ClearCancelCallback();
        delete call_opts;
        delete wrapped_request;
        delete trace;
        if (call->request.store_errors_in_response_body() && !status.ok()) {
          call->response.set_status_code(status.code());
          call->response.set_status_error_message(status.error_message());
          call->SendResponse(ToGrpcStatus(Status::OK()));
        } else {
          call->SendResponse(ToGrpcStatus(status));
        }
      });
  ENQUEUE_REQUEST(RunStep, true);
}

 

6. 业务实现 Master 类

 

6.1 创建

 

前面提到了,GrpcServer 之中建立的是 Master 类的实例。

 

std::unique_ptr<Master> GrpcServer::CreateMaster(MasterEnv* master_env) {
  return std::unique_ptr<Master>(new Master(master_env, 0.0));
}

 

这样,在收到 Client 的消息后,在具体消息响应之中,GrpcMasterService 的线程会调用 master_impl_ 进行处理,就是把业务逻辑委托给 Master 类来实现。所以我们接下来就看看 Master 如何处理。

 

// RPC handler for creating a session.
void CreateSessionHandler(
    MasterCall<CreateSessionRequest, CreateSessionResponse>* call) {
  CreateSessionRequest* rewritten_req = new CreateSessionRequest;
  rewritten_req->mutable_config()->MergeFrom(default_session_config_);
  rewritten_req->MergeFrom(call->request);
  master_impl_->CreateSession(rewritten_req, &call->response,
                              [call, rewritten_req](const Status& status) {
                                call->SendResponse(ToGrpcStatus(status));
                                delete rewritten_req;
                              });
  ENQUEUE_REQUEST(CreateSession, true);
}

 

6.2 定义

 

Master 其实不是 MasterInterface 的派生类,其定义在tensorflow/core/distributed_runtime/master.cc。可以从成员变量 sessions_ 上看出来,主要就是管理 MasterSession。

 

class Master {
 private:
  typedef Master ME;
  // Not owned.
  MasterEnv* env_ = nullptr;
  // Owned.
  mutex mu_;
  // shutdown_ is set to true by the dtor.
  condition_variable shutdown_cv_;
  bool shutdown_ TF_GUARDED_BY(mu_) = false;
  Thread* gc_thread_;
  // Maps session handles to sessions.
  std::unordered_map<string, MasterSession*> sessions_ TF_GUARDED_BY(mu_);
  // Moving average of step times.
  MovingAverage last_1000_steps_ TF_GUARDED_BY(mu_);
  // Cumulative number of steps executed.
  int64 step_count_ TF_GUARDED_BY(mu_);
  // If a session is not active for this many seconds, it will be
  // closed automatically.
  const double session_gc_seconds_;
  // Used to track ids for incoming requests so we can detect duplicates.
  RecentRequestIds recent_request_ids_;
};

 

6.3 功能

 

我们回忆一下之前提到的。

 

分布式运行的核心是如何操作计算图,但是计算功能被拆分为 Client,Master 和 Worker 三个角色。

 

Client 负责构造计算图,Worker 负责执行具体计算,但是 Worker 怎幺知道应该计算什幺?TensorFlow 在两者之间插入了一个 Master 角色来负责协调,调度。

 

虽然 Master 不是 MasterInterface 的派生类,但时其实现了 MasterService 的具体业务。Master 具体负责:

Master 预先知道本地有哪些设备可以作为客户使用的设备,也会发现远程设备,并跟踪这些远程设备的统计数据。
一个 Master 包含多个 “主会话(master sessions)”。每个 master sessions 封装了一个计算图及其相关状态。
主会话将:

精简优化计算图,比如剪枝/分割/插入发送和接受算子。
协调/调度资源。比如哪个计算应该在哪个设备运行,具体就是按照 graph -> Partition -> Device 这个策略把子图划分到硬件设备之上。
把分割之后的各个子图发送给各个 worker,具体每一个子图对应一个 MasterSession。并最终通过在工作者上启动 RunGraph 来驱动图的计算。

Master 维护实时图计算会话的状态。

至此,Master 的静态结构我们已经介绍完毕,具体 Master 功能我们将在后文 Session 部分进行具体介绍。

 

最后,强烈推荐两个大神:

[TensorFlow Internals] ( https://github.com/horance-liu/tensorflow-internals ),虽然其分析的不是最新代码,但是建议对 TF 内部实现机制有兴趣的朋友都去阅读一下,绝对大有收获。
https://home.cnblogs.com/u/deep-learning-stacks/ 西门宇少,不仅仅是 TensorFlow,其公共号还有更多其他领域,业界前沿。

0xFF 参考

 

TensorFlow Internals

 

TensorFlow架构与设计:概述

 

TensorFlow内核剖析

 

TensorFlow架构与设计:OP本质论

 

[译] TensorFlow 白皮书

 

2017TensorFlow开发者峰会

 

https://jcf94.com/2018/02/28/2018-02-28-tfunpacking3/

 

TensorFlow 拆包(五):Distributed

 

TensorFlow Architecture

 

『深度长文』Tensorflow代码解析(五)

 

什幺是in-graph replication和between-graph replication?

 

[腾讯机智] TensorFlow源码解析(1): 创建会话

 

05tensorflow分布式会话

 

第八节,配置分布式TensorFlow

 

TensorFlow 分布式(Distributed TensorFlow)

 

tensorflow源码解析之distributed_runtime

 

Distributed TensorFlow: A Gentle Introduction

 

一文说清楚Tensorflow分布式训练必备知识

 

TensorFlow中的Placement启发式算法模块——Placer

 

TensorFlow的图切割模块——Graph Partitioner

 

TensorFlow中的通信机制——Rendezvous(一)本地传输

 

TensorFlow分布式采坑记

 

TensorFlow技术内幕(九):模型优化之分布式执行

 

Tensorflow架构流程]

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注