REST API プロキシを用いたAmazon SageMaker MLflowへの外部アクセスの簡素化

機械学習 (ML) チームは、MLflow を用いて ML のライフサイクルを効果的に管理しています。Amazon SageMaker MLflow は、包括的な ML 実験の追跡およびモデル管理機能を提供します。しかし、多くの企業では、SDK の直接使用ではなく、HTTPS ベースの統合を必要とする既存のインフラ要件を持っています。

多くの組織は、セキュリティやインフラのパターンを維持しながら、Amazon SageMaker MLflow を既存システムと統合する必要があります。この統合課題は、企業のセキュリティポリシー、ネットワーク制限、またはレガシーシステムの制約により SDK を直接使用できないチームに影響を与えます。

本稿では、MLflow SDK の使用を必要とせずに Amazon SageMaker MLflow への HTTPS アクセスを提供する、安全な Flask ベースの MLflow プロキシサービスの構築方法を実演します。このソリューションは、既存の ML ワークフローを維持しつつクラウドネイティブサービスを採用したい、クラウド変革を進めている組織向けです。

本稿では以下のトピックを取り上げます:

• MLflow の HTTPS リクエスト用の MLflow プロキシサービスの実装。

• 安全なアクセスのための AWS Identity and Access Management (IAM) 認証の構成。

• URL プレサインおよびリクエスト変換の管理。

このソリューションを実装することで、以下が可能になります:

• 標準的な HTTPS エンドポイントを通じて、SageMaker MLflow を安全にアクセスできます。

• お使いの組織のセキュリティ要件への準拠を維持します。

• MLflow を既存のエンタープライズシステムと統合します。

• 実装の複雑さとメンテナンスオーバーヘッドを削減します。

ソリューション概要

軽量な Flask ベースの MLflow プロキシアーキテクチャは、3 つの主要コンポーネントを通じて、エンタープライズシステムと Amazon SageMaker MLflow の間の安全な統合を提供します。

コンポーネント 1: アプリケーションロードバランサー (ALB)

AWS Application Load Balancer はアップストリームルーターとして機能し、以下の機能を提供します:

• MLflow UI および REST API リクエストのトラフィック分散。

• 初期リクエストの処理とルーティング。

• カスタムドメイン名および SSL 終端のサポート。

注: この実装では ALB を使用していますが、要件に応じて Nginx ベースなどの他のルーティングソリューションを使用することも可能です。

コンポーネント 2: Flask MLflow プロキシサービス

アーキテクチャの中核となる Python ベースの Flask アプリケーションは、以下の処理を担当します:

• 着信 HTTPS リクエストのインターセプトと処理。

• AWS 認証およびリクエスト署名の管理。

• セキュアな MLflow エンドポイントアクセスのための URL 変換。

• クライアントへのレスポンスルーティングの処理。

コンポーネント 3: Amazon SageMaker MLflow

AWS マネージドの SageMaker MLflow サービスは、以下の機能を提供します:

• MLflow の 2 つのデプロイモードへの対応:

MLflow Tracking Server – 管理された MLflow トラッキングサーバー。

• MLflowApp – サーバーレス MLflow アプリケーション。

• 追跡情報のためのバックエンドメタデータストア。

• モデルファイルおよびデータのストレージ。

このアーキテクチャは、既存のエンタープライズシステムとの互換性を維持しつつ、安全な通信を提供します。プロキシサービスは橋渡し役として機能し、標準的な HTTPS リクエストを認証された AWS API 呼び出しに変換して、SageMaker MLflow と対話できるようにします。

アーキテクチャとリクエストワークフロー

以下の図は、Flask プロキシサービスが外部クライアントと Amazon SageMaker MLflow の間で安全な通信を提供する方法を示しています。

アーキテクチャ図には、3 つの主要コンポーネントが表示されています:

• 着信トラフィックを処理する ALB (Application Load Balancer)。

• 認証とリクエスト変換を管理する Flask プロキシサービス。

• ML 操作を処理する Amazon SageMaker MLflow。

リクエストワークフロー

このアーキテクチャを通じて、安全な MLflow アクセスがどのように提供されるか、リクエストの流れを見てみましょう。

クライアントが HTTPS リクエストを開始すると、まずすべての着信トラフィックのエントリーポイントとして機能する ALB に到達します。その後、ALB はこれらのリクエストを MLflow プロキシサービスにルーティングします。

リクエストを受信すると、MLflow プロキシサービスは以下の重要な機能を実行します:

• AWS IAM 連携による認証処理を行います。
• URL を変換し、安全なアクセスのために事前署名を付与します。
• 必要に応じて MLflow REST API エンドポイントを処理します。

MLflow プロキシサービスは、SageMaker MLflow の REST エンドポイントに対して API コールを行う前に、受信したリクエストを認証付きの AWS リクエストに変換します。SageMaker MLflow がリクエストを処理した後、応答が返されますが、これは MLflow プロキシサービスによって再度処理され、元のクライアントへルーティングされます。

このワークフローは、エンタープライズシステムと SageMaker MLflow の間の統合を提供しつつ、セキュリティを維持します。

前提条件

このウォークスルーを進めるには、以下の準備が必要です:

• AWS アカウント。

• 以下のツールがインストールされたワークステーション:

以下を作成する権限を設定した AWS Command Line Interface (AWS CLI)：

Amazon Virtual Private Cloud (Amazon VPC) および関連するネットワークコンポーネント。

• Amazon Elastic Compute Cloud (Amazon EC2) インスタンス。

• Amazon SageMaker AI リソース。

• Amazon Simple Storage Service (Amazon S3) バケット。

• AWS Identity and Access Management (IAM) ロールおよびポリシー。

• AWS CloudFormation スacks。

• AWS Application Load Balancers。

• Node.js バージョン 18.0.0 以降。

• NPM。

• AWS Cloud Development Kit (AWS CDK) CLI バージョン 2.100.0 以降。

• pip または pip3 を備えた Python 3.x。

• 必要な知識：

AWS サービスおよび IAM 権限に関する基本的な理解。

• Python および Flask アプリケーションへの慣れ。

• MLflow の概念と運用に関する理解。

• コストに関する考慮事項：

このソリューションは、コストが発生する可能性のある AWS リソースを作成します。

• 主要なコスト要因となるリソースには以下が含まれます：

Amazon EC2 インスタンス。

• Application Load Balancer。

• Amazon SageMaker AI リソース。

• Amazon S3 ストレージ。

AWS サービスの価格に関する情報は、AWS Pricing Calculator をご覧ください。

ソリューションのデプロイ

このセクションでは、AWS アカウント内でソリューションをデプロイし、検証する方法について順を追って説明します。デプロイプロセスには約 40 分かかります。

ステップ 1: AWS CDK を使用してインフラストラクチャをデプロイする

• ソリューションコードをダウンロードし、依存関係をインストールします:

# リポジトリをクローン

git clone https://github.com/aws-samples/sample-sagemaker-mlflow-rest-apis.git

プロジェクトディレクトリに移動して依存関係をインストール

cd sample-sagemaker-mlflow-rest-apis

npm ci

• AWS CDK 用に環境をブートストラップします。AWS アカウントとリージョンがすでに AWS CDK 用にブートストラップされている場合は、このステップはスキップしてください。

CDK のために AWS アカウントとリージョンをブートストラップ:

npx cdk bootstrap aws:///

• AWS アカウントに必要なリソースをデプロイします。本ソリューションは 4 つの CDK スタックで構成されています:

ネットワーキングスタック — VPC およびネットワークコンポーネントを作成します。

SageMaker AI ドメインスタック — SageMaker ドメインを設定します。

SageMaker MLflow スタック — MLflow 追跡サーバーまたは MLflow サーバーレスアプリをデプロイします。

Flask アプリケーションスタック — MLflow プロキシサービス（proxy service）をデプロイします。

以下のいずれかのコマンドを使用して、すべてのスタックをデプロイしてください。

追跡サーバーベースのデプロイの場合:

npx cdk deploy --all --require-approval=never -c mlflowType=tracking

サーバーレスアプリベースのデプロイの場合:

npx cdk deploy --all --require-approval=never -c mlflowType=serverless

ステップ 2: Flask MLflow プロキシサービスのインストールと設定

• EC2 インスタンスへの接続:

CDK の出力、または sagemaker-infra-flaskapp-{mlflowType} AWS CloudFormation スタックの出力セクションから Amazon EC2 インスタンス ID を確認してください。

• AWS Systems Manager Session Manager を使用して接続します。Session Manager 接続ガイドに従ってください。

• Python 3.13 と依存関係、および必要なパッケージをインストールします:

# ルートユーザーに切り替え

sudo su -

cd /root

Python と依存関係をインストール

chmod +x install_python13.sh

./install_python13.sh

注: このスクリプトは Ubuntu ベースのシステム向けに設計されています。他の Linux ディストリビューションの場合は、システムのパッケージマネージャーを使用して Python 3.12 以上、PIP3、および Virtualenv をインストールしてください。

• MLflow プロキシサービスをインストールして起動します:

chmod +x setup_mlflow_proxy_app.sh

./setup_mlflow_proxy_app.sh

• Flask MLflow プロキシサービスのステータスを確認します:

systemctl status mlflowproxy

注: サービスが実行されていない場合は、以下のコマンドでログを確認してください:

journalctl -u mlflowproxy

ステップ 3: MLflow REST API アクセスの検証

このセクションでは、ALB を介して MLflow REST APIs と対話する方法を示します。

*注意：これらの例では HTTP（非暗号化）プロトコルを使用しています。本番環境では HTTPS の使用を推奨します。本記事では API リクエストに curl を使用していますが、お好みのツールをご利用いただけます。提供された curl コマンドは、トラッキングサーバーモードとサーバーレスモードの両方で同じように機能し、プロキシサービスが違いを透過的に処理します。*

• ご自身のワークステーションで以下のコマンドを実行して、ALB の DNS 名を取得してください:

aws cloudformation describe-stacks --stack-name sagemaker-infra-flaskapp-{mlflowType} --query 'Stacks[0].Outputs[?OutputKey==ALBUrl].OutputValue' --output text

• ご自身のワークステーションで以下のコマンドを実行して、MLflow API エンドポイントのテストを行ってください。 , , , および を適切な値に置き換えてください。

実験の作成:

curl -X POST http:///ajax-api/2.0/mlflow/experiments/create -H "Content-Type: application/json" -d '{"name": "mlflow-experiment"}'

• 実験の検索:

curl -X POST http:///ajax-api/2.0/mlflow/experiments/search -H "Content-Type: application/json" -d '{"max_results": 5}'

• 実験の取得:

curl -X GET 'http:///ajax-api/2.0/mlflow/experiments/get?experiment_id=0'

• 実験内のランの実行:

curl -X POST http:///ajax-api/2.0/mlflow/runs/create -H "Content-Type: application/json" -d '{"experiment_id": , "run_name": ""}'

• ランからのアーティファクトのリスト表示:

curl -X GET "http:///ajax-api/2.0/mlflow/artifacts/list?run_id="

• ランへのタグ設定:

curl -X POST "http:///ajax-api/2.0/mlflow/runs/set-tag" -H "Content-Type: application/json" -d '{"run_id": "", "key": "model_type","value": "api-test"}'

• ランの削除:

curl -X POST http:///ajax-api/2.0/mlflow/runs/delete -H "Content-Type: application/json" -d '{"run_id": ""}'

*注意：MLflow UI を開き、前述の curl コマンドを使用して行った変更を確認することもできます。MLflow UI の起動方法については、プレサイン URL を使用した MLflow UI の起動を参照してください。*

クリーンアップ

継続的な課金を避け、このソリューションによって作成されたリソースを削除するには、以下のクリーンアップ手順に従ってください:

• CDK 管理リソースの削除。ワークステーション上のクローンしたリポジトリのルートディレクトリに移動し、以下を実行してください。

追跡サーバーベースのデプロイの場合:

npx cdk destroy --all -c mlflowType=tracking

サーバーレスアプリベースのデプロイの場合:

npx cdk destroy --all -c mlflowType=serverless

注: ネットワーキングおよび SageMaker ドメインスタックは、両方のデプロイモードで共有されています。AWS CDK は、最後の MLflow または Flask アプリスタックペアが削除された場合にのみ、これらを削除します。

• 手動リソースのクリーンアップ。保持ポリシーや依存関係のため、一部のリソースを手動で削除する必要がある場合があります:

Amazon S3 バケット:

Amazon S3 コンソールに移動してください。

• このソリューションによって作成されたバケットを特定します。

• 各バケットの中身を空にし、削除します。

• Amazon CloudWatch ロググループ:

CloudWatch コンソールで、このソリューションに関連付けられたロググループを検索します。

• これらのロググループを削除します。

• Flask ベースのプロキシサービスに対して Amazon CloudWatch の監視を構成し、アプリケーションの健全性を追跡し、異常を検出し、不審な活動に対するアラートを設定します。

• 潜在的なサービス拒否 (DoS) アタックから保護し、個々のクライアントからのリクエスト数を制御するために、Flask ベースのプロキシサービスに対してレート制限を実装します。AWS WAF（ウェブアプリケーションファイアウォール）を ALB と併用して、レートベースルールを実装できます。

• プロキシアクセスをプライベートネットワークに制限するため、内部（インターネット非公開）ALB をデプロイします。この設定により、VPC 内または接続されたネットワークからのトラフィックのみがサービスに到達できるようになります。VPC ピアリングまたは AWS Transit Gateway を通じて接続します。

• クライアントとアプリケーション間の安全な通信のために、ALB レベルで HTTPS 終端を有効化します。AWS Certificate Manager (ACM) を使用して、アプリケーション用の SSL/TLS 証明書のプロビジョニングおよび管理を行うことができます。HTTPS リスナーの構成に関する手順については、Application Load Balancer の HTTPS リスナードキュメントを参照してください。

これらのセキュリティ対策は、一般的な Web 脆弱性から Flask アプリケーションを保護し、コンポーネント間の安全な通信を提供します。

結論

本記事では、Amazon SageMaker MLflow への HTTPS アクセスを提供する安全な Flask ベースのプロキシサービスの構築方法について紹介しました。このソリューションは、既存のインフラストラクチャと AWS マネージド MLflow の機能を橋渡ししながら、エンタープライズのセキュリティ要件を維持することを組織に支援します。

ソリューションのメリット:

• 既存のエンタープライズセキュリティコントロールとの統合。
• 既存の ML ワークフローへの最小限の変更。
• デプロイメントの複雑さの低減。
• REST API の統合。
• エンタープライズプロキシサービスとの互換性。

次のステップ

Amazon SageMaker MLflow および関連トピックについてさらに詳しく知りたい場合は、以下を実行できます:

• Amazon SageMaker MLflow のドキュメントを参照する。
• MLflow 追跡サーバーおよび MLflow アプリケーションに関する記事を読む。

このソリューションをご自身の環境でお試しください。その経験をコメント欄でお知らせください。

執筆者について

Manish Garg

Manish は AWS Professional Services のデリバリーコンサルタントであり、AWS Cloud 上での顧客ワークロードの移行と近代化を専門としています。彼は技術に対する深い情熱を持ち、DevOps プラクティスの分野にも強い関心を寄せています。