Mosaic AI Vector Search: Budget policies

Preview

This feature is in Public Preview.

This article describes how to use budget policies to track vector search costs.

Budget policies enable administrators to group and filter billing records across all Databricks serverless products, and provide a dedicated UI for tracking spending. Budget policies are created by workspace admins and then assigned to users. For general information and details about how to create and manage budget policies, see Attribute usage with serverless budget policies.

Apply a budget policy

You can apply a budget policy when you create an endpoint, or you can edit the endpoint later to apply a policy.

Apply budget policy at endpoint creation

Databricks UI

Follow these steps to create a vector search endpoint and apply a budget policy:

1. In the left sidebar, click Compute.

2. Click the Vector Search tab and click Create.

3. Click the arrow to the right of Advanced Settings to open that section of the dialog.

   

4. Select a policy from the Budget Policy dropdown menu.

5. Click Confirm.

Python SDK

Python

client.create_endpoint(
   name="vector_search_endpoint_name",
   endpoint_type="STANDARD",
   budget_policy_id="12345678-1234-1234-1234-1234567890ab",
)

client.update_endpoint_budget_policy(
   name="vector_search_endpoint_name",
   budget_policy_id="12345678-1234-1234-1234-1234567890ab",
)

REST API

JSON

{
	name="vector_search_endpoint_name",
	budget_policy_id=”12345678-1234-1234-1234-1234567890ab”
}

Edit a budget policy for an endpoint

To edit a budget policy for an endpoint using the Databricks UI:

1. On the endpoint page, click the pencil icon.

   

2. Select a policy from the dropdown menu and click Save.

To edit a policy using the Python SDK or the REST API, see the tabs under Apply budget policy at endpoint creation.

Query spending by budget policy

To query spending by budget policy, you must have the budget policy ID. To get the budget policy ID using the UI, follow these steps as an admin user:

1. Click your username in the top bar of the Databricks workspace and click Settings.
2. Click Compute.
3. Next to Serverless budget policies, click Manage.
4. Select a serverless budget policy. The budget policy ID appears in the upper-right section, About this policy.

To get the budget policy ID using the REST API, see List policies.

The following query aggregates spending by the budget policy ID. Replace <INSERT BUDGET POLICY ID> with the budget policy ID.

SQL

WITH vector_search_usage_by_budget_policy (
 SELECT *,
        CASE WHEN usage_type = "STORAGE_SPACE" THEN 'storage'
             ELSE 'serving'
       END as workload_type
   FROM system.billing.usage
  WHERE billing_origin_product = 'VECTOR_SEARCH'
  AND usage_metadata.endpoint_name IS NOT NULL
  AND usage_metadata.budget_policy_id = '<INSERT BUDGET POLICY ID>'
),
daily_usage_by_budget_policy AS (
 SELECT workspace_id,
      cloud,
      usage_date,
      workload_type,
      usage_metadata.endpoint_name as vector_search_endpoint,
      CASE WHEN workload_type = 'serving' THEN SUM(usage_quantity)
           ELSE null
           END as dbus,
      CASE WHEN workload_type = 'storage' THEN SUM(usage_quantity)
           ELSE null
           END as dsus
FROM vector_search_usage_by_budget_policy
GROUP BY all
ORDER BY 1,2,3,4,5,6 DESC
)
SELECT * FROM daily_usage_by_budget_policy

Limitations

For Vector Search, budget policies have the following limitations:

• Budget policies do not enforce spending limits for Vector Search.
• Budget policies are applied only to endpoints to track serving costs. They are not applied to vector indexes.