いつの間にか MLflow Tracking Server が Artifact のプロキシに対応していた

以前の MLflow Tracking Server では、アーティファクトを保存する場所については URI としてクライアントに伝えるだけだった。クライアントは、サーバから教えてもらった URI に自分でつなぎにいく。この形では、アクセスするためのクレデンシャルがそれぞれのクライアントで必要になるなど、利便性にやや欠ける面があった。

そんな折、どうやら MLflow v1.24 から Tracking Server がアーティファクトをストレージとの間でプロキシする機能が追加されたらしい。ドキュメントにも、通信のシナリオとして以前は存在しなかった 5 と 6 が追加されている。

www.mlflow.org

今回は、上記の追加された機能を試してみよう。なお、通信のシナリオでいうと 5 に該当する。

使った環境は次のとおり。

$ sw_vers
ProductName:    macOS
ProductVersion: 12.4
BuildVersion:   21F79
$ uname -srm                     
Darwin 21.5.0 arm64
$ python -V
Python 3.9.13
$ pip list | grep -i mlflow
mlflow                        1.27.0

下準備

下準備として、あらかじめ必要なパッケージ類をインストールする。今回は Artifact を保存するバックエンドに S3 互換のオブジェクトストレージを使いたい。そのために、OSS の S3 互換オブジェクトストレージの実装である MinIO をインストールしておく。また、S3 を操作するために AWS CLI も入れておこう。

$ brew install minio awscli

そして、肝心の MLflow と、バックエンドに S3 互換のオブジェクトストレージを使うために boto3 も入れておく。

$ pip install mlflow boto3

インストールが終わったら MinIO のサーバを起動する。

$ mkdir -p /tmp/minio
$ minio server /tmp/minio

続いて、MinIO にアクセスできることを確認する。デフォルトのクレデンシャルを使ってバケットの一覧を確認する。この時点では特に何も表示されなければ大丈夫。

$ export AWS_ACCESS_KEY_ID=minioadmin
$ export AWS_SECRET_ACCESS_KEY=minioadmin
$ aws --endpoint-url http://localhost:9000 s3 ls

MLflow Tracking Server が使うためのバケットを作成する。ここでは mlflow-example という名前にした。

$ aws s3 --endpoint-url http://localhost:9000 mb s3://mlflow-example
make_bucket: mlflow-example

プロキシが有効な MLflow Tracking Server を起動する

以上で準備が終わった。次に Artifact のプロキシを有効にした MLflow Tracking Server を起動する。そのためには mlflow server サブコマンドを実行するときに --serve-artifacts オプションをつける。また、同時にアーティファクトを保存するバックエンドを --artifacts-destination で指定する。バックエンドが S3 互換ストレージであれば、エンドポイントやクレデンシャルの情報を環境変数で渡せる。

$ export MLFLOW_S3_ENDPOINT_URL=http://127.0.0.1:9000
$ export AWS_ACCESS_KEY_ID=minioadmin
$ export AWS_SECRET_ACCESS_KEY=minioadmin
$ mlflow server \
    --backend-store-uri sqlite:///tracking.db \
    --artifacts-destination s3://mlflow-example/artifacts \
    --serve-artifacts

このように、Tracking Server がストレージとの間でアーティファクトをプロキシするパターンでは、クレデンシャルなどの情報がサーバ側に集約される。

プロキシが有効な MLflow Tracking Server を利用する

MLflow Tracking Server が起動したら、それを使ってみよう。

いくつかのやり方があるけど、ここでは環境変数 MLFLOW_TRACKING_URI を使って MLflow Tracking Server の URI を指定する。

$ export MLFLOW_TRACKING_URI=http://127.0.0.1:5000

そして Python のインタプリタを起動しよう。

$ python

mlflow パッケージをインポートする。

>>> import mlflow

mlflow.get_tracking_uri() 関数で、Tracking URI が先ほど環境変数で指定したものになっていることを確認する。

>>> mlflow.get_tracking_uri()
'http://127.0.0.1:5000'

また、mlflow.get_artifact_uri() で、Artifact URI が mlflow-artifacts: から始まっていることを確認する。こうなっていれば MLflow Tracking Server が Artifact をプロキシしてくれることを示している。この形では、クライアント側にアーティファクトを保存するためのクレデンシャルは必要ない。

>>> mlflow.get_artifact_uri()
'mlflow-artifacts:/0/b91a91c68cbf4928b3069c6222479f03/artifacts'

試しに適当なテキストファイルをアーティファクトとして記録してみよう。

>>> with tempfile.TemporaryDirectory() as d:
...    filename = 'test-artifact'
...    artifact_path = pathlib.Path(d) / filename
...    with open(artifact_path, 'w') as fp:
...        print('Hello, World!', file=fp)
...    mlflow.log_artifact(artifact_path)
...

実行できたら mlflow.end_run() で実験を完了する。

>>> mlflow.end_run()

記録したアーティファクトが MinIO のストレージに記録されているか確認しよう。

$ aws --endpoint-url http://localhost:9000 s3 ls --recursive s3://mlflow-example
2022-07-03 17:41:01         14 artifacts/0/b91a91c68cbf4928b3069c6222479f03/artifacts/test-artifact

たしかに、なにかファイルができている。

中身を表示してみよう。

$ aws --endpoint-url http://localhost:9000 s3 cp s3://mlflow-example/artifacts/0/b91a91c68cbf4928b3069c6222479f03/artifacts/test-artifact -
Hello, World!

ちゃんと、先ほど記録したメッセージが表示された。どうやら、ちゃんとプロキシが機能しているようだ。

まとめ

今回は MLflow v1.24 で追加された Tracking Server のアーティファクトをプロキシする機能を試してみた。このパターンでは、クライアント側にストレージにアクセスするためのクレデンシャルが必要なくなる。また、クライアントからストレージに直接疎通がないような構成でも問題がない。ただし、逆に言えばクレデンシャルなしでクライアントがストレージにアクセスできてしまう点にはアクセスコントロールの観点で注意が必要となる。