初めに
備忘録として久しぶりに記事を書きます（２回目の投稿もVPC関連...）。
転職してからは機械学習や生成AI周りの案件を社内でひそひそと取り組んでおります。
本記事では機械学習周りのAWSサービスであるSageMaker StudioをVPCオンリーでMLFlow(SageMaker managed)と接続させるまでに起きた事象と対処手順について記載します。（簡単な内容ではあるのですが、一部AWSサポートの方にもお手伝いいただいたため）
状況
とある機械学習案件でSageMaker Studio及びMLFlowをVPCオンリーの環境構築をする必要が生じたため、下図のような環境を構築したい。
ただし、当初は上図の？の部分であるSageMaker managedのMLFlow側のネットワークについてよく理解できていませんでした。
要件整理
先ほどの？の箇所を除いて要件を整理すると次のようになります。
- SageMaker Studioに付けるロールとしてはSageMaker系（MLFlowは別途あるので注意）とS3,Dynamo等のデータ系とECRあたりが必要
- VPCエンドポイントも S3（Gateway Endpoint）、SageMaker API / Runtime、MLflow（Interface Endpoint）等が必要
- セキュリティグループ（SG）のインバウンドについては、それぞれのVPCのSGからのhttpsのポート443を開ける必要がある
といった具合で当初は整理していました。
実際に起きた事象
上記で整理したリソースについて作成し、SageMaker Studioの起動には成功しました。
またSageMaker Studio内のJupyterLabについても、起動時間はかかるものの起動には成功している状態でした。
しかし、下記のコマンドでMLFlowとの接続を試そうとしたときにハングアウトしてしまいました。
import mlflow
# Tracking URI を設定
# 今回はSageMaker Managed MLflow のため以下形式
tracking_uri = "arn:aws:sagemaker:<region>:<account-id>:mlflow-tracking-server/<tracking-server-name>"
mlflow.set_tracking_uri(tracking_uri)
# 実験（Experiment）を指定
experiment_name = "mlflow-connection-test"
mlflow.set_experiment(experiment_name)
# ③ Run を作ってログを投げる
with mlflow.start_run(run_name="connectivity-test") as run:
print(f"Run ID: {run.info.run_id}")
# パラメータ
mlflow.log_param("param1", 123)
# メトリクス
mlflow.log_metric("accuracy", 0.95)
print("Logged successfully!")
しかし、MLFlowのGUIも確認できており、各種VPCエンドポイントの名前解決もできていたため解決までに沼ってしまいました。。
原因
結論から記載すると、先ほどの構成図にはSageMaker Studio自身から自身への疎通の考慮が抜けていることが原因でした。
具体的にはSageMaker Studio が利用する ENI から、同一セキュリティグループに属する ENI（VPC Endpoint 等）への通信許可が不足していたことが原因でした。
そもそも公式のドキュメントをよく読んでみると12、SageMaker managedのMLFlowはSageMaker管理のVPC内で展開されることが分かりました。（自環境のVPCエンドポイントとの疎通はよしなにしてくれるらしい、そのため厳密には自環境のVPCのみには収まっていない）
また、そのSageMaker管理のVPC内で同一リソースでの疎通が発生することをサポートの方から教わりました。
ちなみにMLFlowのアーティファクト等はちゃんと指定した自環境のS3に配置されます、サーバーのみがSageMaker管理のVPCに存在する形です。
対処
対処方法としてはいたってシンプルで、SageMaker StudioのSGに自身のSGからのトラフィックをインバウンドに追加した後、SageMaker Studioのドメインを修正したSGに紐づけて作成しなおし、下記のコマンドでドメインを更新したところ先ほどのハングアウトが解消し、MLFlowのGUIからExperimentの登録が確認できました。（根本すぎて恥ずかしかったです。。）
aws sagemaker update-domain \
--domain-id d-xxxxxxxxxxxx \
--default-user-settings '{
"SecurityGroups": ["sg-xxxxxxx"]
}' \
--subnet-ids subnet-xxxx subnet-yyyy
最後に
内容的には初歩の初歩なのですが、そこまで古いサービスではないからか日本語での開設等がなかったため今回記事を書いてみました。
実際に構築してみて、ある程度ネットワークやインスタンスについては良しなにやってくれるので便利ではあったのですが、維持してみるとちょっと値段がお高い印象でした。。。
もっと安く維持できる方法があればコメントお願いします。