使用@remote装饰器可以访问Amazon SageMaker训练工作负载的私有存储库

使用@remote装饰器访问私有存储库

随着越来越多的客户希望将机器学习(ML)工作负载投入到生产中,组织机构中推动缩短ML代码开发生命周期的力度也越来越大。许多组织更喜欢以Python方法和类的形式编写他们的ML代码,而不是以探索性的方式编写代码(不使用方法或类),因为这有助于更快地发布生产就绪的代码。

通过使用Amazon SageMaker,您可以使用@remote装饰器来运行SageMaker训练作业,只需使用@remote装饰器对Python代码进行注释即可。SageMaker Python SDK会自动将您现有的工作空间环境以及任何关联的数据处理代码和数据集转换为在SageMaker训练平台上运行的SageMaker训练作业。

在本地运行Python函数通常需要安装几个依赖项,这些依赖项可能不在本地Python运行环境中。您可以通过pip或conda等软件包和依赖管理工具来安装它们。

然而,在银行、保险和医疗保健等受监管行业运营的组织在其环境中设置了严格的数据隐私和网络控制措施。这些控制措施通常要求在其任何环境中都不能访问互联网。出于这样的限制的原因是为了对出入流量具有完全控制,以便减少不正当行为者通过其网络发送或接收未经验证的信息的机会。通常还规定在审计和工业合规规定中要求具有这种网络隔离。就机器学习而言,这限制了数据科学家无法从公共存储库(如PyPI、Anaconda或Conda-Forge)下载任何软件包。

为了使数据科学家能够使用他们选择的工具,并同时遵守环境的限制,组织通常在其自己的环境中设置自己的私有软件包存储库。您可以通过多种方式在AWS上设置私有软件包存储库:

  • 使用AWS CodeArtifact
  • 使用Amazon Simple Storage(Amazon S3)
  • 在Amazon Elastic Compute Cloud(Amazon EC2)上托管存储库

在本文中,我们将重点介绍第一种选择:使用CodeArtifact。

解决方案概述

下面的架构图显示了解决方案的架构。

实施解决方案的高级步骤如下:

  • 使用AWS CloudFormation模板设置一个没有互联网访问权限的虚拟私有云(VPC)。
  • 使用第二个CloudFormation模板将CodeArtifact设置为私有PyPI存储库,并为VPC提供连接性,并设置Amazon SageMaker Studio环境以使用私有PyPI存储库。
  • 使用开源SageMaker Python SDK中的@remote装饰器基于MNIST数据集训练一个分类模型。所有依赖项将从私有PyPI存储库中下载。

请注意,本文中的使用SageMaker Studio是可选的。您可以选择在任何集成开发环境(IDE)中工作。您只需要正确设置AWS命令行界面(AWS CLI)凭证。有关更多信息,请参阅配置AWS CLI。

先决条件

您需要一个具有AWS Identity and Access Management(IAM)角色权限的AWS帐户,以管理作为解决方案一部分创建的资源。有关详细信息,请参阅创建AWS帐户。

设置没有互联网连接的VPC

使用vpc.yaml模板创建一个新的CloudFormation堆栈。此模板将创建以下资源:

  • 一个具有两个私有子网的VPC,跨两个可用区,并且没有互联网连接
  • 用于访问Amazon S3的网关VPC端点
  • 为SageMaker、CodeArtifact和其他一些服务设置接口VPC端点,以允许VPC中的资源通过AWS PrivateLink连接到AWS服务

提供一个堆栈名称,例如No-Internet,然后完成堆栈创建过程。

等待堆栈创建过程完成。

使用VPC设置私有存储库和SageMaker Studio

下一步是使用sagemaker_studio_codeartifact.yaml模板部署另一个CloudFormation堆栈。此模板创建以下资源:

  • 连接到先前步骤创建的VPC的SageMaker域
  • CodeArtifact域
  • 连接到上游公共PyPI存储库的CodeArtifact私有存储库

提供堆栈名称并保留默认值或调整以下参数:CodeArtifact域名私有存储库名称SageMaker Studio的用户配置文件名称,以及上游公共PyPI存储库的名称。您还需要提供先前步骤创建的VPC堆栈名称

堆栈创建完成后,SageMaker域应在SageMaker控制台上可见。

为了验证SageMaker Studio中没有可用的互联网连接,请启动SageMaker Studio。选择FileNewTerminal以启动一个终端并尝试curl任何互联网资源。它应该连接失败,如下图所示。

使用私有PyPI存储库的@remote装饰器训练图像分类器

在本节中,我们使用@remote装饰器运行一个PyTorch训练作业,生成一个MNIST图像分类模型。为了实现这一目标,我们设置了一个配置文件,开发了训练脚本,并运行了训练代码。

设置配置文件

我们设置了一个config.yaml文件,并提供了所需的配置,以执行以下操作:

  • 在先前创建的无互联网VPC中运行SageMaker训练作业
  • 通过连接到先前创建的私有PyPI存储库下载所需的软件包

文件的示例如下:

SchemaVersion: '1.0'
SageMaker:
  PythonSDK:
    Modules:
      RemoteFunction:
        Dependencies: '../config/requirements.txt'
        InstanceType: 'ml.m5.xlarge'
        PreExecutionCommands:
            - 'aws codeartifact login --tool pip --domain <domain-name> --domain-owner <AWS account number> --repository <private repository name> --endpoint-url <VPC-endpoint-url-prefixed with https://>
        RoleArn: '<execution role ARN for running training job>'
        S3RootUri: '<s3 bucket to store the job output>'
        VpcConfig:
            SecurityGroupIds: 
            - '<security group id used by SageMaker Studio>'
            Subnets: 
            - '<VPC subnet id 1>'
            - '<VPC subnet id 2>'

Dependencies字段包含指向requirements.txt的路径,该文件包含所有所需的依赖项。请注意,所有依赖项都将从私有存储库下载。 requirements.txt文件包含以下代码:

torch
torchvision
sagemaker>=2.156.0,<3

PreExecutionCommands部分包含连接到私有PyPI存储库的命令。要获取CodeArtifact VPC端点URL,请使用以下代码:

response = ec2.describe_vpc_endpoints(
    Filters=[
        {
            'Name': 'service-name',
            'Values': [
                f'com.amazonaws.{boto3_session.region_name}.codeartifact.api'
            ]
        },
    ]
)

code_artifact_api_vpc_endpoint = response['VpcEndpoints'][0]['DnsEntries'][0]['DnsName']

endpoint_url = f'https://{code_artifact_api_vpc_endpoint}'
endpoint_url

通常,我们可以获取两个CodeArtifact的VPC端点,并且可以在连接命令中使用其中任何一个。有关更多详细信息,请参阅在VPC中使用CodeArtifact。

此外,配置文件中还提供了诸如execution roleoutput locationVPC configurations等配置。这些配置是运行SageMaker训练作业所需的。要了解支持的所有配置的更多信息,请参阅配置文件。

为了使用@remote装饰器与config.yaml文件一起工作并不是强制的。这只是以更清晰的方式向@remote装饰器提供所有配置的方法。所有的配置也可以直接在装饰器的参数中提供,但这会降低长期来看更改的可读性和可维护性。此外,配置文件可以由管理员创建并与环境中的所有用户共享。

开发训练脚本

接下来,我们用简单的Python文件准备训练代码。我们将代码分为三个文件:

  • load_data.py – 包含下载MNIST数据集的代码
  • model.py – 包含用于模型的神经网络架构的代码
  • train.py – 包含使用load_data.py和model.py训练模型的代码

train.py中,我们需要将主要的训练函数装饰如下:

@remote(include_local_workdir=True)
def perform_train(train_data,
                  test_data,
                  *,
                  batch_size: int = 64,
                  test_batch_size: int = 1000,
                  epochs: int = 3,
                  lr: float = 1.0,
                  gamma: float = 0.7,
                  no_cuda: bool = True,
                  no_mps: bool = True,
                  dry_run: bool = False,
                  seed: int = 1,
                  log_interval: int = 10,
                  ):
    # pytorch本地训练代码........

现在我们准备好运行训练代码了。

使用@remote装饰器运行训练代码

我们可以从终端或任何可执行提示符中运行代码。在本文中,我们使用SageMaker Studio笔记本单元格来演示:

!python ./train.py

运行上述命令会触发训练作业。在日志中,我们可以看到它正在从私有PyPI存储库下载软件包。

这样就完成了在没有互联网访问权限的环境中,使用@remote装饰器与私有存储库一起工作的实现。

清理

要清理资源,请按照CLEANUP.md中的说明进行操作。

结论

在本文中,我们学习了如何在没有互联网访问权限的限制性环境中有效使用@remote装饰器的功能。我们还学习了如何使用SageMaker中的配置文件支持将CodeArtifact私有存储库功能集成起来。这个解决方案使得迭代开发变得更加简单快速。另一个附加的优点是,您仍然可以以更自然、面向对象的方式编写训练代码,并且仍然可以使用SageMaker的能力在远程集群上运行训练作业,而对代码的更改很小。本文中展示的所有代码都可以在GitHub存储库中找到。

作为下一步,我们鼓励您查看@remote装饰器功能和Python SDK API,并在您选择的环境和IDE中使用它。amazon-sagemaker-examples存储库中还有其他示例,可以帮助您快速入门。您还可以查看文章Run your local machine learning code as Amazon SageMaker Training jobs with minimal code changes以获取更多详细信息。