[Feature] Integrate Databricks SDK with Model Serving Auth Provider

[aravind-segu]

Changes

This PR introduces a new model serving auth method to Databricks SDK.

• If the correct environment variables are set to identify a model serving environment
• Check to see if there is an oauth file written by the serving environment
• If this file exists use the token here for authentication

Tests

Added Unit tests

• make test run locally
• make fmt applied
• relevant integration tests applied

[smurching]

A few comments, thanks @aravind-segu !

[mgyucht]

Mostly looking good, but we shouldn't throw exceptions from the credential strategy, otherwise DefaultCredentials will stop trying later auth mechanisms. This will be problematic if new strategies are added later. Otherwise, a few nits.

Were you able to test this out with a manual test?

[mgyucht]

Just return None here if you're unable to authenticate. DefaultCredentials doesn't catch these exceptions, so this will interrupt the fallback flow.

[aravind-segu]

I am raising an error here, but then catching it immediately and returning None after. I wanted to add an extra message to differentiate whether reading the token failed or whether the token read was itself none.

[mgyucht]

This method is called with every request. Do we need to check the file every time? We incur the cost of opening a file & reading its contents on every request. Is there an expiry included in the token file? We could use that to determine whether we need to refresh the token.

I tested a simple benchmark:

import json
import timeit
import os

with open('sample.json', 'w') as f:
    json.dump([{'key': 'value'} for _ in range(100)], f)

def benchmark():
    with open('sample.json', 'r') as f:
        data = json.load(f)

print(f"Average time to open, read, and parse file: {timeit.timeit(benchmark, number=1000):.6f} seconds")

This came out to 0.07 seconds on my laptop. It might be nice to avoid this overhead, especially for latency-sensitive use-cases.

[aravind-segu]

Good call, added a 5 min cache

[aravind-segu]

Mostly looking good, but we shouldn't throw exceptions from the credential strategy, otherwise DefaultCredentials will stop trying later auth mechanisms. This will be problematic if new strategies are added later. Otherwise, a few nits.

Were you able to test this out with a manual test?

Manual Test Endpoint Here: https://e2-dogfood.staging.cloud.databricks.com/ml/endpoints/agents_main-default-test_agent_09_16?o=6051921418418893

[aravind-segu]

@mgyucht We discussed in the call about making this a oauth-credential-provider. However that requires a refresht_token we do not have. We can set that to None but what happens then when the expiry is hit.

In terms of direct ingress, I do not think we are going to be hitting Model Serving Direct Ingress Endpoints right now. The only thing that we could do is talk to a Foundational Model Serving Endpoint through the workspace client. Currently we already do that in our Mlflow code by sending a request directly.

WIth this change we will now start to use databricks sdk workspace client to make this call. We will use the code here. As it already worked before, I am assuming that it will work again as this directly calls the api client. Do we still need to integrate oauth-credential-provider now.

[smurching]

We shouldn't raise here right?

[smurching]

Oh I see we catch it further down

[smurching]

Last comments, then LGTM!

[smurching]

LGTM!

[mgyucht]

Close, just a few small comments.

[aravind-segu]

jenkins merge