使用AAD令牌调用Databricks Model Serving Endpoint失败求解决方案

问题场景 #

我尝试使用AAD令牌查询Databricks Model Serving Endpoint，操作及报错情况如下：

生成AAD令牌的代码 #

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"]

可正常运行的请求（列出端点） #

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)

报错的请求（调用模型端点） #

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

payload = ... # 构造模型输入数据
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()

返回的错误信息 #

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.

补充说明：使用个人访问令牌执行相同查询可正常工作。

解决方向 #

问题根源在于模型服务的invocations接口对JWT令牌的aud（受众）声明要求与Databricks REST API不同：

• 当前生成的令牌aud声明是全局Databricks服务ID（2ff814a6-3304-4ab8-85cb-cd0e6f879c1d），符合REST API验证逻辑，但不符合模型服务端点的要求。
• 模型服务端点期望aud声明指向你的Databricks实例URL。

修正方法 #

修改AAD令牌生成时的scope参数，将全局服务ID替换为你的Databricks实例URL：

import requests 
import json

# 替换<databricks-instance>为你的实际实例域名
databricks_instance = "<databricks-instance>"
url = f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
payload = {
  "client_id": client_id,
  "grant_type": "client_credentials",
  "scope": f"{databricks_instance}/.default",
  "client_secret": client_secret
}

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

若使用旧版AAD令牌端点（v1），也可通过resource参数指定实例URL：

payload = {
  "client_id": client_id,
  "grant_type": "client_credentials",
  "resource": databricks_instance,
  "client_secret": client_secret
}

修正后生成的令牌aud声明会匹配你的Databricks实例，即可正常调用invocations接口。

内容的提问来源于stack exchange，提问作者kanimbla