开始使用开源的Amazon SageMaker分布式

数据科学家需要一个一致和可重现的环境来进行机器学习(ML)和数据科学工作负载,使其能够管理依赖项并且安全。AWS Deep Learning Containers已经为常见框架(如TensorFlow、PyTorch和MXNet)的训练和服务模型提供了预构建的Docker镜像。为了改善这一体验,我们在2023 JupyterCon上宣布了SageMaker开源分发的公共测试版。这为不同专业水平的ML开发人员提供了统一的端到端ML体验。开发人员无需再在不同的框架容器之间进行实验,也无需在从本地JupyterLab环境和SageMaker笔记本转移到SageMaker上的生产作业时切换。开源SageMaker分发支持数据科学、ML和可视化的最常见的软件包和库,如TensorFlow、PyTorch、Scikit-learn、Pandas和Matplotlib。您可以从Amazon ECR公共画廊开始使用容器。

在本文中,我们将向您展示如何使用SageMaker开源分发快速在本地环境进行实验并轻松将其提升到SageMaker上的作业。

解决方案概述

在我们的示例中,我们展示如何使用PyTorch训练图像分类模型。我们使用公开可用的PyTorch KMNIST数据集。我们训练一个神经网络模型,测试模型的性能,并最终打印训练和测试损失。此示例的完整笔记本可在SageMaker Studio Lab示例存储库中找到。我们使用开源分发在本地笔记本上进行实验,将其移动到Amazon SageMaker Studio以使用更大的实例,然后将笔记本调度为笔记本作业。

先决条件

您需要以下先决条件:

  • 已安装Docker。
  • 具有管理员权限的活动AWS帐户。
  • 已安装AWS命令行界面(AWS CLI)和Docker的环境。
  • 现有的SageMaker域。要创建域,请参阅Onboard to Amazon SageMaker Domain。

设置本地环境

您可以直接在本地笔记本上开始使用开源分发。要启动JupyterLab,请在终端上运行以下命令:

export ECR_IMAGE_ID='public.ecr.aws/sagemaker/sagemaker-distribution:latest-cpu'
docker run -it \
    -p 8888:8888 \
    --user `id -u`:`id -g` \
    -v `pwd`/sample-notebooks:/home/sagemaker-user/sample-notebooks \
    $ECR_IMAGE_ID jupyter-lab --no-browser --ip=0.0.0.0

您可以将ECR_IMAGE_ID替换为Amazon ECR公共画廊中可用的任何图像标签,或者如果您使用支持GPU的机器,则选择latest-gpu标签。

此命令将启动JupyterLab并在终端上提供一个URL,例如http://127.0.0.1:8888/lab?token= <token>。请复制链接并在您喜欢的浏览器中输入以启动JupyterLab。

设置Studio

Studio是一个端到端的集成开发环境(IDE),可让开发人员和数据科学家在规模上构建、训练、部署和监控ML模型。Studio提供了一个广泛的第一方图像列表,其中包含常见的框架和软件包,如Data Science、TensorFlow、PyTorch和Spark。这些图像使得数据科学家可以通过选择他们选择的计算框架和实例类型来简单地开始使用ML。

现在,您可以使用Studio的“自行添加图像”功能在Studio上使用SageMaker开源分发。要将开源分发添加到SageMaker域中,请完成以下步骤:

  1. 通过在终端上运行以下命令将开源分发添加到您的帐户的Amazon Elastic Container Registry(Amazon ECR)存储库中:

    # 根据您的要求使用latest-cpu或latest-gpu标记
    export ECR_GALLERY_IMAGE_ID='sagemaker-distribution:latest-cpu'
    export SAGEMAKER_IMAGE_NAME='sagemaker-runtime'
    export SAGEMAKER_STUDIO_DOMAIN_ID='d-xxxx'
    export SAGEMAKER_STUDIO_IAM_ROLE_ARN='<studio-default-execution-role-arn>'
    
    docker pull public.ecr.aws/sagemaker/$ECR_GALLERY_IMAGE_ID
    
    export ECR_PRIVATE_REPOSITORY_NAME='sm-distribution'
    export ECR_IMAGE_TAG='sagemaker-runtime-cpu'
    export AWS_ACCOUNT_ID='0123456789'
    export AWS_ECR_REPOSITORY_REGION='us-east-1'
    
    # 创建存储库
    aws --region ${AWS_ECR_REPOSITORY_REGION} ecr create-repository --repository-name $ECR_PRIVATE_REPOSITORY_NAME
    aws --region ${AWS_ECR_REPOSITORY_REGION} ecr get-login-password | docker login --username AWS --password-stdin ${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_ECR_REPOSITORY_REGION}.amazonaws.com
    export ECR_IMAGE_URI=$AWS_ACCOUNT_ID.dkr.ecr.$AWS_ECR_REPOSITORY_REGION.amazonaws.com/$ECR_PRIVATE_REPOSITORY_NAME:$ECR_IMAGE_TAG
    
    # 标记
    docker tag public.ecr.aws/sagemaker/$ECR_GALLERY_IMAGE_ID $ECR_IMAGE_URI
    # 将镜像推送到您的私有存储库
    docker push $ECR_IMAGE_URI
  2. 创建SageMaker图像并将图像附加到Studio域:

    # 创建SageMaker图像
    aws sagemaker create-image \
        --image-name $SAGEMAKER_IMAGE_NAME \
        --role-arn $SAGEMAKER_STUDIO_IAM_ROLE_ARN
    # 创建SageMaker图像版本。
    aws sagemaker create-image-version \
        --image-name $SAGEMAKER_IMAGE_NAME \
        --base-image $ECR_IMAGE_URI
    
    # 可选地,描述图像版本以确保成功创建
    aws sagemaker describe-image-version \
        --image-name $SAGEMAKER_IMAGE_NAME \
        --version-number 1
    
    # 创建应用程序图像配置文件
    cat > /tmp/app-config.json << EOF
    {
       "AppImageConfigName": "app-image-config-$SAGEMAKER_IMAGE_NAME",
       "KernelGatewayImageConfig": { 
          "FileSystemConfig": { 
             "DefaultGid": 100,
             "DefaultUid": 1000,
             "MountPath": "/home/sagemaker-user"
          },
          "KernelSpecs": [ 
             { 
                "DisplayName": "Python 3 (ipykernel)",
                "Name": "python3"
             }
          ]
       }
    }
    EOF
    
    # 创建Amazon SageMaker应用程序图像配置。
    aws sagemaker create-app-image-config \
        --cli-input-json file:///tmp/app-config.json
    
    # 创建默认用户设置文件
    # 如果您有其他自定义映像,请更新文件以使用现有设置
    cat > /tmp/default-user-settings.json << EOF
    {
        "DefaultUserSettings": {
            "KernelGatewayAppSettings": {
                "CustomImages": [
                    {
                        "ImageName": "$SAGEMAKER_IMAGE_NAME",
                        "AppImageConfigName": "app-image-config-$SAGEMAKER_IMAGE_NAME",
                        "ImageVersionNumber": 1
                    }
                ]
            }
        }
    }
    EOF
    
    # 使用新的默认用户设置更新Amazon SageMaker域。
    aws sagemaker update-domain \
        --domain-id $SAGEMAKER_STUDIO_DOMAIN_ID \
        --cli-input-json file:///tmp/default-user-settings.json
  3. 在SageMaker控制台上,通过选择您的域和现有的用户配置文件启动Studio。

  4. 可选地,按照关闭和更新SageMaker Studio中的步骤重新启动Studio。

下载笔记本

从 GitHub 仓库本地下载示例笔记本。

在您选择的 IDE 中打开笔记本并添加一个单元格以安装 torchsummarytorchsummary 包不是发行版的一部分,将其安装在笔记本上将确保笔记本端对端运行。我们建议使用 condamicromamba 管理环境和依赖项。 将以下单元格添加到笔记本中并保存笔记本:

%pip install torchsummary

在本地笔记本上进行实验

通过选择上传图标将笔记本上传到您启动的 JupyterLab UI,如以下屏幕截图所示。

上传后,启动 cv-kmnist.ipynb 笔记本。 您可以立即开始运行单元格,无需安装任何依赖项,如 torch、matplotlib 或 ipywidgets。

如果您遵循了前面的步骤,则可以看到您可以在本地计算机上使用发行版。 在下一步中,我们将在 Studio 上使用相同的发行版,以利用 Studio 的功能。

将实验移至 Studio(可选)

可选地,让我们将实验提升到 Studio。 Studio 的优点之一是底层计算资源完全弹性,因此您可以轻松地将可用资源增加或减少,并且更改会在后台自动进行而不会中断您的工作。 如果您想在较大的数据集和计算实例上运行相同的笔记本,则可以迁移到 Studio。

导航到您之前启动的 Studio UI 并选择上传图标以上传笔记本。

启动笔记本后,将提示您选择映像和实例类型。 在内核启动器上,选择 sagemaker-runtime 作为映像和一个 ml.t3.medium 实例,然后选择 选择

现在,您可以在 Studio 笔记本中运行笔记本,无需在本地开发环境中对笔记本进行任何更改!

将笔记本安排为作业

完成实验后,SageMaker 提供了多种选项来将您的笔记本生产化,例如培训作业和 SageMaker 流水线。 其中一个选项是使用 SageMaker 笔记本作业直接运行笔记本本身作为非交互式、定期笔记本作业。 例如,您可能希望定期重新训练模型,或定期对传入数据进行推断并生成供利益相关者使用的报告。

从 Studio 中选择笔记本作业图标以启动笔记本作业。 如果您已在本地计算机上安装了笔记本作业扩展程序,则还可以直接从本地计算机安排笔记本。 有关设置笔记本作业扩展程序的说明,请参见安装指南。

笔记本作业会自动使用开源发行版的 ECR 映像 URI,因此您可以直接安排笔记本作业。

选择按计划运行,选择一个计划,例如每周六,然后选择创建。如果您想立即查看结果,还可以选择立即运行

当第一个笔记本作业完成后,您可以通过在输出文件下选择笔记本来直接查看Studio UI中的笔记本输出。

其他考虑因素

除了直接使用公开可用的ECR映像以进行ML工作负载外,开源分发还提供以下优点:

  • 用于构建映像的Dockerfile是公开可用的,供开发人员探索和构建自己的映像。您还可以将此映像作为基础映像继承,并安装自定义库以获得可重现的环境。
  • 如果您不习惯Docker,而更喜欢在JupyterLab环境中使用Conda环境,我们为每个发布版本提供了一个env.out文件。您可以使用文件中的说明创建自己的Conda环境,该环境将模仿相同的环境。例如,查看CPU环境文件cpu.env.out。
  • 您可以使用映像的GPU版本来运行与GPU兼容的工作负载,例如深度学习和图像处理。

清理

完成以下步骤来清理您的资源:

  1. 如果您已安排笔记本按计划运行,请暂停或删除笔记本作业定义选项卡上的计划,以避免为未来的作业付费。
  2. 关闭所有Studio应用程序,以避免支付未使用的计算使用费用。请参阅关闭和更新Studio应用程序的说明。
  3. 可选地,删除您创建的Studio域。

结论

在ML生命周期的不同阶段中保持可重现的环境是数据科学家和开发人员面临的最大挑战之一。使用SageMaker开源分发,我们提供了一个带有最常见的ML框架和软件包的相互兼容版本的映像。分发也是开源的,为开发人员提供了对软件包和构建过程的透明度,从而更容易地自定义自己的分发。

在本文中,我们向您展示了如何在本地环境、Studio中以及作为培训作业的容器中使用分发。此功能目前处于公共测试版。我们鼓励您尝试并在公共GitHub存储库上分享您的反馈和问题!