Configure AI Gateway on model serving endpoints

In this article, you learn how to configure Mosaic AI Gateway on a model serving endpoint.

Requirements

• A Databricks workspace in one of the following regions:
  • An external models supported region
  • A Foundation Model APIs provisioned throughput supported region
• A model serving endpoint. You can do one of the following:
  • To create an endpoint for external models, complete steps 1 and 2 of Create an external model serving endpoint.
  • To create an endpoint for provisioned throughput, see Provisioned throughput Foundation Model APIs.

Configure AI Gateway using the UI

In the AI Gateway section of the endpoint creation page, you can individually configure AI Gateway features. See Supported features for which features are available on external model serving endpoints and provisioned throughput endpoints.

The following table summarizes how to configure AI Gateway during endpoint creation using the Serving UI. If you prefer to do this programmatically, see the Notebook example.

Feature	How to enable	Details
Usage tracking	Select Enable usage tracking to enable tracking and monitoring of data usage metrics. This feature is enabled by default for pay-per-token endpoints.	You must have Unity Catalog enabled. Account admins must enable the serving system table schema before using the system tables: system.serving.endpoint_usage, which captures token counts for each request to the endpoint and system.serving.served_entities, which stores metadata for each foundation model. See Usage tracking table schemas Only account admins have permission to view or query the served_entities table or endpoint_usage table, even though the user that manages the endpoint must enable usage tracking. See Grant access to system tables The input and output token count are estimated as (text_length+1)/4 if the token count is not returned by the model.
Payload logging	Select Enable inference tables to automatically log requests and responses from your endpoint into Delta tables managed by Unity Catalog.	You must have Unity Catalog enabled and CREATE_TABLE access in the specified catalog schema. Inference tables enabled by AI Gateway have a different schema than inference tables created for model serving endpoints that serve custom models. See AI Gateway-enabled inference table schema. Payload logging data populates these tables less than hour after querying the endpoint. Payloads larger than 1 MB are not logged. The response payload aggregates the response of all of the returned chunks. Streaming is supported. In streaming scenarios, the response payload aggregates the response of returned chunks.
AI Guardrails	See Configure AI Guardrails in the UI.	Guardrails prevent the model from interacting with unsafe and harmful content that is detected in model inputs and outputs. Output guardrails are not supported for embeddings models or for streaming.
Rate limits	Select Rate limits to enforce request rate limits that manage traffic for your endpoint on a per user and per endpoint basis.	Rate limits are defined in queries per minute (QPM). The default is No limit for both per user and per endpoint.
Traffic splitting	In the Served entities section, specify the percentage of traffic you want to be routed to specific models. To configure traffic splitting on your endpoint programmatically, see Serve multiple external models to an endpoint.	To route all traffic to a specific model, set it to 100%. If you want to specify a fallback-only model, add that model to the endpoint and set its percentage of traffic to 0%. To load balance traffic across models and set up fallbacks, you can expect the following behavior: Requests are randomly split across the entities based on the assigned traffic percentages. If the request hits the first entity and fails, it falls back to the next entity in the order that the served entities were listed during endpoint creation or the most recent endpoint update. The traffic split does not influence the order of fallback attempts.
Fallbacks	Select Enable fallbacks in the AI Gateway section to send your request to other served models on the endpoint as a fallback.	If the initial request routed to a certain entity returns a 429 or 5XX error, the request falls back to the next entity listed on the endpoint. The order in which requests are redirected to fallback served entities is based on the order the models are listed during endpoint creation or the most recent endpoint update. Traffic percentage does not influence the order of fallback attempts sent to served entities. Fallbacks are only supported for external models. You must assign traffic percentages to other models served on the endpoint before you can enable fallbacks to external models. Any external model assigned 0% traffic functions exclusively as a fallback model. You can have a maximum of two fallbacks. Each entity is tried once in sequential order until the request succeeds. If all listed entites have been tried without success the request fails. The first successful or last failed request attempt and response are logged in both the usage tracking and payload logging tables.

The following diagram shows an example where,

• Three served entities are served on a model serving endpoint.
• The request is originally routed to Served entity 3.
• If the request returns a 200 response, the request was successful on Served entity 3 and the request and its response are logged to the usage tracking and payload logging tables of the endpoint.
• If the request returns a 429 or 5xx error on Served entity 3, the request falls back to the next served entity on the endpoint, Served entity 1.
  • If the request returns a 429 or 5xx error on Served entity 1, the request falls back to the next served entity on the endpoint, Served entity 2.
  • If the request returns a 429 or 5xx error on Served entity 2, the request fails since this is the maximum number of fall back entities. The failed request and the response error are logged to the usage tracking and payload logging tables.

Configure AI Guardrails in the UI

The following table shows how to configure supported guardrails.

Guardrail	How to enable	Details
Safety	Select Safety to enable safeguards to prevent your model from interacting with unsafe and harmful content.
Personally identifiable information (PII) detection	Select PII detection to detect PII data such as names, addresses, credit card numbers.
Valid topics	You can type topics directly into this field. If you have multiple entries, be sure to press enter after each topic. Alternatively, you can upload a .csv or .txt file.	A maximum of 50 valid topics can be specified. Each topic cannot exceed 100 characters
Invalid keywords	You can type topics directly into this field. If you have multiple entries, be sure to press enter after each topic. Alternatively, you can upload a .csv or .txt file.	A maximum of 50 invalid keywords can be specified. Each keyword cannot exceed 100 characters.

Usage tracking table schemas

The following sections summarize the usage tracking table schemas for the system.serving.served_entities and system.serving.endpoint_usage system tables.

system.serving.served_entities usage tracking table schema

The system.serving.served_entities usage tracking system table has the following schema:

Column name	Description	Type
served_entity_id	The unique ID of the served entity.	STRING
account_id	The customer account ID for Delta Sharing.	STRING
workspace_id	The customer workspace ID of the serving endpoint.	STRING
created_by	The ID of the creator.	STRING
endpoint_name	The name of the serving endpoint.	STRING
endpoint_id	The unique ID of the serving endpoint.	STRING
served_entity_name	The name of the served entity.	STRING
entity_type	Type of the entity that is served. Can be FEATURE_SPEC, EXTERNAL_MODEL, FOUNDATION_MODEL, or CUSTOM_MODEL	STRING
entity_name	The underlying name of the entity. Different from the served_entity_name which is a user provided name. For example, entity_name is the name of the Unity Catalog model.	STRING
entity_version	The version of the served entity.	STRING
endpoint_config_version	The version of the endpoint configuration.	INT
task	The task type. Can be llm/v1/chat, llm/v1/completions, or llm/v1/embeddings.	STRING
external_model_config	Configurations for external models. For example, {Provider: OpenAI}	STRUCT
foundation_model_config	Configurations for foundation models. For example,{min_provisioned_throughput: 2200, max_provisioned_throughput: 4400}	STRUCT
custom_model_config	Configurations for custom models. For example,{ min_concurrency: 0, max_concurrency: 4, compute_type: CPU }	STRUCT
feature_spec_config	Configurations for feature specifications. For example, { min_concurrency: 0, max_concurrency: 4, compute_type: CPU }	STRUCT
change_time	Timestamp of change for the served entity.	TIMESTAMP
endpoint_delete_time	Timestamp of entity deletion. The endpoint is the container for the served entity. After the endpoint is deleted, the served entity is also deleted.	TIMESTAMP

system.serving.endpoint_usage usage tracking table schema

The system.serving.endpoint_usage usage tracking system table has the following schema:

Column name	Description	Type
account_id	The customer account ID.	STRING
workspace_id	The customer workspace id of the serving endpoint.	STRING
client_request_id	The user provided request identifier that can be specified in the model serving request body.	STRING
databricks_request_id	A Databricks generated request identifier attached to all model serving requests.	STRING
requester	The ID of the user or service principal whose permissions are used for the invocation request of the serving endpoint.	STRING
status_code	The HTTP status code that was returned from the model.	INTEGER
request_time	The timestamp at which the request is received.	TIMESTAMP
input_token_count	The token count of the input.	LONG
output_token_count	The token count of the output.	LONG
input_character_count	The character count of the input string or prompt.	LONG
output_character_count	The character count of the output string of the response.	LONG
usage_context	The user provided map containing identifiers of the end user or the customer application that makes the call to the endpoint. See Further define usage with usage_context.	MAP
request_streaming	Whether the request is in stream mode.	BOOLEAN
served_entity_id	The unique ID used to join with the system.serving.served_entities dimension table to lookup information about the endpoint and served entity.	STRING

Further define usage with usage_context

When you query an external model with usage tracking enabled, you can provide the usage_context parameter with type Map[String, String]. The usage context mapping appears in the usage tracking table in the usage_context column. The usage_context map size cannot exceed 10 KiB.

Account admins can aggregate different rows based on the usage context to get insights and can join this information with the information in the payload logging table. For example, you can add end_user_to_charge to the usage_context for tracking cost attribution for end users.

Bash

{
  "messages": [
    {
      "role": "user",
      "content": "What is Databricks?"
    }
  ],
  "max_tokens": 128,
  "usage_context":
    {
      "use_case": "external",
      "project": "project1",
      "priority": "high",
      "end_user_to_charge": "abcde12345",
      "a_b_test_group": "group_a"
    }
}

Update AI Gateway features on endpoints

You can update AI Gateway features on model serving endpoints that had them previously enabled and endpoints that did not. Updates to AI Gateway configurations take about 20-40 seconds to be applied, however rate limiting updates can take up to 60 seconds.

The following shows how to update AI Gateway features on a model serving endpoint using the Serving UI.

In the Gateway section of the endpoint page, you can see which features are enabled. To update these features, click Edit AI Gateway.

Notebook example

The following notebook shows how to programmatically enable and use Databricks Mosaic AI Gateway features to manage and govern models from providers. See the PUT /api/2.0/serving-endpoints/{name}/ai-gateway for REST API details.

Enable Databricks Mosaic AI Gateway features notebook

Open notebook in new tab

Additional resources

• Mosaic AI Gateway introduction.
• Monitor served models using AI Gateway-enabled inference tables.