稿 深度学习平台团队

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

01 背景介绍

TensorFlow Serving(以下简称 TF Serving) 是谷歌开源、用来部署** TensorFlow模型** 的高性能推理系统。它主要具备如下特点:

  • 同时支持 gRPC 和 HTTP 接口
  • 支持多模型、多版本
  • 支持模型热更新和版本切换

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

总体来讲,使用 TF Serving 来部署推理服务的稳定性和性能都比较好,特别是CTR 预估这种对服务延迟和稳定性要求高的业务。不过,在 TFServing Github 的 [ issue 列表]上,经常有人报告运行中出现内存不断增长导致 OOM 的问题。一个典型的内存 issue «Sharpincrease in memory usage -> server is killed»

在 2018 年就有人提出,目前为止还是 open 状态。

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

02 模型特征 Raw Serving Tensor输入

先来介绍一下 TF 模型 SavedModel 特征输入的两种方式。 一种是模型输入只有一个 ,该输入的 placeholder 是一个 String,String 的内容是tf.Examples; 另一种是模型输入有多个 ,每个输入placeholder 分别对应特征的Tensor。以tf.estimator的API为例子,两者的API分别为 tf.estimator.export.build_parsing_serving_input_receiver_fntf.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 文件中截取部分出来

可以看到

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; 一个 是在 TF Serving 里面总是对 inputs 做一次排序,在TF Serving 里面排序可以省去一次在 TF里面的strings::StrCat 和查找。

两个PR地址👇:

https://github.com/tensorflow/tensorflow/pull/39743

https://github.com/tensorflow/serving/pull/1638

03服务突增高并发请求

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

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

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

这个现象很让人困惑, 到底是什么进程(或线程)突然增加了? 我们之前分析过TFServing的线程模型,主要的工作线程是由模型的 Intra 和 Inter OP 配置来控制的,不