はじめに
みなさんこんにちは、サービスイノベーション部の内村です。 普段の業務では機械学習モデルの開発運用を行ったり、社内のデータ活用推進業務に携わっています。
今回は機械学習モデルの実験管理を行うMLFlow環境をcdkでECS上に構築することを試してみたのでその内容の紹介と 直近のAWS re:Inventで紹介されたSagemaker Studioの新機能Code Editorを使ってSagemaker Studio内でMLFlow UIを表示する方法について紹介します。
- はじめに
- MLflowとは？
- MLflowで実験管理を行う際に必要な環境
- AWS上でのMLflow環境構築
- MLflow環境構築のまとめ
- MLflow UIをSagemaker Studio上で表示する
- Sagemaker Studioとは?
- Sagemaker Studio 環境の構築
- まとめ
- 参考
MLflowとは？
MLflowは機械学習のライフサイクル(実験管理からモデルのデプロイ、モニタリングまで)を管理するためのオープンソースのプラットフォームです。 機械学習モデルの管理、実験管理、モデルのデプロイなど機械学習のライフサイクルを管理、実現するための様々な機能が提供されています。 久々にリリースノートを確認したのですが、最近ではLLM(large language model)に関する機能も追加されていて、開発も活発に行われているようです。
様々な機能が提供されていますが、特に実験管理の機能については使ったことのある方も多いのではないでしょうか。 今回はそんなMLflowをチームで使うために必要なMLflowのリモートサーバー構築を行います。
MLflowで実験管理を行う際に必要な環境
公式ドキュメントを確認したところMLflowで実験管理を行う際に利用できる環境は6パターンに分けられていました。細かい違いはあるものの、大きく分けてローカル環境のみで実行するパターンと、ローカル環境とリモートトラッキングサーバーの両方を利用するパターンに分けられます。
今回は冒頭でも説明した通り、チームでの機械学習モデルを開発するのに利用できるよう、後者のリモートトラッキングサーバーを利用するパターンの環境構築をAWS上で行いました。
1. ローカル環境のみのパターン
ローカル環境でMLflowを実行し、ローカルファイルシステムにログ、アーティファクトを保存するパターンです。 試しにMLflowを使ってみたいという場合にはこのパターンが一番簡単に試せます。
2. ローカル環境 + リモートトラッキングサーバー
ローカル環境でMLflowを実行し、リモートトラッキングサーバーにログ、アーティファクトを保存するパターンです。 リモートトラッキングサーバを用意することで、複数人で実験結果を共有することができるようになります。 また、モデルやデータセットなどのアーティファクトや実験ログについてもリモートに保存することで、 再現性の担保や統一環境での実験管理を行うことができます。
AWS上でのMLflow環境構築
必要なリソースが分かったので、ここから先は実際にAWS上に環境構築をしていきたいと思います。
参考となった公式リポジトリ
MLflow環境を構築する方法を調べたところ、awsが公式のサンプルがいくつか公開されていました。 今回はその中から目的に近かった3つのサンプルについて軽く紹介します。
- amazon-Sagemaker-mlflow-fargate
- pythonのcdkを利用して構築
- 3つの中で一番シンプルだがMLflowアクセス時の認証機能がない
- Sagemaker Studioの環境についても別途自身で構築する必要がある
- aws-Sagemaker-mlflow-cdk
- typescriptのcdkを利用して構築
- パスワードによるbasic認証機能がある
- sagemaker studioの環境も構築できる
- Sagemaker-studio-mlflow-integration
- typescriptのcdkを利用して構築
- cognitoによる認証機能がある
- mlflowのユーザー管理機能もある
- Sagemaker Studioの環境も構築できる
- Amplifyも活用していてリソースも多く複雑
構築するリソース
普段からtypescriptのcdkを利用しており、3はMLflow環境を試すだけにしては複雑過ぎたので、今回は2番目のサンプルを利用して構築を行うことにします。 図ではマルチAZの表記を省略していますが、実際は各リソースがマルチAZで構築されます。 MLflowの説明で記したTracking/UI ServerはFargate上で実行し、filestoreはS3を利用しています。 また、メトリクスやparameterなどのログはAurora Serverlessに保存されます。
サーバレスで構築されているため、必要以上にコストがかからないのも良いところですね。
AWS上にMLflow環境をデプロイする
環境構築はAWSのcloud9上で、READMEに記載されている手順に従って進めました。 特にエラーが出ることもなく進んだのと、必要なコマンドはREADME内に全て記載されているのでここでは割愛します。 普段の業務ではcdkをl1で書いているので、 作成しているリソースに対してl2で書かれたコードの記述量が少なく驚きました。
作成されるリソースの詳細についても参考にしたREADMEで説明されているので、ここでは気になった点をいくつか紹介します。
Task Container
タスク定義ではnginxコンテナとmlflowコンテナの2つのコンテナが定義されています。 nginxコンテナの用途としては、mlflowコンテナに対してbasic認証を行うために利用されています。
Auroa Serverless
Aurora Serverlessは、MLflowの実験ログを保存するためだけに利用されています。 そのため、MLflow Server以外のリソースからは直接アクセスすることはできません。
構築したMLflow環境の確認
作成したMLflow環境の動作確認を行ってみます。 MLflowの利用についてもサンプルリポジトリ内に実験用のノートブックが用意されているので、今回はそちらをそのまま利用しました。
実行後mlflowのUIをWebブラウザ上から確認してみると、以下のように実験結果が表示されていました。 これで無事、モデル学習時のログ、メトリクスの保存とUIの表示ができていることが確認できました。
MLflow環境構築のまとめ
ここまでで、MLflow環境の構築と実験の実行、結果の確認まで行うことができました。
ただ、MLflow serverへのアクセス時にbasic認証を利用しているためセキュリティ面が気になります。 そこで、MLflowのUI表示をSagemaker Studio環境から行えるようにすることで、 UI表示のために開けていたWebからのアクセスを制限してセキュリティを向上させることができないかと考え試行錯誤しました。
結論から言うと今回期待した挙動とは違ったものとなってしまたのですが、 最近発表されたSagemaker Studioの新機能Code Editorを利用してMLflow UIの表示を行なったのでその内容を紹介します。
MLflow UIをSagemaker Studio上で表示する
ウェブブラウザ上で動作しているnotebookの中でさらにwebページを表示するというのは、 ニッチな要望ということもあり良さそうなJupyterlabの拡張機能も見つけることができませんでした。
少し悩んでいたところ、AWS re:Invent 2023でSagemaker Studioの新機能Code Editorが発表され Sagemaker 環境上でVSCodeの様々な拡張機能を利用できるようになりました。 せっかくなので、今回はリリースされたばかりのこの機能を利用してMLflow UIを表示してみます。
Sagemaker Studioとは?
Sagemaker Studioは機械学習モデルの開発、学習、デプロイまで通貫して行うことができる統合開発環境です。 統合開発環境と聞くと分かりづらいかもしれませんが、 機械学習に関連した様々な機能が提供されていて、Jupyter Notebookを使った実験だけではなく、モデルのデプロイ、モニタリング結果の可視化といったことまで行うことができる環境になっています。
Sagemaker Studio 環境の構築
前述のCDKで作成したSagemaker Studio環境ではCode Editor機能が有効になっていないため、こちらの記事を参考に新たにSagemaker Studio環境を構築します。
1.Sagemaker Studio ドメインの作成
CDKのドキュメント等はまだ見つけられなかったのでメインの作成はAWSコンソールから行いました。 Sagemaker Studioのドメインを作成しようとすると最初にシングルユーザー向けのドメインを作成するか、 組織向けのドメインを作成するかを選択する画面が表示されます。 Code Editorを利用するためには組織向けのドメインを作成する必要があるため、組織向けのドメインを作成します。
アップデートでこれまでのSagemaker Studioの環境はclassicと呼ばれるようになっているようで、 今回の目的はCode Editorを利用するためにnewの方を選択します。
以降はこれまで通り、ドメイン名やIAMロールなどを設定していきました。
2.Code Editorの利用
テストユーザを作成した後、起動してみると以下のようなhomeが表示されました。 最初にjupyterlabかCode Editorを選ぶようにな形となっています。
Code Editorを選択すると見慣れたVSCodeのような画面が表示されました。 拡張機能等も利用できるようになっているので、 これまでは一時的な分析用途でのみnotebookを利用していたのですが、 しっかりと環境を構築していくことで、Code-Editor上での開発も行えるようになりそうです。
3.Sagemaker Studio からMLflow UIを表示する
ここから先は簡単でVS Codeの基本機能としてSimple browserというブラウザ表示機能が提供されているので、そちらを利用してMLflow UIを表示していきます。
手順としては、Code Editorn内でCommand + Shift + Pを押すとコマンドパレットが表示されるので、
そこでSimple Browserと入力しクリックすることで、URLを入力することができます。
特に追加の設定などは必要なく、MLflow UIを表示することができました。
4. アクセス経路を確認してみる
最後に確認としてsimple browserでIPアドレスを表示してくれるWebページを開くと、自宅のIPアドレスが表示されていました。
図に書いたローカルブラウザ → Sagemaker Studio Instance → MLflow サーバーの経路でアクセスすることを期待していたのですが
ローカルブラウザでMLflow サーバーにアクセスして、その結果をSagemaker Studioで表示しているようです。
別途確認としてSagemaker Studio上で実行したMLflow uiサーバーへのアクセスも試したのですが、そちらもsimple browserからはアクセスできませんでした。
個人的にはSagemaker環境経由で表示できた方が便利なので今後のアップデート、活用事例の情報が増えてくることに期待したいと思います！
まとめ
今回はMLflow環境の構築とSagemaker Studio環境からMLflow UIを表示する方法について紹介しました。
MLflow UIをSagemaker 環境内で表示する方法については期待した挙動とは違いましたが、 新規機能ということもありアップデートで挙動が変わるかもしれません。 また、今回はあまり触れられませんでしたがCode Editor機能の追加で、 Sagemaker Studio上の開発環境も大きく変わるので、今後も様々な機能を試して開発環境改善に取り組んでいきたいと思います。
