亚马逊AWS官方博客
带你SSH到Amazon SageMaker 训练实例一探究竟
2022.10.25 更新,您现在可以借助 SageMaker SSH Helper,安全地连接到 Amazon SageMaker 的训练作业、处理作业、实时推理终端节点和 SageMaker Studio 笔记本容器,以进行快速交互式实验、远程调试和高级故障排除。详细内容参见:https://github.com/aws-samples/sagemaker-ssh-helper。
1.前言
Amazon SageMaker 是一项托管的专门针对机器学习而构建的端到端机器学习开发平台,Amazon SageMaker 在运行脚本、训练算法和部署模型时大量使用容器镜像来管理运行环境,常见有2种模式:
- BYOS(Bring Your Own Script)自定义脚本,使用预先构建好的镜像,用户只需传入自己的代码来运行程序,不需要自己调试镜像,适用于较简单情景。
- BYOC(Bring Your Own Container)自定义容器,用户需要自己创建镜像来运行程序,适用于用户对任务自定义程度较高的情景。
相较而言, BYOC模式较BYOS有着一定复杂度。如果用户本地已经是容器化环境,有固定的机器学习平台/框架,包含约定俗成的流程、任务提交方式、参数传递方法等,当往SageMaker BYOC 模式适配时,需要花费时间进行调试,例如修改启动脚本等。在实际跟用户交流过程中,有的用户提出了希望能够SSH 到SageMaker training 实例上以加快适配,如:
- BYOC 快速集成
- 直接进入容器排障
- 某些分布式训练框架需要借助 SSH调试,例如 DeepSpeed
- 快捷查看容器资源信息,例如查看 p4d.24xlarge 实例中本地盘的挂载路径
- 训练过程中某些偶发性问题的深度分析
本文将演示一种通过BYOC的方式实现 SSH登陆到SageMaker 训练实例的方法。
2.环境准备
2.1 网络配置
要连通到SageMaker里的训练容器,前提条件是网络可达。我们需要将训练任务设置为VPC模式运行。
由于全球GPU资源紧张,我们尽量把选定区域的所有可用区都利用起来。为方便快速构建符合要求的隔离网络环境,这里提供3个CloudFormation模板:
以美东为例,我们选择6AZ模板。
注意:默认每个区域只有5个弹性IP,可以到 Service Quotas 申请提额 https://us-east-1.console.aws.amazon.com/servicequotas/home/services/ec2/quotas。
找到EC2服务组,输入“IP”进行搜索,定位到 “EC2-VPC Elastic IPs”:
这里演示我们先提升到10,申请后需要等待一段时间。提额生效后,打开 Amazon CloudFormation 控制台创建堆栈页面 https://us-east-1.console.aws.amazon.com/cloudformation/home?region=us-east-1#/stacks/create/template。
上传模板,美东区域我们选择 sagemaker-vpc-6azs。
2.2 安全配置
打开 IAM 控制台角色页面 https://us-east-1.console.aws.amazon.com/iamv2/home?region=us-east-1#/roles,创建新角色:
默认会带AmazonSageMakerFullAccess权限:
点击下一步,输入角色名,例如“SageMaker-Studio-Admin”,然后创建。
创建后,添加内联策略,增加Amazon S3、Amazon ECR、Amazon CodeBuild等权限。
切换到JSON Tab 页:
替换为以下JSON定义:
策略名称保存为:SageMakerCodeBuildPolicy。
后续步骤中,我们将借助 sagemaker-studio-image-build-cli 工具来自动构建 Docker 镜像,参考:https://github.com/aws-samples/sagemaker-studio-image-build-cli。
工具需要用到 Amazon CodeBuild,所以需要修改角色信任关系,添加对 CodeBuild 服务的信任策略。使得 CodeBuild 可以 Assume SageMaker 的执行角色。
替换为以下内容:
2.3 准备SageMaker Studio
为了演示更多的场景,我们设计了松耦合模式,拆开了“提交训练环境”和“远程调试终端”。接下来我们将分别准备:
- SageMaker Studio (打包镜像、运行 Jupyter 笔记本)
- Cloud9 (SSH 可选终端一)
- SageMaker Notebook Instance (SSH 可选终端二)
您实际使用时完全可以根据自身情况只选择一种:
- 如果您已经在用 Notebook Instance,则可以在同一个控制台页面完成命令行和笔记本操作。
- 如果您平时采用EC2机器(例如G5),基于 SageMaker Python SDK 直接提交训练任务,可以只参考 Cloud9 的操作流程。
打开 SageMaker 控制台 Studio 页面:https://us-east-1.console.aws.amazon.com/sagemaker/home?region=us-east-1#/studio-landing。
首先启用配置:
设置Profile名称,执行角色选择前面创建的 SageMaker-Studio-Admin:
其他保持默认,提交创建。等待约5分钟左右创建完成(可以先准备 Cloud9 和 SageMaker Notebook 实例)。
2.4 准备Cloud9
打开 Cloud9 控制台:
https://us-east-1.console.aws.amazon.com/cloud9/home?region=us-east-1。
创建一个Cloud9环境,设置名称例如 “byoc-debug-ide”:
下一步环境配置,注意以下几点:
- 如果只是用来SSH到容器,micro即可。
- 如果有编译代码等其他需求,建议调整为C系或者M系。
- 为节约成本,可以设置空闲时间,例如空闲1小时后自动停止。
- 网络配置,选择任意一个私有子网,例如这里选择“PrivateSubnet01”。
- 因为是放在私有子网,访问方式需要选择“Systems Manager”。
其他保持默认,点击创建。约等待 1~2 分钟即可打开Cloud9控制台。 如果出现确实无法访问的情况(不排除是出口网络问题),可以把刚创建的 Cloud9 删除,重新创建一个。
创建完成后,注意需要绑定安全组,后续才能访问 SageMaker 拉起的容器。
首先跳转到EC2控制台:
然后修改安全组配置:
搜索InternalSecurityGroup,附加到实例:
2.5 准备Notebook
打开 SageMaker 控制台 Notebook 实例页面 https://us-east-1.console.aws.amazon.com/sagemaker/home?region=us-east-1#/notebook-instances。
创建新Notebook,参数如下:
- 名字:byoc-debug-notebook
- 平台:notebook-al2-v1
- 角色:SageMaker-Studio-Admin
- 网络:可以选择任意一个私有子网,为了与前面的Cloud9区分,这里选择“PrivateSubnet02”
- 安全组:InternalSecurityGroup
- 互联网访问:注意一定要选择“Disable”,从而走VPC的网络跟外部通信
3.集成准备
我们在官方原版 SageMaker Immersion Day 的基础上,对BYOC场景进行了扩展,添加了SSH支持。
原版地址:https://github.com/aws-samples/amazon-sagemaker-immersion-day。
扩展版地址:https://github.com/AIMLTOP/amazon-sagemaker-immersion-day。
3.1 拉取代码
打开 SageMaker Studio
第一次打开需要初始化,等待约2分钟左右。进入后,首先开一个终端。
有2种方式:
- 通过File菜单
- 通过Launcher页面的快捷入口
拉取代码:
3.2 打开笔记本
从左边文件浏览器窗口,依次进入 byoc-debug/byoc 目录。
双击打开 byoc.ipynb,Kernel 选择 Python 3 (Data Science) :
3.3 安装依赖包
执行 “cell 00”,安装依赖包,特别是 sagemaker-studio-image-build 工具,用于自动构建 Docker 镜像并推送到 ECR 镜像仓库。
3.4 准备 Dockerfile
在 scikit_container 目录下,有一个准备好的示例Dockefile:
注意:
- SageMaker Immersion Day原版的基础镜像用的是DockerHub的ubuntu:18.04,免费的DockerHub账号,镜像拉取有频率限制。这里替换成了AWS的免费公开镜像仓库。
- 在原版Dockerfile的基础上,增加了SSH的配置。
Dockerfile 有一行关键配置如下:
COPY ssh/id_rsa.pub /root/.ssh/authorized_keys
首先将允许外部访问 SageMaker 训练容器的机器(Cloud9/EC2/SageMaker Notebook 实例等)公钥保存在 ssh目录下的 id_rsa.pub 文件。
接着生成 Docker 镜像时,会复制到容器环境root用户下的.ssh目录。
3.5 生成SSH Key
这里我们演示 2 个场景:
- 从 Cloud9 SSH 到 SageMaker 训练容器
- 从 SageMaker Notebook 实例 SSH 到 SageMaker 训练容器
首先进入 Cloud9 命令行界面,生成 SSH Key:
ssh-keygen -q -t rsa -N '' -f ~/.ssh/id_rsa
.ssh 目录默认不可见,复制到另一个目录方便下载:
下载 id_rsa_cloud9.pub 到本地。
然后进入SageMaker 控制台 Notebook 实例页面 https://us-east-1.console.aws.amazon.com/sagemaker/home?region=us-east-1#/notebook-instances。
找到前面准备好的 byoc-debug-notebook 笔记本,点击如下图所示的 “Open JupyterLab”:
打开后,默认里面是空的,可以通过单击如下图所示的 Terminal 打开终端 shell 环境:
生成 SSH Key:
ssh-keygen -q -t rsa -N '' -f ~/.ssh/id_rsa
.ssh 目录默认不可见,复制到 SageMaker 文件浏览器根目录的下级目录方便下载:
下载 id_rsa_notebook.pub 到本地。
在本地将 id_rsa_cloud9.pub 和 id_rsa_notebook.pub合并成一个文件 id_rsa.pub,类似如下:
如果有更多机器需要 SSH 到训练容器,例如 EC2 机器,流程是一样的:
- 登录到调试机器生成 SSH Key
- 将公钥加到pub
- 重新打包镜像
4.模型训练
4.1 生成镜像
回到 SageMaker Studio 控制台打开的 Jupyter 笔记本页面。
将前面准备好的 id_rsa.pub 上传到 /byoc-debug/byoc/scikit_container/ssh/ 目录,然后执行 “cell 01”
这里注意,如果您修改了仓库名称,不是 byoc 或者 sagemaker-studio 的前缀,您需要在前面的步骤中添加权限。
执行成功,最后输出类似如下:
4.2 准备测试数据
依次执行 “cell 02”,“cell 03” 和 “cell 04”
4.3 准备启动脚本
打开 byoc/scikit_container/decision_trees 目录下的 train 脚本:
需要在启动脚本里:
- 启动SSH后台服务
- 打印容器IP地址
- 设置睡眠等待时间,例如这里的5分钟(因为数据集很小,训练任务很快完成,只有容器还在运行状态才能SSH)
4.4 配置训练容器网络
SSH 的前提是网络可达,提交任务时需要指定在 VPC 中运行。需要配置:
- 指定子网信息,这里替换成6个私有子网的ID。
- 指定安全组信息,这里替换成 InternalSecurityGroup 安全组的ID。
配置完成后,执行“cell 05”。
4.5 启动训练
执行“cell 06”,注意这里的配置,增加了 subnets 和 security_group_ids 参数。
5.开发调试
前面睡眠时间是5分钟,SSH动作要快一些(必须赶在训练任务结束,容器关闭之前,否则不能建连)。
如果前期调试的时候,觉得5分钟太短,可以酌情调整。
5.1 获取容器IP
从日志中获取打印出来的容器IP地址,例如 10.0.159.150:
5.2 Cloud9 SSH
打开 Cloud9 控制台命令行界面,通过 ssh root@container_ip 进行登录,注意替换容器IP,例如:
ssh root@10.0.159.150
5.3 Notebook SSH
打开 SageMaker Notebook 实例命令行界面,通过 ssh root@container_ip 进行登录,注意替换容器IP,例如:
ssh root@10.0.159.150
总结
通过前面的演示,您可以选择任意一种方式使用Notebook或者EC2实例SSH到Amazon SageMaker 训练容器里进行探索。除了从外部环境 SSH 到训练容器,您还可以进一步修改镜像配置,允许容器之间可以互相访问,例如 DeepSpeed 多机之间需要通过SSH免密互相登录。
参考文档
参考1:SageMaker Immersion Day:
https://github.com/aws-samples/amazon-sagemaker-immersion-day
参考2: 了解 Amazon SageMaker Notebook实例网络配置和高级路由选项:
参考3: 从本地公司网络访问 Amazon SageMaker Studio 笔记本: