MLFlow Serving Custom Models in production environment

 

Serve Word2vec pretrained model using pyfunc flavor through docker.

After completing development, if you want to serve your trained model, then it has to be either logged or saved via log_model()/save_model() function using one of the built-in flavors.

MLflow does support multiple built-in flavors such as sklearn, keras etc and the flavor type is saved in the MLmodel file. Below is the example to save the model using sklearn flavor.

mlflow.sklearn.log_model(sk_model, "sk_models")

Flavors are the key concepts which makes MLflow more powerful. Flavors are a convention that deployment tools can use to understand the model, which makes it possible to write tools that work with models from any ML library without having to integrate each tool with each library.

And the below command helps to expose the logged model as a service

mlflow models serve 
-m ‘runs:/616e1e58d5814be58f613f13dfd8e9de/’
-h 0.0.0.0 -p 8000

And you can access the exposed model in the postman like this.

http://127.0.01:8000/invocations
Content-Type: application/json; format=pandas-split
Body: {"columns": ["col1", "col2"], "data": [["v1", "v2"]]}

Problem Statement

Sometimes, you need to expose the model which is not supported by built-in flavor like the use case which we came across, this blog helps in addressing this problem. In those cases, python_function flavor helps to achieve this problem.

The python_function model flavor serves as a default model interface for MLflow Python models. Any MLflow Python model is expected to be loadable as a python_function model.

Here we consider google pretrained Word2vec model, used for NLP use-cases. By default, it doesn’t supported by MLflow built-in flavors, we are going to create a custom function for them.

mlflow.pyfunc

All the built-in flavors defines predict() function, we are going to achieve the same using the Python class which inherits from PythonModel, defining predict() and, optionally, load_context().

#word2vec_wrapper.py
import mlflow.pyfuncclass Wrod2VecWrapper(mlflow.pyfunc.PythonModel):    def load_context(self, context):        
        import gensim        
        self.model = ensim.models.KeyedVectors.load(
            context.artifacts["model_path"], 
            mmap='r'
        )
        
    def predict(self, context, model_input):        
        # your custom code goes here
        return self.model.predict(model_input)

We can log this custom model using pyfunc.log_model() as below

with mlflow.start_run():
    mlflow.pyfunc.log_model(
        python_model=Wrod2VecWrapper(),
        artifact_path="chatbot",
        pip_requirements=["-r requirement.txt"],
        code_path=["nlp.py", "word_to_vec_wrapper.py"],
        artifacts={
            "model_path": "word2vecmodel.mod",
        }
    )

Here code_path helps to package necessary related code files and place it in code folder. In the same way model related files are saved in the artifact folder and it will be look like the below format.

log_model() folder structure

Once the custom models is created and saved in the ml tracking server, you can server them using the mlflow models server command.

Or you can package them using docker which will copy the required files in to the container and install the necessary packages which you have mentioned in the requirements and expose your model as server. MLFlow helps to achieve all the functions with one simple command which is given below

mlflow models build-docker -m 'runs:/98531e8ad3c644a68a83fcdf9e856aab/chatbot' -n chatbot --install-mlflowdocker run -e DISABLE_NGINX=true -e GUNICORN_CMD_ARGS="-b 0.0.0.0" 
-ip 8000:8000 chatbot

Now you can access your model in the postman at port 8000.

Hope this helps Someone.