用 Python resource 模組找出尖峰記憶體用量 一文介紹如何透過 Python resource 模組了解尖峰記憶體用量(peak memory usage), 不過該模組並無法提供更詳細的記憶體統計資料,無法得知具體哪部分的 Python 程式消耗大量記憶體,因此我們需要透過工具剖析(profiling)詳細的記憶體用量,以幫助定位問題之所在。

本文將介紹如何使用 Fil 剖析 Python 程式的記憶體用量,並透過 1 個簡單的範例,實際定位程式中耗用記憶體的部分,並進行改善優化。

本文環境

  • macOS 11.6
  • Python 3.7
  • pandas 1.0.3
  • Fil
$ pip install pandas==1.0.3 filprofiler

測試資料

本文需要使用以下資料集進行測試,該資料集包含至少 200MB 以上的財經新聞 CSV 檔,我們用該資料來模擬實際處理大量資料時,記憶體被耗用的情況:

Daily Financial News for 6000+ Stocks

Fil (aka Fil profiler)

Fil 是 1 套十分易於使用的 Python 記憶體剖析工具,該工具運作原理為透過作業系統所提供的 LD_PRELOAD / DYLD_INSERT_LIBRARIES 機制,在程式啟動之前載入開發者指定的共享函式庫(shared library), 該函式庫可以覆寫(override)某些重要的系統 API, 使得開發者可透過該共享函式庫攔截(intercept)某些 API 的呼叫,例如 Fil 透過覆寫記憶體相關的底層 API, 使得 Fil 得以攔截並統計記憶體分配以及使用情況,達成記憶體剖析的功能。

如果想了解 Fil 會攔截哪些 API, 可以參考 What Fil tracks 一文。

Fil 初體驗

以下範例載入測試資料集中的 raw_partner_headlines.csv, 並查詢多少新聞標題中含有 J.P. Morgan 以及該新聞中所提到的股票代號,並且匯總股票代號出現的總次數,最後只顯示超過 5 次的股票代號:

import pandas as pd


def load_data():
    return pd.read_csv('<path to>/raw_partner_headlines.csv')


def main():
    df = load_data()
    df = df.query('headline.str.contains("J.P. Morgan")', engine='python')
    df = df.groupby(by='stock').count()
    for r in df.sort_values(by='headline', ascending=False).itertuples():
        if r.url > 5:
            print(r.Index, r.url)


main()

上述範例執行結果如下,可以看到被 J.P. Morgan 提到最多次的股票代號為 JPM, 其次為 ARNA:

$ python test.py
JPM 15
ARNA 10
ZIOP 8
BMRN 8
AVEO 8
TEX 6
MDVN 6
NKTR 6
REGN 6
VVUS 6

接著我們實際用 Fil 剖析上述範例會用到多少記憶體,以下是呼叫 fil-profiletest.py 進行記憶體用量剖析的指令:

$ fil-profile run test.py
=fil-profile= Memory usage will be written out at exit, and opened automatically in a browser.
...(略)...

一旦 fil-profile 指令完成記憶體剖析的工作, Fil 就會打開 1 個瀏覽器視窗,並且將分析報告以網頁方式呈現:

上圖顯示前文的範例用了 732.6 MiB 的記憶體,相較於 raw_partner_headlines.csv 檔案大小 400MB 而言,多耗用將近 1 倍的記憶體,理想情況是載入資料之後只佔用 400MB的記憶體空間,但由於 Python 的物件化設計,記憶體耗用的情況是預料之中的事,可以想見一旦 raw_partner_headlines.csv 檔案越大,就需要記憶體越大的機器進行運算,最終勢必對成本(例如機器租用成本)造成影響,所以我們需要優化記憶體的使用。

接著,我們可以循著報告往下找到記憶體耗用的癥結點:

上圖可以發現 main() 函式佔了約 84.59% 的記憶體用量,但具體在 main() 函式的哪部分,可以繼續往下追蹤:

上圖顯示 main() 函式中的 load_data() 函式佔用約 66.38% 的記憶體用量,看來 load_data() 的記憶體用量優化,是 1 個可以進行的方向。

實戰優化

pandas.read_csv() 預設會載入整份資料到記憶體之中,如果能夠分批載入計算將可以有效降低記憶體用量,所以我們可以為 read_csv() 設定 chunksize 參數,讓其分批載入資料,不過分批載入資料的缺點是無法一次計算好最終結果,必須將每批資料的計算結果加總儲存,也就是下列範例中的 counter[r.Index] += r.url 部分,最終再將加總後的結果列印:

import pandas as pd

from collections import Counter


def load_data():
    return pd.read_csv(
        '<path to>/raw_partner_headlines.csv',
        chunksize=1000
    )


def main():
    counter = Counter()
    for df in load_data():
        df = df.query('headline.str.contains("J.P. Morgan")', engine='python')
        df = df.groupby(by='stock').count()
        for r in df.itertuples():
            counter[r.Index] += r.url

    for stock, count in counter.items():
        if count > 5:
            print(stock, count)


main()

上述範例執行結果如下:

$ python test.py
ARNA 10
AVEO 8
BMRN 8
JPM 15
MDVN 6
NKTR 6
REGN 6
TEX 6
VVUS 6
ZIOP 8

其記憶體剖析報告如下:

上述報告可以看到,僅僅只是加個參數,就能夠讓記憶體用量從 732.6 MiB 下降至 150.2 MiB, 下降幅度相當可觀,對於更大的資料也不須擔心記憶體不足的問題,如果是租用雲端機器,也只需要相當低規格的機器即可達成。

以上就是透過 Fil 實際改善記憶體用量的實戰過程。

Happy coding!

References

https://pythonspeed.com/fil/docs/index.html