Databricks GenAI Engineer: Coding pyfunc Chains and Registering Them to Unity Catalog

A chain is the code that turns a user request into an answer: it formats a prompt, optionally retrieves context, calls an LLM and shapes the reply. On Databricks you package that code as an MLflow model so it can be versioned in Unity Catalog and deployed to Model Serving. The flavor you log it with decides how it loads back and how much wrapper code you write.

Databricks recommends ResponsesAgent for new agents, because Agent Framework, the Review App, AI Playground and Model Serving all understand its message format out of the box and you do not write a signature by hand. A plain PythonModel still serves fine, but you define the input format and signature yourself. Any pyfunc wrapper around a LangChain chain loads back as pyfunc, so if other teams need the runnable itself, log with the LangChain flavor.

A custom pyfunc chain is a subclass of mlflow.pyfunc.PythonModel with two methods that matter: load_context runs once each time the model is loaded, and predict runs on every request. Put request-independent setup in the first and request-specific work in the second.

import mlflow
from mlflow.deployments import get_deploy_client

class SupportChain(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # once per load: clients, tokenizers, files packaged with the model
        self.client = get_deploy_client("databricks")
        with open(context.artifacts["glossary"]) as f:
            self.glossary = load_glossary(f)

    def predict(self, context, model_input, params=None):
        question = model_input["question"][0]
        question = mask_account_numbers(question)        # pre-processing
        reply = self.client.predict(endpoint=LLM_ENDPOINT,
                                    inputs=build_request(question))
        return clean_and_validate(reply)                 # post-processing

The general rule: work that does not depend on the request belongs in load_context. Anything built in predict is paid for on every call, and anything built in __init__ has to survive pickling when the instance is logged, which live network clients usually cannot. Deferring setup until a request arrives only moves the cost onto whichever request comes first, whereas load_context finishes before the model accepts traffic.

Pre-processing is code in predict that runs on the request before the LLM call; post-processing is code in predict that runs on the LLM's response before it is returned. Because both live inside the model, every caller of the serving endpoint gets the same behaviour: a web app, a Slack bot, a REST client and SQL users calling it through ai_query().

• Pre-processing examples: masking account numbers or other sensitive values so they never leave the chain, expanding acronyms before retrieval, normalising or rewriting the question, trimming retrieved chunks to fit the request's token budget. Masking has to happen before the endpoint call; doing it afterwards is too late, because the value has already been sent.
• Post-processing examples: stripping Markdown code fences or trailing "Sources:" lines, parsing and validating JSON, reshaping the reply into fixed fields such as answer and sources, removing internal chunk IDs, adding a fallback when validation fails.

Rule of thumb: a requirement that must hold for every caller, including ones you do not control, belongs in code that ships with the model version rather than in each client or in the prompt alone.

A simple chain in LangChain Expression Language (LCEL) is a pipeline of runnables joined with |, where each step receives the previous step's output. The basic generation chain is prompt, then chat model, then output parser: format, generate, parse.

from databricks_langchain import ChatDatabricks
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

llm = ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct")

# prompt -> chat model -> plain string
simple_chain = prompt | llm | StrOutputParser()

# RAG: the dict runs each branch on the same input (the question string)
rag_chain = (
    {"context": retriever | format_docs,
     "question": RunnablePassthrough()}
    | rag_prompt | llm | StrOutputParser()
)

• Order matters. The template must come first to turn variables into messages; the chat model generates; StrOutputParser pulls the text out of the chat message. Putting the parser before the model, or the model before the prompt, breaks the data flow.
• Every template variable needs a branch. In a RAG chain whose template has {context} and {question}, the leading dictionary feeds the question to the retriever (then format_docs joins the documents' page_content into one string) and forwards the question unchanged with RunnablePassthrough. Leaving out a key, or swapping the branches, gives the template missing or wrong values.
• Use the right model class. ChatDatabricks calls a chat endpoint (a Foundation Model APIs endpoint or your own serving endpoint); DatabricksEmbeddings produces vectors and has no generation step.

Coding "according to requirements" means mapping each stated need to a step: a required output format to a parser or post-processing, a required grounding source to a retriever branch, a required model to the endpoint name.

Models from code is the logging method Databricks recommends for chains and agents: the chain lives in a Python file that ends with mlflow.models.set_model(...), and you pass that file's path to log_model instead of a Python object. MLflow stores the source file and runs it again when the model is loaded to rebuild the object.

# chain.py
import mlflow
from mlflow.models import ModelConfig

config = ModelConfig(development_config="config.yaml")
llm = ChatDatabricks(endpoint=config.get("llm_endpoint"))
retriever = build_retriever(config.get("index_name"))   # built in the file
chain = ...                                            # compose the chain
mlflow.models.set_model(chain)

# notebook
with mlflow.start_run():
    mlflow.langchain.log_model(
        lc_model="chain.py",
        name="chain",
        model_config={"llm_endpoint": "prod-llm",
                      "index_name": "ml.support.docs_index"},
        input_example={"messages": [{"role": "user",
                                     "content": "How do I reset MFA?"}]},
    )

• Why from code: nothing is pickled, so chains that hold unserialisable clients log cleanly, and reviewers can read the logged source. Passing an object, wrapped in a PythonModel or not, still serialises it.
• The file must be self-contained. Because it is re-run at load time in a fresh process, variables defined in earlier notebook cells do not exist there; a chain that uses a notebook-defined retriever fails on load with NameError. Build every object inside the file.
• ModelConfig carries values, not objects. development_config is used while you iterate in the notebook but is not used when the model is logged. The model_config argument to log_model supplies (and overrides) the configuration for the logged model; if neither is available at load time, loading raises an error. Logging the same file once per environment with its own model_config (staging endpoint, production endpoint) keeps one codebase.

A logged model must carry everything the serving container needs that your notebook happened to have: Python packages, your own local modules and any data files. Each has its own log_model argument, and mixing them up is a frequent cause of ModuleNotFoundError or FileNotFoundError on the endpoint.

• Inference can miss lazy imports. MLflow infers requirements from what the model imports when it is logged, so a package imported only inside a request-time helper may be left out. Add it with extra_pip_requirements to keep everything else that was inferred.
• Use only one environment argument. conda_env, pip_requirements and extra_pip_requirements are mutually exclusive; passing any two raises an error before anything is logged. code_paths, artifacts and input_example are separate concerns and combine with any one of them.
• Workspace and Volume paths are not the model. A serving container cannot read the notebook's /Workspace/Users/... path. Log the file as an artifact and read context.artifacts[...], ideally once in load_context; reading the original path bypasses the packaged copy. code_paths makes modules importable but does not change which path your code opens.
• Pin versions that matter for reproducibility (the LangChain, MLflow and Vector Search client versions) so a rebuild months later gets the same behaviour.

A model signature is the schema of a model's inputs and outputs, and Unity Catalog requires one on every model version. An input example is a sample request saved with the model. The easiest way to satisfy both needs is to log with an input_example: when no signature is passed, MLflow runs inference on the example, infers the signature from it, and stores the example alongside the model.

• Without a signature, Unity Catalog registration fails. Logging an instance with no optional arguments produces a model that cannot become a Unity Catalog model version. Supplying an input_example when you log lets MLflow infer the signature from a real request, and the same example is stored with the model as documentation of the expected input.
• Writing a signature by hand (mlflow.models.infer_signature or a ModelSignature) also works, but it does not save a sample request.
• Agent interfaces need neither. A ResponsesAgent has a fixed request/response schema, so MLflow infers its signature at logging time. Recent MLflow versions can also infer a pyfunc signature from type hints on predict.
• What a signature does not do: it does not validate business rules, clean output or install anything. Output shaping is post-processing in predict.

The example should look like a real request in the chain's input format, for instance {"question": ["How do I reset my password?"]} for a DataFrame-style pyfunc chain, or a messages list for a chat-style agent.

The basic elements of a RAG application are a model flavor, an embedding model, a retriever, dependencies, an input example and a signature; the two that most often go wrong are the embedding model and the retriever settings. The rule that decides most cases: queries must be embedded with the same model that embedded the indexed chunks. Similarity search compares a query vector with stored vectors, which only means something when both come from one model. A different model, even one with the same dimension count, maps text into a different space, and results degrade silently rather than failing.

The embedding model's input window matters too. databricks-bge-large-en embeds up to 512 tokens and databricks-gte-large-en up to 8,192 (both produce 1,024-dimension vectors). If chunks are much longer than the window, only the beginning of each chunk shapes its vector, so details near the end are hard to retrieve. When the chunk size must stay, switch to a model whose window covers the chunk, and rebuild (re-embed) the index with it; changing only the query model would mix vector spaces, and returning more results or adding hybrid keyword search leaves the truncation in place.

How the retriever is configured depends on who computes the embeddings:

from databricks_langchain import VectorSearchRetrieverTool, DatabricksEmbeddings

retriever_tool = VectorSearchRetrieverTool(
    index_name="ml.support.docs_index",
    text_column="chunk",                     # the text, not the vector column
    embedding=DatabricksEmbeddings(endpoint="databricks-gte-large-en"),
)

Creating and sizing indexes and endpoints (standard vs storage-optimized, sync modes, hybrid search) is covered in the Mosaic AI Vector Search lesson; here you only need to wire a correct retriever into the chain.

Registering to Unity Catalog means creating a governed model version under a three-level name, catalog.schema.model, in the Unity Catalog model registry. Two things are checked when a version is registered: the three-level name and a signature.

mlflow.set_registry_uri("databricks-uc")      # default in MLflow 3

# option 1: register while logging
mlflow.pyfunc.log_model(name="chain", python_model="chain.py",
                        input_example=example,
                        registered_model_name="ml.support.support_rag")

# option 2: register an already-logged model
mlflow.register_model(model_info.model_uri, "ml.support.support_rag")

• Registry URI vs tracking URI. set_registry_uri("databricks-uc") sends registrations to Unity Catalog; "databricks" as a registry URI is the legacy workspace registry. The tracking URI is a different setting that only decides where runs and experiments go. MLflow 3 uses Unity Catalog by default; with MLflow 2.x in a workspace whose default catalog is not a Unity Catalog catalog (for example hive_metastore), you must set it.
• Name notations. Register with catalog.schema.model. Load with models:/catalog.schema.model@alias or models:/catalog.schema.model/<version>. Slash-separated names are not valid registered model names.
• Each log with registered_model_name creates a new version of that model (version 1, 2, 3 …).

Promoting with aliases

In Unity Catalog, you promote a model version by moving an alias, because stages (Staging, Production, Archived) are not supported. An alias such as champion is a named, movable pointer to one version; consumers load models:/ml.support.support_rag@champion and get whatever version the alias points at, so promotion needs no code change.

from mlflow import MlflowClient
MlflowClient().set_registered_model_alias(
    name="ml.support.support_rag", alias="champion", version=6)

model = mlflow.pyfunc.load_model("models:/ml.support.support_rag@champion")

An alias moves only when someone sets it. Registering version 6 leaves champion on version 5 until you reassign it (with the client or in Catalog Explorer). Tags are metadata you can filter on, but you cannot load a model by tag, and copying a model per release (support_rag_prod) forces consumers to change what they load.

Privileges for models

Table privileges do not carry over: SELECT, MODIFY and CREATE TABLE apply to tables, not models. CREATE MODEL on the schema does not let you add versions to someone else's model; under least privilege, grant CREATE MODEL VERSION on that model rather than transferring ownership. Access control for serving endpoints and the resources an endpoint calls is covered in the model serving lesson.

The reliable pre-deployment check is mlflow.models.predict run in an isolated environment built from the model's logged requirements. It loads the logged model and runs it on input data inside a fresh environment, which is what the serving container will do, so a missing dependency fails before any endpoint exists.

# isolated environment from the logged requirements (virtualenv is the default)
mlflow.models.predict(model_uri, input_data=example, env_manager="uv")

# prove a fix without changing the logged model
mlflow.models.predict(model_uri, input_data=example, env_manager="uv",
                      pip_requirements_override=["databricks-vectorsearch"])

# persist the fix once it is proven
mlflow.models.update_model_requirements(model_uri, "add",
                                        ["databricks-vectorsearch"])

• env_manager="local" is not a dependency test. It runs in the current environment, so the notebook's extra libraries hide anything the logged requirements miss. The same is true of mlflow.pyfunc.load_model(uri).predict(...) in the notebook.
• pip_requirements_override adds or changes packages for that one run, so you can confirm that a single package fixes a failure before touching the model. install_mlflow only controls installing MLflow itself.
• update_model_requirements edits the logged model's requirements; use it (or re-log) after the fix is proven.
• infer_signature creates a schema; it tests nothing about the environment.