tensorflow源码解析之distributed_runtime

本文主要介绍TF的分布式运行时的基本概念。为了对TF的分布式运行机制有一个大致的了解，我们先根据//core/中的文件对TF的分布式集群做一个初步的认识，然后在//core//core路径下介绍核心概念/。

TF分布式集群集群定义及理解

在研究TF的分布式运行代码之前，我们需要先了解一下TF分布式运行的基本架构。TF的()由jobs组成，jobs由tasks组成。例如一个集群由两个job组成，job 1命名为“”，包含3个task，job 2命名为“ps”，包含2个task，如下：

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 tasks = 2;
}
message ClusterDef {
    repeated JobDef job = 1;
}

下面我们将分别介绍服务和服务。请注意，服务是由客户提供和使用的，而服务是由客户提供和使用的。

先说服务吧。服务是客户端用来与分布式 TF 计算交互的服务。

一个服务通常包含多个会话，每个会话包含一个计算图及其相关状态，这些会话通常对应同一个会话。

会话的职责包括：

通常，客户端可以通过 RPC 维护一个交互式计算。客户端首先建立一个客户端会话，连接到一个特定的会话，该会话反过来创建一个相应的会话并维护客户端调用之间的状态。

会话创建后，会返回一个句柄给客户端，该句柄可用于与会话交互。

客户端可以将初始图传递给调用，并使用 .

一方面，最常用的操作是它实现了一个 ::Run() API。它支持提供输入、执行图形计算和返回输出。

最后，当客户端不再需要会话时，它需要通过关闭会话来回收与会话相关的资源。由于关闭会话期间的垃圾收集，可以休眠一段时间。

让我们总结一下包含的内容：

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);
}

代码中提到的sum有对应的结构，详见//core//.proto。

沃克

服务定义了一个 TF 服务，可以在某个本地设备上表示和执行数据流图。

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

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);
}

以上内容来自//core/，主要讲解TF中聚类的基本概念和操作流程。以下内容来自//core/，介绍了TF中分布式运行环境中的核心概念。

表示执行计算的实体，对应于 和 。下面是相关类的示意图：

图 ce–>|用于生成|–>|用于生成|–>|派生|–>|派生|ger–>|用于日志服务|–>–>

包含类。该类的作用是当一个RPC返回数据时，可以通过该类将返回结果中的数据解析成张量等元数据信息。

包含存在于 上的类，用于管理会话，包括会话的生成和销毁，并且还维护当前会话句柄到会话的映射。

class SessionMgr {
  public:
    Status CreateSession(...);

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

在TF中，可以表示成两种形式，一种是，另一种是，可以认为两者都对外提供“服务”，但只是两种不同的形式。为它们提供统一的接口：

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

并且都必须由其对应的工厂类生成。工厂类还为其子类提供了注册接口：

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);
};

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

脸

一个类管理本地对象的集合。所有当前发送的张量在接收之前都存储在其中。每个全局对应一个托管本地实例。

包含一个发现可用设备的功能。

保存了未完成的部分运行的要求，这保证了只有在相应的执行程序完成运行时才会将其标记为完成。

在 TF 中，执行器异步执行，直到完成所需的输出（返回张量的操作）或目标（不返回张量的操作）。也就是说，计算图中有两种节点可以作为执行的目标，一种是返回张量的操作对应的节点，另一种是不返回张量的操作对应的节点返回一个张量。部分运行包括两个步骤，首先，设置所有必需的输出和目标，其次，获取输出。在第二步中，可能会出现计算图中需求的输出已经计算完毕，但需求的目标仍在计算中的情况。这时候就派上用场了，虽然这个时候理论上可以返回，因为所有的需求输出都已经计算过了，剩下的需求目标不影响返回的结果。但是 TF 仍然要求我们必须等到所有目标都计算完，因为我们不知道中间输出是否会改变，直到目标完成计算。

/ 相互通信的包装类。

// 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 {}

与单机案例相对应，分布式案例中的会话包括图计算的基本步骤，如资源分配、节点放置、图执行等。

用于与 TF 服务通信的虚拟接口。此接口支持基于 RPC 的实现和进程内实现。

在 TF 中实现服务。对应服务。

环境类包含一个指向必要环境资源的指针。请注意，这些指针是不拥有的。

本地实施。本地的意思是与并且不是跨设备的通信，而是直接在进程内。这样做是为了在同一流程中提供更高效的服务。

包含注册到某个图形的集合。每个注册的图都由一个句柄标识，该句柄由调用者生成并返回给调用者。注册成功后，调用者通过图句柄执行图。每个执行都由一个全局“”唯一标识。在同一张图上，可以独立重复执行多次，只要每次执行的“”不同。

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

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

图面–>|派生|–>|派生|–>|派生|

地址