封裝模型服務的自訂工件

本文說明如何確保模型檔案和工件依賴能在使用 Mosaic AI 模型服務的 部署過程中於 端點取得。

需求

MLflow 1.29 和更新版本

將工件與模型一起封裝

當您的模型在推斷過程中需要檔案或資源時，您可以在記錄模型時將它們封裝到模型資源包中。

如果您使用 Azure Databricks 筆記本，常見的作法是讓這些檔案位於 Unity 目錄磁碟區。 模型有時也會配置為從網路下載數位文件（例如 HuggingFace Tokenizers）。 在部署時靜態擷取所有必要的相依性，即時大規模工作負載才能達到最佳效能。 為此，模型服務要求將 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"])
    ...

在將檔案和產物封裝入您的模型產物後，您可以將模型部署到 模型服務端點。