Issue in creating endpoint for quantized gguf model with llama-cpp-python

[lucaperes]

Hello, Databricks Community,

I am experiencing an issue while trying to serve a quantized model in gguf format using Databricks serving with the llama-cpp-python library.

The model is registered using MLflow and pyfunc on Unity. The model loads without any issues using the load function, indicating that the registration and initial configuration are correct.

The problem arises during the creation of the inference endpoint. Although the model is registered and loaded, I am unable to create the endpoint necessary to perform predictions. The following logs are produced:

[4cmcn] [2025-02-03 15:02:53 +0000] [12033] [INFO] Booting worker with pid: 12033
[4cmcn] [2025-02-03 15:02:55 +0000] [9] [ERROR] Worker (pid:12018) was sent code 132!
[4cmcn] [2025-02-03 15:02:55 +0000] [12045] [INFO] Booting worker with pid: 12045
[4cmcn] [2025-02-03 15:02:55 +0000] [9] [ERROR] Worker (pid:12027) was sent code 132!
[4cmcn] [2025-02-03 15:02:55 +0000] [12049] [INFO] Booting worker with pid: 12049
[4cmcn] [2025-02-03 15:02:56 +0000] [9] [ERROR] Worker (pid:12030) was sent code 132!
[4cmcn] [2025-02-03 15:02:56 +0000] [12062] [INFO] Booting worker with pid: 12062
[4cmcn] [2025-02-03 15:02:56 +0000] [9] [ERROR] Worker (pid:12033) was sent code 132!

Code:

%pip install tkmacosx>=1.0.5
%pip install pynput>=1.7.7
%pip install llama-cpp-python>=0.3.6
%pip install pyperclip>=1.9.0
%pip install transformers>=4.46.2
%pip install pygments>=2.19.1
%pip install cloudpickle>=3.1.1
%pip install mlflow>=2.20.1

from mlflow.models.signature import infer_signature
import mlflow
from typing import Generator, List, Dict, Any, Union, Tuple
from llama_cpp import Llama
from collections import deque
import os
from pathlib import Path

class ChatModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.model_path = None
        self.model = None
    
    def load_context(self, context):
        self.model_path = "/Volumes/ml_lab/generativo/models/granite-3.1-3b-a800m-instruct-Q6_K.gguf"
        self.model = Llama(
            self.model_path,
            n_ctx=8192,
            verbose=False,
            n_threads=8
         )
        
    def create_chat_completion(
                self,
                messages: List[Dict[str, str]], 
                temperature: float = 0.4,
                top_p: float = 0.9,
                top_k: int = 50,
                repeat_penalty: float = 1.2,
                max_tokens: int = 256
    ) -> Generator[str, None, None]:
        """Helper method to create chat completions with standard parameters"""
        if self.model is None:
            raise ValueError("O modelo Llama deve ser fornecido como argumento.")

        output = ""
        for chunk in self.model.create_chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repeat_penalty=repeat_penalty,
            stream=True
        ):

            content = chunk['choices'][0]['delta'].get('content')
            if content:
                if content in ["<end_action>", "<|endoftext|>"]:
                    break
                output += content
                yield output

    def process(self,messages):
        """Processa as mensagens e gera a resposta."""
        response = ""
        for chunk in self.create_chat_completion(messages, max_tokens=2024):
            response = chunk
            yield chunk

    def get_answer(self,messages):
        """Retorna a resposta final da conversa."""
        try:
            return deque(self.process(messages), maxlen=1).pop()
        except IndexError:
            return ""
        
    def predict(self, context, model_input: List[Dict[str,str]]) -> Dict[str,str]:
        """Gera respostas para multiplas entradas."""
        return {"answer": self.get_answer(model_input)}

# Carregando o modelo Llama fora da função de predição
# Isso evita recarregar o modelo a cada predição e economiza recursos
messages = [{'role': 'system',
  'content': 'Você é um assistente que fala português e responde perguntas do usuario baseado no conteudo fornecido.'},
 {'role': 'user', 'content': 'Oi'}]
signature = infer_signature(messages, {"answer": "Olá, tudo bem?"})


mlflow.set_registry_uri("databricks-uc")
with mlflow.start_run():
   model_info = mlflow.pyfunc.log_model(
       python_model=ChatModelWrapper(),
       artifact_path="model",
       registered_model_name="ml_lab.generativo.granite",
       pip_requirements=["tkmacosx>=1.0.5",
                        "pynput>=1.7.7",
                        "llama-cpp-python>=0.3.6",
                        "pyperclip>=1.9.0",
                        "transformers>=4.46.2",
                        "pygments>=2.19.1"],
    )

model = mlflow.pyfunc.load_model(model_info.model_uri)
print(model.predict(messages))

I would like to know if there are any specific guidelines or known adjustments that could resolve this issue. Any help in diagnosing and resolving this matter would be greatly appreciated.