爱奇艺 TensorFlow Serving 内存泄漏优化实践

2023-10-19 11:09:28

文章目录

一、背景介绍
二、模型特征 Raw Serving Tensor输入
三、服务突增高并发请求
四、总结
五、参考文献

TensorFlow Serving由于其便捷稳定的特点在CTR(Click-through Rate，点击率)预估业务场景被广泛的使用，但是其运行时会出现内存不断增长的问题，也不断有相关issue被提交到Github社区，且目前都是Open状态。本文分享了爱奇艺深度学习平台在实践中发现的两个TensorFlow Serving内存泄漏问题，并修复和提交了 PR 到社区，这里将详细介绍问题的背景以及解决的过程，希望能够有所帮助。

一、背景介绍

[TensorFlow Serving]（以下简称 TF Serving）是谷歌开源，用来部署 TensorFlow 模型的高性能推理系统。它主要具备如下特点：
同时支持 gRPC 和 HTTP 接口
支持多模型、多版本
支持模型热更新和版本切换

另外，爱奇艺在 TF Serving的基础上，开源了 [XGBoost Serving] 来支持 GBDT 模型的推理服务，也同样继承了 TF Serving以上特性。

总体来讲，使用 TF Serving 来部署推理服务的稳定性和性能都比较好，特别是 CTR 预估这种对服务延迟和稳定性要求高的业务。不过，在 TF Serving Github 的 [issue 列表]上，经常有人报告运行中出现内存不断增长导致 OOM 的问题。一个典型的内存issue [<<Sharp increase in memory usage -> server is killed>>]在 2018 年就有人提出，目前为止还是 open 状态。

爱奇艺深度学习平台在实践中也遇到了两个类似的问题，我们业务的线上推理服务是通过 Docker 容器进行部署的，但是发现 TF Serving 在某些时候内存会不断增长导致容器 OOM。下面将逐个介绍这两个问题的背景是如何被解决的。

二、模型特征 Raw Serving Tensor输入

先来介绍一下 TF 模型 Saved Model特征输入的两种方式。一种是模型输入只有一个，该输入的 placeholder 是一个 String，String 的内容是tf.Examples；另一种是模型输入有多个，每个输入placeholder 分别对应特征的 Tensor。以 tf.estimator 的 API 为例子，两者的 API 分别为tf.estimator.export.build_parsing_serving_input_receiver_fn 和tf.estimator.export.build_raw_serving_input_receiver_fn。使用 saved_model_cli命令可以明显看到这两种模型输入的不同：

这两种方式在客户端和TF Serving 服务端的处理都有些不同，如下图所示，使用 TF Examples的输入方式需要在客户端先序列化成 String，然后在TF Serving 服务端的模型里使用 parse example op 将 String 反序列化成输入特征的 Tensor，再执行模型的前向计算。而使用 Raw Serving Tensor 就少了这两部分的操作。TF Examples的序列化和反序列化操作会给端到端的推理服务性能带来一定的损耗，因此我们的CTR 类业务基本都在使用第二种模型输入的方式。

介绍完模型输入的背景，下面看一下问题的现象是什么。有一天，业务在上线的时候发现 TF Serving 容器监控中的内存突然开始不断增加，并且丝毫没有要停止的迹象。

我们迅速通过离线复现了问题，并用 gperftools 做了内存的 profiling，下图是从profiling 的 PDF 文件中截取部分出来，如果要看更详细的请看这个 [Pull Request]里面的附件。

可以看到 DirectSession::GetOrCreateExecutors 这个函数通过 hash table的 emplace创建了大量的 String 导致内存不断增加。通过查看TF的源代码可以发现DirectSession 这个类里面有一个成员变量 executors_ 是一个 unordered_map，里面保存了模型输入的 signature 到 ExecutorsAndKeys 的映射。

DirectSession::GetOrCreateExecutors函数里面的逻辑如下：

根据上面的逻辑，如果 input 的数量有 10 个，那么特征输入的组合就有 10!= 3,628,800 种，如果一个 key 有100 Byte，那就需要约 350MB 内存，如果超过 10 个就需要 3G 以上的内存了，这就是内存泄漏的原因。

不过内存泄漏的源头应该是发送给 TF Serving的 PredictRequest 里面inputs 特征顺序不断在变化，从而导致DirectSession::GetOrCreateExecutors函数一直查找不到匹配的输入，然后创建新的字符串插入到 executors_ 中。而发送给 TF Serving的 PredictRequest 是通过 protobuf 定义生成的，inputs 是一个 string 到 TensorProto 的 map。

再进一步查看 Protocol Buffers 关于 map 的文档定义，里面明确说明了map 迭代的顺序是没有定义的，代码不能依赖于map里面的 item 是某种特定的顺序。

回过头看最开始介绍的两种模型输入，使用 tf.Examples 方式的模型输入只会有一个，不会出现这个问题。我们的模型使用的是多个 Raw Tensor 输入，input 特征至少有 10 个以上，因此才会出现这个问题。

找到了内存泄漏的原因，但是为什么会突然出现内存不断上涨呢，业务之前线上跑了很久从未出现过。业务反馈，发送请求的客户端修改了构造请求的逻辑，之前代码里面的顺序是固定的，现在的顺序不是固定的。虽然Protocol Buffers 的定义里面说 map 的迭代顺序是没有定义的，但是实现上如果插入的顺序一样，那么迭代的顺序估计是固定的，但是我们还是不能依赖这种未定义的实现。这样就把整个问题的来龙去脉弄清楚了。

最后，我们分别提了两个 PR，一个是修改 TF 里面的函数GetOrCreateExecutors：https://github.com/tensorflow/tensorflow/pull/39743，一个是在 TF Serving 里面总是对 inputs 做一次排序：https://github.com/tensorflow/serving/pull/1638 ，在 TF Serving 里面排序可以省去一次在 TF 里面的 strings::StrCat 和查找。

三、服务突增高并发请求

业务反馈在流量的高峰期 TF Serving 容器持续出现 OOM 的情况，在平台监控上可以看到 OOM 的事件次数。

在日志里也可以看到容器的 TF Serving 进程被 Kill 掉的现象。

通过平台监控，我们还发现了一个奇怪的现象，TF Serving 容器里面的进程（或线程）数量突然不断增加，在进程（或线程）数量增加到一定程度后，容器就被 OOM kill 了。

这个现象很让人困惑，到底是什么进程（或线程）突然增加了？我们之前分析过TF Serving的线程模型，主要的工作线程是由模型的 Intra 和 Inter OP 配置来控制的，不应该出现增加的情况。于是我们查看了物理机上的 message 日志，发现了被 kill 的线程，居然是 grpcpp_sync_ser。从日志上可以看出，grpcpp_sync_ser申请使用物理内存，这时出现缺页中断，进入到page_fault，但是容器的内存已经被分配光了，于是就被 cgroup kill 了。

进入到容器内部查看，果然发现 TF Serving 进程生成了很多 grpcpp_sync_ser 线程。

根据以上对日志和监控的分析，我们大概还原了事件的过程如下：

很明显问题的根源在于 GRPC Server，于是我们就去查看 GRPC 源码里是如何生成 grpcpp_sync_ser 线程来处理 GRPC 消息的。通过搜索线程名，很快找到了对应线程创建的地方，这里有一个很好的编写代码实践，就是一定要给线程起一个特定的名字，这样出了问题方便定位。

继续分析 MainWorkLoop 函数里的逻辑，总体的逻辑比较复杂，大概相关的逻辑如下。

从上面的逻辑可以看出，如果没有到 GRPC 的 Resource Quota，GRPC Server 就会不断创建新的 WorkerThread 来处理队列里面的消息。通常情况下，如果负载不高，WorkerThread 在处理完消息后就很快退出，但是当服务器负载高的情况下，突发收到大量请求，就会出现同时有很多WorkerThread并存，申请了大量内存，最后引发 OOM。