Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Introduce new Credential Strategies for Agents #882

Merged

Conversation

aravind-segu
Copy link
Contributor

@aravind-segu aravind-segu commented Jan 31, 2025

What changes are proposed in this pull request?

This PR introduces two new credential strategies for Agents, (AgentEmbeddedCredentials, AgentUserCredentials).

Agents currently use the databricks.sdk in order to interact with databricks resources. However the authentication method for these resources is a little unique where we store the token for the authentication in a Credential File on the Kubernetes Container. Therefore in the past we added the Model Serving Credential Strategy to the defaultCredentials list to read this file.

Now we want to introduce a new authentication where the user's token is instead stored in a thread local variable. Agent users will initialize clients as follows:

from databricks.sdk.credentials_provider import ModelServingUserCredentials

invokers_client = WorkspaceClient(credential_strategy = ModelServingUserCredentials())
definers_client = WorkspaceClient()

Then the users can use the invoker_client to interact with resources with the invokers token or the definers_client to interact with resources using the old method of authentication.

Additionally as the users will be using these clients to test their code locally in Databricks Notebooks, if the code is not being run on model serving environments, users need to be able to authenticate using the DefaultCredential strategies.

More details: https://docs.google.com/document/d/14qLVjyxIAk581w287TWElstIeh8-DR30ab9Z6B_Vydg/edit?usp=sharing

How is this tested?

Added unit tests

@aravind-segu aravind-segu changed the title Introduce new Credential Strategies for Agents [Feature] Introduce new Credential Strategies for Agents Feb 6, 2025
@renaudhartert-db renaudhartert-db self-requested a review February 7, 2025 07:29
@renaudhartert-db
Copy link
Contributor

I've started the integration test for this change, result should appear soon.

Meanwhile, could you please run make fmt and commit the changes?

Copy link
Contributor

@renaudhartert-db renaudhartert-db left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM assuming resolution of my last comment. Thanks @aravind-segu!

@aravind-segu aravind-segu force-pushed the model_serving_invokers_token branch from 5fb5b16 to 037d295 Compare February 13, 2025 00:38
@aravind-segu aravind-segu force-pushed the model_serving_invokers_token branch from 037d295 to 89739b7 Compare February 13, 2025 00:40
Signed-off-by: aravind-segu <aravind.segu@databricks.com>
Signed-off-by: aravind-segu <aravind.segu@databricks.com>
Signed-off-by: aravind-segu <aravind.segu@databricks.com>
@aravind-segu aravind-segu force-pushed the model_serving_invokers_token branch from 89739b7 to fd0e4db Compare February 13, 2025 00:41
Copy link

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-py

Inputs:

  • PR number: 882
  • Commit SHA: fd0e4dbcac8616a975f50e86dcd276296e20f12a

Checks will be approved automatically on success.

@renaudhartert-db renaudhartert-db added this pull request to the merge queue Feb 13, 2025
Merged via the queue into databricks:main with commit 41f5f4b Feb 13, 2025
17 of 18 checks passed
hectorcast-db added a commit that referenced this pull request Feb 13, 2025
### New Features and Improvements

 * Introduce new Credential Strategies for Agents ([#882](#882)).

### Internal Changes

 * GetRun logic paginates more arrays ([#867](#867)).
hectorcast-db added a commit that referenced this pull request Feb 13, 2025
 * Introduce new Credential Strategies for Agents ([#882](#882)).

 * GetRun logic paginates more arrays ([#867](#867)).
github-merge-queue bot pushed a commit that referenced this pull request Feb 13, 2025
### New Features and Improvements

* Introduce new Credential Strategies for Agents
([#882](#882)).

### Bug Fixes

* Fix public documentation

### Internal Changes

* GetRun logic paginates more arrays
([#867](#867)).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants