モデルサービングのためにカスタムアーティファクトをパッケージ化する

この記事では、モデルのファイルとアーティファクトの依存関係が、Databricksエンドポイントを使用してモデルサービングで使用できることを確認する方法について説明します。

要件

MLflow 1.29 以降

モデルにアーティファクトをパッケージする

推論中にモデルでファイルまたはアーティファクトが必要な場合は、モデルのログを記録するときにそれらをモデルアーティファクトにパッケージ化できます。

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

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 エンドポイントに提供できます。