将模型服务的自定义项目打包

本文介绍如何确保模型的文件和项目依赖项可用于 Azure Databricks 终结点的模型服务。

要求

MLflow 1.29 和更高版本

将项目和模型一起打包

如果模型在推理过程中需要文件或项目,你可以在记录模型时将其打包到模型项目中。

如果使用的是 Azure Databricks 笔记本,一种常见做法是让这些文件驻留在 Unity 目录卷。 模型有时还配置为从 Internet 下载项目(例如 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"])
    ...

在模型项目内打包文件和项目后,可以将模型服务到 模型服务终结点