Query serving endpoints for custom models

In this page, learn how to format scoring requests for your served model, and how to send those requests to the model serving endpoint. The guidance is relevant to serving custom models, which Databricks defines as traditional ML models or customized Python models packaged in the MLflow format. Register the models in Unity Catalog or in the workspace model registry. Examples include scikit-learn, XGBoost, PyTorch, and Hugging Face transformer models. See Deploy models using Model Serving for more information about this functionality and supported model categories.

For query requests for AI and LLM workloads, see Use foundation models.

Requirements

• A model serving endpoint.
• For the MLflow Deployment SDK, MLflow 2.9 or above is required.
• Scoring request in an accepted format.
• To send a scoring request through the REST API or MLflow Deployment SDK, you must have a Databricks API token.

Querying methods and examples

Model Serving provides the following options for sending scoring requests to served models:

Pandas DataFrame scoring example

The following example assumes a ENDPOINT_INVOCATION_URL like https://<databricks-instance>/serving-endpoints/<endpoint-name>/invocations, where <databricks-instance> is the name of your Databricks instance, and a Databricks REST API token called DATABRICKS_API_TOKEN.

See Supported scoring formats.

REST API

Score a model accepting dataframe split input format.

curl -X POST -u token:$DATABRICKS_API_TOKEN $ENDPOINT_INVOCATION_URL \
  -H 'Content-Type: application/json' \
  -d '{"dataframe_split": [{
    "columns": ["sepal length (cm)", "sepal width (cm)", "petal length (cm)", "petal width (cm)"],
    "data": [[5.1, 3.5, 1.4, 0.2], [4.9, 3.0, 1.4, 0.2]]
    }]
  }'

Score a model accepting tensor inputs. Tensor inputs should be formatted as described in TensorFlow Serving's API documentation.

curl -X POST -u token:$DATABRICKS_API_TOKEN $ENDPOINT_INVOCATION_URL \
  -H 'Content-Type: application/json' \
  -d '{"inputs": [[5.1, 3.5, 1.4, 0.2]]}'

MLflow Deployments SDK

import os
import mlflow.deployments

os.environ["DATABRICKS_HOST"] = "https://<workspace_host>.databricks.com"
os.environ["DATABRICKS_TOKEN"] = "dapi-your-databricks-token"

client = mlflow.deployments.get_deploy_client("databricks")

response = client.predict(
            endpoint="test-model-endpoint",
            inputs={"dataframe_split": {
                    "index": [0, 1],
                    "columns": ["sepal length (cm)", "sepal width (cm)", "petal length (cm)", "petal width (cm)"],
                    "data": [[5.1, 3.5, 1.4, 0.2], [4.9, 3.0, 1.4, 0.2]]
                    }
                }
          )

PowerBI

You can score a dataset in Power BI Desktop using the following steps:

1. Open dataset you want to score.

2. Go to Transform Data.

3. Right-click in the left panel and select Create New Query.

4. Go to View > Advanced Editor.

5. Replace the query body with the code snippet below, after filling in an appropriate DATABRICKS_API_TOKEN and ENDPOINT_INVOCATION_URL.

   (dataset as table ) as table =>
   let
     call_predict = (dataset as table ) as list =>
     let
       apiToken = DATABRICKS_API_TOKEN,
       modelUri = ENDPOINT_INVOCATION_URL,
       responseList = Json.Document(Web.Contents(modelUri,
         [
           Headers = [
             #"Content-Type" = "application/json",
             #"Authorization" = Text.Format("Bearer #{0}", {apiToken})
           ],
           Content = {"dataframe_records": Json.FromValue(dataset)}
         ]
       ))
     in
       responseList,
     predictionList = List.Combine(List.Transform(Table.Split(dataset, 256), (x) => call_predict(x))),
     predictionsTable = Table.FromList(predictionList, (x) => {x}, {"Prediction"}),
     datasetWithPrediction = Table.Join(
       Table.AddIndexColumn(predictionsTable, "index"), "index",
       Table.AddIndexColumn(dataset, "index"), "index")
   in
     datasetWithPrediction
   

6. Name the query with your desired model name.

7. Open the advanced query editor for your dataset and apply the model function.

Tensor input example

The following example scores a model accepting tensor inputs. Tensor inputs should be formatted as described in TensorFlow Serving's API docs. This example assumes a ENDPOINT_INVOCATION_URL like https://<databricks-instance>/serving-endpoints/<endpoint-name>/invocations, where <databricks-instance> is the name of your Databricks instance, and a Databricks REST API token called DATABRICKS_API_TOKEN.

curl -X POST -u token:$DATABRICKS_API_TOKEN $ENDPOINT_INVOCATION_URL \
    -H 'Content-Type: application/json' \
    -d '{"inputs": [[5.1, 3.5, 1.4, 0.2]]}'

Supported scoring formats

For custom models, Model Serving supports scoring requests in Pandas DataFrame (JSON) or Tensor input (JSON or Protobuf).

Pandas DataFrame

Requests should be sent by constructing a JSON-serialized Pandas DataFrame with one of the supported keys and a JSON object corresponding to the input format.

• (Recommended)dataframe_split format is a JSON-serialized Pandas DataFrame in the split orientation.

  {
    "dataframe_split": {
      "index": [0, 1],
      "columns": ["sepal length (cm)", "sepal width (cm)", "petal length (cm)", "petal width (cm)"],
      "data": [
        [5.1, 3.5, 1.4, 0.2],
        [4.9, 3.0, 1.4, 0.2]
      ]
    }
  }
  

• dataframe_records is JSON-serialized Pandas DataFrame in the records orientation.

  Note

  This format does not guarantee the preservation of column ordering, and the split format is preferred over the records format.

  {
    "dataframe_records": [
      {
        "sepal length (cm)": 5.1,
        "sepal width (cm)": 3.5,
        "petal length (cm)": 1.4,
        "petal width (cm)": 0.2
      },
      {
        "sepal length (cm)": 4.9,
        "sepal width (cm)": 3,
        "petal length (cm)": 1.4,
        "petal width (cm)": 0.2
      },
      {
        "sepal length (cm)": 4.7,
        "sepal width (cm)": 3.2,
        "petal length (cm)": 1.3,
        "petal width (cm)": 0.2
      }
    ]
  }
  

The response from the endpoint contains the output from your model, serialized with JSON, wrapped in a predictions key.

{
  "predictions": [0, 1, 1, 1, 0]
}

Tensor input

When your model expects tensors, like a TensorFlow or Pytorch model, there are two supported format options for sending requests: instances and inputs.

If you have multiple named tensors per row, then you have to have one of each tensor for every row.

• instances is a tensors-based format that accepts tensors in row format. Use this format if all the input tensors have the same 0-th dimension. Conceptually, each tensor in the instances list could be joined with the other tensors of the same name in the rest of the list to construct the full input tensor for the model, which would only be possible if all of the tensors have the same 0-th dimension.

  { "instances": [1, 2, 3] }
  

  The following example shows how to specify multiple named tensors.

  {
    "instances": [
      {
        "t1": "a",
        "t2": [1, 2, 3, 4, 5],
        "t3": [
          [1, 2],
          [3, 4],
          [5, 6]
        ]
      },
      {
        "t1": "b",
        "t2": [6, 7, 8, 9, 10],
        "t3": [
          [7, 8],
          [9, 10],
          [11, 12]
        ]
      }
    ]
  }
  

• inputs send queries with tensors in columnar format. This request is different because there are actually a different number of tensor instances of t2 (3) than t1 and t3, so it is not possible to represent this input in the instances format.

  {
    "inputs": {
      "t1": ["a", "b"],
      "t2": [
        [1, 2, 3, 4, 5],
        [6, 7, 8, 9, 10]
      ],
      "t3": [
        [
          [1, 2],
          [3, 4],
          [5, 6]
        ],
        [
          [7, 8],
          [9, 10],
          [11, 12]
        ]
      ]
    }
  }
  

The response from the endpoint is in the following format.

{
  "predictions": [0, 1, 1, 1, 0]
}

Protobuf tensor input

For models with a tensor-based input signature, you can send the request body as a serialized KServe v2 ModelInferRequest instead of JSON, using Content-Type: application/x-protobuf or application/vnd.kserve.protobuf. This avoids parsing numeric tensors from JSON text, so it is faster, especially for large payloads. Endpoint URL and authentication are unchanged.

Supported tensor datatypes are BOOL, UINT8, UINT16, UINT32, UINT64, INT8, INT16, INT32, INT64, FP16, FP32, FP64, and BYTES. For text and binary data in BYTES tensors, see Text and binary inputs in protobuf requests.

Send the data for each input in raw_input_contents, in the same order as inputs. Requests that put tensor data in the typed contents fields, such as fp32_contents, fail with HTTP 400.

The following example shows how to build a protobuf request body.

import numpy as np
from tritonclient.grpc.service_pb2 import ModelInferRequest, ModelInferResponse  # or any other module that provides KServe v2's ModelInferRequest and ModelInferResponse classes

# FP32 tensor of shape [2, 4]
tensor = np.array([[5.1, 3.5, 1.4, 0.2], [4.9, 3.0, 1.4, 0.2]], dtype=np.float32)

request = ModelInferRequest()
request.inputs.add(name="input", datatype="FP32", shape=tensor.shape)
request.raw_input_contents.append(tensor.tobytes())

body = request.SerializeToString()

The response is a serialized ModelInferResponse with tensor data in raw_output_contents.

response = ModelInferResponse()
response.ParseFromString(response_bytes)

values = np.frombuffer(response.raw_output_contents[0], dtype=np.float32).reshape(response.outputs[0].shape)

Text and binary inputs in protobuf requests

A BYTES tensor holds variable-length values, such as text or the contents of a file. A request can mix BYTES tensors with numeric tensors, for example a title and a feature vector for each row. In raw_input_contents, each BYTES value is a 4-byte little-endian length followed by the bytes of the value, as defined by the KServe v2 protocol.

The model's input signature determines what your model receives for each BYTES input:

The following example logs a model that takes a text input and a numeric input. Declare text inputs as object so that one long value doesn't pad every other value in the request.

import mlflow
import numpy as np
from mlflow.models import ModelSignature
from mlflow.types import Schema, TensorSpec


class TitleTagger(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        titles = model_input["title"]  # numpy array of Python str
        features = model_input["features"]  # float32 array of shape (rows, 600)
        genres = ["Sci-Fi" if "space" in title.lower() else "Drama" for title in titles]
        return {"genre": np.array(genres, dtype=object), "score": features.mean(axis=1)}


signature = ModelSignature(
    inputs=Schema(
        [
            TensorSpec(np.dtype(object), (-1,), "title"),  # text: use object, not str
            TensorSpec(np.dtype(np.float32), (-1, 600), "features"),
        ]
    ),
    outputs=Schema(
        [
            TensorSpec(np.dtype(object), (-1,), "genre"),
            TensorSpec(np.dtype(np.float32), (-1,), "score"),
        ]
    ),
)

with mlflow.start_run():
    mlflow.pyfunc.log_model(
        name="model",
        python_model=TitleTagger(),
        signature=signature,
        registered_model_name="main.default.title_tagger",
    )

To build a request, put text and binary values in a NumPy array with dtype=object, and encode it with serialize_byte_tensor from tritonclient. Other KServe v2 clients provide equivalent helpers.

import numpy as np
from tritonclient.grpc.service_pb2 import ModelInferRequest, ModelInferResponse
from tritonclient.utils import deserialize_bytes_tensor, serialize_byte_tensor, triton_to_np_dtype

titles = np.array(["Lost in Space", "Love, Actually"], dtype=object)
features = np.random.rand(2, 600).astype(np.float32)

request = ModelInferRequest()
request.inputs.add(name="title", datatype="BYTES", shape=titles.shape)
request.raw_input_contents.append(serialize_byte_tensor(titles).item())
request.inputs.add(name="features", datatype="FP32", shape=features.shape)
request.raw_input_contents.append(features.tobytes())

body = request.SerializeToString()

Text and binary outputs, such as str, bytes, or object arrays and lists of strings, are returned as BYTES tensors. Each value comes back as Python bytes.

response = ModelInferResponse()
response.ParseFromString(response_bytes)

outputs = {}
for tensor, raw in zip(response.outputs, response.raw_output_contents):
    if tensor.datatype == "BYTES":
        values = deserialize_bytes_tensor(raw)
    else:
        values = np.frombuffer(raw, dtype=triton_to_np_dtype(tensor.datatype))
    outputs[tensor.name] = values.reshape(tensor.shape)

genres = [genre.decode("utf-8") for genre in outputs["genre"]]

A model that returns a pandas DataFrame with both text and numeric columns fails with HTTP 500 for protobuf requests. Return a dictionary of NumPy arrays instead, as in the preceding example.

Text and binary inputs have the following limits:

• Decoded size: The text and binary values in one request can take up at most 256 MiB after decoding. With a str or bytes signature, every value is padded to the longest value in the request, so one long value can push a request over the limit. Requests over the limit fail with HTTP 400. Use an object signature for free text, and send files of similar size in the same request.
• Trailing zero bytes: With a str or bytes signature, a value that ends in a zero byte (0x00) fails with HTTP 400, because fixed-width NumPy arrays drop trailing zero bytes. To send files that end in a zero byte, see Send files in protobuf requests.
• Request size: Protobuf requests have the same 16 MB size limit as JSON requests. See Model Serving limits and regions.
• Performance: Numeric tensors are copied directly into NumPy arrays, but text values are decoded one value at a time, at about the speed of JSON parsing. Requests that are mostly text see little speedup over JSON.

Send files in protobuf requests

BYTES tensors can carry files whose size varies from row to row, such as images, audio, and documents. Send one file per row as raw bytes, with a bytes input signature. The model receives the exact bytes of each file and decodes it. A JSON request can carry binary data only as base64 text, which is about 33% larger.

Common uses include the following:

• Images: Send JPEG or PNG files to image classification, detection, or embedding models, which decode and resize them. On a GPU endpoint, the model can decode the images on the GPU.
• Audio: Send WAV, FLAC, or MP3 files to speech recognition or audio classification models.
• Documents: Send PDF files or other documents to parsing and extraction models.
• Serialized records: Send Parquet files or other serialized data that the model reads.

If your data has a fixed size for each row, such as decoded pixels or audio samples, send it as a numeric tensor such as UINT8 or FP32 instead.

The following example logs a model that reads image files and returns the size of each image.

import io

import mlflow
import numpy as np
from mlflow.models import ModelSignature
from mlflow.types import Schema, TensorSpec


class ImageInfo(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        from PIL import Image

        widths, heights = [], []
        for data in model_input["image"]:  # one encoded image file per row
            image = Image.open(io.BytesIO(bytes(data)))
            widths.append(image.width)
            heights.append(image.height)
        return {"width": np.array(widths, dtype=np.int32), "height": np.array(heights, dtype=np.int32)}


signature = ModelSignature(
    inputs=Schema([TensorSpec(np.dtype(bytes), (-1,), "image")]),  # files: use bytes
    outputs=Schema(
        [
            TensorSpec(np.dtype(np.int32), (-1,), "width"),
            TensorSpec(np.dtype(np.int32), (-1,), "height"),
        ]
    ),
)

with mlflow.start_run():
    mlflow.pyfunc.log_model(
        name="model",
        python_model=ImageInfo(),
        signature=signature,
        registered_model_name="main.default.image_info",
        extra_pip_requirements=["pillow"],
    )

To send files, read each file into a NumPy array with dtype=object. Don't use a fixed-width NumPy bytes array, which drops trailing zero bytes from each value.

from pathlib import Path

paths = ["cat.jpg", "dog.png"]
images = np.array([Path(path).read_bytes() for path in paths], dtype=object)

request = ModelInferRequest()
request.inputs.add(name="image", datatype="BYTES", shape=images.shape)
request.raw_input_contents.append(serialize_byte_tensor(images).item())

body = request.SerializeToString()

With a bytes signature, a file that ends in a zero byte fails with HTTP 400. JPEG, PNG, PDF, and Parquet files end with a fixed marker, so you can send them as is. ZIP-based files such as DOCX and XLSX, gzip files, and some audio and video files can end in a zero byte. To send these files, encode each file on the client and decode it in the model. Encoding with base64 adds about 33% to the size of each file. Appending one nonzero byte adds 1 byte.

import base64

# Option 1: base64. In the model, decode each value with base64.b64decode(bytes(data)).
encoded = np.array([base64.b64encode(Path(path).read_bytes()) for path in paths], dtype=object)

# Option 2: append one byte. In the model, remove it with bytes(data)[:-1].
padded = np.array([Path(path).read_bytes() + b"\x01" for path in paths], dtype=object)

Notebook example

See the following notebook for an example of how to test your Model Serving endpoint with a Python model:

Test Model Serving endpoint notebook

Get notebook

Additional resources

• Log requests and responses for serving endpoints (legacy).
• Use foundation models.
• Debugging guide for Model Serving.