Azure Machine Learning CLI v2 Job の input に追加された mlflow_model について

はじめに

AzureML CLI v2 では yaml 形式でジョブの設定を記述することで機械学習ジョブを実行することができます。

各ジョブはinputsとして何らかの入力を取ることができ、機械学習においてはuri_fileやuri_folderを設定し、Datastore として AzureML に登録したストレージ内のどこかを指定したり、既に Data として登録した何らかのデータセットを指定することで、ジョブを実行するコンピューティングに対象のディレクトリをマウントしつつマウントパスをスクリプトの引数として渡すことができ、ジョブのスクリプト内部からデータセットを読み込むことが可能になります。

CLI (v2) pipeline job YAML schema - Azure Machine Learning | Microsoft Learn

Reference documentation for the CLI (v2) pipeline job YAML schema.

https://learn.microsoft.com/en-us/azure/machine-learning/reference-yaml-job-pipeline#job-inputs

CLI (v2) pipeline job YAML schema - Azure Machine Learning | Microsoft Learn-image

そんなジョブのinputsのtypeの 1 つとして、mlflow_modelというものが増えていることに気が付きました。

追加日時は 2022/8/8 となっていますが、このmlflow_modelについての詳細をドキュメント上から見つけられていません。名前から察するにモデルをinputsとして渡すことができる仕組みのようですが、使い方と挙動を実際に試して探った結果をこちらにメモします。

https://github.com/MicrosoftDocs/azure-docs/commit/60cff4fc04a7145c123155bbaad91b0830cb489d

コード準備

手元にあった Component と Pipeline Job の定義ファイルをベースに以下のように書き換えました。

pipeline.yml

$schema: https://azuremlschemas.azureedge.net/latest/pipelineJob.schema.json
type: pipeline
display_name: mlow-nyc-taxi-regression-pipeline
experiment_name: mlow-nyc-taxi-regression-pipeline
compute: azureml:cpu-cluster
jobs:
  train_job:
    type: command
    component: file:./component.yml
    inputs:
      train_data:
        type: uri_file
        path: azureml:mlow-nyc-taxi-train@latest
      valid_data:
        type: uri_file
        path: azureml:mlow-nyc-taxi-valid@latest
      model:
        type: mlflow_model
        path: <何らかのモデル識別子>

component.yml

$schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
type: command
name: mlow-nyc-taxi-train-job
display_name: mlow-nyc-taxi-train
code: ../02_python_script
command: >-
  python build_model.py  --input_train_data ${{inputs.train_data}} --input_valid_data ${{inputs.valid_data}} --model ${{inputs.model}} --mode remote
inputs:
  train_data:
    type: uri_file
  valid_data:
    type: uri_file
  model:
    type: mlflow_model
environment:
  image: mcr.microsoft.com/azureml/openmpi4.1.0-ubuntu20.04:20220504.v1
  conda_file: ../environment.yml

ジョブの本体である Python スクリプトには以下の記述を仕込み、ジョブが実行されたとき--modelに何が渡されるのか、uri_folder等とのアナロジーとしてパスが渡されることが期待されるが、そのパスには何があるのかをprintで確認します。

import os
print(args.model)
files = os.listdir(args.model)
for file in files:
    print(file)

実験

<何らかのモデル識別子> の部分をモデルを表現しそうな文字列で置換して挙動を確認します。

結論から言って、Asset ID、azureml: から始まるモデル名とバージョンの指定の 2 つの方式について、ジョブにモデルを渡すことに成功しました。

Asset ID

登録済みモデルの Asset ID を指定してジョブを実行します。

Asset ID は以下の図の黄色マーカー部分から取得することができます。

ジョブのデプロイおよび実行いずれも成功し、以下のような出力を得ることができました。

/mnt/azureml/cr/j/414192a55936424c8760b57aab1cc6d2/cap/data-capability/wd/INPUT_model
MLmodel
conda.yaml
model.lgb
python_env.yaml
requirements.txt

1 行目を見ると、やはりモデルが含まれるディレクトリをマウントしてそのパスを代入する仕様のようです。この辺りの挙動は Data や Pipeline を構成する Component 間の入出力と同様です。

2 行目以下は対象ディレクトリに含まれるファイルのリストですが、これは指定したモデル (MLflow Models に沿った形式) を構成するファイル群と一致しています。

Pipeline Job の全体としては以下の図のようになり、モデルを入力として取れていることが分かります。

azureml:<model_name>:<version>

Endpoint をデプロイする場合、既に登録してあるモデルの名前とバージョンを使用して

path: azureml:nyc_taxi_lgb_model:2

のような形式でモデルを指定することができます。

この形式によるモデル指定の場合も、同様にジョブのデプロイと実行に成功し、Asset ID を指定した場合と同様の結果を得ることができました。結果は同じなので割愛します。

また

path: azureml:nyc_taxi_lgb_model@latest

のように@latestで同名モデルのうち最新のものを指定できるか試したところ、こちらも同様に成功しました。

結論

inputsのtype: mlflow_model は Asset ID かモデル名およびバージョン指定により登録済みモデルをジョブの入力として渡すことができるようです。

このとき、モデルを構成するファイル群を含むディレクトリがジョブの実行環境にマウントされ、${{inputs.model}}にはマウント後のパスが代入されて渡される様子です。

おわりに

バッチ推論の仕組みとして Batch Endpoint というものがありますが、これは複数ファイルに対して並列に推論結果を付与するもので、複雑な前処理や集計処理を要求するケースでは使用できないか他の仕組みを組み合わせるを必要があります。

このtype: mlflow_modelを前提とすれば、複雑な前処理もしくは集計処理を必要とするバッチ推論実装が可能になりそうです。

Command (単発のスクリプト)、Sweep (パラメーターチューニング)、Pipeline (Component を連結したジョブ) の 3 種類についてはtype: mlflow_modelをサポートしている様子ですので、これらのジョブ実行方式については学習済みモデルを利用したバッチ処理の実装に利用できます。とりわけ Pipeline Job が有望で、Component の単位で複数の処理を連結することが可能なので、これを使用すれば複雑な前処理・集計処理を伴う推論ジョブを実行することが可能そうです。

Parallel (並列処理) ジョブについてはinputsにtype: mlflow_modelの記載がなく、モデルを入力として取れるか不明です。とはいえ並列処理が必要な場合でも Pipeline Job を構成する Component は MpiConfiguration を設定することで分散実行も可能ですので、この辺りを利用すれば並列の推論結果付与も可能そうです。

CLI (v2) parallel job YAML schema - Azure Machine Learning | Microsoft Learn

Reference documentation for the CLI (v2) parallel job YAML schema.