Tensorflow Serving源码详解

时间:2024-04-03 11:37:21

Tensorflow Serving的启动过程分析

在main.cc中

有一个设置开启batching的选项
Tensorflow Serving源码详解

还有一个设置batching参数的文件路径
Tensorflow Serving源码详解
这些flags参数都会被解析到options这个对象中,接下来将这个options传入了开启server的函数,接着开启server
Tensorflow Serving源码详解
接着进入server.cc,可以发现有一系列对于batching的设定,这里可以看到主要是对session_bundle_config这个变量进行了一些设置
Tensorflow Serving源码详解

接下来根据session_bundle_config设置了option中的platform_config_map
Tensorflow Serving源码详解

这个函数在platform_config_util.cc中,可以发现session_bundle_config被用来设置了source_adapter_config
Tensorflow Serving源码详解

再回到server.cc中,可以看到创建流程进入了ServerCore的Create函数
Tensorflow Serving源码详解

接下来来到server_core.cc,在它的Create函数中主要调用了Initialize函数和ReloadConfig函数,其中Initialize中创建了AspiredVersionManager,而ReloadConfig主要用来创建其他的一些组件如Adapter,Router,Source
Tensorflow Serving源码详解
在ReloadConfig中可以看到如下代码,该段代码创建了Adapter,Router,Source等组件并将其连接了起来
Tensorflow Serving源码详解
Tensorflow Serving源码详解

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

接下来先回到server.cc,在从ServerCore返回后,主要是建立和启动了Grpc Server和Http Server(如下图)
Tensorflow Serving源码详解
Tensorflow Serving源码详解

Tensorflow Serving的模型加载过程分析

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

之前说过实际执行模型加载卸载管理的是AspiredVersionsManager,在aspired_versions_manager.cc里面,如下
Tensorflow Serving源码详解

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

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

这里同时要注意传递给WrapSessionForBatching函数的batch_scheduler_参数,该参数定义如下,可以看到tensorflow serving默认使用的是SharedBatchScheduler,且共享同一个SharedBatchScheduler,也就是说所有模型的batch管理都由一个SharedBatchScheduler来统一处理
Tensorflow Serving源码详解
在WrapSessionForBatching函数中可以看到其实是对于模型中的每个TensorfSignature都创建一个相应的队列,这里的TensorfSignature表示的是模型的输入输出节点
Tensorflow Serving源码详解

关于SignatureDef建议阅读下述文档

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

TensorfSignature其实就是把SignatureDef中的Tensor部分提取出来了,相关代码如下,一目了然
Tensorflow Serving源码详解

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

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

Tensorflow Serving的Batch执行过程分析

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

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

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

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

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

  1. 如下图可以看到ThreadLogic每次从当前Queue的指针获取该Queue所对应的Batch来处理
    Tensorflow Serving源码详解

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

  3. 在session执行之前,该函数会把一个batch的请求的tensor进行合并,最后生成一个merged_inputs来进行处理,如下图
    Tensorflow Serving源码详解

  4. 接着对合并完成的merged_inputs进行前向传播
    Tensorflow Serving源码详解

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

  6. 在Http Server创建过程中创建了一个线程池,这个线程池就是用来处理Http请求的,线程池的大小是可调度CPU个数的四倍(如下图)
    Tensorflow Serving源码详解
    Tensorflow Serving源码详解

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

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

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

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

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