Empaquetar artefactos personalizados para Model Serving

En este artículo se describe cómo garantizar que las dependencias de archivos y artefactos de su modelo estén disponibles en su punto de conexión Implementar modelos mediante Model Serving.

Requisitos

MLflow 1.29 y versiones posteriores

Empaquetado de artefactos con modelos

Cuando el modelo requiere archivos o artefactos durante la inferencia, puede empaquetarlos en el artefacto del modelo al registrarlo.

Si trabaja con notebooks de Azure Databricks, una práctica habitual es que estos archivos se almacenen en volúmenes de Unity Catalog. A veces, los modelos también se configuran para descargar artefactos desde Internet (como HuggingFace Tokenizers). Las cargas de trabajo en tiempo real a escala funcionan mejor cuando todas las dependencias necesarias se capturan estáticamente en tiempo de implementación. Por este motivo, Model Serving requiere que los artefactos de los volúmenes de Unity Catalog se empaqueten en el propio artefacto del modelo mediante interfaces de MLflow. Los artefactos de red cargados con el modelo se deben empaquetar con el modelo siempre que sea posible.

Con el comando log_model(), de MLflow puede registrar un modelo y sus artefactos dependientes con el parámetro artifacts.

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

En los modelos pyFunc, se puede acceder a las rutas de acceso de estos artefactos desde el context objeto en context.artifactsy se pueden cargar de la manera estándar para ese tipo de archivo.

Por ejemplo, en un 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"])
    ...

Después de empaquetar sus archivos y artefactos dentro del artefacto del modelo, puede implementar su modelo en un endpoint de servicio de modelos.