亚马逊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 实例上以加快适配,如:

  1. BYOC 快速集成
  2. 直接进入容器排障
  3. 某些分布式训练框架需要借助 SSH调试,例如 DeepSpeed
  4. 快捷查看容器资源信息,例如查看 p4d.24xlarge 实例中本地盘的挂载路径
  5. 训练过程中某些偶发性问题的深度分析

本文将演示一种通过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 S3Amazon ECRAmazon 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”:

下一步环境配置,注意以下几点:

  1. 如果只是用来SSH到容器,micro即可。
  2. 如果有编译代码等其他需求,建议调整为C系或者M系。
  3. 为节约成本,可以设置空闲时间,例如空闲1小时后自动停止。
  4. 网络配置,选择任意一个私有子网,例如这里选择“PrivateSubnet01”。
  5. 因为是放在私有子网,访问方式需要选择“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种方式:

  1. 通过File菜单
  2. 通过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:

注意:

  1. SageMaker Immersion Day原版的基础镜像用的是DockerHub的ubuntu:18.04,免费的DockerHub账号,镜像拉取有频率限制。这里替换成了AWS的免费公开镜像仓库。
  2. 在原版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 个场景:

  1. 从 Cloud9 SSH 到 SageMaker 训练容器
  2. 从 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 脚本:

需要在启动脚本里:

  1. 启动SSH后台服务
  2. 打印容器IP地址
  3. 设置睡眠等待时间,例如这里的5分钟(因为数据集很小,训练任务很快完成,只有容器还在运行状态才能SSH)

4.4 配置训练容器网络

SSH 的前提是网络可达,提交任务时需要指定在 VPC 中运行。需要配置:

  1. 指定子网信息,这里替换成6个私有子网的ID。
  2. 指定安全组信息,这里替换成 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实例网络配置和高级路由选项:

https://thinkwithwp.com/blogs/machine-learning/understanding-amazon-sagemaker-notebook-instance-networking-configurations-and-advanced-routing-options/

参考3: 从本地公司网络访问 Amazon SageMaker Studio 笔记本:

https://thinkwithwp.com/blogs/machine-learning/access-an-amazon-sagemaker-studio-notebook-from-a-corporate-network/

本篇作者

龙斌

亚马逊云科技解决方案架构师,负责协助客户业务系统上云的解决方案架构设计和咨询,现致力于容器和机器学习相关领域的研究。

王世帅

亚马逊云科技机器学习产品技术专家,负责基于亚马逊云科技的机器学习方案的咨询与设计,专注于机器学习的推广与应用,对于云端深度学习模型分布式训练,NLP等领域有丰富经验,多次面向开发者进行云端机器学习产品的介绍与最佳实践经验分享。