MLflow+AWS+SageMakerの環境を立ち上げる

機械学習の実験管理ツールの一つに、MLflowがあります。MLflowを使うと、学習時のハイパーパラメータや精度を記録することができ、効率的にモデルの開発を進めることができるようになります。機械学習の実験管理を検討されている方は、一度は聞いたことがあるツールかと思います。
MLflowをローカル環境で動かしたあと、これをクラウド上で利用するため、簡単な構成のサービスの動きを確認したいときがあるかと思います。しかし、MLflowをローカルではなくクラウド(特にAWS)上にデプロイし、クライアントが学習した結果が保存されるまでの流れを、通しで説明している情報は少ないと感じています。そこで本記事では、MLflowをAWS上にデプロイし、クライアントPCまたはSageMakerで実行された学習結果を確認するまでの手順を説明したいと思います。

はじめに

MLflowをAWS上で動かすサンプルを探すと、AWSのブログで紹介されている以下の記事がヒットします。本投稿では、この記事で紹介されているサンプルについて、AWSへのMLflowのデプロイ方法とMLflowの使い方について説明したいと思います。
https://aws.amazon.com/jp/blogs/news/machine-learning-managing-your-machine-learning-lifecycle-with-mlflow-and-amazon-sagemaker/

事前準備

上の記事で紹介されている方法でAWSにMLflowをデプロイするためには、デプロイを実行するPCに下のライブラリを事前にインストールしておく必要があります

• Python
• node
• AWS CLI
• AWS CDK

私は以下のバージョンのでAWSへデプロイできることを確認できました。動作確認した環境の構築方法を紹介したいと思います。すでにインストールができている場合は、「AWSにMLflowをデプロイする」までスキップしてください。

モジュール	バージョン
OS	Ubuntu 20.04 (Windows 10上のVirtualBox)
Python	v.3.9.7
node	v.14.18.1
AWS CLI	v.2.3.0
AWS CDK	v.1.110.0

Pythonのインストール

Pythonのインストールは、以下のサイトに書かれている方法でインストールできます。
https://www.python.jp/install/ubuntu/index.html

以下のコマンドを入力すると、Pythonのバージョン3.9.7をインストールできます。

$ sudo apt update
$ sudo apt install build-essential libbz2-dev libdb-dev \
  libreadline-dev libffi-dev libgdbm-dev liblzma-dev \
  libncursesw5-dev libsqlite3-dev libssl-dev \
  zlib1g-dev uuid-dev tk-dev
$ wget "https://www.python.org/ftp/python/3.9.7/Python-3.9.7.tar.xz"
$ tar xJf Python-3.9.7.tar.xz
$ cd Python-3.9.7
$ ./configure
$ make
$ sudo make install

nodeのインストール

nodeのインストールは、以下のサイトに書かれている方法でインストールしました。
https://qiita.com/seibe/items/36cef7df85fe2cefa3ea

以下のコマンドを入力すると、nodeのバージョン14.18.1をインストールできます。

$ sudo apt install -y nodejs npm
$ sudo npm install n -g
$ sudo n 14.18.1                 % 安定版はsudo n stable
$ sudo apt purge -y nodejs npm
$ exec $SHELL -l

AWS CLIのインストール

AWS CLIのインストールは、AWS公式のドキュメントに書かれている方法でインストールできます。
https://docs.aws.amazon.com/ja_jp/cli/latest/userguide/install-cliv2-linux.html

以下のコマンドを入力すると、AWS CLIの最新版バージョン2.3.0をインストールできます。

$ curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
$ unzip awscliv2.zip
$ sudo ./aws/install

AWS CDKのインストール

AWS CDKのインストールは、AWS公式のドキュメントに書かれている方法でインストールできます。
https://docs.aws.amazon.com/cdk/latest/guide/getting_started.html

まず、CLIからAWSに接続するためのIAMユーザを作成します。AWSコンソールのIAMからユーザーを追加を選択してください。

ユーザー名を"mlflow-cli"、アクセスキープログラムによるアクセスにチェックを入れます。

サンプルなので、ここではユーザーにAdministrator Access権限を付与して、ユーザーを作成します。

完了画面で表示されるアクセスキーIDとシークレットアクセスキーをメモしておきます。

Ubuntuのコンソールに戻り、AWSと接続するための接続情報を~/.awsに保存します。作成するファイルは、configとcredentialsの2つです。configファイルには接続するリージョンを、credentialsファイルにはメモしておいたアクセスキーとシークレットアクセスキーを記述してください。

$ ls ~/.aws
config credentials
$ cat ~/.aws/config
[default]
region = ap-northeast-1
$ cat ~/.aws/credentials
[default]
aws_access_key_id = XXXXX
aws_secret_access_key = XXXXX

続けて、AWS CDKをインストールします。以下のコマンドを入力すると、AWS CDKの最新版バージョン1.110.0をインストールできます。

$ sudo npm install -g aws-cdk

AWSにMLflowをデプロイする

デプロイ

AWSにMLflowをデプロイするサンプルは、以下のgithubにアップされています。ここのREADMEに書かれているコマンドを打つと、AWS上にMLflowと関連モジュールをデプロイできます。
https://github.com/aws-samples/amazon-sagemaker-mlflow-fargate

$ git clone https://github.com/aws-samples/amazon-sagemaker-mlflow-fargate.git
$ cd amazon-sagemaker-mlflow-fargate

$ python3 -m venv .venv
$ source .venv/bin/activate
$ pip3 install -r requirements.txt

$ ACCOUNT_ID=$(aws sts get-caller-identity --query Account | tr -d '"')
$ AWS_REGION=$(aws configure get region)
$ cdk bootstrap aws://${ACCOUNT_ID}/${AWS_REGION}
$ cdk deploy --parameters ProjectName=mlflow --require-approval never

$ deactivate  % venv環境を抜ける

ここから、デプロイされたAWSのアーキテクチャについて説明します。手っ取り早くMLflowを動かしたい方は、次の章「SageMakerのセットアップ」までスキップしてください。

構成図

このサンプルのアーキテクチャは上のとおりです。この図は、先述のAWSのブログから引用しています。

基本構成

Load balancer
ネットワークロードバランサ(NLB)であり、本構成においてはMLflowへアクセスするためゲートウェイ的な要素として使われています。このAuto ScalingのURI(TCP 80番ポート)にアクセスすると、Fargateの5000番ポートに接続しMLflowが動きます。
Fargate
MLflowサーバの実体です。このサンプルでは、MLflowサーバがインストールされたコンテナをDockerfileを使って作成し、Fargateで起動することでAWS上で動作します。MLflowサーバは5000番ポートを開いて、NLBからリクエストを受け付ける設定となっています。
RDS
MLflowに登録された学習時のハイパーパラメータやメトリクスが保存されます。パラメータ保存用のDBが一つ追加されます。
S3
MLflowに登録されたモデルなどのファイルが保存されます。ファイルは"mlflow-artifacts-XXXXX"バケットに保存されます。

MLflowを触ってみる

ローカルPCで学習

まずは、ローカルのPCで学習した結果が、AWS上に構築したMLflowに登録できることを確認したいと思います。今回は、MLflowのチュートリアルに紹介されている学習スクリプトを使って動作を確認します。まずは、以下のコマンドを実行してGitHubから学習プログラムをチェックアウトして、実行環境を作ってください。

$ pip3 install mlflow scikit-learn boto3
$ git clone https://github.com/mlflow/mlflow
$ cd mlflow/examples
$ cp -p sklearn_elasticnet_wine/train.py sklearn_elasticnet_wine/train.py.bk

学習した結果をAWS上のMLflowに送るため、トラッキングサーバのURIを設定する必要があります。train.pyをエディタで開き、"def eval_metrics"の前に下の一行を追加してください。

mlflow.set_tracking_uri("http://[ロードバランサのDNS名]")

ここで、次のコマンドを実行すると学習が実行されます。

$ python3 sklearn_elasticnet_wine/train.py

学習された結果がMLflowに登録されているか、確認してみましょう。ブラウザでロードバランサのDNS名のページを開いてください。正しく実行されていれば、下のキャプチャのように学習結果が登録されていることが分かります。

詳細画面を見ると、学習時のハイパーパラメータと評価メトリクスの値を確認できます。また、Artifactsに学習によって生成されたモデルなどが登録されていることが分かります。

SageMakerで学習

続けて、Amazon SageMakerで学習した結果が、ローカルPCと同じくMLflowに登録されることを確認します。SageMakerで実行するスクリプトは、MLflowをデプロイするのに使用したGitHubにあるサンプルノートブックを使用します。
まず、サンプルを実行するSageMakerの環境を作ります。SageMakerのノートブックインスタンスから「ノートブックインスタンスの作成」を実行します。

作成するノートブックインスタンスの情報として、以下を入力してください。
ノートブックインスタンス名：mlflow-sample
ノートブックインスタンスのタイプ：ml.t2.medium

また、新しくIAMロールを作成してインスタンスに付与してください。このとき、デプロイにより作成されたMLflowのバケット(mlflow-artifacts-XXXXX)を、アクセスを許可する特定のS3バケットとして追加して、ロールを作成します。

続けて、学習で作成されたモデルファイルをS3に保存するため、IAMロールにアクセス権限を与えます。コンソールのIAMから先ほど作成したIAMロールを選択し、"AmazonS3FullAccess"をアタッチしてください。

ノートブックインスタンスが作成されたら、「JupyterLabを開く」を押してください。起動すると下のような画面が表示されます。amazon-sagemaker-mlflow-fargate/labに格納されているsource_dirと1_track_experiments.ipynbを、D&Dで画面左側のペインに追加します。

1_track_experiment.ipynbを開きます。カーネルは"conda_python3"を選択します。tracking_uriをロードバランサのDNS名に変更してください。この時、URIは"http://"付きで指定します。そして、ノートブックを実行して学習を開始してください。

ブラウザからMLflowのUIにアクセスすると、正しく実行されていれば学習結果が登録されていることが分かります。ローカルPCの時と同様に、パラメータとメトリクス、モデルファイル等が登録されています。

後片付け

このサンプルを試したあとに、デプロイしたAWSのリソースを削除する方法を説明します。まず、"mlflow-artifacts-XXXXX"バケットを空にして、サンプルの実行によりS3に追加されたファイルをAWSコンソールから削除します(バケット本体は削除しない)。

続けて、AWSコンソールのCloudFormationから、デプロイされた"DeploymentStack"スタックを削除します。

なお、コマンドを実行する前にデプロイされたRDSやS3のバケットを手動で削除しないようにしてください。削除すると、CloudFormationの定義と実際のAWSリソースの間で不整合が起こってしまい、スタックを削除するときにエラーが発生してしまいます。

おわりに

本記事では、AWS上にMLflowをデプロイし、ローカルPCまたはSageMakerで学習した結果を記録するまでの手順を紹介しました。このサンプルは簡易的なものではありますが、開発の実運用での構成を考えるときに参考になりそうです。また、ここで紹介したのはMLflowの機能の一部なので、この環境をベースに他のサンプルも動かして、ぜひ色々と試してみてください。