4.4 Deploying a Custom Model to a Model Endpoint

4.4 Deploying a Custom Model to a Model Endpoint

Not every deployable model is a plain scikit-learn estimator. A production model frequently needs pre-processing, post-processing, business rules, an ensemble of several estimators, or a lookup table bundled alongside the weights. MLflow's pyfunc custom model is the mechanism that packages all of that behind one uniform predict interface, and Databricks Model Serving deploys it exactly like any built-in flavour.

Step 1 — Wrap the Logic in mlflow.pyfunc.PythonModel

import mlflow
import pandas as pd

class RiskScorer(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        """Runs ONCE per serving replica when the container starts."""
        import joblib, json
        self.model = joblib.load(context.artifacts["classifier"])
        with open(context.artifacts["thresholds"]) as f:
            self.thresholds = json.load(f)

    def predict(self, context, model_input: pd.DataFrame) -> pd.DataFrame:
        """Runs on EVERY request."""
        # Pre-processing that must match training exactly
        features = model_input.assign(
            spend_ratio=model_input["monthly_spend"] / model_input["credit_limit"]
        )[self.model.feature_names_in_]

        proba = self.model.predict_proba(features)[:, 1]

        # Post-processing: business banding the raw estimator cannot express
        band = pd.cut(proba,
                      bins=[0, self.thresholds["low"], self.thresholds["high"], 1],
                      labels=["low", "medium", "high"])

        return pd.DataFrame({"probability": proba, "risk_band": band.astype(str)})

The split between the two methods is the part exam questions probe:

Loading the model inside predict would repeat expensive deserialisation on every call and destroy the endpoint's latency profile.

Step 2 — Log It with Artifacts, Dependencies, and a Signature

from mlflow.models import infer_signature

signature = infer_signature(sample_input_df, sample_output_df)

with mlflow.start_run(run_name="risk_scorer_custom"):
    mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=RiskScorer(),
        artifacts={
            "classifier": "/dbfs/tmp/risk_clf.joblib",
            "thresholds": "/dbfs/tmp/thresholds.json",
        },
        pip_requirements=["scikit-learn==1.4.2", "pandas==2.2.1", "joblib==1.4.0"],
        signature=signature,
        input_example=sample_input_df.head(3),
        registered_model_name="prod_ml.risk.risk_scorer",
    )

• artifacts is a dictionary of files copied into the model package. Their paths are resolved through context.artifacts at load time — never hard-code a DBFS path inside predict, because the serving container may not have that mount.
• pip_requirements pins the environment the endpoint rebuilds. Omitting a transitive dependency is the single most common cause of an endpoint that fails to start.
• signature declares the input and output schema. Unity Catalog registration requires it, and the endpoint validates incoming requests against it, so a malformed payload returns a clear error instead of a silent misprediction.

Step 3 — Create the Serving Endpoint

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import EndpointCoreConfigInput, ServedEntityInput

w = WorkspaceClient()

w.serving_endpoints.create_and_wait(
    name="risk-scorer-endpoint",
    config=EndpointCoreConfigInput(
        served_entities=[
            ServedEntityInput(
                name="risk-scorer-v1",
                entity_name="prod_ml.risk.risk_scorer",   # 3-level UC name
                entity_version="1",
                workload_size="Small",
                scale_to_zero_enabled=True,
            )
        ]
    ),
)

The same endpoint can be created from the Serving UI or the REST API. Whichever route is used, three things must be true: the model is registered in Unity Catalog, the serving identity can read and execute it, and the logged environment resolves.

Sizing and Access Control

Two parts of that configuration decide whether the endpoint survives production:

Permissions are validated at create and update time, and the requirement is three grants rather than one: USE CATALOG on the catalog, USE SCHEMA on the schema, and EXECUTE on the model. A custom pyfunc that reads a second governed table at request time needs its own grants on that table too.

Shipping Helper Code with the Model

RiskScorer above is self-contained, but a real custom model usually imports project modules. Those modules do not exist inside the serving container unless you send them:

• code_paths=["./src/features"] copies the listed files or directories into the model package and puts them on sys.path at load time.
• mlflow.models.set_model(RiskScorer()) is the models-from-code alternative: point python_model at a .py file ending in that call and MLflow records the source instead of pickling the object, which sidesteps the class-not-found errors a pickled custom model raises once its defining notebook is gone.

Changing an existing endpoint uses the update-config path rather than create. The new served entity is built and brought up before traffic moves, so a failed build leaves the running version serving.

Step 4 — Query It

import requests

response = requests.post(
    f"{host}/serving-endpoints/risk-scorer-endpoint/invocations",
    headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
    json={"dataframe_records": [
        {"monthly_spend": 2400.0, "credit_limit": 8000.0, "tenure_months": 36}
    ]},
)
print(response.json()["predictions"])

Because the custom model returns a DataFrame with probability and risk_band, the endpoint's response carries both fields — the post-processing lives inside the deployed artifact rather than in every calling application.

Failure Modes Specific to Custom Models