Databricks Community

damselfly20 · yesterday

For my RAG use case, I've registered my langchain chain as a model to Unity Catalog. When I'm trying to serve the model, container image creation fails with the following error in the build log:

[...]
#16 178.1 Downloading langchain_core-0.3.17-py3-none-any.whl (409 kB)
#16 178.1 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 409.3/409.3 kB 48.1 MB/s eta 0:00:00
#16 178.1 Downloading langchain_text_splitters-0.3.2-py3-none-any.whl (25 kB)
#16 178.1 Downloading langsmith-0.1.142-py3-none-any.whl (306 kB)
#16 178.1 ━━━━━━━━━━━Pip subprocess error:
#16 178.1 error: subprocess-exited-with-error
#16 178.1
#16 178.1 × python setup.py bdist_wheel did not run successfully.
#16 178.1 │ exit code: 1
#16 178.1 ╰─> [23 lines of output]
#16 178.1 /opt/conda/envs/mlflow-env/lib/python3.11/site-packages/setuptools/_distutils/dist.py:261: UserWarning: Unknown distribution option: 'cffi_modules'
#16 178.1 warnings.warn(msg)
#16 178.1 running bdist_wheel
#16 178.1 running build
#16 178.1 running build_py
#16 178.1 creating build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 copying src/snappy/snappy_formats.py -> build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 copying src/snappy/hadoop_snappy.py -> build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 copying src/snappy/snappy_cffi.py -> build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 copying src/snappy/__main__.py -> build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 copying src/snappy/snappy.py -> build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 copying src/snappy/__init__.py -> build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 copying src/snappy/snappy_cffi_builder.py -> build/lib.linux-x86_64-cpython-311/snappy
#16 178.1 running build_ext
#16 178.1 building 'snappy._snappy' extension
#16 178.1 creating build/temp.linux-x86_64-cpython-311/src/snappy
#16 178.1 gcc -pthread -B /opt/conda/envs/mlflow-env/compiler_compat -DNDEBUG -fwrapv -O2 -Wall -fPIC -O2 -isystem /opt/conda/envs/mlflow-env/include -fPIC -O2 -isystem /opt/conda/envs/mlflow-env/include -fPIC -I/opt/conda/envs/mlflow-env/include/python3.11 -c src/snappy/crc32c.c -o build/temp.linux-x86_64-cpython-311/src/snappy/crc32c.o
#16 178.1 g++ -pthread -B /opt/conda/envs/mlflow-env/compiler_compat -DNDEBUG -fwrapv -O2 -Wall -fPIC -O2 -isystem /opt/conda/envs/mlflow-env/include -fPIC -O2 -isystem /opt/conda/envs/mlflow-env/include -fPIC -I/opt/conda/envs/mlflow-env/include/python3.11 -c src/snappy/snappymodule.cc -o build/temp.linux-x86_64-cpython-311/src/snappy/snappymodule.o
#16 178.1 src/snappy/snappymodule.cc:33:10: fatal error: snappy-c.h: No such file or directory
#16 178.1 33 | #include <snappy-c.h>
#16 178.1 | ^~~~~~~~~~~~
#16 178.1 compilation terminated.


#16 178.1 error: command '/usr/bin/g++' failed with exit code 1
#16 178.1 [end of output]
#16 178.1
#16 178.1 note: This error originates from a subprocess, and is likely not a problem with pip.
#16 178.1 ERROR: Failed building wheel for python-snappy
#16 178.1 ERROR: Could not build wheels for python-snappy, which is required to install pyproject.toml-based projects
#16 178.1
#16 178.1 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 306.7/306.7 kB 49.2 MB/s eta 0:00:00
#16 178.1 Downloading Markdown-3.7-py3-none-any.whl (106 kB)
#16 178.1 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 106.3/106.3 kB 21.6 MB/s eta 0:00:00
[...]

The latest RAG version contains a cross encoder for reranking. Before the reranking was included, everything worked just fine. This is my chain code before:

import os
import mlflow
from operator import itemgetter
from databricks.vector_search.client import VectorSearchClient
from langchain_community.chat_models import ChatDatabricks
from langchain_community.vectorstores import DatabricksVectorSearch
from langchain_core.runnables import RunnableLambda
from langchain_core.runnables import RunnableParallel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough

# Enable MLflow Tracing
mlflow.langchain.autolog()

# Return the string contents of the most recent message from the user
def extract_user_query_string(chat_messages_array):
    return chat_messages_array[-1]["content"]

# Get the conf from the local conf file
model_config = mlflow.models.ModelConfig(
    development_config="chain_config.yaml"
)

databricks_resources = model_config.get("databricks_resources")
retriever_config = model_config.get("retriever_config")
llm_config = model_config.get("llm_config")

# Connect to the Vector Search Index
vsc = VectorSearchClient(disable_notice=True)
vs_index = vsc.get_index(
    endpoint_name=databricks_resources.get("vector_search_endpoint_name"),
    index_name=retriever_config.get("vector_search_index"),
)
vector_search_schema = retriever_config.get("schema")

# Turn the Vector Search index into a LangChain retriever
vector_search_as_retriever = DatabricksVectorSearch(
    vs_index,
    text_column=vector_search_schema.get("chunk_text"),
    columns=[
        vector_search_schema.get("primary_key"),
        vector_search_schema.get("chunk_text"),
        vector_search_schema.get("document_uri"),
    ],
).as_retriever(search_kwargs=retriever_config.get("parameters"))

# Required to:
# 1. Enable the RAG Studio Review App to properly display retrieved chunks
# 2. Enable evaluation suite to measure the retriever
mlflow.models.set_retriever_schema(
    primary_key=vector_search_schema.get("primary_key"),
    text_column=vector_search_schema.get("chunk_text"),
    doc_uri=vector_search_schema.get("document_uri"),
)

# Method to format the docs returned by the retriever into the prompt
def format_context(docs):
    sources = []
    [
        sources.append(
            {"content": doc.page_content, "url": doc.metadata["url"]}
        )
        for doc in docs
    ]
    return sources

# Prompt Template for generation
prompt = PromptTemplate(
    template=llm_config.get("llm_prompt_template"),
    input_variables=llm_config.get("llm_prompt_template_variables"),
)

# FM for generation
model = ChatDatabricks(
    endpoint=databricks_resources.get("llm_endpoint_name"),
    extra_params=llm_config.get("llm_parameters")
)

# RAG Chain
# The "|" syntax is powered by the LangChain Expression Language (LCEL)
# To learn more about LCEL, read the documentation: https://python.langchain.com/v0.1/docs/expression_language/

second_chain_part = prompt | model | StrOutputParser()

chain = RunnableParallel(
    {
        "question": itemgetter("messages")
        | RunnableLambda(extract_user_query_string),
        "history": itemgetter("messages") | RunnablePassthrough(),
        "context": itemgetter("messages")
        | RunnableLambda(extract_user_query_string)
        | vector_search_as_retriever
        | RunnableLambda(format_context),
    }
).assign(answer=second_chain_part)

# Tell MLflow logging where to find your chain
mlflow.models.set_model(model=chain)

... and this is my chain code after:

import os
import mlflow
from operator import itemgetter
from databricks.vector_search.client import VectorSearchClient
from langchain_community.chat_models import ChatDatabricks
from langchain_community.vectorstores import DatabricksVectorSearch
from langchain_core.runnables import RunnableLambda
from langchain_core.runnables import RunnableParallel
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import CrossEncoderReranker
from langchain_community.cross_encoders import HuggingFaceCrossEncoder


# Enable MLflow Tracing
mlflow.langchain.autolog()

# Return the string contents of the most recent message from the user
def extract_user_query_string(chat_messages_array):
    return chat_messages_array[-1]["content"]

# Get the conf from the local conf file
model_config = mlflow.models.ModelConfig(
    development_config="chain_config.yaml"
)

databricks_resources = model_config.get("databricks_resources")
retriever_config = model_config.get("retriever_config")
llm_config = model_config.get("llm_config")

# Connect to the Vector Search Index
vsc = VectorSearchClient(disable_notice=True)
vs_index = vsc.get_index(
    endpoint_name=databricks_resources.get("vector_search_endpoint_name"),
    index_name=retriever_config.get("vector_search_index"),
)
vector_search_schema = retriever_config.get("schema")

# Turn the Vector Search index into a LangChain retriever
vector_search_as_retriever = DatabricksVectorSearch(
    vs_index,
    text_column=vector_search_schema.get("chunk_text"),
    columns=[
        vector_search_schema.get("primary_key"),
        vector_search_schema.get("chunk_text"),
        vector_search_schema.get("document_uri"),
    ],
).as_retriever(search_kwargs=retriever_config.get("parameters"))

# Required to:
# 1. Enable the RAG Studio Review App to properly display retrieved chunks
# 2. Enable evaluation suite to measure the retriever
mlflow.models.set_retriever_schema(
    primary_key=vector_search_schema.get("primary_key"),
    text_column=vector_search_schema.get("chunk_text"),
    doc_uri=vector_search_schema.get("document_uri"),
)


# Load CrossEncoder model for reranking
reranking_model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
compressor = CrossEncoderReranker(model=reranking_model, top_n=5)
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor, base_retriever=vector_search_as_retriever
)


# Method to format the docs returned by the retriever into the prompt
def format_context(docs):
    sources = []
    [
        sources.append(
            {"content": doc.page_content, "url": doc.metadata["url"]}
        )
        for doc in docs
    ]
    return sources

# Prompt Template for generation
prompt = PromptTemplate(
    template=llm_config.get("llm_prompt_template"),
    input_variables=llm_config.get("llm_prompt_template_variables"),
)

# FM for generation
model = ChatDatabricks(
    endpoint=databricks_resources.get("llm_endpoint_name"),
    extra_params=llm_config.get("llm_parameters"),
)

# RAG Chain
# The "|" syntax is powered by the LangChain Expression Language (LCEL)
# To learn more about LCEL, read the documentation: https://python.langchain.com/v0.1/docs/expression_language/

second_chain_part = prompt | model | StrOutputParser()

chain = RunnableParallel(
    {
        "question": itemgetter("messages")
        | RunnableLambda(extract_user_query_string),
        "history": itemgetter("messages") | RunnablePassthrough(),
        "context": itemgetter("messages")
        | RunnableLambda(extract_user_query_string)
        | compression_retriever
        | RunnableLambda(format_context),
    }
).assign(answer=second_chain_part)

# Tell MLflow logging where to find your chain
mlflow.models.set_model(model=chain)

I do realize that basically the same issue is already discussed here (https://community.databricks.com/t5/machine-learning/serving-endpoint-container-image-creation-fails...), but downgrading my ML cluster from DBR 15.4 LTS ML to DBR 14.3 LTS ML didn't solve the problem for me.

Databricks Community

Serving Endpoint: Container Image Creation Fails

Connect with Databricks Users in Your Area

Now Hiring: Databricks Community Technical Moderator

Data + AI Summit: Call for Presentations

Season's Speedings: Databricks SQL Delivers 4x Performance Boost Over Two Years

Databricks Community Champion - October 2024 - Filip Niziol

Become Our Next Monthly Community Champion!