Use authentication credential secrets in Azure Machine Learning jobs

Article
09/27/2024

APPLIES TO: Python SDK azure-ai-ml v2 (current)

Authentication information such as your user name and password are secrets. For example, if you connect to an external database in order to query training data, you would need to pass your username and password to the remote job context. Coding such values into training scripts in clear text is insecure as it would potentially expose the secret.

The Azure Key Vault allows you to securely store and retrieve secrets. In this article, learn how you can retrieve secrets stored in a key vault from a training job running on a compute cluster.

Important

The Azure Machine Learning Python SDK v2 and Azure CLI extension v2 for machine learning do not provide the capability to set or get secrets. Instead, the information in this article uses the Azure Key Vault Secrets client library for Python.

Prerequisites

Before following the steps in this article, make sure you have the following prerequisites:

Tip

Many of the prerequisites in this section require Contributor, Owner, or equivalent access to your Azure subscription, or the Azure Resource Group that contains the resources. You may need to contact your Azure administrator and have them perform these actions.

An Azure subscription. If you don't have an Azure subscription, create a trial subscription before you begin. Try the trial subscription.
An Azure Machine Learning workspace. If you don't have one, use the steps in the Create resources to get started article to create one.
An Azure Key Vault. If you used the Create resources to get started article to create your workspace, a key vault was created for you. You can also create a separate key vault instance using the information in the Quickstart: Create a key vault article.

Tip

You do not have to use same key vault as the workspace.
(Optional) An Azure Machine Learning compute cluster configured to use a managed identity. The cluster can be configured for either a system-assigned or user-assigned managed identity.
If your job runs on a compute cluster, grant the managed identity for the compute cluster access to the secrets stored in key vault. Or, if the job runs on serverless compute, grant the managed identity specified for the job access to the secrets. The method used to grant access depends on how your key vault is configured:
- Azure role-based access control (Azure RBAC): When configured for Azure RBAC, add the managed identity to the Key Vault Secrets User role on your key vault.
- Azure Key Vault access policy: When configured to use access policies, add a new policy that grants the get operation for secrets and assign it to the managed identity.
A stored secret value in the key vault. This value can then be retrieved using a key. For more information, see Quickstart: Set and retrieve a secret from Azure Key Vault.

Tip

The quickstart link is to the steps for using the Azure Key Vault Python SDK. In the table of contents in the left navigation area are links to other ways to set a key.

Get secrets

There are two ways to get secrets during training:

Using a managed identity associated with the compute resource the training job runs on.
Using your identity by having the compute run the job on your behalf.

Managed identity
Your identity

Add the azure-keyvault-secrets and azure-identity packages to the Azure Machine Learning environment used when training the model. For example, by adding them to the conda file used to build the environment.

The environment is used to build the Docker image that the training job runs in on the compute cluster.

From your training code, use the Azure Identity SDK and Key Vault client library to get the managed identity credentials and authenticate to key vault:

from azure.identity import DefaultAzureCredential
from azure.keyvault.secrets import SecretClient

credential = DefaultAzureCredential()

secret_client = SecretClient(vault_url="https://my-key-vault.vault.azure.cn/", credential=credential)

After authenticating, use the Key Vault client library to retrieve a secret by providing the associated key:
```
secret = secret_client.get_secret("secret-name")
print(secret.value)
```

Add the azure-keyvault-secrets, azure-identity, and azure-ai-ml packages to the Azure Machine Learning environment used when training the model. For example, by adding them to the conda file used to build the environment.

The environment is used to build the Docker image that the training job runs in on the compute cluster.
From your training code, use the Azure Machine Learning SDK and Key Vault client library to get the managed identity credentials and authenticate to key vault. The AzureMLOnBehalfOfCredential class is used to authenticate on behalf of your user identity:
```
from azure.ai.ml.identity import AzureMLOnBehalfOfCredential
from azure.keyvault.secrets import SecretClient

credential = AzureMLOnBehalfOfCredential()
secret_client = SecretClient(vault_url="https://my-key-vault.vault.azure.cn/", credential=credential)
```
After authenticating, use the Key Vault client library to retrieve a secret by providing the associated key:
```
secret = secret_client.get_secret("secret-name")
print(secret.value)
```

When you submit the training job, you must specify that it runs on behalf of your identity by using identity=UserIdentityConfiguration(). The following example submits a job using this parameter:

from azure.ai.ml import Input, command
from azure.ai.ml.constants import AssetTypes
from azure.ai.ml.entities import UserIdentityConfiguration

job = command(
    code="./sdk/ml/azure-ai-ml/samples/src",
    command="python read_data.py --input_data ${{inputs.input_data}}",
    inputs={"input_data": Input(type=AssetTypes.MLTABLE, path="./sample_data")},
    environment="AzureML-sklearn-1.0-ubuntu20.04-py38-cpu:1",
    compute="cpu-cluster",
    identity=UserIdentityConfiguration(),
)

For an example of using the Azure CLI to submit a job that uses your identity, visit Https://github.com/Azure/azureml-examples/blob/d4c90eead3c1fd97393d0657f7a78831490adf1c/cli/jobs/single-step/on-behalf-of/README.md.

For an example of submitting a training job using the Azure Machine Learning Python SDK v2, see Train models with the Python SDK v2.

Use authentication credential secrets in Azure Machine Learning jobs

Prerequisites

Get secrets

Related content

Additional resources