tensorflow源码解析之distributed_runtime

本篇主要介绍TF的分布式运行时的基本概念。为了对TF的分布式运行机制有一个大致的了解，我们先结合/tensorflow/core/protobuf中的文件给出对TF分布式集群的初步理解，然后介绍/tensorflow/core/distributed_runtime路径下的核心概念。

TF分布式集群

集群定义和理解

在研读TF的分布式运行时代码之前，我们需要先看下TF分布式运行的基本架构。TF的集群（cluster）由作业（job）构成，作业由任务（task）构成。举个例子，一个由两个作业构成的集群，作业1名为“worker”，包含了3个任务，作业2名为“ps”，包含了2个任务，如下：

Cluster:
    job { name:'worker'
            tasks {key:0 value:'worker1:2222'}
            tasks {key:1 value:'worker2:2222'}
            tasks {key:2 value:'worker3:2222'}
    }
    job { name:'ps'
            tasks {key:0 value:'ps0:2222'}
            tasks {key:1 value:'ps1:2222'}
    }

下面再看TF对于集群的定义，就一目了然了：

message JobDef {
    string name = 1;//作业的名称
    
    //作业包含的任务id到hostname:port字符串的映射，也就是任务的编号到任务的数据传输接口
    map<int32, string> tasks = 2;
}

message ClusterDef {
    repeated JobDef job = 1;
}

以下我们会分别介绍Master服务和Worker服务，注意，Master服务是由Master提供，供客户端使用的，而Worker服务是由Worker提供，供Master使用的。

master

先来讲Master服务。Master服务是一种被客户端用来与分布式的TF计算交互的服务。

一个Master服务通常会包含了多个master会话，每一个会话包含了一张计算图以及与之相关的状态，这些master会话通常会对应同一个client会话。

一个Master会话的职责包括：

节点放置；
插入恰当的节点以实现跨设备和跨进程的数据流和资源管理；
发布命令给worker，使之运行分配给它的计算子图；

通常，客户端可以通过RPC的形式与一个Master之间保持一个交互式的计算。客户端首先建立一个客户端的会话，连接到一个特定的Master，这个Master接着创建一个对应的Master会话，并且在客户端的调用之间维持状态。

Master会话创建之后，Master会返回一个句柄给客户端，这个句柄可以被用来进行客户端和Master会话之间的交互。

客户端可以在CreateSession调用中传递一个初始的图给Master，并且使用ExtendSession向图中添加节点。

对于一个Master来说，最常用的操作是RunStep，它实现了一个Session::Run()的API。它支持提供输入，执行图计算，返回输出。

最后，当客户端不再需要Master会话的时候，它需要通过CloseSession关闭这个会话，Master可以回收跟会话相关的资源。Master在关闭会话期间可以会因为垃圾回收而休眠一段时间。

我们来总结下MasterService包含的内容：

service MasterService {
    rpc CreateSession(CreateSessionRequest) returns (CreateSessionResponse);
    rpc ExtendSession(ExtendSessionRequest) returns (ExtendSessionResponse);
    rpc PartialRunStep(PartialRunStepRequest) returns (PartialRunSetupResponse);
    rpc RunStep(RunStepRequest) returns (RunStepResponse);
    rpc CloseSession(CloseSessionRequest) returns (CloseSessionResponse);
    rpc ListDevices(ListDevicesRequest) returns (ListDeviceResponse);
    rpc Reset(ResetRequest) returns ( ResetResponse);
}

代码中提到的xxxRequest和xxxResponse，都有对应的结构，详见/tensorflow/core/protobuf/master.proto。

woker

Worker服务定义了一种TF的服务，它可以代表MasterService，在一些局部的设备上执行数据流图。

一个Worker服务保留了多个注册图，每一个注册图都是客户端完整图的一个子图，包含了仅需要在当前worker上计算的节点。

service WorkerService {
    rpc GetStatus(GetStatusRequest) returns (GetStatusResponse);
    rpc CreateWorkerSession(CreateWorkerSessionRequest) returns (CreateWorkerSessionResponse);
    rpc RegisterGraph(RegisterGraphRequest) returns (RegisterGraphResponse);
    rpc DeregisterGraph(DeregisterGraphRequest) returns (DeregisterGraphResponse);
    rpc RunGraph(RunGraphRequest) returns (RunGraphResponse);
    rpc CleanupGraph(CleanupGraphRequest) returns (CleanupGraphResponse);
    rpc CleanupAll(CleanupAllRequest) returns (CleanupAllResponse);
    rpc RecvTensor(RecvTensorRequest) returns (RecvTensorResponse) {}
    rpc Logging(LoggingRequest) returns (LoggingResponse);
    rpc Tracing(TracingRequest) returns (TracingResponse);
}

以上内容来自/tensorflow/core/protobuf，主要为了讲解TF中集群的基本概念和运行过程，以下内容来自/tensorflow/core/distributed_runtime，介绍TF中分布式运行时环境中的核心概念。

worker

Worker代表了执行计算的实体，与Client和Master相对应。以下是相关类的关系图：

tensor_coding

包含了TensorResponse类，这个类的作用是，当一个RPC返回了数据时，通过这个类可以把返回结果中的数据解析为张量，以及其它的元数据信息。

session_mgr

包含了SessionMgr类，它存在于Worker上，为Worker管理会话，包括了会话的产生和销毁，同时还维护了一个当前Worker上的会话句柄到会话的映射。

class SessionMgr {
  public:
    Status CreateSession(...);
    Status DeleteSession(...);
  private:
    const WorkerEnv* const worker_env_;
    const WorkerCacheFactory worker_cache_factory_;
    std::map<string, std::unique_ptr<WorkerSession>> sessions_ GUARDED_BY(mu_);
};

server_lib

TF中的server，可以表现为两种形式，一种是Worker，一种是Master，可以认为，两者都是对外提供了“服务”，只不过是两种不同的形式。ServerInterface为它们提供了统一的接口：

class ServerInterface {
  public:
    virtual Status Start() = 0;
    virtual Status Stop() = 0;
    virtual Status Join() = 0;
};

而所有的Server必须由其对应的工厂类产生，工厂类还提供了对其子类的注册接口：

class ServerFactory {
  public:
    virtual Status NewServer(...);
    
    //任何一个工厂类的子类，都必须用这个方法将其一个对象注册到这里
    static void Register(const string& server_type, ServerFactory* factory);
    
    //根据server_def，寻找一个能产生指定server的工厂
    static Status GetFactory(const ServerDef& server_def, ServerFactory** out_factory);
};

scheduler

根据Graph和CostModel的信息，计算不同调度策略下，每个节点的最早开始时间和最晚开始时间，三个类SlackAnalysis，GreedyScheduler，PriorityScheduler分别代表了松弛策略、贪心调度策略和优先级调度策略。

rendezvous_mgr_interface

类RendezvousMgr管理着一个局部rendezvous对象的集合。所有被当前的Worker发送的张量，在接收之前都在这个RendezvousMgr中保存着。每一个全局的step_id都对应着一个被RendezvousMgr管理的一个局部的rendezvous实例。

remote_device

包含了一个函数，NewRemoteDevices，它可以发现remote_worker上的可用设备。

partial_run_mgr

PartialRunMgr保存了未完成的局部运行的需求，它保证只有当对应的执行器完成运行时，它才会被标记为完成。

在TF的worker中，执行器会异步的执行，直到需求的输出（能够返回张量的操作）或者目标（不会返回张量的操作）完成。也就是说，计算图中有两类节点都可以作为worker执行的目标，一类是返回张量的操作对应的节点，一类是不返张量的操作对应的节点。一个局部运行包含两步，第一，设置所有需要的输出和目标，第二，获得输出。在第二步时，可能存在一种情况，即计算图中需求的输出已经计算完成，但需求的目标仍在计算。这时候，PartialRunMgr就发挥作用了，虽然这时理论上可以返回了，因为所有需求的输出都计算完成了，剩余的需求目标并不影响返回的结果。但TF仍然要求必须等到所有的目标都完成计算才行，因为在目标完成计算之前，我们并不知道中间的输出是否会发生变化。

message_wrappers

在Master和Worker之间相互通信的Request/Response的包装类。

// Wrapper classes for the `MasterService.RunStep` request message.
class RunStepRequestWrapper {}
class MutableRunStepRequestWrapper : public RunStepRequestWrapper {}
class InMemoryRunStepRequest : public MutableRunStepRequestWrapper {}
class MutableProtoRunStepRequest : public MutableRunStepRequestWrapper {}
class ProtoRunStepRequest : public RunStepRequestWrapper {}

// Wrapper classes for the `WorkerService.RunGraph` request message.
class RunGraphRequestWrapper {}
class MutableRunGraphRequestWrapper : public RunGraphRequestWrapper {}
class InMemoryRunGraphRequest : public MutableRunGraphRequestWrapper {}
class MutableProtoRunGraphRequest : public MutableRunGraphRequestWrapper {}
class ProtoRunGraphRequest : public RunGraphRequestWrapper {}

// Wrapper classes for the `WorkerService.RunGraph` response message.
class MutableRunGraphResponseWrapper {}
class InMemoryRunGraphResponse : public MutableRunGraphResponseWrapper {}
class OwnedProtoRunGraphResponse : public MutableRunGraphResponseWrapper {}
class NonOwnedProtoRunGraphResponse : public MutableRunGraphResponseWrapper {}

// Wrapper classes for the `MasterService.RunStep` response message.
class MutableRunStepResponseWrapper {}
class InMemoryRunStepResponse : public MutableRunStepResponseWrapper {}
class OwnedProtoRunStepResponse : public MutableRunStepResponseWrapper {}
class NonOwnedProtoRunStepResponse : public MutableRunStepResponseWrapper {}

master_session

与单机情况下的DirectSession对应的，分布式情况下的Master会话，它包含了图计算的基本步骤，比如资源分配、节点放置、图执行等。

master_interface

用于与TF的Master服务通信的虚拟接口。这个接口既支持基于RPC的master实现，也支持进程内部的master实现。

master

TF中Master服务的实现。与Worker服务对应。

master_env

Master的环境类，包含了一个Master所必须的环境资源指针。注意Master并不拥有这些指针。

local_master

局部Master的实现。局部Master的含义是，与Client的通信不是跨设备的，而是直接在进程内部进行的。这个Master的实现，是为了给同进程内部的Client提供更高效的Master服务。

graph_mgr

GraphMgr包含了注册到某个worker的图的集合。每一个注册的图都会被一个句柄标识，这个句柄由GraphMgr产生，并且返回给调用者。在注册成功之后，调用者通过一个图句柄来执行一张图。每一次的执行都被一个全局的"step_id"唯一标识。在同一张图上，可以重复和独立的执行多次，只要每一次执行的"step_id"都是不同的。

call_options

为不同的RPC系统提供了可插拔的调用接口。

base_rendezvous_mgr

为RendezvousMgrInterface提供了不同的实现，具体框架图如下：

graph TB RendezvousMgrInterface-->|派生|BaseRendezvousMgr RemoteRendezvous-->|派生|BaseRemoteRendezvous Rendezvous-->|派生|RemoteRendezvous

github地址

posted on 2018-09-17 23:35 jicanghai 阅读(2300) 评论(0) 编辑收藏举报

刷新页面返回顶部

jicanghai