「Amazon SageMakerトレーニングワークロードで@remoteデコレータを使用してプライベートリポジトリにアクセスする」

「Amazon SageMakerトレーニングワークロードで@remoteデコレータを使用してプライベートリポジトリにアクセスする」' - Access private repository using @remote decorator in Amazon SageMaker training workload.

機械学習（ML）のワークロードを本番環境に展開する顧客が増えるにつれて、組織内でMLコードの開発ライフサイクルを短縮するための取り組みが進んでいます。多くの組織では、Pythonのメソッドやクラスの形式でMLコードを本番用に書くことを好むことがあります。これにより、本番用のコードをより迅速に提供することができます。

Amazon SageMakerを使用すると、Pythonコードに@remoteデコレータを付けることで、SageMakerトレーニングジョブを簡単に実行できます。SageMaker Python SDKは、既存のワークスペース環境と関連するデータ処理コード、データセットをSageMakerトレーニングジョブに自動的に変換して実行します。

Python関数をローカルで実行する場合、通常は複数の依存関係が必要ですが、これらはローカルのPythonランタイム環境に含まれていない場合もあります。パッケージや依存関係の管理ツール（pipやcondaなど）を使用してインストールすることができます。

ただし、銀行、保険、医療などの規制業界で運営されている組織では、厳格なデータプライバシーやネットワーキングの制御が行われています。これらの制御は、環境内でのインターネットアクセスを許可しないことがよく求められます。この制限の理由は、出入りトラフィックを完全に制御して、不正なアクターがネットワークを介して検証されていない情報を送受信する可能性を低減するためです。また、監査や産業規制の一環として、このようなネットワークの分離を必要とする場合もあります。機械学習の場合、これによりデータサイエンティストはPyPI、Anaconda、Conda-Forgeなどの公開リポジトリからパッケージをダウンロードすることができません。

• 「Amazon SageMakerのトレーニングワークロードで、@remoteデコレータを使用してプライベートリポジトリにアクセスする」
• 「AIとの親交を深める」
• AIに親しむ

データサイエンティストが環境の制約を尊重しながら、選択したツールにアクセスできるようにするために、組織は独自のプライベートパッケージリポジトリを設定することがあります。AWS上でプライベートパッケージリポジトリを設定する方法はいくつかあります:

• AWS CodeArtifactを使用する
• Amazon Simple Storage（Amazon S3）を使用する
• Amazon Elastic Compute Cloud（Amazon EC2）上にリポジトリをホストする

この記事では、最初のオプションであるCodeArtifactの使用に焦点を当てます。

ソリューションの概要

以下のアーキテクチャ図は、ソリューションのアーキテクチャを示しています。

このソリューションの実装手順は次のとおりです。

• AWS CloudFormationテンプレートを使用して、インターネット接続のない仮想プライベートクラウド（VPC）をセットアップします。
• 2番目のCloudFormationテンプレートを使用して、CodeArtifactをプライベートなPyPIリポジトリとしてセットアップし、VPCへの接続性を提供し、Amazon SageMaker Studio環境がプライベートなPyPIリポジトリを使用できるようにします。
• オープンソースのSageMaker Python SDKの@remoteデコレータを使用して、MNISTデータセットに基づく分類モデルをトレーニングします。すべての依存関係はプライベートなPyPIリポジトリからダウンロードされます。

この記事では、SageMaker Studioの使用はオプションです。お好みの統合開発環境（IDE）で作業することもできます。ただし、AWS CLIの資格情報を正しく設定する必要があります。詳細については、「AWS CLIの設定」を参照してください。

事前準備

ソリューションの一部として作成されるリソースを管理する権限を持つAWS Identity and Access Management（IAM）ロールを持つAWSアカウントが必要です。詳細については、「AWSアカウントの作成」を参照してください。

インターネット接続のないVPCのセットアップ

vpc.yamlテンプレートを使用して新しいCloudFormationスタックを作成します。このテンプレートでは、次のリソースが作成されます:

• インターネット接続のない2つのアベイラビリティーゾーンにまたがる2つのプライベートサブネットを持つVPC
• Amazon S3にアクセスするためのゲートウェイVPCエンドポイント
• SageMaker、CodeArtifact、その他のサービスへのリソースがAWS PrivateLinkを介してAWSサービスに接続できるようにするためのインターフェースVPCエンドポイント

「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を起動します。[ファイル]、[新規作成]、[ターミナル]を選択してターミナルを起動し、インターネットリソースに対してcurlコマンドを実行してみてください。次のスクリーンショットに示すように、接続に失敗するはずです。

プライベートPyPIリポジトリを使用して@remoteデコレータを使って画像分類器をトレーニングする

このセクションでは、@remoteデコレータを使用して、MNIST画像分類モデルを生成するPyTorchトレーニングジョブを実行します。これを実現するために、設定ファイルを設定し、トレーニングスクリプトを開発し、トレーニングコードを実行します。

設定ファイルの設定

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には2つのVPCエンドポイントがあり、接続コマンドでどちらでも使用できます。詳細については、「VPCからCodeArtifactを使用する」を参照してください。

また、execution role、output location、およびVPC configurationsなどの設定は、設定ファイルで提供されます。これらの設定は、SageMakerトレーニングジョブを実行するために必要です。サポートされるすべての設定についての詳細は、「Configuration file」を参照してください。

@remoteデコレーターを使用するためには、必ずしもconfig.yamlファイルを使用する必要はありません。これは、すべての設定を@remoteデコレーターに直接指定するためのよりクリーンな方法です。すべての設定をデコレーターの引数に直接指定することもできますが、これは可読性や保守性を損なう可能性があります。また、管理者が設定ファイルを作成し、環境のすべてのユーザーと共有することもできます。

トレーニングスクリプトの開発

次に、シンプルなPythonファイルでトレーニングコードを準備します。コードは3つのファイルに分割されています：

• 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プライベートリポジトリの機能を統合する方法も学びました。このソリューションにより、反復的な開発がより簡単かつ高速になります。もう一つの利点は、トレーニングコードをより自然でオブジェクト指向的な方法で記述し、コードの変更を最小限に抑えながらリモートクラスタでトレーニングジョブを実行することができることです。この投稿のすべてのコードはGitHubリポジトリで利用できます。

次のステップとして、@remoteデコレーターの機能とPython SDK APIを確認し、選択した環境とIDEで使用してみることをお勧めします。迅速に始めるために、amazon-sagemaker-examplesリポジトリに追加の例も用意されています。詳細については、「Run your local machine learning code as Amazon SageMaker Training jobs with minimal code changes」の投稿もご覧ください。

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!