Custom PyFunc Models: Log Artifacts and Register in Unity Catalog

1.Why pyfunc, and the two methods it requires

Databricks uses the term *custom model* for any Python model or custom code that Model Serving deploys as a production API. You can log a custom model with a built-in MLflow flavor such as sklearn. You need pyfunc when a built-in flavor can't express what you want. Typical cases are preprocessing that has to run before the model's predict function, a framework that MLflow doesn't support natively, outputs that need post-processing, per-request branching logic, or code that is entirely custom.

You write a pyfunc model as a subclass of mlflow.pyfunc.PythonModel. It has two required methods, and the split between them affects performance.

load_context runs once. Load anything the model needs only once to operate here, such as weights, tokenizers and lookup objects. Databricks notes that doing so minimises what has to be loaded during predict, which speeds up inference. predict holds the logic that runs every time a request arrives. Helper methods for formatting input and output are ordinary class methods that predict calls. In the example below, load_context reads the weights and tokenizer paths from context.artifacts, and predict chains format, inference and post-processing.

class CustomModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        self.model = torch.load(context.artifacts["model-weights"])
        from preprocessing_utils.my_custom_tokenizer import CustomTokenizer
        self.tokenizer = CustomTokenizer(context.artifacts["tokenizer_cache"])

    def format_inputs(self, model_input):
        # insert some code that formats your inputs
        pass

    def format_outputs(self, outputs):
        predictions = (torch.sigmoid(outputs)).data.numpy()
        return predictions

    def predict(self, context, model_input):
        model_input = self.format_inputs(model_input)
        outputs = self.model.predict(model_input)
        return self.format_outputs(outputs)

Sources12

2.Logging custom artifacts with the model

The context.artifacts lookups in that class only work if the files were packaged when the model was logged. In a Databricks notebook those files often live in Unity Catalog volumes. Some models also download items from the internet at load time, for example Hugging Face tokenizers. Databricks says real-time workloads at scale perform best when all dependencies are captured statically at deployment time. For that reason, Model Serving does not let a model read from a volume path at serving time. The files have to be copied into the model artifact itself. Network artifacts should also be packaged whenever possible.

You do this with the artifacts parameter of log_model. It's a dictionary that maps a name of your choice to a source path.

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

At load time the dictionary keys become your handles. context.artifacts["model-weights"] resolves to the packaged copy of the file, and you open it the usual way for that file type. Because each key in artifacts must match a lookup in load_context, keep the names identical in both places.

Sources3

3.Declaring code and library dependencies

Artifacts cover files. Libraries and shared code are declared separately. During deployment, Databricks builds a production container from the logged model. The base image may include some system-level dependencies, but application-level dependencies must be listed explicitly in the MLflow model. Native flavors capture their packages automatically. For custom pyfunc models, you declare them yourself.

code_path is how the earlier CustomModel can import preprocessing_utils.my_custom_tokenizer on the endpoint. The folder is logged with the model and added to the import path when the model loads.

One trap is specific to Databricks Runtime ML. Those runtimes ship mlflow-skinny rather than the full mlflow package. If you log a pyfunc model there without pip_requirements, MLflow records mlflow-skinny in conda.yaml. Model Serving then can't build the container image, because it requires mlflow. The fix is to pin mlflow explicitly.

mlflow.pyfunc.log_model(
    name="model",
    python_model=your_model,
     ? =["mlflow==3.8.1"],  # use mlflow, not mlflow-skinny
    registered_model_name="catalog.schema.model_name",
)

Sources21

4.Registering in Unity Catalog and validating before you deploy

Once the model is logged, register it. Unity Catalog is recommended over the legacy workspace registry. Unity Catalog adds four requirements:

- A signature. New model versions in Unity Catalog must have a model signature. You can create one with infer_signature or by using Databricks autologging. With MLflow 2.5.0 and above, you can also pass an input_example to log_model, and MLflow infers the signature from it. - The right registry. In MLflow 3 the default registry URI is databricks-uc. On older setups, call mlflow.set_registry_uri("databricks-uc") first. - A three-level name in the form <catalog>.<schema>.<model>. You can pass it as registered_model_name in log_model (as in the snippet above), or call register_model afterwards. - Privileges. To create a new registered model you need CREATE MODEL and USE SCHEMA on the schema, plus USE CATALOG on the catalog. To add versions to an existing model, you must own it or hold CREATE MODEL VERSION on it, plus USE SCHEMA and USE CATALOG.

mlflow.register_model("models:/{model_id}", "{registered_model_name}")

Deploying a new version takes time. Databricks quotes about 10 minutes for packaging and provisioning, and longer for complex models. Before you deploy, Databricks recommends checking that the model can be served by validating it with mlflow.models.predict. If deployment does fail with dependency errors, Databricks recommends testing the model locally.

Sources41

Exam traps

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

1. 1.Logging a pyfunc model on Databricks Runtime ML without pip_requirements is fine, because MLflow captures its own package automatically.Why is that wrong?

   DBR ML ships mlflow-skinny, which MLflow then records in conda.yaml. Model Serving needs the full mlflow package and can't build the image without it, so pin mlflow==<version> in pip_requirements.

   Covered in Declaring code and library dependencies

2. 2.A served model can keep reading its weights from a Unity Catalog volume path at inference time.Why is that wrong?

   Files from volumes must be packaged into the model artifact with the artifacts parameter and read through context.artifacts.

   Covered in Logging custom artifacts with the model