Query Databricks Model Serving Endpoint with AAD token fails

Question

I am trying to query databricks model serving endpoint with AAD token.

I first generate the AAD token as follows:

import requests 
import json

url = f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
payload = {
  "client_id": client_id,
  "grant_type": "client_credentials",
  "scope": "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d/.default",
  "client_secret": client_secret
}

response = requests.post(url, data=payload)
token = json.loads(response.text)["access_token"]

For example following command to list endpoints works perfectly fine:

url = "https://<databricks-instance>/api/2.0/serving-endpoints"
headers = {'Authorization': f'Bearer {token}', 'Content-Type': 'application/json'}

response = requests.request(method='GET', headers=headers, url=url)

The problem arises when I try to query the endpoint with some data as follows:

url = "https://<databricks-instance>/serving-endpoints/test-model/invocations"
headers = {'Authorization': f'Bearer {token}', 'Content-Type': 'application/json'}


payload = ... # Create model input data
data_json = json.dumps(payload)
response = requests.request(method='POST', headers=headers, url=url, data=data_json)
if response.status_code != 200:
  raise Exception(f'Request failed with status {response.status_code}, {response.text}')
else:
  response.json()

This returns an error message:

Request failed with status 400

Problem accessing /serving-endpoints/test-model/invocations. Reason:

    io.jsonwebtoken.IncorrectClaimException: Expected aud claim to be: , but was: 2ff814a6-3304-4ab8-85cb-cd0e6f879c1d.

Can someone point me in the right direction of how to resolve this problem?

FYI: the same query works with personal access token

Answer 1

Model serving doesn't support AAD authentication yet, only personal access tokens (PATs). But it should be supported in the future.