はじめに
この記事はMLOps Advent Calendar 2024の14日目の記事です。
Databricksという会社でソフトウェアエンジニアをしている渡辺と申します。普段は業務でMLflowのCore maintainerをしています。
2024年6月に、Amazon SageMakerでのフルマネージド型MLflowが提供開始されました。MLflowを開発している身ながら、SageMaker上でのセットアップ方法やプロダクト全体についてはあまり把握できていなかったので、今回実際にサービスを使って理解を深めることにしました。
この記事ではあくまで１ユーザとしてサービスを体験するもので、OSS版やDatabricksのマネージドサービスとの機能比較は行っていません。実際にサービスを検討する際は、公式のドキュメンテーションや比較記事をご参考ください。
MLflowとは
MLflowはオープンソースの機械学習プラットフォームで、主に実験やモデルの管理、評価、デプロイメント、トレーシングなどの機能を備えています。
これまでのMLflowの使い方
MLflowでは、「トラッキングサーバー」というサーバーを中心に、実験、モデル、その他さまざまなアセットの管理が行われます。これまでは、このトラッキングサーバーに関しては、
- 自分でホストする
- Databricks Managed MLflow を使う
の2つの選択肢が主流でした。今回はこれにSageMaker上でマネージドを使うという選択肢が加わった形になります。
これまでもAWS FargateでMLflowのサーバーを立ちあげて、SageMaker上からアクセスするという方法は可能でしたが(AWS公式ガイド)、複数のAWSサービスを跨いた構成となるため一癖ある運用でした。正式なマネージドサービスがサポートされることで、こういったセットアップを行わなくても、AWSでMLflowを用いた実験管理が行えるようになりました。
なお、 SageMaker内にも”Experiment”という実験管理のメカニズムが存在していますが、2024年現在AWS公式でMagaed MLflowへの移行が推奨されています。
Managed MLflow on SageMakerの構成
SageMaker Managed MLflowは以下のようなシステム構成で成り立っています。
それでは実際にManaged MLflowを SageMakerで立ち上げて、実際にモデルの学習とデプロイまで行ってみます。内容としてはこちらのAWS公式ブログをなぞりながら、少し違うモデルを使ってみます。
SageMaker Studioのセットアップ
SageMakerでマネージドMLflowを使うためには、SageMaker Studioを立ち上げる必要があります。 SageMaker Studioはノートブックやモデルレジストリなどを含んた機械学習用の統合開発環境(IDE)で、学習からデプロイまで一連のサイクルをStudio内で行うことができます。
つい最近のre:Inventで Amazon SageMaker Unified Studioという新しいプロダクトが発表されました。Glue, Athena, EMR, あるいはBedrockのような、従来のSageMakerでは範囲外だったものの、機械学習プロジェクトで必要になりやすいサービスを改めて統合環境として提供するプロダクトのようです。
軽く試してみたところ、AI OpsツールとしてMLflowも既に統合されていたのですが、Unified Studioはまだプレビュー段階なので、今回は現行のStudioで試しました。
はじめにドメインを作成します。ドメインは SageMaker Studioのリソース管理における最上位のコンテナになります。今回はシングルユーザー向けのクイックセットアップを利用します。
ドメインは数分で立ち上がりました。作成される構成については、こちらのクラスメソッドさんの記事で詳細に解説してくれています。
横の”Studio”タブからトップページに戻ると、実際にSageMaker Studioにログインすることができます。
左上の”Applications＂欄にMLflowが表示されています。
MLflowトラッキングサーバーの立ち上げ
MLflowアイコンをクリックすると、Managed MLflowのコンソールが開きます。一番上に現在走っているトラッキングサーバーのリストがあり、下の方にスクロールすると、MLflowや SageMakerのドキュメントやサンプロノートブックへのリンクがあります。
まずはMLflowトラッキングサーバーを1台立ててみます。真ん中の”Create"をクリックすると、サーバー名とストレージとなるS3バケットを聞かれます。先ほどの全体構成にあったように、このS3バケットにはモデルの重みや画像といったArtifactsファイルのみが保存されます。
インスタンスサイズはSmall, Medium, Largeの3種類が用意されていて、デフォルトはSmallです。参考のチームサイズやTPSは公式ドキュメントに載っているのですが、一般的なチーム規模ではSmallで十分そうです。
また、MLflowのバージョンは2024年12月現在では、2.13と2.16の2つがサポートされています。現在の最新バージョンは2.19なのでややタイムラグはあるようですが、サーバ側は必ずしも最新でなくても、ノートブック側では最新のMLflowを使うこともできるので、実務上大きな問題はなさそうです。
適当な名前とS3バケットを指定してサーバーの作成を開始します。
トラッキングサーバーの起動には20分ほどかかりました。ここはもう少し早くなってくれればという印象です。
MLflowアクセス用の権限を設定
MLflowサーバーがスタートしたところで、左上の3点マークに”Open UI”というタブがあったので開こうとしたのですが、権限エラーとなってしまいました。
User: arn:aws:sts::[ACOUNT_ID]:assumed-role/AmazonSageMaker-ExecutionRole-[TIMESTAMP]/SageMaker is not authorized to perform: sagemaker-mlflow:AccessUI on resource: arn:aws:sagemaker:ap-northeast-1:[ACOUNT_ID]:mlflow-tracking-server/mlflow-test-server because no identity-based policy allows the sagemaker-mlflow:AccessUI action
どうやらSageMaker Studioの実行IAMロールにMLflow関係の権限が足りていないようです。エラーに表示されているIAMロールをチェックすると、MLflowに関する権限が何も付与されていませんでした。
今回は簡単のため、デフォルトのSageMakerのExecution Roleにそのまま権限を付与します。実際にはプロジェクトの役割ごとに個別のロールを作成して、 MLflowトラッキングサーバーに対しても最小権限を付与する運用が適切と思われます。
Built-inのポリシーはなかったので、インラインポリシーを作成します。まずインラインポリシーの対象サービスとして” SageMaker MLflow”を選択します。
Policy Statementではトラッキングサーバーに対する権限を、非常に細かい粒度で設定することができます。現実では、例えばサイエンティストにモデルの作成や削除権限を与え、MLエンジニアにはデプロイ用にモデルの読み取りのみを許可するなどといった運用が考えられるでしょうか。
もちろんアクションとなる対象のサーバーも細かく設定することができます。
また、このロールは SageMaker側からトラッキングサーバーにアクセスするための権限ですが、トラッキングサーバー側のアクセス権限も指定できるので、例えばMLflow経由で誤って他のチームのS3バケットにアクセスするような心配もありません。
ただ今回は簡単なテストなので、とりあえず全ての権限を付与しておきます。ポリシーを作成して SageMakerに戻ると、MLflow UIが開けるようになっていました。 UI自体はOSS版と変わりなく見慣れたものでした。
2024年12月現在の東京リージョンでは、 Sagemaker Managed MLflowはSmallインスタンスで1時間あたり0.789 USDが課金されます（料金表）。機能を試し終わった後や中断する場合は、必ずサーバーを停止するか削除するようにしてください。
Experimentの作成
ひとまずUIが開けたので新しいMLflow Experimentを作成してみます。
作成したMLflow Experimentは SageMakerの”Experiments"ページからも見ることができました。リンクをクリックするとMLflowが開く、という仕様になっています。そのため、旧式の SageMaker Experimentをそのまま置き換える形で利用できそうです。
SageMakerノートブックでモデルを学習
さて一通りセットアップが完了したので、実際にモデルを開発していきます。
左上のJupyterLabアイコンからノートブック環境をスタートすることができます。SageMaker Studioではノートブック環境はSpaceという単位で管理されます。適当な名前で新しいSpaceを作成し、数分でSpaceが起動すると、JupyterLabを開くことができるようになりました。
また、実際にコードを書いていく前に、MLflowコンソール内のサンプルノートブックを見てみました。
簡単なクライアント側の設定とScikit-learnベースの決定木モデルを学習するサンプルコードが含まれています。ただ、SageMaker独自のscikit-learnラッパーを用いた実装になっていたので、今回はこちらは用いず、代わりにXGBoostを使って簡単な分類器を学習していこうと思います。
MLflowのインストール
基本的にはサーバーと同じバージョンをインストールするのが安全ですが、2.16.2以降でサーバ側に大きな変更はないので、とりあえず最新のMLflow(2.19.0)をインストールします。
また、モデルを SageMakerのサービングにデプロイする際、sagemaker-mlflowというプラグインが必要になるみたいです。
pip install mlflow sagemaker-mlflow==0.1.0
また作成したトラッキングサーバーのARNとExperimentの名前を設定します。
mlflow.set_tracking_uri("arn:aws:sagemaker:<REGION>:<ACCOUNT_ID>:mlflow-tracking-server/mlflow-test-server")
mlflow.set_experiment("MLflow SageMaker Test")
モデルの学習
それでは実際にモデルを学習していきます。今回はXGBoostで簡単なIris分類器を学習します。
# 1. データセットのダウンロードと分割
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
iris = load_iris(as_frame=True)
df = pd.DataFrame(data=iris.data, columns=iris.feature_names)
X_train, X_test, y_train, y_test = train_test_split(iris.data, iris.target, test_size=0.2)
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.25)
# 2. Classifierクラスの初期化
xgb_classifier = XGBClassifier(
n_estimators=10,
max_depth=3,
learning_rate=1,
objective="binary:logistic",
random_state=123,
)
学習を開始する前にMLflowの自動ロギングを有効化します。手動でMLflowに記録することもできるのですが、autolog()の一行だけでパラメータやモデルなど自動で記録してくれるので便利です。
mlflow.xgboost.autolog()
xgb_classifier.fit(X_train, y_train, eval_set=[(X_val, y_val)])
学習が終わった後、MLflow UIをもう一度開いてみます。新しいRunが作成され、データセットやパラメータ、 Validationメトリックなどが自動で記録されています。
モデルをModel Registryに登録
MLflowに記録されたモデルは、UI（あるいはAPI）からモデルレジストリに登録することができます。 新しいモデル”mlflow-sagemaker-test"を作成して登録します。モデルのVersion 1が作成されました。
嬉しいことに、SageMaker Managed MLflowの場合はMLflowに登録したモデルを自動で SageMakerのModel Registryにも反映してくれました。Lossの情報やハイパラもMLflow側からこちらに反映されています。
本当は SageMakerの推論エンドポイントへのデプロイメントまで試したかったのですが、時間の都合上間に合いませんでした。後ほど試して追記しようと思っています...
所感
- SageMaker Studioの中で、これまでのような複雑なセットアップなしでMLflowが利用できるのはかなり便利だなと感じました。また特にレジストリ周りで SageMakerとの連携が取れているのは良いなと感じました。
- 公式のドキュメントやサンプルが充実していたのと、ノートブックでの開発はOSS版と全く同じように行えたので、特に大きく詰まらずにモデルの登録まで行うことができました。
- 今回は現行のStudioで体験してみましたが、新たに発表された SageMaker Unified Studioの中ではどのような使用感になるのかは、また試してみようと思います。