PrivateLinkによる低レイテンシー推論

SageMaker 公式ドキュメントで、SageMaker Endpointの低レイテンシー推論のベストプラクティスとして、PrivateLinkの利用が紹介されています。
低レイテンシーのリアルタイム推論AWS PrivateLink - Amazon SageMaker

リアルタイム推論のレイテンシーは、2つの要素で構成されています。

オーバーヘッドレイテンシー：ネットワーク通信など、推論処理以外でかかる時間
推論レイテンシー：MLモデルの推論処理にかかる時間

オーバヘッドレイテンシーの原因として以下の説明がされています。

デフォルトでは、 SageMaker 2 つ以上のインスタンスを持つエンドポイントが少なくとも 2 つにデプロイされているAWS アベイラビリティーゾーン (AZ) とどの AZ のインスタンスでも呼び出しを処理できます。
その結果、1 つ以上の AZ の「ホップ」が発生し、オーバーヘッド遅延の原因となります。

つまり、デフォルトのSageMaker Endpointでは「冗長化された複数のAZのインスタンスにリクエストを割り振るため、AZ間通信が発生しネットワーク通信時間が増える」ということです。

上記の問題に対し、PrivateLinkの「プライベート DNS 名の有効化」を利用することで以下の状況を実現できると紹介されてます。

SageMaker Endpointへのトラフィックを同一VPC 内に保持する
呼び出しトラフィックと発信元のクライアントを同じAZに保持する。AZ間通信が回避され、オーバーヘッドレイテンシーが減少する。

本記事では、PrivateLinkを介して、閉鎖網VPC内のEC2からSageMaker Endpointへリクエストを投げる手順をまとめます。

手順

VPCエンドポイントの作成

AWS PrivateLinkを使用するには、VPCエンドポイントを作成し、SageMaker Endpointへ接続を行います。

PrivateLinkとVPCエンドポイントの説明はこちらの記事が参考になります。
【AWS】PrivateLinkとVPCエンドポイントについて、使い方と注意点の話 - BFT名古屋 TECH BLOG

マネジメントコンソールのVPCサービスページから、VPCエンドポイントの作成を行います。

サービスに、推論エンドポイントに対して推論処理を実行するAPIで必要になる「com.amazonaws.ap-northeast-1.sagemaker.runtime」を選択します。

VPCとサブネットを選択します。ap-northeast-1a, ap-northeast-1cのプライベートサブネットを選択しました。

「DNS名を有効化」にチェックを入れます。

VPCエンドポイントに関連付けるセキュリティグループを選択します。

デフォルトのセキュリティグループを選択しました。

推論処理を実行するためのVPCエンドポイントができました。

SageMaker Endpointの作成

VPCエンドポイント作成時に指定したサブネットに、SageMaker Endpointをデプロイします。

VPCエンドポイントに関連付けたセキュリティグループを指定します。

モデル作成は本記事では扱いませんが、こちらにSageMakerでのモデル作成手順をまとめています。
SageMaker で学習ジョブを実行する ~組み込みアルゴリズム~ - 肉球でキーボード

import sagemaker

sess = sagemaker.Session()

model_name = "privatelink-model"
endpoint_name = 'privatelink-endpoint'

subnet_a = "subnet-02d49088de4cb6ac2"
subnet_c = "subnet-0b59cdf939f464244"
security_group = "sg-07c022fcdbd50acbd"

sess.create_model(
    model_name,
    role,
    model.prepare_container_def(),
    vpc_config={
        "SecurityGroupIds": [security_group],
        "Subnets": [subnet_a, subnet_c],
    },
)

predictor = model.deploy(
    initial_instance_count=2,
    instance_type='ml.t2.medium',
    endpoint_name=endpoint_name,
)

SageMaker Endpointが作成され、AZ指定ができていることが確認できます。

閉鎖網VPC内のEC2インスタンスからリクエスト

SageMaker Endpointと同一VPC内のプライベートサブネットのEC2から、VPC Endpointを介して推論リクエストを送ります。

以下のような構成を作成します。

SageMaker Endpointと同一のVPCにEC2インスタンスを起動します。

同一VPCのap-northeast-1aサブネットに起動するEC2のネットワーク設定は以下のようになります。

本記事ではprivate subnet内のEC2インスタンスの接続方法は省略しますが、以下の方法があります。

踏み台サーバーの利用: EC2踏み台サーバを経由してssh接続 | ITサービス指向エンジニアブログ
AWS Systems Manager (SSM) のSession Managerの利用: 閉域網で Amazon SageMaker を利用する際のポイントと手順 | Amazon Web Services ブログ

EC2から推論リクエストを送ります。

import boto3

runtime = boto3.Session().client("sagemaker-runtime")

endpoint_name = 'privatelink-endpoint'
response = runtime.invoke_endpoint(
        EndpointName=endpoint_name,
        ContentType="application/json",
        Accept="application/json",
        Body="0",
)

閉鎖網VPC内でSageMaker Endopointへの推論リクエストを送る手順は以上になります。