Packaging a Pyfunc Chain: Artifacts, Code Paths, Dependencies, and Signatures

1.Packaging the files your chain loads

A pyfunc chain usually needs files when it runs, such as model weights or a tokenizer cache. In notebooks these often sit in Unity Catalog volumes, and some models download pieces from the internet, for example HuggingFace tokenizers. Databricks notes that "Real-time workloads at scale perform best when all required dependencies are statically captured at deployment time." Files in volumes therefore must be packaged, and network artifacts should be packaged whenever possible.

You package them with the artifacts parameter of log_model(). It is a dictionary that maps a name to a path. At load time, each path is available inside the model under context.artifacts[<name>], which is why the load_context in a pyfunc model reads context.artifacts["model-weights"]. The names you pick when logging are the keys your code uses when loading.

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

Sources1

2.Shipping shared preprocessing code

Preprocessing logic often lives in your own .py modules, like a custom tokenizer or prompt-formatting helpers, and %pip install can't install those. To include them, you pass the paths when you log. Authors "can log full code references that load into the path," so a model logged with code_path = ["preprocessing_utils/"] can import from preprocessing_utils inside its own methods. MLflow stores these files in a code directory next to the model. When the model loads, "MLflow adds these files or directories to the Python path." Custom Python wheel files can be shipped the same way.

The parameter name depends on the MLflow version, and so does the name of the argument that identifies the model:

Sources23

3.Pinning dependencies so serving can build the image

When you log a model, MLflow automatically writes requirements.txt and conda.yaml files. For mlflow.pyfunc.log_model, MLflow tries to work out the dependencies itself with mlflow.models.infer_pip_requirements. A custom chain imports libraries that MLflow has no flavor for, so check what got captured, and log requirements with exact versions (f"nltk=={nltk.__version__}", not just nltk).

You can control the result in two ways:

- **extra_pip_requirements adds packages that inference missed. - pip_requirements or conda_env** replaces the whole set. This is "generally discouraged because this overrides the dependencies which MLflow picks up automatically."

One case makes overriding mandatory. Databricks Runtime ML ships mlflow-skinny instead of the full mlflow package. If you log a pyfunc model there without pip_requirements, mlflow-skinny is what ends up in conda.yaml. Model Serving requires mlflow and can't build the container image without it.

# DBR ML ships with mlflow-skinny by default, so specify mlflow explicitly
# to ensure Model Serving compatibility.
mlflow.pyfunc.log_model(
    name="model",
    python_model=your_model,
    pip_requirements=["mlflow==3.8.1"],  # use mlflow, not mlflow-skinny
    registered_model_name="catalog.schema.model_name",
)

Sources32

4.Signature, input example, and testing the chain before serving

Databricks recommends adding a signature and an input example. The signature isn't optional if you want the chain in Unity Catalog, because "Signatures are necessary for logging models to the Unity Catalog." infer_signature works it out from sample inputs and the model's outputs on them:

from mlflow.models.signature import infer_signature

signature = infer_signature(training_data, model.predict(training_data))
mlflow.sklearn.log_model(model, "model", signature=signature)

These examples use the sklearn flavor, but mlflow.pyfunc.log_model takes the same signature and input_example keywords. The input_example is a sample request, for example {"feature1": 0.5, "feature2": 3}.

Next, test the logged chain the same way a consumer would call it. You can load any Python MLflow model with mlflow.pyfunc.load_model() as a generic Python function and call .predict(). That runs load_context and then your pre-process, model, post-process path. Databricks also suggests checking that the model can actually be served before you deploy it, using mlflow.models.predict. Once it passes, you can register it to Unity Catalog or the Workspace Registry and serve it from a Model Serving endpoint. Those steps belong to a separate objective.

model = mlflow.pyfunc.load_model(model_path)
model.predict(model_input)

Sources452

Exam traps

Each one states something that sounds right. Open it to see what is actually true.

1. 1.MLflow's automatic dependency capture on Databricks Runtime ML is always good enough for Model Serving.Why is that wrong?

   Databricks Runtime ML includes mlflow-skinny, which is what conda.yaml records. Model Serving needs the full mlflow package, so pin it in pip_requirements.

   Covered in Pinning dependencies so serving can build the image

2. 2.A served chain can keep reading its weights and tokenizer files straight from a Unity Catalog volume path.Why is that wrong?

   Model Serving requires files in volumes to be packaged into the model with the artifacts parameter. load_context then reads them through context.artifacts.

   Covered in Packaging the files your chain loads

3. 3.The parameter for shipping local .py modules is code_path in every MLflow version.Why is that wrong?

   MLflow 3 uses code_paths. code_path is the MLflow 2.x name.

   Covered in Shipping shared preprocessing code

Practise it for real

Log a pyfunc chain with preprocessing and postprocessing so that it loads and predicts the way a serving endpoint would call it.

1. 1.Subclass mlflow.pyfunc.PythonModel. Load the files you need in load_context, and in predict call format_inputs, then the model, then format_outputs.

   Why: One-time loading goes in load_context. Per-request pre- and post-processing goes in predict.

   You should see: A class whose predict returns post-processed output rather than raw model output.

2. 2.Call mlflow.pyfunc.log_model with python_model set to your class, artifacts pointing at your weights or tokenizer files, and code_paths listing any local helper modules.

   Why: Model Serving needs files and local code packaged inside the model artifact.

   You should see: A logged model whose artifacts include your files and a code directory.

3. 3.In the same log_model call, set pip_requirements with an exact mlflow==<version> pin, and add a signature and input_example.

   Why: Databricks Runtime ML would otherwise record mlflow-skinny, and Unity Catalog requires a signature.

   You should see: A conda.yaml that lists mlflow rather than mlflow-skinny, plus a recorded signature.

4. 4.Run model = mlflow.pyfunc.load_model(model_path), then model.predict(model_input) on a sample request.

   Why: This runs load_context and the full pre-process, model, post-process path the way a consumer would.

   You should see: Formatted predictions come back with no import or missing-file errors.

5. 5.Validate the logged model with mlflow.models.predict before deploying it.

   Why: Databricks recommends checking that the model can be served before you deploy.

   You should see: Validation completes without dependency or environment errors.