モデル提供用のカスタム成果物をパッケージ化する

この記事では、Azure Databricks エンドポイントでサービスを提供する Model でモデルのファイルと成果物の依存関係を使用できるようにする方法について説明します。

要件

MLflow 1.29 以降

モデルを使用して成果物をパッケージ化する

推論中にモデルにファイルまたは成果物が必要な場合は、モデルをログに記録するときにモデル成果物にパッケージ化できます。

Azure Databricks ノートブックを使用している場合、一般的な方法は、これらのファイルを Unity カタログ ボリュームに配置することです。 モデルは、インターネット (HuggingFace Tokenizers など) から成果物をダウンロードするように構成される場合もあります。 大規模なリアルタイム ワークロードでは、デプロイ時に必要なすべての依存関係が静的に取得された場合に最高のパフォーマンスが得られます。 このため、モデル サービスでは、MLflow インターフェイスを使用して、Unity カタログ ボリューム成果物をモデル成果物自体にパッケージ化する必要があります。 モデルと共に読み込まれたネットワーク成果物は、可能な限りモデルと共にパッケージ化する必要があります。

MLflow コマンド log_model() を使用すると、artifacts パラメーターを使用してモデルとその依存成果物をログに記録できます。

mlflow.pyfunc.log_model(
    ...
    artifacts={'model-weights': "/Volumes/catalog/schema/volume/path/to/file", "tokenizer_cache": "./tokenizer_cache"},
    ...
)

PyFunc モデルでは、これらの成果物のパスは、context.artifacts の下の context オブジェクトからアクセス可能であり、そのファイルの種類に対する標準の方法で読み込むことができます。

たとえば、カスタム MLflow モデルでは次のようになります。

class ModelPyfunc(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        self.model = torch.load(context.artifacts["model-weights"])
        self.tokenizer = transformers.BertweetTokenizer.from_pretrained("model-base", local_files_only=True, cache_dir=context.artifacts["tokenizer_cache"])
    ...

ファイルと成果物がモデル成果物内にパッケージ化されたら、モデルを Model Serving エンドポイントに提供できます。