tensorflow中一个环境变量引发的内存泄漏血案

起因

最近信息流推荐的业务方在使用tensorflow进行分布式训练时，反馈说程序有内存泄露的情况。详细了解之后，现场情况是这样的：

• 数据从hdfs读取，checkpoint也保存到hdfs
• 对于推荐模型，查表多，计算少，所以跑在了CPU上。而由于单个进程的CPU利用率上不去，所以每个机器上都起了多个tensorflow进程(当然，如何提高单进程时的程序性能，是另一个话题了，这里先不谈)。
• 随着运行时间的增加，系统的空闲物理内存在逐渐减少，最终会引发Linux的OOM Killer杀掉某个进程。而由于PS占用的物理内存最大，所以基本上都是PS被杀掉。

初次分析

尽管ps被kill，但内存消耗却不一定是ps引起的。为了进一步确定问题，我首先观察了下各进程virtual memory和res memory的使用情况：

while true; do
    # 打印virtual memory和res memory
    ps ux | grep 'job_name=ps\|job_name=worker' | grep -v grep | awk '{print $5,$6}' 
    sleep 30
done

通过对内存使用的观察，我大致总结了一些现象：

1. 无论是ps还是worker，virtual memory都要比res memory大出不少来。这应该是比较正常的现象。
2. ps的virtual memory和res memory都还维持在一个比较稳定的状态，不像是有内存泄漏的样子。
3. worker的virtual memory有缓慢的上涨，而res memory则比较快的进行增长，但也还远远没有达到virtual memory的大小。由于一台物理机上起了好几个worker进程，所以物理内存会消耗的比较快。

通过这几点，我开始脑补问题的原因：

• ps的virtual和res memory都比较平稳，所以应该不像是内存泄漏的样子
• worker的virtual memory增长远没有res memory快，这种情况有点像申请了一个大内存池，然后池里面的内存在发生着泄露。

出于这个原因，我就没急着上gperftools这种内存检测工具。考虑到tensorflow进程里由于hdfs的使用而嵌入了一个jvm，所以我觉得这搞不好是java的问题：申请了一大坨堆内存，然后开始慢慢的把它们都用满。

验证

为了验证猜想，我先把代码改成了读本地，非常幸运的是：问题消失了。所以很自然的，我认为应该是jvm申请了太大的堆内存，总体造成了物理内存的浪费。

于是我设置了下hdfs c接口的jvm参数，将堆内存限制为1G：

export LIBHDFS_OPTS=-Xmx1g -Xms256m -Xmn128m

运行了一段时间后，java开始报OutOfMemory：

java_oom.png

再来

虽然问题没解决，但java OOM的异常堆栈给提供了一个很有用的信息：程序在创建hadoop Filesystem对象时出错了。翻一下tensorflow的代码，从注释中你就会发现这其实是不符合程序本意的：

tensorflow_hadoop_source.png

tensorflow希望只有一个FileSystem的对象，但它依赖于hdfs的cache层来保证这一点。所以，程序在运行一段时间后，还会去创建新的FileSystem对象，非常不合理。

无奈只好开始啃hadoop的代码。发现c接口可以通过设置一个开关参数来打印FileSystem的泄漏情况：

hadoop_source.png

在tensorflow调用hdfs的代码中加上这个参数后，FileSystem对象的创建过程得到了更加清楚的展示：

leak_detail.png

至此，已经开始逐渐浮出水面了：

• 由于某种原因，tensorflow在调用hdfs进行数据读写时，每次都会创建一个新的FileSystem对象。而这些对象很有可能是一直在cache中存放而没有进行删除的。

后来，hdfs的同事通过对java dump文件的一些分析，也证实了这个结论。有关jvm的内存分析工具，不再详细展开，大家可以用jmap为关键字进行搜索。

root cause

找到内存泄漏的来源后，就开始从代码层面进行分析，大体流程如下：

• 我们在使用tensorflow的时候，会通过KERB_TICKET_CACHE_PATH环境变量来指定hdfs kerberos的ticket cache(如果你对kerberos不熟悉，可以看我的这篇文章)。而一旦设置了该环境变量后，访问hadoop就会创建一个新的UserGroupInformation，从而创建一个新的FileSystem。
• 通过和hdfs的同学沟通，可以在程序外部执行kinit，并且将ticket cache放到默认位置后，不设置该环境变量也可以访问带安全的hdfs。
• 我们之所以使用了该环境变量，可能是由于受老的tensorflow文档的影响。但就hdfs而言，使用自定义路径ticket cache接口的行为，的确也略微有些不太清晰。

所以，最后通过去除这个环境变量，这个问题得以解决。

写在最后

这么简单的一个问题，花了我其实有将近一周的时间调试，回想下还是有些啼笑皆非的。整个调试过程的感慨如下：

• 找bug有时候也是个运气活。想从风马牛不相及的现象回溯到原因中去，能够找到怀疑的方向是非常重要的。而为了提高自己在方向判断上的敏锐度，还是得努力扩充自己在系统层面的知识面才行。
• 对于一个复杂的系统而言，把接口行为定义清楚，把文档写清楚，真的相当重要。很多看似在开发阶段节省下来的少量时间，在维护阶段很有可能都得加倍偿还回去。