Tensorflow Serving源码详解

Tensorflow Serving的启动过程分析

在main.cc中

有一个设置开启batching的选项

还有一个设置batching参数的文件路径

这些flags参数都会被解析到options这个对象中，接下来将这个options传入了开启server的函数，接着开启server

接着进入server.cc，可以发现有一系列对于batching的设定，这里可以看到主要是对session_bundle_config这个变量进行了一些设置

接下来根据session_bundle_config设置了option中的platform_config_map

这个函数在platform_config_util.cc中，可以发现session_bundle_config被用来设置了source_adapter_config

再回到server.cc中，可以看到创建流程进入了ServerCore的Create函数

接下来来到server_core.cc，在它的Create函数中主要调用了Initialize函数和ReloadConfig函数，其中Initialize中创建了AspiredVersionManager，而ReloadConfig主要用来创建其他的一些组件如Adapter，Router，Source

在ReloadConfig中可以看到如下代码，该段代码创建了Adapter，Router，Source等组件并将其连接了起来

这里的逻辑是Source监听文件系统，如果检测到模型，则通过Router找到合适的Adapter，tensorflow对应的就是SavedModelSourceAdapter，接着Adapter将Source传递过来的模型信息传递到AspiredVersionManager，接着AspiredVersionManager接受到请求来根据策略进行模型加载卸载。具体AspiredVersionManager是怎么进行模型加载卸载的将在后面介绍。

接下来先回到server.cc，在从ServerCore返回后，主要是建立和启动了Grpc Server和Http Server（如下图）

Tensorflow Serving的模型加载过程分析

这里主要从代码看一下tensorflow serving的模型加载过程，也就是从一个model文件到一个可执行session的过程，而着重看一下batching模式下的session加载过程，因为在batching模式下使用的是不是传统的session，而是封装后的BatchingSession，主要看一下这个BatchingSession封装的过程。

之前说过实际执行模型加载卸载管理的是AspiredVersionsManager，在aspired_versions_manager.cc里面，如下

其中可以看到一个PeriodicFunction，这就是一个定时执行的线程，它主要执行了三个函数，其中实际执行模型加载卸载的是第三个函数，即InvokePolicyAndExecuteAction，而在InvokePolicyAndExecuteAction中就可以看到具体的模型加载卸载的过程了，如下。

接下来撇去中间的层层调用，最后来到的是saved_model_bundle_factory.cc，这个文件里包含了加载saved_model格式模型的具体操作。如下可以看到如果设置了batching的话，tensorflow serving会执行一个WrapSessionForBatching的函数来封装原有的session。其中的SignatureDef记录了计算图的输入输出节点的情况。

这里同时要注意传递给WrapSessionForBatching函数的batch_scheduler_参数，该参数定义如下，可以看到tensorflow serving默认使用的是SharedBatchScheduler，且共享同一个SharedBatchScheduler，也就是说所有模型的batch管理都由一个SharedBatchScheduler来统一处理

在WrapSessionForBatching函数中可以看到其实是对于模型中的每个TensorfSignature都创建一个相应的队列，这里的TensorfSignature表示的是模型的输入输出节点

关于SignatureDef建议阅读下述文档

https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/signature_defs.md

TensorfSignature其实就是把SignatureDef中的Tensor部分提取出来了，相关代码如下，一目了然

所以这里要注意的是，在Tensorflow Serving中不是每个Model建立一个队列，而是对每个Model的不同Signature建立一个队列！

接下来会把原本的session封装成BatchingSession，如下

再接下来顺着函数调用链，可以看到如下代码，这里就用到了上述WrapSessionForBatching函数传递进来的TensorSignature和queue_creator，且可以发现在如下代码中传入了一个lambda函数来作为process_batch_callback，这个callback函数会在batch处理的时候被调用，可以发现在这个lambda表达式中调用了一个ProcessBatch函数，这个函数就是具体执行Batch前向传播的。

Tensorflow Serving的Batch执行过程分析

前面说到Tensorflow Serving的Batch管理用的是SharedBatchScheduler，那么先看一看SharedBatchScheduler相关的程序源代码，这部分代码在Tensorflow的代码仓库里。

在SharedBatchScheduler中，根据前述分析，在模型加载的时候会根据TensorflowSignature创建相应的队列，这些队列都会被一个SharedBatchSchduler所管理，在SharedBatchScheduler中有一系列工作线程ThreadLogic（如下图所示），在创建SharedBatchScheduler创建的时候就会自动创建默认是可调度cpu个数的工作线程，这些线程的处理逻辑就在ThreadLogic函数里

下面进入ThreadLogic函数，它里面主要有几点需要关注

1. 如下图可以看到ThreadLogic是以一种round-robin的方式在处理每个Queue

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-1waWJ626-1594887779634)(file:///C:/Users/ljint/AppData/Local/Temp/msohtmlclip1/01/clip_image044.jpg)]

2. 如下图可以看到ThreadLogic每次从当前Queue的指针获取该Queue所对应的Batch来处理
   

3. 如下图可以看到ThreadLogic最终是调用对应Queue的Batch处理函数来对Batch进行处理，事实上Queue的ProcessBatch函数里面调用的就是在模型加载过程中所传进去callback函数（如下图）
   
   
   接下来看一下这个具体的Batch处理流程，也就是BatchingSession传入的process_batch_callback的处理流程。主要注意如下几点。

4. 在session执行之前，该函数会把一个batch的请求的tensor进行合并，最后生成一个merged_inputs来进行处理，如下图
   

5. 接着对合并完成的merged_inputs进行前向传播
   

6. 最后对处理完成的Tensor进行分拆（这里Tensorflow Serving假设输出Tensor的dim 0大小也是Batch大小B），经过这一步操作之后每个请求的output字段也都被填上了值可以各自返回了
   
   那么新来的请求是怎么进入Queue的呢？这就是下面接着分析请求的路径。这里只分析Http的接口。下面先从Http Server的启动函数开始，代码如下，主要关注如下几点。

7. 在Http Server创建过程中创建了一个线程池，这个线程池就是用来处理Http请求的，线程池的大小是可调度CPU个数的四倍（如下图）
   
   

8. Tensorflow Serving使用的Http Server是一个事件驱动的服务器，底层基于libevent，如下图可以看到创建了一个EvHttpServer，Ev就是Event的意思。处理逻辑就是libevent负责处理IO事件，而上述创建的线程池来执行具体操作。
   

9. 从如下代码中可以看到这里注册了请求的处理函数，也就是线程池中的线程所要执行的函数。
   
   接下来再看请求的具体处理过程，先看到如下代码，这里是按照请求的不同进行了分类，我们主要关注Predict的操作，下面进入Predict请求的相关处理。
   

顺着函数调用链最后进入了RunPredict函数，如下图可以看到在此函数中先是解析了请求的Signature，接着分为三部进行处理，即PreProcess，Predict，PostProcess，我们主要关注Predict的过程。可以看到Predict的过程也就是调用session->Run的过程，而由于在Batching模式下，session已经被封装成了BatchingSession，所以这里调用的其实是BatchingSession的Run函数！

下面进入BatchingSession的Run函数，主要发现要注意以下几点

1. 在BatchingSession中先会对Signature进行验证，查找是否已经有相应的Signature所对应的队列（如下图），而如果没找到的话就会直接执行！
   
   
2. 而如果找到了相应的队列，则就会把请求加入队列，并且陷入等待状态，在当Batch处理完成以后，会通知相应请求的等待线程（也就是Http Server线程池里处理请求的线程）