[原始碼解析] TensorFlow 分散式環境(2)---Master 靜態邏輯

羅西的思考發表於2022-03-19

[原始碼解析] 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 部分進行具體介紹。

最後,強烈推薦兩個大神:

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架構流程]

相關文章