亚马逊AWS官方博客

面向 GPU 服务器的 SageMaker 无痛使用指南(一)—SageMaker Notebook 笔记本实例

随着生成式 AI(GenAI)的兴起,包括大型语言模型(LLM)和扩散模型(Diffusion Model)等,导致了对 GPU 服务器资源的巨大需求。由于训练和推理这些模型需要大量的计算能力,因此 GPU 服务器资源变得紧缺。 Amazon SageMaker 作为一个端到端的机器学习操作(MLOps)平台,为 GenAI 模型的训练和推理提供了一系列工具和开发集成框架。它预配置了充足的 GPU 服务器资源池,并在训练和部署过程中提供算力资源的自动拉起、弹性伸缩和释放等功能。

然而,SageMaker 对 GPU 服务器算力资源的 SDK 封装增加了使用的难度,需要一定的学习周期。 对于许多技术背景的客户而言,他们已经拥有算法团队,熟悉 GenAI 的底层技术,不需要使用 SageMaker 的 SDK,而希望像使用原生 EC2 实例一样登录、远程调试 GPU 服务器,并灵活地安装部署自己需要的业界开源训练和推理框架,以保持原有的使用习惯和灵活的扩展性。

因此,本系列博客旨在整理面向原生 GPU 服务器的 SageMaker 使用方法,包括 SageMaker Notebook 实例、Training Job 训练任务、Inference Endpoint 推理端点、SageMaker Hyperpod HPC 高性能集群等各种 SageMaker 功能组件。通过这种方式,用户可以按照自己以往的习惯及使用基础设施的方式,从底层操控和使用 SageMaker 的 GPU 服务器,而无需关注 SageMaker SDK 及框架使用的约定和技术细节。

本文我们先介绍 SageMaker Notebook 笔记本实例上 GPU 服务器的使用方法。

SageMaker Notebook 笔记本实例简介

SageMaker Notebook 笔记本实例是一种托管服务器环境,主要面向开发人员构建、训练、调试和部署机器学习模型,笔记本实例提供了可配置的各种 GPU 的服务器,用户可以根据工作负载选择合适的实例类型。在笔记本实例的 GPU 服务器中预装了常用的机器学习框架和库,如 TensorFlow、PyTorch、scikit-learn、XGBoost 等,也预置了集成开发环境 Jupyter Notebook,可以在笔记本中编写、运行和调试代码。Notebook 笔记本实例可以直接访问 Amazon S3、Amazon EFS 等存储服务,便于存取训练数据和模型工件,也可以连接到 Amazon RDS、Redshift 等数据库服务。Notebook 笔记本实例可以像 EC2 实例一样,根据需求启动、停止和终止。

像 EC2 一样使用 Notebook 笔记本实例进行开发测试的主要场景有以下几种:

  • SSH 登陆 SageMaker Notebook 实例 GPU 服务器
  • IDE 远程调试 SageMaker Notebook 实例 GPU 服务器上的代码
  • Notebook 实例 GPU 服务器上 Web 访问

以下我们逐个进行讲解。

SSH 登陆 SageMaker Notebook instance

虽然有 juypter lab 的可视化开发环境,但算法工程师通常需要 SSH 登陆 GPU 服务器,SageMaker Notebook 笔记本实例是托管的 GPU 服务器资源,它没有绑定公网 IP,不能像 EC2 一样直接 ssh 登陆,但我们可以通过一台和该笔记本实例在同一个 VPC 网络中的跳板机进行 ssh 的转发,即可方便地登陆到 SageMaker notebook 实例,其架构图如下:

如上 bastion 即为同一 VPC 内的跳板机,bastion-private-key.pem 是登陆需要的密钥文件,与 EC2 实例登陆一样的方式。

我们详细介绍下具体的实施步骤。

前置条件:

1. sagemaker notebook 创建时关联 vpc subnet 设置安全组
2. 安全组需要对跳板机开放 22 端口
3. 在此 vpc 中这台跳板机可以通过 ssh 远程连接

登陆堡垒机,在堡垒机中获取 ssh public key:

sh-4.2$ cat ~/.ssh/authorized_keys
ssh-rsa AAAAB3NzaC1ycIr86.....kLdQjgZosRllpcoV83JVMfZ 203.pem

登陆 notbook instance,将上述 publick key 粘贴添加到 notbook instance 的 ssh 目录下:

vim ~/.ssh/authorized_keys

在笔记本实例生命周期中,选择创建笔记本时,加入如下脚本:

PEM_FILE="<你的ssh私钥>.pem" 
S3_BUCKET="your-s3-bucket-name" 
S3_KEY="path/to/<你的ssh私钥>.pem" 
 aws s3 cp s3://$S3_BUCKET/$S3_KEY $PEM_FILE 
cat $PEM_FILE >> /home/ubuntu/.ssh/authorized_keys 

注意选择 scope 为创建笔记本,以便每次 create 新笔记本实例时能自动配置密钥:

在 notbook instance 中获取链接 vpc 的网络,跳板机上 ssh 转发的时候,是用该 vpc 内的 ip 地址,来进行连接的:

ip -4 addr show eth2 | grep -oP '(?<=inet\s)\d+(.\d+){3}'
可以看到接口为: notebook-eth2-ip
对应的ip地址为:10.0.15.195

在本地终端打 ssh 隧道:

ssh -i /Users/tangqy/AMAZON/keypairs/ec203.pem -N -L 8422:10.0.15.195:22 ec2-user@ec2-52-25-118-245.us-west-2.compute.amazonaws.com

其中 10.0.15.195 为笔记本实例的私网 ip
ec2-52-25-118-245.us-west-2.compute.amazonaws.com 为跳板机公网域名

现在在本地终端,已经可以像登陆 EC2 一样,通过本地转发端口直接链接到远程的 notbook 笔记本实例:

ssh -i /Users/tangqy/AMAZON/keypairs/ec203.pem -p 8422 ec2-user@localhost

IDE 远程 Debug

既然已经可以像登陆 EC2 一样 SSH 到 SageMaker Notebook 笔记本实例,当然我们也可以在 IDE 中使用远程调试连接到该实例,满足客户算法工程师开发调试 GPU 服务器上代码的需求。

我们以 vscode 为例,在 vscode 的 hosts 中配置如下:

Host notebook_forward
HostName localhost
IdentityFile /Users/tangqy/AMAZON/keypairs/ec203.pem
User ec2-user
Port 8422

在 vscode 里面点击连接到远程主机,选择刚才配置的 notebook_forward hosts 配置:

连接上 notebook 笔记本实例之后,即可以使用 notebook 的远程 kernel,进行单步调试:

vscode 里面可以 python intereptor 远程 debug python 环境:

与在 EC2 上进行远程调试类似,在 vscode 中打开的远程 jupyter/python 代码是自动同步的,并且可以单步,跳入跳出等调试方法:

SageMaker Notebook 笔记本实例上 Web 应用调试

目前业界很多 GenAI 的开源框架和模型,都提供了 GUI 的开发和调试交互界面。在应用构建阶段,在 SageMaker Notebook 实例启用这类开源框架或者模型的 Web 应用程序是一种非常有用的方式,它可以帮助开发人员或者内部测试人员更高效地交互式测试和对比 GenAI 模型的效果。

我们以 SVC Sovit 的开源模型的 demo ui 为例,来看看如何在笔记本实例上开启 web 应用界面。

用以上步骤 ssh 登陆到笔记本实例上,我们可以拉起 sovit svc 的 demo ui 应用,如下所示:

(python3) [ec2-user@ip-172-16-33-176 so-vits-svc]$ pwd
/home/ec2-user/SageMaker/So-vit-tts/so-vits-svc
(base) [ec2-user@ip-172-16-33-176 so-vits-svc]$ conda activate python3
(python3) [ec2-user@ip-172-16-33-176 so-vits-svc]$ python3 webUI.py
(python3) [ec2-user@ip-172-16-33-176 so-vits-svc]$ lsof -i:8080
COMMAND   PID     USER   FD   TYPE  DEVICE SIZE/OFF NODE NAME
python  23131 ec2-user   24u  IPv4 5461098      0t0  TCP *:webcache (LISTEN)

从上文章节中同样的方式,我们从本地通过跳板机打 ssh 隧道,通过该跳板机转发到笔记本实例:

ssh -i /Users/tangqy/AMAZON/keypairs/<你的私钥>.pem -N -L 8422:10.0.15.195:8080 ec2-user@ec2-52-25-118-245.us-west-2.compute.amazonaws.com
  • 8080 端口为笔记本实例上 sovit tts svc 的 gradio web ui 应用监听端口
  • 8422 为本地端口

这样本地浏览器上访问 localhost:8422 即可以访问到笔记本实例上启动的 sovit svc 的 web 应用:

总结

本文介绍了如何在 SageMaker Notebook 笔记本实例上如何通过一种方便的方式来访问和使用 GPU 服务器资源,并调试和开发自己需要的业界开源模型及框架。

在 SageMaker notebook 笔记本实例上,可以通过 SSH 登陆 SageMaker Notebook 实例,可以通过 IDE 远程调试笔记本实例上的代码,也可以在笔记本实例上启动开源模型的 Web 应用进行方便的调试。

通过上述方法,用户可以像使用原生 EC2 实例一样登录、远程调试 SageMaker Notebook 实例的 GPU 服务器,并灵活地安装部署所需的框架和模型,满足了算法团队的使用习惯和扩展性需求。

参考资料

SageMaker Notebook 实例:https://docs.aws.amazon.com/sagemaker/latest/dg/nbi.html

VSCode 远程开发:https://code.visualstudio.com/docs/remote/ssh

OpenSSH 远程端口转发:https://www.ssh.com/academy/ssh/tunneling/

本篇作者

唐清原

亚马逊云科技高级解决方案架构师,负责 Data Analytic & AIML 产品服务架构设计以及解决方案。10+数据领域研发及架构设计经验,历任 IBM 咨询顾问,Oracle 高级咨询顾问,澳新银行数据部领域架构师职务。在大数据 BI,数据湖,推荐系统,MLOps 等平台项目有丰富实战经验。