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

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的异常堆栈给提供了一个很有用的信息:程序在创建hadoop Filesystem对象时出错了。翻一下tensorflow的代码,从注释中你就会发现这其实是不符合程序本意的:

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

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

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

至此,已经开始逐渐浮出水面了:

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

文章被以下专栏收录