机器学习模型API多进程内存共享

最近在做机器学习系统,模型训练好了,API也开发完了,目前使用的API服务套件是:fastapi + gunicorn + docker。

但是有一个问题一直没有解决,就是当API服务启用多个worker时,一个worker会加载一次模型到内存中,内存资源的消耗非常大,如果直接这样使用,会极大的提升硬件成本,随着模型的增多,API服务的资源消耗是不可想象的,所以暂停了模型优化工作,在API上线生产环境前,这个问题需要着重解决一下。

目标很清晰,就是找到一个方法使API服务的workers共享内存,也就是模型只加载一次到内存,内存资源占用量也只是所有模型的大小总和,而不会随着worker数量的增加而成倍的占用内存。

docker方面应该没什么可以做的,所以解决问题得从fastapi和gunicorn入手。

fastapi方面,能想到的就是代码和服务架构两个方向,代码方面尝试了mmap,没成功,主要是模型加载时,启用mmap出问题了,没有深究,因为经过思考,就算代码层面成功应用了mmap,workers可能也不会共享内存,因为workers是gunicorn管理的;服务架构方面,可以将模型加载独立成一个服务,但是可能会有性能问题,暂时没动,但是想到gunicorn也许能完成类似的事情。所有的思考结果都指向了gunicorn,所以直接研究gunicorn。

研究gunicorn,当然主要还是参数和配置的研究,先研究参数,发现了 –preload,官方说明如下:

Load application code before the worker processes are forked.

By preloading an application you can save some RAM resources as well as speed up server boot times. Although, if you defer application loading to each worker process, you can reload your application code easily by restarting workers.

主要意思就是在worker进程被派生之前加载应用代码,这样可以节约内存资源。

看到这个,我心中即将熄灭的希望之火顿时形成燎原之势,经过尝试,这个参数果然达到了预期效果。

由于使用的是fastapi官方提供的gunicorn docker镜像,经过研究镜像用法,发现了增加gunicorn命令行参数的方法

示例启用命令:

docker run -d --name bai-api -p 8001:80 -v /home/xwd/projects/bai_models:/models -e MAX_WORKERS="2" -e TIMEOUT="360" -e GUNICORN_CMD_ARGS="--preload" bai-api

其中GUNICORN_CMD_ARGS可以指定额外的gunicorn命令行参数。

Guess you like

Origin blog.csdn.net/xwd127429/article/details/116738258