带你 SSH 到 Amazon SageMaker 训练实例一探究竟

亚马逊云开发者

于 2023-03-06 19:59:54 发布

阅读量282

点赞数

文章标签： ssh 运维

原文链接：https://mp.weixin.qq.com/s?__biz=Mzg4NjU5NDUxNg==&mid=2247534639&idx=1&sn=c5e4a999eefeb2eb305fe01a266f6941&chksm=cf954694f8e2cf8285306f7a7cd24800aa22d109688b8e0273a7e63692ab7ef09fce9b408fd5&scene=126&sessionid=0

版权

本文介绍了如何使用AmazonSageMakerSSHHelper安全地连接到SageMaker的训练和处理作业，以便进行快速交互式实验、远程调试和故障排除。通过BYOC模式，用户可以SSH到SageMaker训练实例，适用于快速集成、容器排障和分布式训练框架的调试。文章详细阐述了环境准备、网络配置、安全配置以及通过Cloud9和SageMakerNotebook实例进行SSH的步骤。

摘要由CSDN通过智能技术生成

您现在可以借助 Amazon SageMaker SSH Helper，安全地连接到 SageMaker 的训练作业、处理作业、实时推理终端节点和 SageMaker Studio 笔记本容器，以进行快速交互式实验、远程调试和高级故障排除。详细内容参见：https://github.com/aws-samples/sagemaker-ssh-helper

前言

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.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 的网络跟外部通信

集成准备

我们在官方原版 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

● 重新打包镜像

模型训练

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分钟，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://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/

本篇作者