株式会社ホクソエムのブログ

R, Python, データ分析, 機械学習

MLflowのデータストアを覗いてみる

(2020/08/14 flavorについての記載を一部修正)

はじめに

こんにちは、ホクソエムサポーターの藤岡です。 最近、MLflowを分析業務で使用しているのですが、お手軽に機械学習のモデルや結果が管理できて重宝しています。 また、特定のライブラリに依存しないなど、使い方の自由度も非常に高いところが魅力的です。

ただ、ザ・分析用のPythonライブラリという感じでとにかく色々なものが隠蔽されており、 サーバにつなぐクライアントさえもプログラマあまりは意識する必要がないという徹底っぷりです。 もちろんマニュアル通りに使う分には問題ないですが、 ちゃんと中身を知っておくと自由度の高さも相まって色々と応用が効くようになり、 様々なシチュエーションで最適な使い方をすることができるようになります。

というわけで、今回はMLflowの記録部分を担う、 Experiment, Run, Artifactについてその正体に迫ってみます。

なお、MLflow自体についての機能や使い方についての解説はすでに良記事がたくさんあり、 ググればすぐにヒットするのでここでは割愛します。

データストアの3要素

MLflowのデータ管理はRun、Experiment、Artifactの三つの要素から構成されています。 大まかには、それぞれの役割は以下の通りです。

  • Run: 一回の試行（e.g. 実験, 学習, ... etc.）
  • Experiment: Runを束ねるグループ
  • Artifact: Runで得られた出力や中間生成物の保管先

これらの関係は、図のようになっています。

f:id:kazuya_fujioka:20200412001402p:plain

なお、Artifact LocationはArtifactの格納先のことであり、 あるExperimentに対して一つのArtifact Locationが紐づいています。 逆に、あるArtifact Locationが複数のExperimentと結びつくようにすることも可能ですが、 経験的には管理の点からしてそのような設計は避けるべきだと思います。

では、これらの三要素について、その役割と実装について説明していきます・

Run

概要

まず、全ての基本となるのはRunです。 文字通り、一回の試行を表すものです。 例えば、データセットをXGBoostに投げ込んで学習させる実験を一回やったとすれば、それに対応するのが一つのRunです。

MLflowを使うときはmlflow.start_run関数から始まることがほとんどだと思いますが、ここでRunオブジェクトが生成されています。

Runの基本要素は、Parameter、Tag、Metricsの三つです。 XGBoostの例で言えば、

  • Parameter: ハイパーパラメータ
  • Tag: モデルの識別名
  • Metrics: テストデータでのAUC

みたいな感じでしょうか。 いずれも、複数のキーバリューを取ることができます。

上では一例を挙げてみましたが、ParameterもTagもMetricsも自由に定義できるので、 個々人の裁量で最適な設計をしてあげるのがいいと思います。 特に、TagにするかParameterにするかはその後のMLflow UIのユーザビリティを大きく左右するので、 よく考えるのがいいと思います。 また、そもそもParameterのうち実験で扱わないようなものは敢えて記録しないなど、 スリム化することも念頭に置いておくといいかと思います。

実装

Runはいくつかのコンポーネントに分かれて定義されています。 低レベルAPIを触る際にキーとなる部分なので、それぞれ細かく解説していきます。

Runオブジェクト

まず、RunオブジェクトはTag、 Parameter、 MetricsとRunに関するメタデータを記録したデータコンテナです*1。メタデータは以下のように階層的に格納されています*2。

Run
 ┝ data
 │  ┝ params
 │  ┝ tags
 │  ┗ metrics
 ┗ info
    ┝ run_id
    ┝ experiment_id
    ┝ user_id
    ┝ status
    ┝ start_time
    ┝ end_time
    ┝ artifact_uri
    ┗ lifecycle_stage

なお、これらの要素は全てイミュータブルオブジェクトとして定義されています。 言い換えると、直接の書き換えは非推奨です。

実際にインスタンス化する場合には、MlflowClientのcreate_runメソッドを使います。これはstart_run関数の内部でも呼ばれていて、実はこのオブジェクト*3が返されています。 Runを記録するときに直接呼び出すことは無いので、サンプルプログラム等ではたいてい捨てられてしまっていますが。

一方、サーバから読み出す場合には、MlflowClient.get_runメソッドを使います。 start_runの場合と違ってrun_idが必要になるので注意しましょう。run_idはMLflow UIからも取得できますが、スクリプト中ではExperimentを通じて取得するのが楽です。 その取得方法についてはExperimentの節で扱います。

Runオブジェクトの内部構造と呼び出し方さえ分かってしまえば、実験結果の検索や操作もスクリプトから思いのままに実現可能です。 MLflow UIはリッチなUIを提供してくれている反面、小回りが利かないこともあるので、困ったらRunを自分で直接触るのがいいと思います。

RunDataオブジェクト

RunのdataプロパティにはRunの基本要素となる3要素が格納されています。このプロパティはRunDataというデータコンテナとして定義されており、Parameterがparamsに、Tagがtagsに、Metricsがmetricsにそれぞれ辞書オブジェクトとして格納されています。

RunInfoオブジェクト

RunのinfoプロパティにはRunのメタデータが格納されています。このメタデータには以下の情報が入っています。

run_id

RunのIDです。 あるrun_idは任意のExperimentに対して一意になります*4。

experiment_id

そのRunの属するExperimentのIDです。

user_id

そのRunを実行したUserのIDです。

status

Runの状態です。以下の5ステータスが定義されています。

  • RUNNING: Runを実行中
  • SCHEDULED: Runの実行がスケジュールされている状態 *5
  • FINISHED: Runが正常終了した状態
  • FAILED: Runが失敗して終了した状態
  • KILLED: RunがKillされて終了した状態
start_time / end_time

Runが開始 / 終了した時間です。

artifact_uri

ArtifactのURIです。

lifecycle_stage

Runの削除判定用フラグです。Runが削除されていればDELETED、そうでなければACTIVEが設定されています。 あるRunをAPI経由で削除した場合、データ自体が削除されるのではなくこのフラグがDELETEDに更新されます。

Experiment

概要

Runを束ねるのがExperimentの役割です。 モデルをたくさん投げ込んだときでもExperimentごとに整理しておけばアクセスしにくくなるのを防げますし、 MLflow UIの可視化機能やMetrics比較等を十全に活かすためにも、適切な単位ごとにExperimentで分けることは重要です。

例えば、あるKaggleコンペでExperimentを作って、その中に作成したモデルをRunとして記録するのはもちろん、 コンペによってはXGBoostやRFのようにモデルの種別で複数のExperimentを作ることも有効かもしれません。

Experimentの実現形式はメタデータの格納方式に依存します。 利用可能な格納方式は

  • ファイル (ローカル)
  • RDB (MySQL, MSSQL, SQLite, PostgreSQL)
  • HTTP サーバー (MLflow Tracking Server)
  • Databricks workspace

のいずれかです。

例えばファイルストア（Run, Experimentをファイルとして保管する形式）の場合、 以下のようなディレクトリがExperimentの実態となります。

<experiment-id>
 ┝ meta.yaml
 ┝ <run-1-id>/
 │  ┝ meta.yaml
 │  ┝ metrics/
 │  ┝ params/
 │  ┗ tags/
 ┝ <run-2-id>/
  ...
 ┗ <run-N-id>/
    ┝ meta.yaml
    ┝ metrics/
    ┝ params/
    ┗ tags/

experimentとrunの親子関係がディレクトリ構成で表現され、各種メタデータがmeta.yamlに記録されています。 metrics, params, tagsの中には各種データがファイルで格納されています。 なお、artifact_uriが指定されなければ、tags等と同じフォルダにartifactsというフォルダが作成されてその中に格納されます。

実装

Experimentオブジェクトの扱いはRunオブジェクトとよく似ているので、相違点だけ述べて詳細な解説はスキップします。

主な相違点として、以下のものが挙げられます。

  • 高レベルAPI経由での作成はmlflow.start_runではなくmlflow.create_experiment
  • クライアント経由での作成はMlflowClient.create_runではなくMlflowClient.create_experiment
  • 取得方法はid経由 (get_experiment関数) だけではなく name経由 (get_experiment_by_name関数) も可能（1つのサーバ内で名前がユニークであるという制約のため）
  • メタデータが以下の4つのみ
    • name: Experimentの名前
    • experiment_id: ID
    • artifact_location: Artifact LocationのURI
    • lifecycle_stage: Runと同様のACTIVE/DELETED
    • tags: Runと同様

Artifact

概要

Artifactとは、Runの結果や途中経過で生じたファイルを格納するためのストレージです。 Artifactを置く先のファイルストレージが扱えるファイルであればなんでも格納が可能です。 もちろん、csv化ができるpandas DataFrame や pickle化ができるPythonオブジェクトも同様です。

Artifactは以下のファイルシステムに対応しています。

  • Amazon S3
  • Azure Blob Storage
  • Google Cloud Storage
  • FTP server
  • SFTP Server
  • NFS
  • HDFS

上記のファイルシステムであればRunの場所に関わらず任意の格納場所 (Artifact Location) をURIで指定できますが、図で示したようにExperiment単位で指定しなければならないことに注意してください。 このURI以下にRunごとのフォルダが切られ、そこにRunの中間生成物等が入ります。

Artifactのもっとも重要な機能としては、モデルオブジェクトの格納が挙げられます。 モデルの格納の場合、単にファイルにして格納するだけでなく、それをあとで読み込んでデプロイできるようにしなければ、あまり意味がありません。 MLflowは多くの外部ライブラリに対してこのデプロイ機能を、それも環境の変化に対しても頑健な形で実装しています。

実装

Artifactに格納されるモデルファイルは、どのような形でセーブ/ロードが実現されているのでしょうか。 答えはシンプルで、セーブの際にはローカルに一回保存してからファイルシステムごとに定められたプロトコルで指定したURIへと転送し、ロードはその逆のプロセスです。

ただし、モデルファイルについては概要の節で述べたとおりのデプロイ機能を実現するために複数のファイルが生成・格納されます。 この格納方法は扱うモデルの実装されたモジュールごとに異なります。 MLflowでは、あるモジュールで生成したモデルを格納・呼び出しするための格納方法をflavorと読んでいます flavorと呼ばれるモジュールの中で定義しています*6。

例えば、sklearn flavor, xgb flavorといった具合です。

flavorには、例えば、sklearn flavor, xgboost flavorといったものがあります。

各flavorはmlflow以下に.pyファイル モジュールとしてそれぞれ実装されています。 例えば、XGBoostのflavorを見てみると、中には呼び出し可能な関数として、

  • get_default_conda_env : デフォルトのconda envを生成
  • save_model : ローカルへのモデルのセーブ
  • log_model : Artifactへのモデルの保存
  • load_model : ローカル/Artifactからのモデルの読み込み

の4つの関数が定義されています*7。 これは他のflavorでも同様です。 とはいっても、多くの場合ではこの中で直接使うのはlog_modelとload_modelだけであり、それ以外は内部的に呼ばれるだけかと思います。 頻度は高くないものの、他の二つも外部から使用する可能性があるものです。

これらの関数を使い、モデル本体のファイル（xgbの場合はmodel.pkl） conda env (conda.env) , MLmodelの三つが入ったフォルダを作成 / 読み込みします。 このフォルダが、MLflowにおけるモデルオブジェクトになります。

MLmodelファイルはモデル作成に使ったflavorの情報等が格納されたコンフィグファイルです。 この内容を元に読み込み方法を決定し、実行します。 また、モデルを実行（mlflow runコマンド）する場合にはconda envファイルを元に実行環境を作成します。

flavorは上述の四つの関数がカギであり、これらを理解することができればオリジナルのflavorを作ることもできます。

(2020/08/17 誤った内容を修正。なお、MLflow公式によるflavorの正確な定義はこちらをどうぞ）。

終わりに

MLflowは本当に便利で知名度も高くて少しずつ記事も増えてきているのですが、Run, Experiment, Artifactの三つについての解説が物足りなかったので本記事を書いてみました。 話題を絞ったぶん少し深いところまで掘り下げてみましたが、いかがだったでしょうか？

MLflowの実装は、flavorのあたりはちょっと無理矢理感がある気もしますが、低レベルのAPIのあたりは実装も参考になるし使いやすいしでぜひ読んで欲しいライブラリの一つです。 これを機に利用はもちろん、実装に興味を持っていただければ幸いです。

では、よきPythonライフを！

おまけ

Runの名は？

MLflowのWeb UIを触ると、Runにも名前 (Run Nameの項目) が設定できるようになっているのが分かります。 しかし、ここまでの説明の通り、RunのメタデータとしてRunの名前に該当する項目はありません。

実は、Run Nameはtagsに"mlflow.runName"というタグ名で記録することで表示させることができます。 他にも、このような特殊なタグはmlflow.utils.mlflow_tags内で定義されています。

ところで、"Artifact"って？

Artifactを英和辞典で調べると、遺物とか人工物とかそんなふんわりとした意味しか出てこなくて、使い始めた当初は具体的に何を入れるものなのかがよく分かりませんでした。

本記事を書くにあたって改めて辞書で調べてみると以下のような定義となっていました。

Something observed in a scientific investigation or experiment that is not naturally present but occurs as a result of the preparative or investigative procedure.

ざっくりと訳すと、科学的調査や実験で生じた人工物、というような感じです。 機械学習の実験で使うことを考えると、やっぱり、中間生成物やモデル等なんでも入れていいというような意図を感じます。

以前に業務で使用していたときには学習に使うデータマートとかもMLflowで管理してみたのですが、案外できてしまった（しかも割と便利だった）ので、本当に何を入れてもいいんだと思います。 もっとも、環境に特別な制約がなくマート管理に特化したものを採用できるのであれば、そちらの方がいいと思いますが。

NoneはTagに入りますか？

入ります。Parameterも受けつけます。 ただし、Metricsだけはエラーを吐きます。

データ活用のための数理モデリング入門

データ活用のための数理モデリング入門

*1:実際にはもう少し別の機能が定義されていますが、コミッターでもない限り使わないと思うので省略しています

*2:v1.7.2時点でdeprecateされたものは省略しています

*3:正確にはその子クラスのActiveRunオブジェクトなのですが、コンテキストマネージャ化されていること以外は同じなのでここでは同一のオブジェクトとして扱います

*4:例えば、ファイル形式のストアの実装を見てみると、uuid.uuid4を使っているので、重複はまず無いです。REST形式のストアの方は不明ですが、sqlalchemyのストアでも同様です

*5:使ったことがないので詳細は不明ですが、おそらくDatabricksやKubernetesと連携した際に使用されるパラメータ

*6:flavorという名前の由来はこの辺りかと思います。

*7:autologはexperimentalかつoptionalなので省略しました