Empacotar artefatos personalizados para o Model Serving

Este artigo descreve como garantir que os arquivos e dependências de artefatos do seu modelo estejam disponíveis no seu ponto de extremidade do Mosaic AI Model Serving para implementação de modelos.

Requisitos

MLflow 1.29 e superior

Empacotar artefatos com modelos

Quando seu modelo requer arquivos ou artefatos durante a inferência, você pode empacotá-los no artefato do modelo quando registrar o modelo.

Se estiveres a trabalhar com cadernos do Azure Databricks, uma prática comum é fazer com que esses ficheiros residam em volumes do Catálogo Unity. Às vezes, os modelos também são configurados para baixar artefatos da internet (como os tokenizadores HuggingFace). As cargas de trabalho em tempo real em escala têm melhor desempenho quando todas as dependências necessárias são capturadas estaticamente no momento da implantação. Por esse motivo, o Model Serving requer que os artefatos dos volumes do Unity Catalog sejam incorporados diretamente no artefato do modelo usando as interfaces do MLflow. Os artefatos de rede carregados com o modelo devem ser empacotados com o modelo sempre que possível.

Com o comando MLflow log_model(), você pode registrar um modelo e seus artefatos dependentes com o artifacts parâmetro.

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

Em modelos PyFunc, os caminhos desses artefatos estão acessíveis no objeto context em context.artifacts, e eles podem ser carregados da maneira padrão para esse tipo de arquivo.

Por exemplo, em um modelo MLflow personalizado:

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"])
    ...

Depois que seus arquivos e artefatos forem empacotados em seu artefato de modelo, você poderá servir seu modelo para um ponto de extremidade de serviço de modelo.