tensorflow中一個環境變數引發的記憶體洩漏血案

TensorFlow HDFS · 發表 2019-01-24 13:15:19

摘要：起因最近資訊流推薦的業務方在使用tensorflow進行分散式訓練時，反饋說程式有記憶體洩露的情況。詳細瞭解之後，現場情況是這樣的：資料從hdfs讀取，checkpoint也儲存到hdfs 對於推薦模型，查表多，計算少，所以跑在了CPU上。而由於單個程序的CPU利用率...

起因

最近資訊流推薦的業務方在使用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

通過對記憶體使用的觀察，我大致總結了一些現象：

無論是ps還是worker，virtual memory都要比res memory大出不少來。這應該是比較正常的現象。
ps的virtual memory和res memory都還維持在一個比較穩定的狀態，不像是有記憶體洩漏的樣子。
worker的virtual memory有緩慢的上漲，而res memory則比較快的進行增長，但也還遠遠沒有達到virtual memory的大小。由於一臺物理機上起了好幾個worker程序，所以實體記憶體會消耗的比較快。

通過這幾點，我開始腦補問題的原因：

出於這個原因，我就沒急著上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為關鍵字進行搜尋。

找到記憶體洩漏的來源後，就開始從程式碼層面進行分析，大體流程如下：

我們在使用tensorflow的時候，會通過 KERB_TICKET_CACHE_PATH 環境變數來指定hdfs kerberos的ticket cache(如果你對kerberos不熟悉，可以看我的這篇文章)。而一旦設定了該環境變數後，訪問hadoop就會建立一個新的UserGroupInformation，從而建立一個新的FileSystem。
通過和hdfs的同學溝通，可以在程式外部執行kinit，並且將ticket cache放到預設位置後，不設定該環境變數也可以訪問帶安全的hdfs。
我們之所以使用了該環境變數，可能是由於受老的tensorflow文件的影響。但就hdfs而言， 使用自定義路徑ticket cache介面的行為，的確也略微有些不太清晰 。

所以，最後通過去除這個環境變數，這個問題得以解決。

這麼簡單的一個問題，花了我其實有將近一週的時間除錯，回想下還是有些啼笑皆非的。整個除錯過程的感慨如下：

找bug有時候也是個運氣活。想從風馬牛不相及的現象回溯到原因中去，能夠找到懷疑的方向是非常重要的。而為了提高自己在方向判斷上的敏銳度，還是得努力擴充自己在系統層面的知識面才行。
對於一個複雜的系統而言，把介面行為定義清楚，把文件寫清楚，真的相當重要。很多看似在開發階段節省下來的少量時間，在維護階段很有可能都得加倍償還回去。
碼農不易，搬磚不易，且行且珍惜。