Use federated identities with SDK

Use identity federation with JSON Web Tokens (JWTs) to authenticate with the W&B SDK and CLI without API keys.

Use identity federation to sign in to the W&B SDK and CLI with your organizational credentials, instead of using a long-lived API key. If your W&B organization admin has configured SSO for your organization, you already use those credentials to sign in to the W&B app UI. Identity federation is like SSO for the W&B SDK, but uses JSON Web Tokens (JWTs) directly. Use identity federation as an alternative to API keys.

This page is for organization admins who configure the JWT issuer for a W&B organization. It’s also for users or service accounts that authenticate to W&B using JWTs.

RFC 7523 forms the underlying basis for identity federation with SDK.

Identity federation is available in preview for Multi-tenant Cloud, Dedicated Cloud, and Self-Managed. An Enterprise license is required. For details or assistance, contact your AISE or support.

This document uses the terms “identity provider” and “JWT issuer” interchangeably. Both refer to the same thing in the context of this capability.

Set up the JWT issuer

Before users can authenticate with JWTs, an organization admin must set up a federation between your W&B organization and a publicly accessible JWT issuer.

1. Navigate to the Settings tab in your organization dashboard.
2. In the Authentication option, click Set up JWT Issuer.
3. Add the JWT issuer URL in the text box and click Create.

W&B automatically looks for an OIDC discovery document at the path ${ISSUER_URL}/.well-known/openid-configuration. From the discovery document, W&B locates the JSON Web Key Set (JWKS) at the relevant URL. W&B uses the JWKS for real-time validation of the JWTs to ensure that the relevant identity provider issued them.

After this step, your W&B organization is federated with the JWT issuer. Users in your organization can then authenticate to W&B using JWTs issued by that provider.

Use the JWT to access W&B

After an organization admin sets up a JWT issuer, users can start accessing W&B projects using JWTs issued by that identity provider. The mechanism for using JWTs is as follows:

1. You must sign in to the identity provider using one of the mechanisms available in your organization. You can access some providers in an automated manner using an API or SDK, while others are only accessible through a relevant UI. Contact your W&B organization admin or the owner of the JWT issuer for details.
2. After you’ve retrieved the JWT by signing in to your identity provider, store it in a file at a secure location. Configure the absolute file path in an environment variable WANDB_IDENTITY_TOKEN_FILE.
3. Access your W&B project using the W&B SDK or CLI. The SDK or CLI automatically detects the JWT and exchanges it for a W&B access token after validating the JWT. The W&B access token grants access to the relevant APIs for enabling your AI workflows, such as logging runs, metrics, and artifacts. By default, the access token is stored at the path ~/.config/wandb/credentials.json. You can change that path by specifying the environment variable WANDB_CREDENTIALS_FILE.

JWTs are short-lived credentials that address the shortcomings of long-lived credentials such as API keys and passwords. The JWT expiry time depends on your identity provider’s configuration. Refresh the JWT before it expires, and ensure that it’s stored in the file referenced by the environment variable WANDB_IDENTITY_TOKEN_FILE.

The W&B access token also has a default expiry duration, after which the SDK or CLI tries to refresh it using your JWT. If the user JWT has also expired by that time and isn’t refreshed, authentication fails. If possible, implement the JWT retrieval and post-expiry refresh mechanism as part of the AI workload that uses the W&B SDK or CLI.

JWT validation

To ensure that only valid tokens grant access, the JWT undergoes the following validations. These validations run when the SDK or CLI exchanges the JWT for a W&B access token and then accesses a project:

• W&B verifies the JWT signature using the JWKS at the W&B organization level. This is the first line of defense, and if this fails, that means there’s a problem with your JWKS or how your JWT is signed.

• The iss claim in the JWT must equal the issuer URL configured at the organization level.

• The sub claim in the JWT must equal the user’s email address as configured in the W&B organization.

• The aud claim in the JWT must equal the name of the W&B organization that houses the project you’re accessing as part of your AI workflow.

  On Dedicated Cloud or Self-Managed instances:

  • To skip audience validation, you can set the environment variable FEDERATED_AUTH_AUDIENCES to wandb.
  • Some organizations have specific requirements for the audience. To customize the aud value, set the environment variable FEDERATED_AUTH_AUDIENCES to a string with a comma-separated list of audience values.

• W&B checks the exp claim in the JWT to determine whether the token is valid or has expired and needs to be refreshed.

External service accounts

W&B has long supported built-in service accounts with long-lived API keys. With identity federation for the SDK and CLI, you can also bring external service accounts that use JWTs for authentication. The organization’s configured issuer must issue those JWTs. A team admin can configure external service accounts within the scope of a team, like the built-in service accounts.

To configure an external service account, a team admin must:

1. Navigate to the Service Accounts tab for your team.
2. Click New service account.
3. Provide a name for the service account, select Federated Identity as the Authentication Method, provide a Subject, and click Create.

After this step, the external service account is registered with the team and can use JWTs issued by the configured identity provider to access W&B. The sub claim in the external service account’s JWT must equal the subject configured by the team admin in the team-level Service Accounts tab. W&B verifies that claim as part of JWT validation. The aud claim requirement is similar to that for human user JWTs.

When using an external service account’s JWT to access W&B, it’s often easier to automate the workflow. Automation generates the initial JWT and refreshes it as needed. To attribute runs logged using an external service account to a human user, configure the environment variables WANDB_USERNAME or WANDB_USER_EMAIL for your AI workflow, similar to built-in service accounts.

W&B recommends using a mix of built-in and external service accounts across your AI workloads with different levels of data sensitivity. This mix strikes a balance between flexibility and simplicity.