利用memray重构项目，提升内存利用效率

人生苦短，快用memray

不知不觉距离上次大的重构 #43 已经过去一年多，不曾想冷启动时间又肉眼可见的慢了起来，坚决不能忍，搞起

因为是某个炼丹项目的配套，代码有好几个人的参与。工程这边还好，组内有明确的 lint 规则以及一直坚持的 codereview，代码质量基本可控；但炼丹组的代码就显得很飘逸了，一些暗坑往往出现在这部分。炼丹我是不懂的，咋搞？上工具呗

这次使用的是：memray，官网介绍很全面了，工具使用起来也是非常简单，我就用了两句

1
2
memray run hello.py
memray flamegraph memray-hello.py.xxx.bin

几行用来测试的代码

1
2
3
4
5
6
7
8
9
import torch  # 实际项目的引用关系要复杂得多


def do_ai_pipe():
    print(torch.__version__)


if __name__ == '__main__':
    print('hello world!')

跑出来的火焰图是这样的

好家伙，我跑个 hello world 就干掉三十来兆内存，明显是import torch搞的鬼，果断改一下import位置

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
def do_ai_pipe():
    # 可能有人会有疑问：改成 local import 的话，岂不是每次调用方法都要重复 import 一次，这多浪费
    # 先说答案：不会，因为 import 是有全局缓存的，第二次 import 直接从缓存里取了
    import torch  # 实际项目的引用关系要复杂得多

    print(torch.__version__)


if __name__ == '__main__':
    print('hello world!')

再跑个图

可以看到内存骤降至 KB 级别，代码几乎是不需要做什么改动的。

当然这只是一个简单的例子：某些情况下一些方法引用了一个比较重的依赖，那么就可以考虑把这个重依赖从global import改为local import，可以极大降低项目整体启动的时间和内存占用。

借助memray可以清晰直观地展现出项目所有调用的内存使用情况，你可以很快定位到最需要优化的func，然后根据实际需要，稍微做一点调整，即可获得不错的性能收益。

算个收益账：

支出：2.5 小时，基本就是按照memray的指示，把该延迟引用的延迟引用，能惰性求值的换成惰性求值，不必要的初始化状态丢到对应实例化方法中，等等
收益：未优化前启动需要 3.5~5s 左右按 4s 算，优化后仅 0.8~1.2s按 1s 算，按开发一天调试启停项目 500 次计算，每天节省 500 * 3s / 60 = 25min 相当于一个番茄钟的时间

算起来 6 个工作日即可回本，这买卖很划算。

甩个实际成果：项目冷启动内存占用减半，速度提升了2~3倍。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
yyyy-mm-dd HH:18:28,729 - [INFO] [MainThread] (server:27) starting remarkable
yyyy-mm-dd HH:18:28,729 - [INFO] [MainThread] (base_handler:54) loading web handlers
yyyy-mm-dd HH:18:29,661 - [INFO] [MainThread] (file_utils:41) PyTorch version 1.4.0 available.
yyyy-mm-dd HH:18:31,789 - [INFO] [MainThread] (base_handler:79) /
yyyy-mm-dd HH:18:31,806 - [INFO] [MainThread] (server:89) Server started...
yyyy-mm-dd HH:18:31,806 - [INFO] [MainThread] (server:90) the web server is listening on http://0.0.0.0:8000


yyyy-mm-dd HH:20:06,448 - [INFO] [MainThread] (server:27) starting remarkable
yyyy-mm-dd HH:20:06,448 - [INFO] [MainThread] (base_handler:54) loading web handlers
yyyy-mm-dd HH:20:06,495 - [INFO] [MainThread] (base_handler:64) loading plugins
yyyy-mm-dd HH:20:07,615 - [INFO] [MainThread] (base_handler:79) /
yyyy-mm-dd HH:20:07,631 - [INFO] [MainThread] (server:89) Server started...
yyyy-mm-dd HH:20:07,631 - [INFO] [MainThread] (server:90) the web server is listening on http://0.0.0.0:8000

# NOTE: I am not responsible for any expired content.
create@2022-05-25T08:13:58+08:00
update@2022-06-05T00:28:06+08:00
comment@https://github.com/ferstar/blog/issues/63

利用memray重构项目，提升内存利用效率

相关文章：