[源码解析] TensorFlow 分布式环境(5) --- Session

会话机制是TensorFlow 分布式运行时的核心，我们接下来按照从 Client 到 worker 的流程，把 Session 机制从前到后走一边。

本系列其他文章是：

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

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

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

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

[源码解析] TensorFlow 分布式环境(3)--- Worker 静态逻辑

[源码解析] TensorFlow 分布式环境(4) --- WorkerCache

1. 概述

1.1 Session 分类

分布式模式由如下 sessions 彼此协作完成了会话控制，其中：

  • GrpcSession 位于 Client 之上，控制 Client 的会话生命周期；
  • MasterSession 位于 Master 之上，可能存在多个 Client 同时接入到同一个 Master，Master 会为每个 Client 构建一个 MasterSession。MasterSession 控制 Master 的会话生命周 期；
  • WorkerSession 位于 Worker 之上，可能存在多个 Master 接入到同一个 Worker，Worker 会为每个 Master 创建一个 WorkerSession。WorkerSession 控制 Worker 的会话生命周期；

如下图所示，这里 Master 和 Worker 都是一个 Server，每个 Server 之上运行一个 MasterService，一个 WorkerService，每个 Server 可能会扮演不同角色，具体取决于用户如何配置计算图和集群。因为存在这种两层一对多关系，为了区别这种不同的数据流和控制关系，有逻辑关系的这三个 session 绑定在同一个 session_handle 之上，每个 session_handle 标示一条完整的数据流。

图 1 Session 关系

1.2 会话流程

我们从 GrpcSession 入手，其基本功能如下：

  • 创建会话

    • 获取远端设备集;
    • 在 Master 之上创建 MasterSession;
    • 在各个 Worker 之上创建 WorkerSession;
  • 迭代执行
    • 启动执行;
    • 图分裂;
    • 注册子图;
    • 运行子图;
  • 关闭会话
    • 关闭 MasterSession
    • 关闭 WorkerSession;

1.2.1 MasterSession 生命周期

在分布式模式下，Master 运行时被 MasterSession 控制，其生命周期如下图所示。

图 2 MasterSession 生命周期

1.2.2 WorkerSession 生命周期

在分布式模式下，Worker 运行时由 WorkerSession 控制，其生命周期如下图所示。

图 3 WorkerSession 生命周期

2. GrpcSession

GrpcSession 是 tensorflow::grpc::MasterService 的简单封装。其使用远程设备集作为计算资源，使用 grpc 作为远端调用机制，让调用者在远端设备上对 TensorFlow 图进行计算。

2.1 定义

我们依然只给出成员变量定义和部分重要函数，其就是利用 master_ 对 tensorflow::grpc::MasterService 进行调用。

  1. class GrpcSession : public Session {
  2.   // 有多种创建方式
  3.   Status Create(const GraphDef& graph) override;
  4.   Status Create(const RunOptions& run_options, const GraphDef& graph) override;
  5.   Status Create(GraphDef&& graph) override;
  6.   Status Create(const RunOptions& run_options, GraphDef&& graph) override;  
  8.  private:
  9.   const SessionOptions options_;
  10.   std::unique_ptr<MasterInterface> master_;
  11.   mutex mu_;
  13.   // handle_ returned by the master to identify this session.
  14.   string handle_ TF_GUARDED_BY(mu_);
  16.   // The current version of the graph.
  17.   int64_t current_graph_version_ TF_GUARDED_BY(mu_);
  19.   bool is_local_ = false;
  20. };

2.2 注册&工厂类

GrpcSession 的使用是通过工厂类完成，比如：

  1. Status NewSession(const SessionOptions& options, Session** out_session) {
  2.   SessionFactory* factory;
  3.   Status s = SessionFactory::GetFactory(options, &factory);
  4.   if (!s.ok()) {
  5.     *out_session = nullptr;
  6.     return s;
  7.   }
  8.   // Starts exporting metrics through a platform-specific monitoring API (if
  9.   // provided). For builds using "tensorflow/core/platform/default", this is
  10.   // currently a no-op.
  11.   session_created->GetCell()->Set(true);
  12.   s = factory->NewSession(options, out_session);
  13.   if (!s.ok()) {
  14.     *out_session = nullptr;
  15.   }
  16.   return s;
  17. }

GrpcSession 由 GrpcSessionFactory 来多态创建，如果 protocal 使用了"grpc://"，就会产生 GrpcSession。而 GrpcSessionFactory 会实现注册到系统之上。

  1. const char* const kSchemePrefix = "grpc://";
  2. const size_t kSchemePrefixLength = strlen(kSchemePrefix);
  4. class GrpcSessionFactory : public SessionFactory {
  5.  public:
  6.   bool AcceptsOptions(const SessionOptions& options) override {
  7.     return absl::StartsWith(options.target, kSchemePrefix);
  8.   }
  10.   Status NewSession(const SessionOptions& options,
  11.                     Session** out_session) override {
  12.     std::unique_ptr<GrpcSession> session;
  13.     TF_RETURN_IF_ERROR(GrpcSession::Create(options, &session));
  14.     *out_session = session.release();
  15.     return Status::OK();
  16.   }
  18.   // Invokes the session specific static method to reset containers.
  19.   Status Reset(const SessionOptions& options,
  20.                const std::vector<string>& containers) override {
  21.     return GrpcSession::Reset(options, containers);
  22.   }
  23. };
  25. class GrpcSessionRegistrar {
  26.  public:
  27.   GrpcSessionRegistrar() {
  28.     SessionFactory::Register("GRPC_SESSION", new GrpcSessionFactory());
  29.   }
  30. };
  31. static GrpcSessionRegistrar registrar;

2.3 创建GrpcSession

GrpcSession::Create 方法完成了获取工作。Client 通过 GrpcSession 调用 Master Service，但是具体如何与 Master Service 交互？则通过 MasterInterface。

所以说，这里最重要的就是如何构建 MasterInterface 实例。我们前文提到过，MasterInterface有两种实现，都是用来和 Master service 进行通信，分别对应了不同的应用场景。

  • LocalMaster 用于进程间的直接通信，此时 Client 和 Master 在同一个进程。
  • GrpcRemoteMaster 则使用 Grpc 来和 Master service 进行通信，此时Client 和 Master 分别部署在两个不同进程。GrpcRemoteMaster 其实就实现了 gRPC 客户端，它通过 Stub 访问远端 Master 上的 MasterService 服务。

图上两个矩形封装的 Master 代表实际的 Master 类，此类实现了具体 Master 功能。

图 1 Master 逻辑关系

从下面代码可以看到，GrpcSession 会依据 options.target 来决定如何创建，options.target 一般就是"grpc://"，如果通过 LocalMaster::Lookup 方法得到 LocalMaster 类，就直接使用，如果没有找到，就使用 NewGrpcMaster 来生成一个 GrpcRemoteMaster。

  1. /* static */
  2. Status GrpcSession::Create(const SessionOptions& options,
  3.                            std::unique_ptr<GrpcSession>* out_session) {
  4.   std::unique_ptr<GrpcSession> session(new GrpcSession(options));
  5.   std::unique_ptr<MasterInterface> master;
  6.   // For testing, we enable the client to disable the use of the local
  7.   // master registry, so that the RPC stack is exercised.
  8.   if (!options.config.rpc_options().use_rpc_for_inprocess_master()) {
  9.     master = LocalMaster::Lookup(options.target);
  10.   }
  11.   if (!master) {
  12.     SharedGrpcChannelPtr master_channel;
  13.     TF_RETURN_IF_ERROR(
  14.         NewHostPortGrpcChannel(options.target.substr(kSchemePrefixLength),
  15.                                &options.config.rpc_options(), &master_channel));
  16.     master.reset(NewGrpcMaster(master_channel));
  17.   } else {
  18.     session->is_local_ = true;
  19.   }
  20.   session->SetRemoteMaster(std::move(master));
  21.   *out_session = std::move(session);
  22.   return Status::OK();
  23. }

2.4 创建MasterSession

在 GrpcSession 创建之后，系统会接着创建 MasterSession，这是通过 GrpcSession::Create(graph_def) 完成的。GrpcSession::Create(graph_def) 会构建 CreateSessionRequst 消息，然后通过 GrpcRemoteMaster 把初始计算图发给 Master。Master 收到 CreateSessionRequst 消息之后就构建相应的 MasterSession，然后返回 CreateSessionResponse 给 GrpcSession，消息包括。

  • 该 MasterSession 的 session_handle。用于标识 Master 侧的 MasterSession 实例
  • 初始计算图的版本号 graph_version。用于后续发起 ExtendSession 操作，比如往原始的计算图中追加新的节点。

图 2 创建MasterSession

具体代码如下，首先是两个 create 方法，其最终调用到 CreateImpl。

  1. Status GrpcSession::Create(const RunOptions& run_options,
  2.                            const GraphDef& graph) {
  3.   return Create(run_options, GraphDef(graph));
  4. }
  6. Status GrpcSession::Create(GraphDef&& graph) {
  7.   CallOptions call_options;
  8.   call_options.SetTimeout(options_.config.operation_timeout_in_ms());
  9.   return CreateImpl(&call_options, std::move(graph));
  10. }

CreateImpl 方法如下：

  1. Status GrpcSession::CreateImpl(CallOptions* call_options, GraphDef graph) {
  2.   {
  3.     mutex_lock l(mu_);
  4.     if (!handle_.empty()) {
  5.       return errors::InvalidArgument("A session is alive.");
  6.     }
  7.   }
  8.   CreateSessionRequest req;
  9.   *req.mutable_config() = options_.config;
  10.   req.mutable_graph_def()->Swap(&graph);
  11.   req.set_target(options_.target);
  12.   ReEncodeConsts(req.mutable_graph_def());
  13.   CreateSessionResponse resp;
  14.   Status s = master_->CreateSession(call_options, &req, &resp);
  15.   if (s.ok()) {
  16.     SetHandleAndGraphVersion(resp.session_handle(), resp.graph_version());
  17.   }
  18.   return s;
  19. }

2.4.1 GrpcRemoteMaster::CreateSession

GrpcRemoteMaster 是位于 Client 的 gRPC 客户端实现，它的 CreateSession 方法只是通过 gRPC stub 来调用 远端服务 MasterService 的 CreateSession 接口，其实就是发送一个 CreateSessionRequest 请求。

  1. Status CreateSession(CallOptions* call_options,
  2.                      const CreateSessionRequest* request,
  3.                      CreateSessionResponse* response) override {
  4.   return CallWithRetry(call_options, request, response,
  5.                        &MasterServiceStub::CreateSession);
  6. }

2.4.2 GrpcMasterService::CreateSessionHandler

GrpcMasterService 是 Master 提供的 gRPC 服务，收到 CreateSessionRequest 消息之后， 服务调用 GrpcMasterService::CreateSessionHandler 来处理消息，而真正业务处理是由 master_impl_（Master 类的实例）来完成，就是调用了 Master::CreateSession。

当 master_impl_ 处理完成后，会向 Client 返回 CreateSessionResponse 响应。

  1. // RPC handler for creating a session.
  2. void CreateSessionHandler(
  3.     MasterCall<CreateSessionRequest, CreateSessionResponse>* call) {
  4.   CreateSessionRequest* rewritten_req = new CreateSessionRequest;
  5.   rewritten_req->mutable_config()->MergeFrom(default_session_config_);
  6.   rewritten_req->MergeFrom(call->request);
  7.   master_impl_->CreateSession(rewritten_req, &call->response,
  8.                               [call, rewritten_req](const Status& status) {
  9.                                 call->SendResponse(ToGrpcStatus(status));
  10.                                 delete rewritten_req;
  11.                               });
  12.   ENQUEUE_REQUEST(CreateSession, true);
  13. }

2.4.3 Master::CreateSession

Master::CreateSession 会从线程池之中拿到一个线程，在线程之中会做如下处理：

  • 如果定义了 clust_spec，则按照配置寻找所有的 worker。
  • 获取远端设备。
  • 获取远端worker。
  • 通过factory 建立 MasterSession。
  • 利用 worker_cache_factory，让 MasterSession 建立 WorkerSession 会话。
  • 通过 sessions_.insert 在 Master 内部的 <session_handle, MasterSession> 二元组之中保存对应关系，这样后续 Master 就可以通过 session_handle 得到对应的 MasterSession。
  1. void Master::CreateSession(const CreateSessionRequest* req,
  2.                            CreateSessionResponse* resp, MyClosure done) {
  3.   SchedClosure([this, req, resp, done]() {
  4.     Status status;
  5.     WorkerCacheFactoryOptions worker_cache_factory_options;
  6.     string grpc_protocol("grpc");
  7.     worker_cache_factory_options.protocol = &grpc_protocol;
  8.     auto call_done = gtl::MakeCleanup([&status, &done] { done(status); });
  9.     status = ValidateExternalGraphDefSyntax(req->graph_def());
  10.     if (!status.ok()) return;
  12.     // The following 4 variables are set differently, depending on whether this
  13.     // session uses a client-provided clusterspec or not.
  14.     WorkerCacheInterface* worker_cache = nullptr;
  15.     // Note: worker_cache_ptr will be null except if this session is using a
  16.     // client-supplied ClusterDef (ClusterSpec propagation).
  17.     std::unique_ptr<WorkerCacheInterface> worker_cache_ptr;
  18.     std::unique_ptr<DeviceSet> device_set;
  19.     // TODO(saeta): Convert to std::make_unique when available.
  20.     std::unique_ptr<std::vector<std::unique_ptr<Device>>> remote_devices(
  21.         new std::vector<std::unique_ptr<Device>>());
  23.     if (req->config().has_cluster_def()) { // 如果定义了集群
  24.       worker_cache_factory_options.cluster_def = &req->config().cluster_def();
  26.       // Set the server_def's job_name and task_index fields.
  27.       string normalized_string;
  28.       string grpc_protocol(kGrpcProtocol);
  29.       if (req->target().compare(0, grpc_protocol.length(), grpc_protocol) ==
  30.           0) {
  31.         normalized_string =
  32.             req->target().substr(grpc_protocol.length(), string::npos);
  33.       } else {
  34.         normalized_string = req->target();
  35.       }
  36.       for (auto&& job : req->config().cluster_def().job()) {
  37.         for (auto&& task : job.tasks()) {
  38.           if (task.second == normalized_string) {
  39.             if (worker_cache_factory_options.job_name != nullptr) {
  40.               return;
  41.             }
  42.             if (env_->local_devices[0]->parsed_name().job == job.name() &&
  43.                 env_->local_devices[0]->parsed_name().task == task.first) {
  44.               return;
  45.             }
  46.             worker_cache_factory_options.job_name = &job.name();
  47.             worker_cache_factory_options.task_index = task.first;
  48.           }
  49.         }
  50.       }
  51.       worker_cache_factory_options.rpc_options = &req->config().rpc_options();
  52.       // Create the worker cache from the computed server_def.
  53.       status = env_->worker_cache_factory(worker_cache_factory_options,
  54.                                           &worker_cache);
  55.       if (!status.ok()) return;
  56.       worker_cache_ptr = std::unique_ptr<WorkerCacheInterface>(worker_cache);
  57.       // Ping all the workers and build the list of devices that the
  58.       // session will use.
  59.       // 获取设备
  60.       status =
  61.           DeviceFinder::GetRemoteDevices(req->config().device_filters(), env_,
  62.                                          worker_cache, remote_devices.get());
  63.       if (!status.ok()) return;
  64.       device_set.reset(new DeviceSet);
  65.       for (auto&& d : *remote_devices) {
  66.         device_set->AddDevice(d.get());
  67.         DeviceNameUtils::ParsedName name = d->parsed_name();
  68.         if (name.job == *worker_cache_factory_options.job_name &&
  69.             name.task == worker_cache_factory_options.task_index &&
  70.             name.type == "CPU" && name.id == 0) {
  71.           device_set->set_client_device(d.get());
  72.         }
  73.       }
  74.     } else { // 没有集群
  75.       worker_cache = env_->worker_cache;
  76.       // Ping all the workers and build the list of devices that the
  77.       // session will use.
  78.       // 获取远端设备
  79.       status =
  80.           DeviceFinder::GetRemoteDevices(req->config().device_filters(), env_,
  81.                                          worker_cache, remote_devices.get());
  82.       if (!status.ok()) return;
  83.       device_set.reset(new DeviceSet);
  84.       for (auto&& d : *remote_devices) {
  85.         device_set->AddDevice(d.get());
  86.       }
  87.       int num_local_devices = 0;
  88.       for (Device* d : env_->local_devices) {
  89.         device_set->AddDevice(d);
  90.         if (num_local_devices == 0) {
  91.           // Uses the first local device as the client device.
  92.           device_set->set_client_device(d);
  93.         }
  94.         num_local_devices++;
  95.       }
  96.     }
  98.     SessionOptions options;
  99.     options.config = req->config();
  101.     // 获取远端worker
  102.     std::vector<string> filtered_worker_list;
  103.     DeviceFinder::GetRemoteWorkers(req->config().device_filters(), env_,
  104.                                    worker_cache, &filtered_worker_list);
  106.     // 通过factory找到会话
  107.     MasterSession* session = env_->master_session_factory(
  108.         options, env_, std::move(remote_devices), std::move(worker_cache_ptr),
  109.         std::move(device_set), std::move(filtered_worker_list));
  111.     GraphDef* gdef =
  112.         const_cast<CreateSessionRequest*>(req)->mutable_graph_def();
  114.     // 建立会话，把图传给会话
  115.     status = session->Create(std::move(*gdef), worker_cache_factory_options);
  116.     if (!status.ok()) {
  117.       session->Close().IgnoreError();
  118.       session->Unref();
  119.       return;
  120.     }
  121.     resp->set_session_handle(session->handle());
  122.     // Insert into the session map, which takes ownership of the session.
  123.     {
  124.       mutex_lock l(mu_);
  125.       CHECK(sessions_.insert({session->handle(), session}).second);
  126.     }
  127.   });
  128. }

3. MasterSession

MasterSession 位于 Master 之上，可能存在多个 Client 同时接入到同一个 Master，Master 会为每个 Client 构建一个 MasterSession。MasterSession 控制 Master 的会话生命周 期。

3.1 定义

MasterSession 的定义如下。

  1. // MasterSession wraps ClientGraph in a reference counted object.
  2. // This way, MasterSession can clear up the cache mapping Run requests to
  3. // compiled graphs while the compiled graph is still being used.
  4. class MasterSession::ReffedClientGraph : public core::RefCounted {
  5.  public:
  6.   ReffedClientGraph(const string& handle, const BuildGraphOptions& bopts,
  7.                     std::unique_ptr<ClientGraph> client_graph,
  8.                     const SessionOptions& session_opts,
  9.                     const StatsPublisherFactory& stats_publisher_factory,
  10.                     bool is_partial, WorkerCacheInterface* worker_cache,
  11.                     bool should_deregister)
  12.       : session_handle_(handle),
  13.         bg_opts_(bopts),
  14.         client_graph_before_register_(std::move(client_graph)),
  15.         session_opts_(session_opts),
  16.         is_partial_(is_partial),
  17.         callable_opts_(bopts.callable_options),
  18.         worker_cache_(worker_cache),
  19.         should_deregister_(should_deregister),
  20.         collective_graph_key_(
  21.             client_graph_before_register_->collective_graph_key) {
  22.     VLOG(1) << "Created ReffedClientGraph for node with "
  23.             << client_graph_before_register_->graph.num_node_ids();
  25.     stats_publisher_ = stats_publisher_factory(handle, bopts, session_opts);
  27.     // Initialize a name to node map for processing device stats.
  28.     for (Node* n : client_graph_before_register_->graph.nodes()) {
  29.       name_to_node_details_.emplace(
  30.           n->name(),
  31.           NodeDetails(n->type_string(),
  32.                       strings::StrCat(
  33.                           "(", absl::StrJoin(n->requested_inputs(), ", "))));
  34.     }
  35.   }
  37.   ~ReffedClientGraph() override {
  38.     if (should_deregister_) {
  39.       DeregisterPartitions();
  40.     } else {
  41.       for (Part& part : partitions_) {
  42.         worker_cache_->ReleaseWorker(part.name, part.worker);
  43.       }
  44.     }
  45.   }
  47.  private:
  48.   const string session_handle_;
  49.   const BuildGraphOptions bg_opts_;
  51.   // NOTE(mrry): This pointer will be null after `RegisterPartitions()` returns.
  52.   std::unique_ptr<ClientGraph> client_graph_before_register_ TF_GUARDED_BY(mu_);
  53.   const SessionOptions session_opts_;
  54.   const bool is_partial_;
  55.   const CallableOptions callable_opts_;
  56.   WorkerCacheInterface* const worker_cache_;  // Not owned.
  58.   struct NodeDetails {
  59.     explicit NodeDetails(string type_string, string detail_text)
  60.         : type_string(std::move(type_string)),
  61.           detail_text(std::move(detail_text)) {}
  62.     const string type_string;
  63.     const string detail_text;
  64.   };
  65.   std::unordered_map<string, NodeDetails> name_to_node_details_;
  67.   const bool should_deregister_;
  68.   const int64_t collective_graph_key_;
  69.   std::atomic<int64_t> execution_count_ = {0};
  71.   // Graph partitioned into per-location subgraphs.
  72.   struct Part {
  73.     // Worker name.
  74.     string name;
  76.     // Maps feed names to rendezvous keys. Empty most of the time.
  77.     std::unordered_map<string, string> feed_key;
  79.     // Maps rendezvous keys to fetch names. Empty most of the time.
  80.     std::unordered_map<string, string> key_fetch;
  82.     // The interface to the worker. Owned.
  83.     WorkerInterface* worker = nullptr;
  85.     // After registration with the worker, graph_handle identifies
  86.     // this partition on the worker.
  87.     string graph_handle;
  89.     Part() : feed_key(3), key_fetch(3) {}
  90.   };
  92.   // partitions_ is immutable after RegisterPartitions() call
  93.   // finishes.  RunPartitions() can access partitions_ safely without
  94.   // acquiring locks.
  95.   std::vector<Part> partitions_;
  97.   mutable mutex mu_;
  99.   // Partition initialization and registration only needs to happen
  100.   // once. `!client_graph_before_register_ && !init_done_.HasBeenNotified()`
  101.   // indicates the initialization is ongoing.
  102.   Notification init_done_;
  104.   // init_result_ remembers the initialization error if any.
  105.   Status init_result_ TF_GUARDED_BY(mu_);
  107.   std::unique_ptr<StatsPublisherInterface> stats_publisher_;
  108. };

3.2 创建

MasterSession::Create(graph_def) 的工作如下：

  • 调用 MakeForBaseGraph 来初始化计算图，并生成 SimpleGraphExecutionState 实例;
  • 调用 CreateWorkerSessions，如果动态配置集群，则广播通知给所有 Worker，让其创建对应的 WorkerSession。
  1. Status MasterSession::Create(GraphDef&& graph_def,
  2.                              const WorkerCacheFactoryOptions& options) {
  3.   if (session_opts_.config.use_per_session_threads() ||
  4.       session_opts_.config.session_inter_op_thread_pool_size() > 0) {
  5.     return errors::InvalidArgument(
  6.         "Distributed session does not support session thread pool options.");
  7.   }
  8.   if (session_opts_.config.graph_options().place_pruned_graph()) {
  9.     session_opts_.config.mutable_graph_options()->set_place_pruned_graph(false);
  10.   }
  12.   GraphExecutionStateOptions execution_options;
  13.   execution_options.device_set = devices_.get();
  14.   execution_options.session_options = &session_opts_;
  15.   {
  16.     mutex_lock l(mu_);
  17.     TF_RETURN_IF_ERROR(GraphExecutionState::MakeForBaseGraph(
  18.         std::move(graph_def), execution_options, &execution_state_));
  19.   }
  20.   should_delete_worker_sessions_ = true;
  21.   return CreateWorkerSessions(options);
  22. }

3.2.1 创建计算图

这里会构建 GraphExecutionState，依据 GraphDef 构建对应的 FullGraph。

GraphDef 是原始图结构，ConvertGraphDefToGraph 完成从 GraphDef 到 Graph 的格式转换，GraphDef 包含了图的元数据，Graph 则包含图结构的其他信息，被运行时系统所使用。

  1. /* static */ Status GraphExecutionState::MakeForBaseGraph(
  2.     GraphDef&& graph_def, const GraphExecutionStateOptions& options,
  3.     std::unique_ptr<GraphExecutionState>* out_state) {
  5.   auto flib_def = absl::make_unique<FunctionLibraryDefinition>(
  6.       OpRegistry::Global(), graph_def.library());
  8.   TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(&graph_def, *flib_def, 0));
  10.   if (options.session_options->config.graph_options().place_pruned_graph() ||
  11.       !options.session_options->config.experimental()
  12.            .optimize_for_static_graph()) {
  13.     auto ret = absl::WrapUnique(new GraphExecutionState(
  14.         absl::make_unique<GraphDef>(std::move(graph_def)), std::move(flib_def),
  15.         options));
  17.     // When place_pruned_graph is true, a different Graph* will be initialized
  18.     // each time we prune the original graph, so there is no need to
  19.     // construct a Graph* in this case.
  20.     if (!options.session_options->config.graph_options().place_pruned_graph()) {
  21.       auto base_graph = absl::make_unique<Graph>(OpRegistry::Global());
  22.       TF_RETURN_IF_ERROR(ConvertGraphDefToGraph({}, *ret->original_graph_def_,
  23.                                                 base_graph.get()));
  24.       TF_RETURN_IF_ERROR(ret->InitBaseGraph(std::move(base_graph)));
  25.     }
  26.     *out_state = std::move(ret);
  27.   } else {
  28.     auto ret = absl::WrapUnique(
  29.         new GraphExecutionState(nullptr, std::move(flib_def), options));
  30.     auto base_graph = absl::make_unique<Graph>(OpRegistry::Global());
  31.     TF_RETURN_IF_ERROR(
  32.         ConvertGraphDefToGraph({}, std::move(graph_def), base_graph.get()));
  33.     TF_RETURN_IF_ERROR(ret->InitBaseGraph(std::move(base_graph)));
  34.     *out_state = std::move(ret);
  35.   }
  36.   return Status::OK();
  37. }

InitBaseGraph 会调用 Placer.run 完成算子编排。就是把计算图之中的算子放到最适合的设备上计算，这样可以最大化效率。Placer 会对 Graph 做分析，并且结合用户的要求对每个Node如何放置进行微调，具体原则有如下四种：

  • 尽量满足用户的要求。用户可以通过 device 信息或者 loc 来制定设备，尽量优先满足。
  • 尽量使用快速设备。TF 系统之中每个设备都有优先级，级别越高计算性能越好，优先选择级别高的设备。
  • 尽量保证程序可运行。如果某个 Node 指定了在某种设备上执行，但是系统之中没有，则会选择一个可用的设备来重写 Placement。
  • 尽量考虑近邻性。比如尽量让 Consumer 和 Producer 在同一个设备上，避免无意义的跨设备拷贝。
  1. Status GraphExecutionState::InitBaseGraph(std::unique_ptr<Graph>&& new_graph) {
  2.   // Save stateful placements before placing.
  3.   RestoreStatefulNodes(new_graph.get());
  5.   GraphOptimizationPassOptions optimization_options;
  6.   optimization_options.session_handle = session_handle_;
  7.   optimization_options.session_options = session_options_;
  8.   optimization_options.graph = &new_graph;
  9.   optimization_options.flib_def = flib_def_.get();
  10.   optimization_options.device_set = device_set_;
  12.   TF_RETURN_IF_ERROR(OptimizationPassRegistry::Global()->RunGrouping(
  13.       OptimizationPassRegistry::PRE_PLACEMENT, optimization_options));
  15.   Placer placer(new_graph.get(), "", flib_def_.get(), device_set_,
  16.                 /* default_local_device= */ nullptr,
  17.                 session_options_ == nullptr ||
  18.                     session_options_->config.allow_soft_placement(),
  19.                 session_options_ != nullptr &&
  20.                     session_options_->config.log_device_placement());
  21.   TF_RETURN_IF_ERROR(placer.Run());
  23.   TF_RETURN_IF_ERROR(OptimizationPassRegistry::Global()->RunGrouping(
  24.       OptimizationPassRegistry::POST_PLACEMENT, optimization_options));
  26.   for (const Node* n : new_graph->nodes()) {
  27.     node_name_to_cost_id_map_[n->name()] = n->cost_id();
  28.   }
  30.   SaveStatefulNodes(new_graph.get());
  31.   graph_ = new_graph.release();
  32.   return Status::OK();
  33. }

3.2.2 创建 WorkerSession

当 MasterSession 创建成功后，如果没有动态配置集群 (默认的分布式配置环境)， 则不会广播所有 Worker 动态地创建 WorkerSession。事实上，每个 Worker 都存在一个 SessionMgr 实例，它持有一个名为 legacy_session_ 的 WorkerSession 实例。因此，每个 Worker 存在一个全局唯一的 WorkerSession 实例。

图 3 创建 WorkerSession

逻辑如下：

  • 首先，调用 ReleaseWorker 来释放已有的 workers。
  • 其次，调用 GetOrCreateWorker 重新在缓存之中获取 Worker，如果没有，缓存自会构建。
  • 最后，遍历 Workers，调用 CreateWorkerSessionAsync 来让每个 Worker 各自创建一个 WorkerSession，每个请求都会用 set_session_handle(handle_) 来把 MasterSession 的 session_handle 设置进入，这样每个 WorkerSession 都和 MasterSession 共享同样的 session_handle，它们都隶属于同一个 MasterSession。

为了收集全部 Workers 返回的消息，这里使用了计数器 BlockingCounter 来等待，其会把初始数值设置为 Worker 数目，当收集全部 Workers 的 CreateWorkerSessionResponse 响应消息之后，计数器会减少为 0，则 BlockingCounter 会被唤醒。

  1. Status MasterSession::CreateWorkerSessions(
  2.     const WorkerCacheFactoryOptions& options) {
  3.   const std::vector<string> worker_names = filtered_worker_list_;
  4.   WorkerCacheInterface* worker_cache = get_worker_cache();
  6.   struct WorkerGroup {
  7.     // The worker name. (Not owned.)
  8.     const string* name;
  10.     // The worker referenced by name. (Not owned.)
  11.     WorkerInterface* worker = nullptr;
  13.     // Request and responses used for a given worker.
  14.     CreateWorkerSessionRequest request;
  15.     CreateWorkerSessionResponse response;
  16.     Status status = Status::OK();
  17.   };
  18.   BlockingCounter done(worker_names.size());
  19.   std::vector<WorkerGroup> workers(worker_names.size());
  21.   // Release the workers.
  22.   auto cleanup = gtl::MakeCleanup([&workers, worker_cache] {
  23.     for (auto&& worker_group : workers) {
  24.       if (worker_group.worker != nullptr) {
  25.         worker_cache->ReleaseWorker(*worker_group.name, worker_group.worker);
  26.       }
  27.     }
  28.   });
  30.   string task_name;
  31.   string local_device_name;
  32.   DeviceNameUtils::SplitDeviceName(devices_->client_device()->name(),
  33.                                    &task_name, &local_device_name);
  34.   const int64_t client_device_incarnation =
  35.       devices_->client_device()->attributes().incarnation();
  37.   Status status = Status::OK();
  38.   // Create all the workers & kick off the computations.
  39.   for (size_t i = 0; i < worker_names.size(); ++i) {
  40.     workers[i].name = &worker_names[i];
  41.     workers[i].worker = worker_cache->GetOrCreateWorker(worker_names[i]);
  42.     workers[i].request.set_session_handle(handle_);
  43.     workers[i].request.set_master_task(task_name);
  44.     workers[i].request.set_master_incarnation(client_device_incarnation);
  45.     if (session_opts_.config.share_cluster_devices_in_session() ||
  46.         session_opts_.config.experimental()
  47.             .share_cluster_devices_in_session()) {
  48.       for (const auto& remote_dev : devices_->devices()) {
  49.         *workers[i].request.add_cluster_device_attributes() =
  50.             remote_dev->attributes();
  51.       }
  53.       if (!session_opts_.config.share_cluster_devices_in_session() &&
  54.           session_opts_.config.experimental()
  55.               .share_cluster_devices_in_session()) {
  56.       }
  57.     }
  59.     DeviceNameUtils::ParsedName name;
  60.     if (!DeviceNameUtils::ParseFullName(worker_names[i], &name)) {
  61.       status = errors::Internal("Could not parse name ", worker_names[i]);
  62.       return status;
  63.     }
  64.     if (!name.has_job || !name.has_task) {
  65.       status = errors::Internal("Incomplete worker name ", worker_names[i]);
  66.       return status;
  67.     }
  69.     if (options.cluster_def) {
  70.       *workers[i].request.mutable_server_def()->mutable_cluster() =
  71.           *options.cluster_def;
  72.       workers[i].request.mutable_server_def()->set_protocol(*options.protocol);
  73.       workers[i].request.mutable_server_def()->set_job_name(name.job);
  74.       workers[i].request.mutable_server_def()->set_task_index(name.task);
  75.       // Session state is always isolated when ClusterSpec propagation
  76.       // is in use.
  77.       workers[i].request.set_isolate_session_state(true);
  78.     } else {
  79.       // NOTE(mrry): Do not set any component of the ServerDef,
  80.       // because the worker will use its local configuration.
  81.       workers[i].request.set_isolate_session_state(
  82.           session_opts_.config.isolate_session_state());
  83.     }
  84.     if (session_opts_.config.experimental()
  85.             .share_session_state_in_clusterspec_propagation()) {
  86.       // In a dynamic cluster, the ClusterSpec info is usually propagated by
  87.       // master sessions. However, in data parallel training with multiple
  88.       // masters
  89.       // ("between-graph replication"), we need to disable isolation for
  90.       // different worker sessions to update the same variables in PS tasks.
  91.       workers[i].request.set_isolate_session_state(false);
  92.     }
  93.   }
  95.   for (size_t i = 0; i < worker_names.size(); ++i) {
  96.     auto cb = [i, &workers, &done](const Status& s) {
  97.       workers[i].status = s;
  98.       done.DecrementCount();
  99.     };
  100.     workers[i].worker->CreateWorkerSessionAsync(&workers[i].request,
  101.                                                 &workers[i].response, cb);
  102.   }
  104.   done.Wait();
  105.   for (size_t i = 0; i < workers.size(); ++i) {
  106.     status.Update(workers[i].status);
  107.   }
  108.   return status;
  109. }
GrpcRemoteWorker

GrpcRemoteWorker 是 gRPC 的客户端，通过 stub 调用远端 WorkerService 相应的服务接口。

  1. void CreateWorkerSessionAsync(const CreateWorkerSessionRequest* request,
  2.                               CreateWorkerSessionResponse* response,
  3.                               StatusCallback done) override {
  4.   IssueRequest(request, response, createworkersession_, std::move(done));
  5. }
GrpcWorkerService

远端 Worker 之中，接收到消息是在 GrpcWorkerService 之中，当收到 CreateWorkerSessionRequest 消息，将 由 CreateWorkerSessionHandler 回调处理，CreateWorkerSessionHandler 是一个宏，其在线程池中启动一个可运行的线程，触发 Worker（就是GrpcWorker） 的 CreateWorkerSession 方法来动态创建 WorkerSession 实例。

  1. #define HANDLE_CALL(method, may_block_on_compute_pool)                        \
  2.   void method##Handler(WorkerCall<method##Request, method##Response>* call) { \
  3.     auto closure = [this, call]() {                                           \
  4.       Status s = worker_->method(&call->request, &call->response);            \
  5.       if (!s.ok()) {                                                          \
  6.         VLOG(3) << "Bad response from " << #method << ": " << s;              \
  7.       }                                                                       \
  8.       call->SendResponse(ToGrpcStatus(s));                                    \
  9.     };                                                                        \
  10.     if ((may_block_on_compute_pool)) {                                        \
  11.       worker_->env()->env->SchedClosure(std::move(closure));                  \
  12.     } else {                                                                  \
  13.       worker_->env()->compute_pool->Schedule(std::move(closure));             \
  14.     }                                                                         \
  15.     ENQUEUE_REQUEST(method, false);                                           \
  16.   }
  18.   HANDLE_CALL(CreateWorkerSession, false);

4. WorkerSession

其实，GrpcWorker 最终调用的是 WorkerInterface.CreateWorkerSession 方法。

  1. Status CreateWorkerSession(const CreateWorkerSessionRequest* request,
  2.                            CreateWorkerSessionResponse* response) {
  3.   return CallAndWait(&ME::CreateWorkerSessionAsync, request, response);
  4. }

CreateWorkerSessionRequest 消息之中携带了 MasterSession 分配的 session_handle，GrpcWorker 将据此创建一个 WorkerSession，session_handle 在这个 Worker 之内唯一标识这个 WorkerSession。

在 GrpcWorker 的 WorkerEnv 上下文之中有一个 SessionMgr，SessionMgr 负责统一管理和维护所有的 WorkerSession 生命周期。SessionMgr 与 WorkerSession 是一对多的关系，每个 WorkerSession 实例使用 session_handle 标识。

  1. void Worker::CreateWorkerSessionAsync(const CreateWorkerSessionRequest* request,
  2.                                       CreateWorkerSessionResponse* response,
  3.                                       StatusCallback done) {
  4.   Status s = env_->session_mgr->CreateSession(
  5.       request->session_handle(), request->server_def(),
  6.       request->cluster_device_attributes(), request->isolate_session_state(),
  7.       request->master_task(), request->master_incarnation());
  8.   done(s);
  9. }

4.1 SessionMgr

4.1.1 定义

重点是如下，维护了 session_handle 和 WorkerSession 之间的对应关系，每个 WorkerSession 由 session_handle 来标识。

  • std::map<string, std::shared_ptr> sessions_ ：维护了对应关系。

  • std::shared_ptr legacy_session_ ：本地 WorkerSession 实例。

图 4 SessionMgr

  1. class SessionMgr {
  2.  public:
  3.   typedef std::function<Status(const ServerDef&, WorkerCacheInterface**)>
  4.       WorkerCacheFactory;
  6.   explicit SessionMgr(
  7.       WorkerEnv* worker_env, const string& default_worker_name,
  8.       std::unique_ptr<WorkerCacheInterface> default_worker_cache,
  9.       WorkerCacheFactory worker_cache_factory);
  10.   ~SessionMgr() {}
  12.   // Allocates state for a new session.
  13.   Status CreateSession(const string& session, const ServerDef& server_def,
  14.                        bool isolate_session_state);
  15.   Status CreateSession(
  16.       const string& session, const ServerDef& server_def,
  17.       const protobuf::RepeatedPtrField<DeviceAttributes>& device_attributes,
  18.       bool isolate_session_state);
  20.   // Create WorkerSession from the master with the given `master_task` and
  21.   // `master_incarnation`. We first look for existing WorkerSessions associated
  22.   // with the specified master task. If there are sessions created by the same
  23.   // master but with a different incarnation, it indicates that the remote
  24.   // master has restarted before deleting the sessions on worker. When it
  25.   // happens, old sessions associated with the master will be automatically
  26.   // removed before the new session is created.
  27.   Status CreateSession(
  28.       const string& session, const ServerDef& server_def,
  29.       const protobuf::RepeatedPtrField<DeviceAttributes>& device_attributes,
  30.       bool isolate_session_state, string master_task,
  31.       int64_t master_incarnation);
  33.   void ResetDefaultWorkerCache(WorkerCacheInterface* worker_cache);
  35.   // Updates state (worker cache, devices) of worker session identified by
  36.   // session name (`session`) based on a new server_def and set of devices.
  37.   Status UpdateSession(const string& session, const ServerDef& server_def,
  38.                        const protobuf::RepeatedPtrField<DeviceAttributes>&
  39.                            cluster_device_attributes,
  40.                        bool isolate_session_state);
  42.   // Locates the worker session for a given session handle
  43.   Status WorkerSessionForSession(const string& session_handle,
  44.                                  std::shared_ptr<WorkerSession>* out_session);
  45.   std::shared_ptr<WorkerSession> LegacySession();
  47.   Status DeleteSession(const string& session);
  49.   static string WorkerNameFromServerDef(const ServerDef& server_def);
  51.   void SetLogging(bool active);
  53.   void RetrieveLogs(int64_t step_id, LoggingResponse* response);
  55.   void ClearLogs();
  57.  private:
  58.   WorkerEnv* const worker_env_;  // Not owned.
  60.   // A note about destruction:
  61.   // We must delete graph_mgr before device_mgr, due to shared
  62.   // ownership of OpKernels in the executors. (The graph_mgr will
  63.   // free all stateless OpKernels, and pass over borrowed stateful
  64.   // OpKernels, which are also held in their respective devices'
  65.   // OpSegments.)
  66.   //
  67.   // legacy_session_ owns the worker_env_.device_mgr, and so we must ensure
  68.   // that sessions_'s WorkerSessions are deleted (which do not own the
  69.   // underlying devices, but instead own RenamedDevices) before
  70.   // legacy_session_ is deleted. Further, we must ensure that WorkerSession's
  71.   // device_mgr is deleted after WorkerSession's graph_mgr.
  73.   std::unique_ptr<WorkerCacheInterface> default_worker_cache_;
  74.   std::shared_ptr<WorkerSession> legacy_session_;
  76.   bool is_logging_active_ = false;
  78.   const WorkerCacheFactory worker_cache_factory_;
  80.   Status WorkerSessionForSessionLocked(
  81.       const string& session_handle, std::shared_ptr<WorkerSession>* out_session)
  82.       TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
  84.   mutex mu_;
  85.   // A map from session identifier to internal session structure.
  86.   std::map<string, std::shared_ptr<WorkerSession>> sessions_ TF_GUARDED_BY(mu_);
  88.   // Incarnation and WorkerSession handle associated with a master task.
  89.   struct MasterAssociatedSession {
  90.     const int64_t master_incarnation;
  91.     const string session_handle;
  92.   };
  93.   // A map from master task name to its associated worker sessions.
  94.   std::unordered_multimap<string, MasterAssociatedSession>
  95.       master_to_associated_sessions_ TF_GUARDED_BY(mu_);
  96. };

4.1.2 建立 Session

CreateSession 方法会创建 WorkerSession 和 GraphMgr。

  1. Status SessionMgr::CreateSession(
  2.     const string& session, const ServerDef& server_def,
  3.     const protobuf::RepeatedPtrField<DeviceAttributes>&
  4.         cluster_device_attributes,
  5.     bool isolate_session_state, string master_task,
  6.     int64_t master_incarnation) {
  7.   mutex_lock l(mu_);
  8.   if (session.empty()) {
  9.     return errors::InvalidArgument("Session must be non-empty.");
  10.   }
  12.   // For given master task name, check if one or more `WorkerSession`s have been
  13.   // created previously on this worker, and if so garbage collect the expired
  14.   // `WorkerSession`s. This happens when the master fails before sending
  15.   // `DeleteSession` requests, which can cause `WorkerSession`s to be leaked.
  16.   if (!master_task.empty()) {
  17.     auto it_range = master_to_associated_sessions_.equal_range(master_task);
  18.     if (it_range.first != it_range.second &&
  19.         it_range.first->second.master_incarnation != master_incarnation) {
  20.       auto it = it_range.first;
  21.       while (it != it_range.second) {
  22.         auto session_it = sessions_.find(it->second.session_handle);
  23.         if (session_it != sessions_.end()) {
  24.           sessions_.erase(session_it);
  25.         }
  26.         it = master_to_associated_sessions_.erase(it);
  27.       }
  28.     }
  29.   }
  31.   WorkerCacheInterface* worker_cache = nullptr;
  32.   string worker_name;
  33.   if (server_def.cluster().job().empty()) {
  34.     worker_cache = new WorkerCacheWrapper(default_worker_cache_.get());
  35.     worker_name = legacy_session_->worker_name();
  36.   } else {
  37.     TF_RETURN_IF_ERROR(worker_cache_factory_(server_def, &worker_cache));
  38.     worker_name = WorkerNameFromServerDef(server_def);
  39.   }
  41.   if (worker_cache != nullptr && default_worker_cache_ != nullptr) {
  42.     worker_cache->SetLogging(this->is_logging_active_);
  43.   }
  45.   std::shared_ptr<WorkerSession> worker_session;
  46.   std::vector<std::unique_ptr<Device>> cluster_devices;
  48.   if (isolate_session_state || server_def.cluster().job_size()) {
  50.     // Create a private copy of the DeviceMgr for the WorkerSession.
  51.     std::vector<std::unique_ptr<Device>> renamed_devices;
  52.     for (Device* d : worker_env_->local_devices) {
  53.       renamed_devices.push_back(RenamedDevice::NewRenamedDevice(
  54.           worker_name, d, false, isolate_session_state));
  55.     }
  56.     auto device_mgr = MakeUnique<StaticDeviceMgr>(std::move(renamed_devices));
  57.     LookupLocalDevice cb = [&device_mgr](StringPiece name, Device** device) {
  58.       return device_mgr->LookupDevice(name, device);
  59.     };
  60.     AsRemoteDevices(worker_env_->env, cluster_device_attributes, cb,
  61.                     &cluster_devices);
  62.     std::unique_ptr<DynamicDeviceMgr> remote_devices;
  63.     if (!cluster_device_attributes.empty()) {
  64.       remote_devices = MakeUnique<DynamicDeviceMgr>();
  65.       TF_RETURN_IF_ERROR(
  66.           remote_devices->AddDevices(std::move(cluster_devices)));
  67.     }
  69.     auto graph_mgr = MakeUnique<GraphMgr>(worker_env_, device_mgr.get());
  70.     worker_session.reset(
  71.         new WorkerSession(session, worker_name,
  72.                           std::unique_ptr<WorkerCacheInterface>(worker_cache),
  73.                           std::move(device_mgr), std::move(graph_mgr),
  74.                           std::move(remote_devices)));
  75.   } else {
  76.     AsRemoteDevices(worker_env_->env, cluster_device_attributes, nullptr,
  77.                     &cluster_devices);
  78.     std::unique_ptr<DynamicDeviceMgr> remote_devices;
  79.     if (!cluster_device_attributes.empty()) {
  80.       remote_devices = MakeUnique<DynamicDeviceMgr>();
  81.       TF_RETURN_IF_ERROR(
  82.           remote_devices->AddDevices(std::move(cluster_devices)));
  83.     }
  84.     // Borrow the WorkerEnv's DeviceMgr for the WorkerSession, so
  85.     // that resources using it can use its devices after the
  86.     // WorkerSession has been deleted.
  87.     auto graph_mgr = MakeUnique<GraphMgr>(worker_env_, worker_env_->device_mgr);
  88.     worker_session = WorkerSession::CreateWithBorrowedDeviceMgr(
  89.         session, worker_name,
  90.         std::unique_ptr<WorkerCacheInterface>(worker_cache),
  91.         worker_env_->device_mgr, std::move(graph_mgr),
  92.         std::move(remote_devices));
  93.   }
  95.   sessions_.insert(std::make_pair(session, std::move(worker_session)));
  96.   if (!master_task.empty()) {
  97.     MasterAssociatedSession s{master_incarnation, session};
  98.     master_to_associated_sessions_.emplace(master_task, s);
  99.   }
  100.   return Status::OK();
  101. }

4.1.3 注册图

我们用 RegisterGraphAsync 为例来看看 worker 内部功能。可以看到其使用 GraphMgr 完成了基础功能。

  1. void Worker::RegisterGraphAsync(const RegisterGraphRequest* request,
  2.                                 RegisterGraphResponse* response,
  3.                                 StatusCallback done) {
  4.   std::shared_ptr<WorkerSession> session;
  5.   Status s;
  6.   if (request->create_worker_session_called()) {
  7.     s = env_->session_mgr->WorkerSessionForSession(request->session_handle(),
  8.                                                    &session);
  9.   } else {
  10.     session = env_->session_mgr->LegacySession();
  11.   }
  12.   if (s.ok()) {
  13.     s = session->graph_mgr()->Register(
  14.         request->session_handle(), request->graph_def(), session.get(),
  15.         request->graph_options(), request->debug_options(),
  16.         request->config_proto(), request->collective_graph_key(),
  17.         session->cluster_flr(), response->mutable_graph_handle());
  18.   }
  19.   done(s);
  20. }

4.2 WorkerSession

4.2.1 定义

WorkerSession 之中比较重要的几个成员变量包括几个管理类 GraphMgr，DeviceMgr，DynamicDeviceMgr：

  • string session_name_ ：Session 名称。

  • string worker_name_ ：Worker 名称，比如 /job:mnist/replica:0/task:1。

  • std::shared_ptr worker_cache_ ：Worker 缓存。

  • std::unique_ptr graph_mgr_ ：本 session 注册的计算图，每个 Worker 可以注册和运行多个计算图，每个计算图使用 graph）handle 标识。

  • std::unique_ptr device_mgr_ ：本地计算设备集合信息。

图 5 WorkerSession 概念

  1. // WorkerSession encapsulates all of the state relating to a given session.
  2. class WorkerSession {
  3.  public:
  4.   // Collection of local devices. These devices are typically
  5.   // RenamedDevices in all except the SessionMgr.legacy_session_ and
  6.   // sessions created with `isolate_session_state == false`. In the
  7.   // those cases, this method returns a pointer to a borrowed
  8.   // DeviceMgr (typically the `worker_env.device_mgr`).
  9.   DeviceMgr* device_mgr() {
  10.     return device_mgr_ ? device_mgr_.get() : borrowed_device_mgr_;
  11.   }
  13.   DynamicDeviceMgr* remote_device_mgr() { return remote_device_mgr_.get(); }
  15.   const string& session_name() const { return session_name_; }
  16.   const string& worker_name() const { return worker_name_; }
  18.   WorkerCacheInterface* worker_cache() const {
  19.     tf_shared_lock l(worker_session_state_mu_);
  20.     return worker_cache_.get();
  21.   }
  22.   GraphMgr* graph_mgr() const { return graph_mgr_.get(); }
  24.   ClusterFunctionLibraryRuntime* cluster_flr() const {
  25.     return cluster_flr_.get();
  26.   }
  28.   WorkerSession(const string& session_name, const string& worker_name,
  29.                 std::unique_ptr<WorkerCacheInterface> worker_cache,
  30.                 std::unique_ptr<DeviceMgr> device_mgr,
  31.                 std::unique_ptr<GraphMgr> graph_mgr,
  32.                 std::unique_ptr<DynamicDeviceMgr> remote_device_mgr);
  34.   static std::shared_ptr<WorkerSession> CreateWithBorrowedDeviceMgr(
  35.       const string& session_name, const string& worker_name,
  36.       std::unique_ptr<WorkerCacheInterface> worker_cache,
  37.       DeviceMgr* borrowed_device_mgr, std::unique_ptr<GraphMgr> graph_mgr,
  38.       std::unique_ptr<DynamicDeviceMgr> remote_device_mgr);
  40.   // In the eager runtime we allow WorkerSession to be updated, where the
  41.   // worker cache will be recreated. If WorkerSession upate is expected and a
  42.   // worker in the cache is used in RPCs, the caller should hold a shared
  43.   // pointer to avoid the workers getting deleted.
  44.   std::shared_ptr<WorkerCacheInterface> GetSharedWorkerCache() {
  45.     tf_shared_lock l(worker_session_state_mu_);
  46.     return worker_cache_;
  47.   }
  49.   // Update an existing worker session with new set of remote workers and
  50.   // devices. Added devices will be owned by the worker session, and removed
  51.   // devices will be freed by their names.
  52.   Status UpdateWorkerCacheAndDevices(
  53.       std::unique_ptr<WorkerCacheInterface> new_worker_cache,
  54.       std::vector<std::unique_ptr<Device>> added_remote_devices,
  55.       const std::vector<Device*>& removed_remote_devices);
  57.   ~WorkerSession();
  59.  private:
  60.   WorkerSession(const string& session_name, const string& worker_name,
  61.                 std::unique_ptr<WorkerCacheInterface> worker_cache,
  62.                 DeviceMgr* borrowed_device_mgr,
  63.                 std::unique_ptr<GraphMgr> graph_mgr,
  64.                 std::unique_ptr<DynamicDeviceMgr> remote_device_mgr);
  66.   // The name of the session.
  67.   const string session_name_;
  69.   // The name of the worker. E.g., /job:mnist/replica:0/task:1.
  70.   const string worker_name_;
  72.   mutable mutex worker_session_state_mu_;
  73.   // Object from which WorkerInterface instances can be obtained.
  74.   std::shared_ptr<WorkerCacheInterface> worker_cache_
  75.       TF_GUARDED_BY(worker_session_state_mu_);
  77.   // graph_mgr keeps track of the registered graphs of this session.
  78.   //
  79.   // Note: graph_mgr must be deleted before rendezvous_mgr!
  80.   // Note: graph_mgr must be deleted before device_mgr!
  81.   const std::unique_ptr<GraphMgr> graph_mgr_;
  83.   std::unique_ptr<ClusterFunctionLibraryRuntime> cluster_flr_;
  85.   const std::unique_ptr<DeviceMgr> device_mgr_;
  86.   DeviceMgr* const borrowed_device_mgr_;  // Not owned.
  87.   std::unique_ptr<DynamicDeviceMgr> remote_device_mgr_;
  88. };

至此，session 基本流程我们梳理完成，下面就会对业务进行详细分析。

0xFF 参考

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

[源码解析] TensorFlow 分布式环境(5) --- Session的更多相关文章

  1. [源码解析] TensorFlow 分布式环境(6) --- Master 动态逻辑

    [源码解析] TensorFlow 分布式环境(6) --- Master 动态逻辑 目录 [源码解析] TensorFlow 分布式环境(6) --- Master 动态逻辑 1. GrpcSess ...

  2. [源码解析] TensorFlow 分布式环境(7) --- Worker 动态逻辑

    [源码解析] TensorFlow 分布式环境(7) --- Worker 动态逻辑 目录 [源码解析] TensorFlow 分布式环境(7) --- Worker 动态逻辑 1. 概述 1.1 温 ...

  3. [源码解析] TensorFlow 分布式环境(8) --- 通信机制

    [源码解析] TensorFlow 分布式环境(8) --- 通信机制 目录 [源码解析] TensorFlow 分布式环境(8) --- 通信机制 1. 机制 1.1 消息标识符 1.1.1 定义 ...

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

    [源码解析] TensorFlow 分布式环境(1) --- 总体架构 目录 [源码解析] TensorFlow 分布式环境(1) --- 总体架构 1. 总体架构 1.1 集群角度 1.1.1 概念 ...

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

    [源码解析] TensorFlow 分布式环境(2)---Master 静态逻辑 目录 [源码解析] TensorFlow 分布式环境(2)---Master 静态逻辑 1. 总述 2. 接口 2.1 ...

  6. [源码解析] TensorFlow 分布式环境(3)--- Worker 静态逻辑

    [源码解析] TensorFlow 分布式环境(3)--- Worker 静态逻辑 目录 [源码解析] TensorFlow 分布式环境(3)--- Worker 静态逻辑 1. 继承关系 1.1 角 ...

  7. [源码解析] TensorFlow 分布式环境(4) --- WorkerCache

    [源码解析] TensorFlow 分布式环境(4) --- WorkerCache 目录 [源码解析] TensorFlow 分布式环境(4) --- WorkerCache 1. WorkerCa ...

  8. [源码解析] TensorFlow 分布式 DistributedStrategy 之基础篇

    [源码解析] TensorFlow 分布式 DistributedStrategy 之基础篇 目录 [源码解析] TensorFlow 分布式 DistributedStrategy 之基础篇 1. ...

  9. [源码解析] TensorFlow 分布式之 MirroredStrategy

    [源码解析] TensorFlow 分布式之 MirroredStrategy 目录 [源码解析] TensorFlow 分布式之 MirroredStrategy 1. 设计&思路 1.1 ...

随机推荐

  1. No compiler is provided in this environment. Perhaps you are running on a JRE rather than a JDK 问题解决

    1. 问题描述 使用idea对Java工程执行mvn compile命令进行编译,出现以下报错: [ERROR] Failed to execute goal org.apache.maven.plu ...

  2. IDEA中使用Docker

    开发环境 IDEA:2020.3.2 Docker:20.10.12 注意,如果没有开启Docker远程连接,请先开启Docker远程连接. 1. 打开或新建一个Web项目 可参考使用IDEA新建一个 ...

  3. 使用传统的三层架构出现的问题.引入Spring底层实现原理来解决(工厂模式+反射+XML配置文件/注解)

    以前写的代码 mapper层 public interface PersonMapper { void selectPersonList(); } public class PersonMapperI ...

  4. 微服务从代码到k8s部署应有尽有系列（八、各种队列）

    我们用一个系列来讲解从需求到上线.从代码到k8s部署.从日志到监控等各个方面的微服务完整实践. 整个项目使用了go-zero开发的微服务,基本包含了go-zero以及相关go-zero作者开发的一些中 ...

  5. java-23种设计模式概述【软件设计模式基本介绍（是什么、作用、优点）、模式的分类和介绍】

    一.设计模式基本介绍(是什么.作用.优点) 1.软件设计模式是什么? 软件设计模式(Software Design Pattern),又称设计模式. 2.设计模式的作用 ★ 提高代码的可复用性.可维护 ...

  6. Smartbi权限安全管理系统_保障数据权限安全

    思迈特软件Smartbi具有完善的安全管理体系,Smartbi权限安全管理系统它可以控制用户功能权限.数据访问权限.资源访问权限.Smartbi权限安全管理系统支持按用户.用户组.角色进行管理:支持多 ...

  7. Java的诞生历史

    Java帝国的诞生 1.1972年C诞生 贴近硬件,运行极快,效率极高 操作系统,编译器,数据库,网络系统等都采用C语言 但是,它的指针和内存管理给程序员安上了"枷锁".它的指针没 ...

  8. 3rd S-curve velocity profile

    3rd S-curve  (1) (2) (3) (4) 完整的三次S曲线包括上面的七个阶段.前面三个阶段为加速阶段,从初始速度Vs加速到Vmax: (5) 整个加速阶段的位移为: (6) 后面三个阶 ...

  9. Qt：Shadow Build

    每个编辑器有Build和Run两个设置界面. 在Build界面上,有一个"Shadow build"复选框.如果勾选此项,编译后将在项目的同级目录下建立一个编译后的文件目录,目录名 ...

  10. Python的内置数据结构

    Python内置数据结构一共有6类: 数字 字符串 列表 元组 字典 文件 一.数字 数字类型就没什么好说的了,大家自行理解 二.字符串 1.字符串的特性(重要): 序列化特性:字符串具有一个很重要的 ...