带你 SSH 到 Amazon SageMaker 训练实例一探究竟
您现在可以借助 Amazon SageMaker SSH Helper，安全地连接到 SageMaker 的训练作业、处理作业、实时推理终端节点和 SageMaker Studio 笔记本容器，以进行快速交互式实验、远程调试和高级故障排除。详细内容参见：https://github.com/aws-samples/sagemaker-ssh-helper
01
前言
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 训练实例的方法。
02
环境准备
2.1 网络配置
要连通到 SageMaker 里的训练容器，前提条件是网络可达。我们需要将训练任务设置为 VPC 模式运行。
由于全球 GPU 资源紧张，我们尽量把选定区域的所有可用区都利用起来。为方便快速构建符合要求的隔离网络环境，这里提供3个 CloudFormation 模板：
● 跨3个 AZ (大部分区域)
https://github.com/AIMLTOP/amazon-sagemaker-immersion-day/blob/master/byoc/infra/sagemaker-vpc-3azs.yaml
● 跨4个 AZ (俄勒冈，东京，首尔)
https://github.com/AIMLTOP/amazon-sagemaker-immersion-day/blob/master/byoc/infra/sagemaker-vpc-4azs.yaml
● 跨6个 AZ (美东)
https://github.com/AIMLTOP/amazon-sagemaker-immersion-day/blob/master/byoc/infra/sagemaker-vpc-6azs.yaml
以美东为例，我们选择 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
Amazon CloudFormation：
https://aws.amazon.com/cn/cloudformation/
上传模板，美东区域我们选择 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 等权限。
Amazon S3：
https://aws.amazon.com/cn/s3/
Amazon ECR：
https://aws.amazon.com/cn/ecr/
Amazon CodeBuild：
https://aws.amazon.com/cn/codebuild/
切换到 JSON Tab 页：
替换为以下 JSON 定义：
JSON
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Action": [
"codebuild:DeleteProject",
"codebuild:CreateProject",
"codebuild:BatchGetBuilds",
"codebuild:StartBuild"
],
"Resource": "arn:aws:codebuild:*:*:project/sagemaker-studio*"
},
{
"Effect": "Allow",
"Action": "logs:CreateLogStream",
"Resource": "arn:aws:logs:*:*:log-group:/aws/codebuild/sagemaker-studio*"
},
{
"Effect": "Allow",
"Action": [
"logs:GetLogEvents",
"logs:PutLogEvents"
],
"Resource": "arn:aws:logs:*:*:log-group:/aws/codebuild/sagemaker-studio*:log-stream:*"
},
{
"Effect": "Allow",
"Action": "logs:CreateLogGroup",
"Resource": "*"
},
{
"Effect": "Allow",
"Action": [
"ecr:CreateRepository",
"ecr:BatchGetImage",
"ecr:CompleteLayerUpload",
"ecr:DescribeImages",
"ecr:DescribeRepositories",
"ecr:UploadLayerPart",
"ecr:ListImages",
"ecr:InitiateLayerUpload",
"ecr:BatchCheckLayerAvailability",
"ecr:PutImage"
],
"Resource": [
"arn:aws:ecr:*:*:repository/sagemaker-studio*",
"arn:aws:ecr:*:*:repository/byoc*"
]
},
{
"Effect": "Allow",
"Action": "ecr:GetAuthorizationToken",
"Resource": "*"
},
{
"Effect": "Allow",
"Action": [
"s3:GetObject",
"s3:DeleteObject",
"s3:PutObject"
],
"Resource": "arn:aws:s3:::sagemaker-*/*"
},
{
"Effect": "Allow",
"Action": [
"s3:CreateBucket"
],
"Resource": "arn:aws:s3:::sagemaker*"
},
{
"Effect": "Allow",
"Action": [
"iam:GetRole",
"iam:ListRoles"
],
"Resource": "*"
},
{
"Effect": "Allow",
"Action": "iam:PassRole",
"Resource": "arn:aws:iam::*:role/*",
"Condition": {
"StringLikeIfExists": {
"iam:PassedToService": "codebuild.amazonaws.com"
}
}
}
]
}左滑查看更多
策略名称保存为：SageMakerCodeBuildPolicy。
后续步骤中，我们将借助 SageMaker-studio-image-build-cli 工具来自动构建 Docker 镜像，参考：https://github.com/aws-samples/sagemaker-studio-image-build-cli
工具需要用到 Amazon CodeBuild，所以需要修改角色信任关系，添加对 CodeBuild 服务的信任策略。使得 CodeBuild 可以 Assume SageMaker 的执行角色。
替换为以下内容：
{
"Version": "2012-10-17",
"Statement": [
{
"Sid": "",
"Effect": "Allow",
"Principal": {
"Service": [
"codebuild.amazonaws.com",
"sagemaker.amazonaws.com"
]
},
"Action": "sts:AssumeRole"
}
]
}左滑查看更多
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 的网络跟外部通信
03
集成准备
我们在官方原版 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 页面的快捷入口
拉取代码：
cd ~/
echo "clone code from github"
git clone https://github.com/AIMLTOP/amazon-sagemaker-immersion-day byoc-debug
cd ~/byoc-debug左滑查看更多
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 账号，镜像拉取有频率限制。这里替换成了 Amazon 的免费公开镜像仓库。
• 在原版 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 目录默认不可见，复制到另一个目录方便下载：
mkdir -p ~/environment/ssh
cp -R ~/.ssh/id_rsa.pub ~/environment/ssh/id_rsa_cloud9.pub左滑查看更多
下载 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 文件浏览器根目录的下级目录方便下载：
mkdir -p ~/SageMaker/ssh
cp -R ~/.ssh/id_rsa.pub ~/SageMaker/ssh/id_rsa_notebook.pub左滑查看更多
下载 id_rsa_notebook.pub 到本地。
在本地将 id_rsa_cloud9.pub 和 id_rsa_notebook.pub 合并成一个文件 id_rsa.pub，类似如下：
如果有更多机器需要 SSH 到训练容器，例如 EC2 机器，流程是一样的：
● 登录到调试机器生成 SSH Key
● 将公钥加到 pub
● 重新打包镜像
04
模型训练
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 参数。
05
开发调试
前面睡眠时间是5分钟，SSH 动作要快一些（必须赶在训练任务结束，容器关闭之前，否则不能建连）。
如果前期调试的时候，觉得5分钟太短，可以酌情调整。
5.1 获取容器 IP
从日志中获取打印出来的容器 IP 地址，例如 10.0.159.150：
5.2 Cloud9 SSH
打开 Cloud9 控制台命令行界面，通过 ssh root@container_ip 进行登录，注意替换容器 IP，例如：
5.3 Notebook SSH
打开 SageMaker Notebook 实例命令行界面，通过 ssh root@container_ip 进行登录，注意替换容器 IP，例如：
06
总结
通过前面的演示，您可以选择任意一种方式使用 Notebook 或者 EC2 实例 SSH 到 Amazon SageMaker 训练容器里进行探索。除了从外部环境 SSH 到训练容器，您还可以进一步修改镜像配置，允许容器之间可以互相访问，例如 DeepSpeed 多机之间需要通过 SSH 免密互相登录。
07
参考文档
● 参考1：
《SageMaker Immersion Day》
https://github.com/aws-samples/amazon-sagemaker-immersion-day
● 参考2：
《了解 Amazon SageMaker Notebook实例网络配置和高级路由选项》
https://aws.amazon.com/blogs/machine-learning/understanding-amazon-sagemaker-notebook-instance-networking-configurations-and-advanced-routing-options/
● 参考3：
《从本地公司网络访问 Amazon SageMaker Studio 笔记本》
https://aws.amazon.com/blogs/machine-learning/access-an-amazon-sagemaker-studio-notebook-from-a-corporate-network/
本篇作者
龙斌
亚马逊云科技解决方案架构师，负责协助客户业务系统上云的解决方案架构设计和咨询，现致力于容器和机器学习相关领域的研究。
王世帅
亚马逊云科技机器学习产品技术专家，负责基于亚马逊云科技的机器学习方案的咨询与设计，专注于机器学习的推广与应用，对于云端深度学习模型分布式训练，NLP 等领域有丰富经验，多次面向开发者进行云端机器学习产品的介绍与最佳实践经验分享。
听说，点完下面4个按钮
就不会碰到bug了！
