What is Databricks Feature Serving?
Databricks Feature Serving makes data in the Databricks platform available to models or applications deployed outside of Azure Databricks. Feature Serving endpoints automatically scale to adjust to real-time traffic and provide a high-availability, low-latency service for serving features. This page describes how to set up and use Feature Serving. For a step-by-step tutorial, see Deploy and query a feature serving endpoint.
When you use Mosaic AI Model Serving to serve a model that was built using features from Databricks, the model automatically looks up and transforms features for inference requests. With Databricks Feature Serving, you can serve structured data for retrieval augmented generation (RAG) applications, as well as features that are required for other applications, such as models served outside of Databricks or any other application that requires features based on data in Unity Catalog.
Why use Feature Serving?
Databricks Feature Serving provides a single interface that serves pre-materialized and on-demand features. It also includes the following benefits:
- Simplicity. Databricks handles the infrastructure. With a single API call, Databricks creates a production-ready serving environment.
- High availability and scalability. Feature Serving endpoints automatically scale up and down to adjust to the volume of serving requests.
- Security. Endpoints are deployed in a secure network boundary and use dedicated compute that terminates when the endpoint is deleted or scaled to zero.
Requirements
- Databricks Runtime 14.2 ML or above.
- To use the Python API, Feature Serving requires
databricks-feature-engineering
version 0.1.2 or above, which is built into Databricks Runtime 14.2 ML. For earlier Databricks Runtime ML versions, manually install the required version using%pip install databricks-feature-engineering>=0.1.2
. If you are using a Databricks notebook, you must then restart the Python kernel by running this command in a new cell:dbutils.library.restartPython()
. - To use the Databricks SDK, Feature Serving requires
databricks-sdk
version 0.18.0 or above. To manually install the required version, use%pip install databricks-sdk>=0.18.0
. If you are using a Databricks notebook, you must then restart the Python kernel by running this command in a new cell:dbutils.library.restartPython()
.
Databricks Feature Serving provides a UI and several programmatic options for creating, updating, querying, and deleting endpoints. This article includes instructions for each of the following options:
- Databricks UI
- REST API
- Python API
- Databricks SDK
To use the REST API or MLflow Deployments SDK, you must have a Databricks API token.
Important
As a security best practice for production scenarios, Databricks recommends that you use machine-to-machine OAuth tokens for authentication during production.
For testing and development, Databricks recommends using a personal access token belonging to service principals instead of workspace users. To create tokens for service principals, see Manage tokens for a service principal.
Authentication for Feature Serving
For information about authentication, see Authentication for Azure Databricks automation - overview.
Create a FeatureSpec
A FeatureSpec
is a user-defined set of features and functions. You can combine features and functions in a FeatureSpec
. FeatureSpecs
are stored in and managed by Unity Catalog and appear in Catalog Explorer.
The tables specified in a FeatureSpec
must be published to an online table or a third-party online store. See Use online tables for real-time feature serving or Third-party online stores.
You must use the databricks-feature-engineering
package to create a FeatureSpec
.
from databricks.feature_engineering import (
FeatureFunction,
FeatureLookup,
FeatureEngineeringClient,
)
fe = FeatureEngineeringClient()
features = [
# Lookup column `average_yearly_spend` and `country` from a table in UC by the input `user_id`.
FeatureLookup(
table_name="main.default.customer_profile",
lookup_key="user_id",
features=["average_yearly_spend", "country"]
),
# Calculate a new feature called `spending_gap` - the difference between `ytd_spend` and `average_yearly_spend`.
FeatureFunction(
udf_name="main.default.difference",
output_name="spending_gap",
# Bind the function parameter with input from other features or from request.
# The function calculates a - b.
input_bindings={"a": "ytd_spend", "b": "average_yearly_spend"},
),
]
# Create a `FeatureSpec` with the features defined above.
# The `FeatureSpec` can be accessed in Unity Catalog as a function.
fe.create_feature_spec(
name="main.default.customer_features",
features=features,
)
Create an endpoint
The FeatureSpec
defines the endpoint. For more information, see Create custom model serving endpoints, the Python API documentation, or the Databricks SDK documentation for details.
Note
For workloads that are latency sensitive or require high queries per second, Model Serving offers route optimization on custom model serving endpoints, see Configure route optimization on serving endpoints.
REST API
curl -X POST -u token:$DATABRICKS_API_TOKEN ${WORKSPACE_URL}/api/2.0/serving-endpoints \
-H 'Content-Type: application/json' \
-d '"name": "customer-features",
"config": {
"served_entities": [
{
"entity_name": "main.default.customer_features",
"workload_size": "Small",
"scale_to_zero_enabled": true
}
]
}'
Databricks SDK - Python
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import EndpointCoreConfigInput, ServedEntityInput
workspace = WorkspaceClient()
# Create endpoint
workspace.serving_endpoints.create(
name="my-serving-endpoint",
config = EndpointCoreConfigInput(
served_entities=[
ServedEntityInput(
entity_name="main.default.customer_features",
scale_to_zero_enabled=True,
workload_size="Small"
)
]
)
)
Python API
from databricks.feature_engineering.entities.feature_serving_endpoint import (
ServedEntity,
EndpointCoreConfig,
)
fe.create_feature_serving_endpoint(
name="customer-features",
config=EndpointCoreConfig(
served_entities=ServedEntity(
feature_spec_name="main.default.customer_features",
workload_size="Small",
scale_to_zero_enabled=True,
instance_profile_arn=None,
)
)
)
To see the endpoint, click Serving in the left sidebar of the Databricks UI. When the state is Ready, the endpoint is ready to respond to queries. To learn more about Mosaic AI Model Serving, see Mosaic AI Model Serving.
Get an endpoint
You can use the Databricks SDK or the Python API to get the metadata and status of an endpoint.
Databricks SDK - Python
endpoint = workspace.serving_endpoints.get(name="customer-features")
# print(endpoint)
Python API
endpoint = fe.get_feature_serving_endpoint(name="customer-features")
# print(endpoint)
Get the schema of an endpoint
You can use the REST API to get the schema of an endpoint. For more information about the endpoint schema, see Get a model serving endpoint schema.
ACCESS_TOKEN=<token>
ENDPOINT_NAME=<endpoint name>
curl "https://example.databricks.com/api/2.0/serving-endpoints/$ENDPOINT_NAME/openapi" -H "Authorization: Bearer $ACCESS_TOKEN" -H "Content-Type: application/json"
Query an endpoint
You can use the REST API, the MLflow Deployments SDK, or the Serving UI to query an endpoint.
The following code shows how to set up credentials and create the client when using the MLflow Deployments SDK.
# Set up credentials
export DATABRICKS_HOST=...
export DATABRICKS_TOKEN=...
# Set up the client
import mlflow.deployments
client = mlflow.deployments.get_deploy_client("databricks")
Note
As a security best practice, when you authenticate with automated tools, systems, scripts, and apps, Databricks recommends that you use personal access tokens belonging to service principals instead of workspace users. To create tokens for service principals, see Manage tokens for a service principal.
Query an endpoint using APIs
This section includes examples of querying an endpoint using the REST API or the MLflow Deployments SDK.
REST API
curl -X POST -u token:$DATABRICKS_API_TOKEN $ENDPOINT_INVOCATION_URL \
-H 'Content-Type: application/json' \
-d '{"dataframe_records": [
{"user_id": 1, "ytd_spend": 598},
{"user_id": 2, "ytd_spend": 280}
]}'
MLflow Deployments SDK
Important
The following example uses the predict()
API from the MLflow Deployments SDK. This API is Experimental and the API definition might change.
import mlflow.deployments
client = mlflow.deployments.get_deploy_client("databricks")
response = client.predict(
endpoint="test-feature-endpoint",
inputs={
"dataframe_records": [
{"user_id": 1, "ytd_spend": 598},
{"user_id": 2, "ytd_spend": 280},
]
},
)
Query an endpoint using the UI
You can query a serving endpoint directly from the Serving UI. The UI includes generated code examples that you can use to query the endpoint.
In the left sidebar of the Azure Databricks workspace, click Serving.
Click the endpoint you want to query.
In the upper-right of the screen, click Query endpoint.
In the Request box, type the request body in JSON format.
Click Send request.
// Example of a request body.
{
"dataframe_records": [
{"user_id": 1, "ytd_spend": 598},
{"user_id": 2, "ytd_spend": 280}
]
}
The Query endpoint dialog includes generated example code in curl, Python, and SQL. Click the tabs to view and copy the example code.
To copy the code, click the copy icon in the upper-right of the text box.
Update an endpoint
You can update an endpoint using the REST API, the Databricks SDK, or the Serving UI.
Update an endpoint using APIs
REST API
curl -X PUT -u token:$DATABRICKS_API_TOKEN ${WORKSPACE_URL}/api/2.0/serving-endpoints/<endpoint_name>/config \
-H 'Content-Type: application/json' \
-d '"served_entities": [
{
"name": "customer-features",
"entity_name": "main.default.customer_features_new",
"workload_size": "Small",
"scale_to_zero_enabled": True
}
]'
Databricks SDK - Python
workspace.serving_endpoints.update_config(
name="my-serving-endpoint",
served_entities=[
ServedEntityInput(
entity_name="main.default.customer_features",
scale_to_zero_enabled=True,
workload_size="Small"
)
]
)
Update an endpoint using the UI
Follow these steps to use the Serving UI:
- In the left sidebar of the Azure Databricks workspace, click Serving.
- In the table, click the name of the endpoint you want to update. The endpoint screen appears.
- In the upper-right of the screen, click Edit endpoint.
- In the Edit serving endpoint dialog, edit the endpoint settings as needed.
- Click Update to save your changes.
Delete an endpoint
Warning
This action is irreversible.
You can delete an endpoint using the REST API, the Databricks SDK, the Python API, or the Serving UI.
Delete an endpoint using APIs
REST API
curl -X DELETE -u token:$DATABRICKS_API_TOKEN ${WORKSPACE_URL}/api/2.0/serving-endpoints/<endpoint_name>
Databricks SDK - Python
workspace.serving_endpoints.delete(name="customer-features")
Python API
fe.delete_feature_serving_endpoint(name="customer-features")
Delete an endpoint using the UI
Follow these steps to delete an endpoint using the Serving UI:
- In the left sidebar of the Azure Databricks workspace, click Serving.
- In the table, click the name of the endpoint you want to delete. The endpoint screen appears.
- In the upper-right of the screen, click the kebab menu
and select Delete.
Monitor the health of an endpoint
For information about the logs and metrics available for Feature Serving endpoints, see Monitor model quality and endpoint health.
Access control
For information about permissions on Feature Serving endpoints, see Manage permissions on your model serving endpoint.
Example notebook
This notebook illustrates how to use the Databricks SDK to create a Feature Serving endpoint using Databricks Online Tables.
Feature Serving example notebook with online tables
Feedback
https://aka.ms/ContentUserFeedback.
Coming soon: Throughout 2024 we will be phasing out GitHub Issues as the feedback mechanism for content and replacing it with a new feedback system. For more information see:Submit and view feedback for