Model Servingのカスタム成果物をパッケージ化する

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

要件

MLflow 1.29 以降

モデルを使用した成果物のパッケージ化

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

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

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

mlflow.pyfunc.log_model(
    ...
    artifacts={'model-weights': "/dbfs/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 エンドポイントに提供できます。