Abilitare la federazione delle identità del carico di lavoro per i carichi di lavoro IAM AWS

I carichi di lavoro AWS possono eseguire l'autenticazione per Azure Databricks senza segreti a lungo termine scambiando un token di identità OIDC firmato da AWS. Il percorso consigliato invoca l'API AWS STS GetWebIdentityToken, che funziona ovunque il carico di lavoro disponga di credenziali AWS.

Casi d'uso

  • Funzioni lambda che chiamano api Azure Databricks (attivazione di processi, esecuzione di query su SQL Warehouse)
  • Pipeline ETL basate su EC2/ECS che eseguono l'autenticazione per Azure Databricks senza segreti
  • Carichi di lavoro di ML basati su EKS che accedono agli endpoint di serving dei modelli
  • Modelli tra account in cui i carichi di lavoro in un account AWS vengono federati in un account Azure Databricks gestito da un altro team
  • Miglioramento della postura di sicurezza eliminando Azure Databricks pat o segreti di lunga durata da AWS Secrets Manager

Prerequisiti di AWS

Per abilitare la federazione delle identità del carico di lavoro per i carichi di lavoro IAM AWS, è necessario completare i passaggi seguenti.

Passaggio 1: Abilitare la federazione delle identità in uscita di AWS IAM

Abilitare la federazione delle identità in uscita nell'account AWS:

import boto3
boto3.client('iam').enable_outbound_web_identity_federation()

È anche possibile abilitare questa opzione nella console IAM in Impostazioni account>Abilita "federazione delle identità web in uscita".

Passaggio 2: Concedere sts:GetWebIdentityToken l'autorizzazione

Concedi l'autorizzazione al ruolo IAM sts:GetWebIdentityToken del carico di lavoro:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": ["sts:GetWebIdentityToken"],
      "Resource": "*",
      "Condition": {
        "ForAllValues:StringEquals": {
          "sts:IdentityTokenAudience": "databricks"
        },
        "NumericLessThanEquals": {
          "sts:DurationSeconds": 300
        }
      }
    }
  ]
}

Annotazioni

La condizione del gruppo di destinatari garantisce che il ruolo possa richiedere solo token destinati a Azure Databricks. La condizione di durata limita la durata del token a 300 secondi. È possibile modificare la durata fino a 3600 secondi in base alle esigenze del carico di lavoro, ma sono consigliate durate più brevi.

Passaggio 3: Annotare l'URL dell'emittente specifico dell'account

Recupera l'URL dell'emittente specifico dell'account:

import boto3
info = boto3.client('iam').get_outbound_web_identity_federation_info()
print(info['IssuerUrl'])  # https://<uuid>.tokens.sts.global.api.aws

Creare una politica di federazione

Importante

Azure Databricks i criteri di federazione vengono creati a livello di account (non a livello di area di lavoro). L'host dell'interfaccia della riga di comando Azure Databricks deve essere impostato su https://accounts.azuredatabricks.net e l'utente deve essere un amministratore dell'account.

Crea un criterio di federazione delle identità dei carichi di lavoro usando Databricks CLI. Impostare l'autorità emittente sull'URL dell'autorità emittente specifica dell'account dal passaggio 3. Per istruzioni dettagliate, vedere Configurare un criterio di federazione di un'entità servizio.

databricks account service-principal-federation-policy create ${SP_ID} --json '{
  "oidc_policy": {
    "issuer": "https://<uuid>.tokens.sts.global.api.aws",
    "audiences": ["databricks"],
    "subject": "arn:aws:iam::<account-id>:role/<workload-role-name>"
  }
}'

Eseguire l'autenticazione in Databricks

Dopo aver creato i criteri di federazione, usare Databricks SDK per autenticare i carichi di lavoro AWS. L'esempio seguente usa il pattern IdTokenSource dell'SDK per recuperare un token AWS STS e scambiarlo per ottenere un token OAuth di Azure Databricks.

import boto3
from databricks.sdk import WorkspaceClient
from databricks.sdk import oidc
from databricks.sdk.core import Config, credentials_strategy, oidc_credentials_provider


class AwsStsTokenSource(oidc.IdTokenSource):
    def __init__(self, audience="databricks", region="us-east-1"):
        self._audience = audience
        self._region = region

    def id_token(self) -> oidc.IdToken:
        sts = boto3.client("sts", region_name=self._region)
        resp = sts.get_web_identity_token(
            Audience=[self._audience],
            SigningAlgorithm="RS256",
            DurationSeconds=300,
        )
        return oidc.IdToken(jwt=resp["WebIdentityToken"])


@credentials_strategy("aws-sts-wif", [])
def aws_sts_wif_strategy(cfg: Config):
    return oidc_credentials_provider(cfg, AwsStsTokenSource())


w = WorkspaceClient(
    host="https://my-workspace.cloud.databricks.com",
    client_id="<service-principal-uuid>",
    credentials_strategy=aws_sts_wif_strategy
)
# No secrets needed
clusters = w.clusters.list()

Annotazioni

È consigliabile usare la durata del token di 300 secondi. È possibile modificare fino a 3600 secondi in base alle esigenze del carico di lavoro.

Per un esempio di scambio di token manuale, vedere Eseguire l'autenticazione con un token del provider di identità.

Informazioni di riferimento sulla documentazione di AWS