Deploy your own docker container for model inference using AWS Sagemaker
AWS Sagemaker是亚马逊推出的专门用于机器学习训练、推理的一个平台。最近探索在sagemaker上面部署自动扩缩容的Stable Diffusion模型,踩了一些坑,这里记录一下。
自动扩缩容的GPU平台选择
如果我们想要部署大模型(比如Stable Diffusion)用于对外提供服务的话,我们的部署方案最好是可以自动扩缩容的,也就是当请求量少的时候自动减少机器数量,节约成本,当请求量增多的时候自动增加机器数量,减少排队人数。
很多云平台也提供了这种自动扩缩容的部署方案,这里大概列举几个供大家参考:
- Huggingface的inference endpoint
只需要将项目文件全部上传到Huggingface仓库,指定好python依赖包,以及写一个调用模型进行推理的函数,Huggingface就会自动进行部署,支持自动扩缩容。优点在于操作简单,缺点在于成本相对较高。之前我输出过一个教程:
- Replicate平台
需要自己在本地先打包好完整的镜像(包含所有项目依赖,以及模型文件),然后将镜像推送到Replicate平台。之后Replicate平台会提供出API,用于模型的自动扩缩容的调用。优点在于成本较低,缺点在于每次更改都需要重新打包镜像。我之前也产出过一个教程:
- Runpod平台
是目前操作最简单也最经济的一个方案,我在教程里有详细描述。Runpod支持自定义扩缩容策略、单独挂载数据卷使得可以随时上传或者更换模型而无需重新打包镜像,同时按使用时长付费,收费标准也是目前最便宜的一个。我的教程:
它就类似于国内的AutoDL平台,不过AutoDL只对企业版账号提供弹性部署的服务。
- AWS Sagemaker
AWS Sagemaker也可以进行模型的弹性部署,类似于AWS lambda(无服务器计算),但是AWS Lambda不支持GPU调度,所以我们想要在AWS平台上弹性部署运行在GPU上的模型的话,还是要选择Sagemaker。
Sagemaker支持的部署方式及定价
从它的官方文档可以看到,对于模型部署,Sagemaker提供了几种不同的方式:
- Real-time inference: 实时推理,也就是能够实时响应请求,这种部署方式就需要至少有一台GPU机器一直在运行当中,因此也会一直产生费用
- Serverless inference: 无服务器推理,也就是自动扩缩容,最少可以缩容到0台实例
- Serverless inference with provisioned concurrency: 无服务器推理,但是支持保留指定数量的机器一直在运行当中。这个看似和上面的real-time inference一样,但其实Sagemaker有差异,也是我踩坑的地方,后面再讲
- Asynchronous inference: 异步推理,所有请求都会放入一个队列,然后进行消费,产生结果之后会将结果放入S3存储当中。适用于耗时较久的推理任务
- Batch Transform: 没研究过文档
这些服务的定价可以从这里看到,比如我们可以看到“按需定价”--“实时推理”的价格,需要注意的是,由于我们需要部署GPU实例,所以要拉到下面“加速计算”那里,才是适用于我们的场景的实例,价格也和其他的GPU平台差不多。
Sagemaker部署自定义镜像
因为我们的需求是部署之后支持自动扩缩容的,所以还是需要用到docker镜像。如果只是在单个服务器上部署出来进行调用,当然很简单,我们可以自己装好依赖,输入启动命令,上传好项目和模型文件,然后就能调用。但是自动扩缩容不可能每次都去装环境,所以就需要用到docker。
要在Sagemaker上部署自定义的docker镜像,需要做好下面的前置准备:
- 注册AWS账号,并且绑定信用卡
- 了解docker相关知识,知道如何安装docker,如何编写Dockerfile,如何使用
docker build
,docker push
等命令 - 查看Sagemaker官方文档
接下来简单介绍一下如何部署。
实现handler和Dockerfile
我们的镜像里需要有一个handler来对模型进行初始化以及接收并处理请求,Sagemaker对于我们实现的handler有一定要求,比如要对某个特定的path进行响应,因此我们不用自己从头实现API,而是可以直接使用Sagemaer inference toolkit,我们在镜像的handler中,基于这个SDK构建出我们的调用模型的逻辑就行了,具体的handler写法和Dockerfile写法可以参考:GitHub example
构建镜像并上传到AWS ECR
上面编写好Dockerfile以及准备好项目文件之后,我们就可以在本地构建镜像。Sagemaker不支持从Dockerhub读取镜像,只支持AWS自己的ECR(Elastic Container Registry),所以我们要将镜像上传到ECR。
首先去AWS网页版的ECR控制台,新建一个私有仓库。注意这里必须是私有仓库而不能是public repository,否则在Sagemaker那边会报错并且还不会提示你是因为是public仓库的原因。
创建好ECR私有仓库之后,可以在网页上看到推送镜像的命令。我们只需要在本地先安装好AWS Cli,运行aws configure
进行验证,然后再按照ECR的推送命令进行推送就行了。
在Sagemaker中关联镜像
来到Sagemaker控制台,左边点击“Inference”,选择“Models”,然后点击创建模型:
在“Location of inference code image”的地方,输入刚刚ECR的镜像URI就行
最后点击创建。
创建Endpoint configuration
点击左边的Endpoint configuration并创建:
这里我们可以选择Provisioned,表示会预留几台机器一直在运行中,也可以选择Serverless。
然后下面创建一个Variant:
这里可以关联上刚刚创建的model,也可以选择要部署的机器类型和预留的机器数量。
创建Endpoint
最后点击Endpoint然后创建一个就行,等它的状态变为“In Service”,就可以查看日志以及发起调用了。调用Endpoint可以参考官方文档,需要在本地装一个python包进行调用。
Sagemaker挂载数据卷
如果我们的模型文件本身比较小的话,上面的流程操作下来不会有什么问题。但是,如果像Stable Diffusion这种模型文件本身就4、5个G的大小,加上镜像文件本身的依赖,一个镜像可能就有12个G,这种情况下我们可能就没法顺利部署。
首先,Sagemaker的serverless inference,本身支持的容量大小是有限的,我没在文档里看到,不过如果我的镜像体积大于10G,就会在部署endpoint的过程中给我报错说镜像太大。
其次,real-time inference的官方文档上说,在部署的时候会挂载EBS数据卷,我一开始以为是额外挂载的数据卷,后来才反应过来是整个镜像用这么多存储空间,不过也比serverless inference要好,至少如果我们把模型文件一起打包到镜像中,也可以进行部署了。
但是这样依然不够灵活,最理想的还是像Runpod平台的方式,可以支持挂载数据卷,这样在打包好镜像之后就不必一次次重新打包了,想要更换模型文件只需要去数据卷中更换就行。
Sagemaker其实也提供了单独存储模型文件的方式,但是它的方式是将模型文件放到S3中,每次启动容器的时候,会从S3将模型文件下载到容器当中,这样做缺点十分明显,第一是增加了容器的冷启动时间,第二是下载下来的模型文件依然会占用容器的存储空间,可能导致空间不够用的情况。
至于Sagemaker是否支持在容器中挂载数据卷(storage volume),比如用AWS自身的EFS(Elastic File System)进行挂载,我Google了一下,发现了一篇文章,但是后来研究了一下才发现,Sagemaker的studio、domains、notebook instance、training job确实可以挂载一个EFS数据卷,使得模型数据的存储变得很方便,但是在inference模型推理中,根本不支持挂载EFS,这个帖子说的很明白了:stackoverflow.com/questions/7…
综上所述,Sagemaker在成本上和其他的GPU平台差不多,但是部署便捷性上远不如Runpod平台,所以还是推荐别折腾了,直接用Runpod进行弹性部署。