Custom pyfunc Models with On-Demand Feature Engineering

1.Custom model objects with mlflow.pyfunc

Sometimes the transformation has to happen inside the model when a request arrives, not in an upstream pipeline. MLflow's Python function flavor, pyfunc, packages any Python code as a model. Databricks lists the cases where you need it: the model needs preprocessing before inputs reach predict, the framework is not natively supported by MLflow, raw outputs need post-processing, the model branches per request, or you want to deploy fully custom code.

A custom model subclasses mlflow.pyfunc.PythonModel and has two required functions with different jobs.

In this example, load_context loads the weights and the tokenizer from context.artifacts. predict runs the per-request feature work (format_inputs), scores the model, and post-processes the result (format_outputs).

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)

Sources1

2.Packaging a pyfunc so it can be served

Feature-engineering code is often shared across models. With code_path, the model carries that code, so the import in load_context above works after loading. For example: mlflow.pyfunc.log_model(CustomModel(), "model", code_path = ["preprocessing_utils/"]).

The runtime also catches people out. Databricks Runtime ML includes mlflow-skinny by default, not the full mlflow package. If you log a pyfunc there without pip_requirements, the model's conda.yaml captures mlflow-skinny, and Model Serving cannot build the container. Pin mlflow explicitly:

# 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",
)

Before deploying, check that the model can actually be served by using mlflow.models.predict. You can then register the model to Unity Catalog and serve it from a Model Serving endpoint.

Sources1

3.On-demand features: computing features at inference time

A hand-written pyfunc is one way to compute features per request. Feature Engineering in Unity Catalog offers a managed alternative: on-demand features. These are features whose values are not known ahead of time and are calculated at inference from request-time inputs. You define the calculation as a Python UDF governed by Unity Catalog. This requires a Unity Catalog-enabled workspace and Databricks Runtime 13.3 LTS ML or above.

During training, you reference the UDF with a FeatureFunction and map its arguments to columns with input_bindings. You pass it to create_training_set in feature_lookups, alongside ordinary FeatureLookups. Here, x comes from the request and y comes from a feature table:

  FeatureFunction(
      udf_name="main.default.example_feature",    # UDF must be in Unity Catalog so uses a three-level namespace
      input_bindings={
        "x": "new_source_input",
        "y": "materialized_feature_value"
      },
      output_name="on_demand_feature",
  ),

The request-side input, new_source_input, must be in the training DataFrame. At inference it must be in the DataFrame you pass to score_batch, or in the real-time request. The step that makes the feature logic travel with the model is logging: you must log with the Feature Store log_model method. After that, score_batch calculates every feature, including on-demand ones, and Model Serving runs the UDF for each request. For Feature Views, the equivalent building block is a CustomUDF with a RequestSource for request-time inputs.

import mlflow
mlflow.set_registry_uri(" ? ")

fe.log_model(
    model=model,
    artifact_path="main.default.model",
    flavor=mlflow.sklearn,
    training_set=training_set,
    registered_model_name="main.default.recommender_model"
)

Sources2

4.Missing values and dependencies: batch vs online

The same UDF runs in two environments, and they do not behave the same way. A UDF that takes a looked-up feature as input receives a different missing value depending on where it runs, and dependencies reach the two environments by different routes.

A robust UDF handles both missing-value forms. You can also give a FeatureLookup a fallback with default_values. If you renamed the feature with rename_outputs, the default has to use the renamed name.

%sql
CREATE OR REPLACE FUNCTION square(x INT)
RETURNS INT
LANGUAGE PYTHON AS
$$
import numpy as np
if x is None or np.isnan(x):
  return 0
return x * x
$$

Sources2

Exam traps

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

1. 1.A model trained on a training set with a FeatureFunction can be logged with plain mlflow.sklearn.log_model and will still compute on-demand features at inference.Why is that wrong?

   Only a model logged with the Feature Store log_model method carries the feature metadata that makes batch scoring and serving evaluate on-demand features automatically.

   Covered in On-demand features: computing features at inference time

2. 2.A pyfunc logged on Databricks Runtime ML without pip_requirements will serve fine, because MLflow is already installed there.Why is that wrong?

   DBR ML ships mlflow-skinny, so conda.yaml captures it, and Model Serving cannot build the image. Pin mlflow==<version> in pip_requirements.

   Covered in Packaging a pyfunc so it can be served

3. 3.Dependencies declared in the UDF's ENVIRONMENT clause are also available when the model is served online.Why is that wrong?

   The ENVIRONMENT clause covers offline SQL execution only. For a served model, list the packages in extra_pip_requirements when calling fe.log_model.

   Covered in Missing values and dependencies: batch vs online