利用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!')

跑出来的火焰图是这样的

image

好家伙,我跑个 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!')

再跑个图

image

可以看到内存骤降至 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
加载评论