带你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.thinkwithwp.com/servicequotas/home/services/ec2/quotas。

找到EC2服务组，输入“IP”进行搜索，定位到 “EC2-VPC Elastic IPs”：

这里演示我们先提升到10，申请后需要等待一段时间。提额生效后，打开 Amazon CloudFormation 控制台创建堆栈页面 https://us-east-1.console.thinkwithwp.com/cloudformation/home?region=us-east-1#/stacks/create/template。

上传模板，美东区域我们选择 sagemaker-vpc-6azs。

2.2 安全配置

打开 IAM 控制台角色页面 https://us-east-1.console.thinkwithwp.com/iamv2/home?region=us-east-1#/roles，创建新角色：

默认会带AmazonSageMakerFullAccess权限：

点击下一步，输入角色名，例如“SageMaker-Studio-Admin”，然后创建。

创建后，添加内联策略，增加Amazon S3、Amazon ECR、Amazon CodeBuild等权限。

切换到JSON Tab 页：

替换为以下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.thinkwithwp.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.thinkwithwp.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.thinkwithwp.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页面的快捷入口

拉取代码：

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账号，镜像拉取有频率限制。这里替换成了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 目录默认不可见，复制到另一个目录方便下载：

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.thinkwithwp.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
重新打包镜像

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，例如：