diff --git a/docs/further.md b/docs/further.md index 8910b82..2494234 100644 --- a/docs/further.md +++ b/docs/further.md @@ -1,6 +1,6 @@ # Azure Batch Authentication -The plugin uses [DefaultAzureCredential](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python) to create and destroy Azure Batch resources. The caller must have Contributor permissions on the Azure Batch account for the plugin to work properly. If you are using the Azure Storage plugin you should also have the Storage Blob Data Contributor role for the storage account(s) you use. +The plugin uses a CustomAzureCredential chain that prefers the use of AzureCliCredential, then falls back to a ManagedIdentityCredential, and finally, an EnvironmentCredential (service principal) to create and destroy Azure Batch resources. The caller must have Contributor permissions on the Azure Batch account for the plugin to work properly. If you are using the Azure Storage plugin you should also have the Storage Blob Data Contributor role for the storage account(s) you use. To run a Snakemake workflow using your azure identity you need to ensure you are logged in using the [Azure CLI](https://learn.microsoft.com/en-us/cli/azure/): diff --git a/snakemake_executor_plugin_azure_batch/__init__.py b/snakemake_executor_plugin_azure_batch/__init__.py index 40a4586..ba28868 100644 --- a/snakemake_executor_plugin_azure_batch/__init__.py +++ b/snakemake_executor_plugin_azure_batch/__init__.py @@ -15,7 +15,6 @@ import azure.batch.models as bm from azure.batch import BatchServiceClient from azure.core.exceptions import HttpResponseError -from azure.identity import DefaultAzureCredential from azure.mgmt.batch import BatchManagementClient from snakemake_interface_common.exceptions import WorkflowError from snakemake_interface_executor_plugins.executors.base import SubmittedJobInfo @@ -30,6 +29,7 @@ from snakemake_executor_plugin_azure_batch.constant import AZURE_BATCH_RESOURCE_ENDPOINT from snakemake_executor_plugin_azure_batch.util import ( AzureIdentityCredentialAdapter, + CustomAzureCredential, read_stream_as_string, unpack_compute_node_errors, unpack_task_failure_information, @@ -278,8 +278,7 @@ def __post_init__(self): def init_batch_client(self): """ - Initialize the BatchServiceClient and BatchManagementClient using - DefaultAzureCredential. + Initialize the BatchServiceClient and BatchManagementClient Sets: - self.batch_client @@ -287,12 +286,9 @@ def init_batch_client(self): """ try: - # initialize BatchServiceClient - default_credential = DefaultAzureCredential( - exclude_managed_identity_credential=True - ) adapted_credential = AzureIdentityCredentialAdapter( - credential=default_credential, resource_id=AZURE_BATCH_RESOURCE_ENDPOINT + credential=CustomAzureCredential(), + resource_id=AZURE_BATCH_RESOURCE_ENDPOINT, ) self.batch_client = BatchServiceClient( adapted_credential, self.settings.account_url @@ -300,7 +296,7 @@ def init_batch_client(self): # initialize BatchManagementClient self.batch_mgmt_client = BatchManagementClient( - credential=default_credential, + credential=CustomAzureCredential(), subscription_id=self.settings.subscription_id, ) diff --git a/snakemake_executor_plugin_azure_batch/util.py b/snakemake_executor_plugin_azure_batch/util.py index e735531..479970e 100644 --- a/snakemake_executor_plugin_azure_batch/util.py +++ b/snakemake_executor_plugin_azure_batch/util.py @@ -6,7 +6,21 @@ from azure.core.pipeline import PipelineContext, PipelineRequest from azure.core.pipeline.policies import BearerTokenCredentialPolicy from azure.core.pipeline.transport import HttpRequest -from azure.identity import DefaultAzureCredential +from azure.identity import ( + AzureCliCredential, + ChainedTokenCredential, + EnvironmentCredential, + ManagedIdentityCredential, +) + + +def CustomAzureCredential() -> ChainedTokenCredential: + credential_chain = ( + AzureCliCredential(), + ManagedIdentityCredential(), + EnvironmentCredential(), + ) + return ChainedTokenCredential(*credential_chain) # The usage of this credential helper is required to authenticate batch with managed @@ -26,13 +40,13 @@ def __init__( azure.common.credentials or msrestazure. Args: - credential: Any azure-identity credential (DefaultAzureCredential by + credential: Any azure-identity credential (CustomAzureCredential by default) resource_id (str): The scope to use to get the token (default ARM) """ super(AzureIdentityCredentialAdapter, self).__init__(None) if credential is None: - credential = DefaultAzureCredential() + credential = CustomAzureCredential() self._policy = BearerTokenCredentialPolicy(credential, resource_id, **kwargs) def _make_request(self):