Add managed identity support for Azure Batch container registries

docs/azure.mdx

@@ -550,6 +550,23 @@ Nextflow uses the following environment variables if the registry credentials ar
 - `AZURE_REGISTRY_PASSWORD`: the password for Azure Container Registry authentication
 :::
 
+Alternatively, you can authenticate to a private registry with a user-assigned [managed identity](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/overview) instead of a username and password:
+
+```groovy
+azure {
+    registry {
+        server = '<registry-server>'     // e.g., 'myregistry.azurecr.io'
+        managedIdentityResourceId = '<managed-identity-resource-id>' // e.g., '/subscriptions/<sub>/resourceGroups/<rg>/providers/Microsoft.ManagedIdentity/userAssignedIdentities/<name>'
+    }
+}
+```
+
+The value is the full ARM resource ID of the managed identity, not its client ID. When set, it takes precedence over `userName` and `password`, and can also be provided via the `AZURE_REGISTRY_MANAGED_RESOURCE_ID` environment variable.
+
+:::note
+The managed identity must already be attached to the pool nodes. Nextflow cannot attach an identity when auto-creating a pool, so this option requires a pre-existing pool (or an identity attached out-of-band).
+:::
+
 Container images from public registries such as Docker Hub can be used without additional configuration, even when a private registry is specified in the `azure.registry` scope. The Docker runtime on Azure Batch VMs automatically uses the appropriate registry based on the container image specified for the task.
 
 ### VM images

docs/reference/config.mdx

@@ -605,6 +605,14 @@ The client ID for an Azure [managed identity](https://learn.microsoft.com/en-us/
 
 When `true`, use the system-assigned [managed identity](https://learn.microsoft.com/en-us/entra/identity/managed-identities-azure-resources/overview) to authenticate Azure resources. Defaults to environment variable `AZURE_MANAGED_IDENTITY_SYSTEM`.
 
+##### `azure.registry.managedIdentityResourceId`
+
+The ARM resource ID of a user-assigned managed identity used to authenticate to a private container registry, e.g. `/subscriptions/<sub>/resourceGroups/<rg>/providers/Microsoft.ManagedIdentity/userAssignedIdentities/<name>`. When set, it takes precedence over `azure.registry.userName` and `azure.registry.password`. Defaults to environment variable `AZURE_REGISTRY_MANAGED_RESOURCE_ID`.
+
+:::note
+The managed identity must already be attached to the pool nodes. Nextflow cannot attach an identity when auto-creating a pool, so this option requires a pre-existing pool (or an identity attached out-of-band).
+:::
+
 ##### `azure.registry.password`
 
 The password to connect to a private container registry.

plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy

@@ -36,6 +36,7 @@ import com.azure.compute.batch.models.BatchJobCreateContent
 import com.azure.compute.batch.models.BatchJobConstraints
 import com.azure.compute.batch.models.BatchJobUpdateContent
 import com.azure.compute.batch.models.BatchNodeFillType
+import com.azure.compute.batch.models.BatchNodeIdentityReference
 import com.azure.compute.batch.models.BatchPool
 import com.azure.compute.batch.models.BatchPoolCreateContent
 import com.azure.compute.batch.models.BatchPoolInfo
@@ -874,13 +875,20 @@ class AzBatchService implements Closeable {
         final registryOpts = config.registry()
 
         if( registryOpts && registryOpts.isConfigured() ) {
-            final containerRegistries = new ArrayList<ContainerRegistryReference>(1)
-            containerRegistries << new ContainerRegistryReference()
+            final registryRef = new ContainerRegistryReference()
                     .setRegistryServer(registryOpts.server)
-                    .setUsername(registryOpts.userName)
-                    .setPassword(registryOpts.password)
+            if( registryOpts.usesManagedIdentity() ) {
+                registryRef.setIdentityReference( new BatchNodeIdentityReference().setResourceId(registryOpts.managedIdentityResourceId) )
+                log.debug "[AZURE BATCH] Connecting Azure Batch pool to Container Registry '$registryOpts.server' using managed identity '$registryOpts.managedIdentityResourceId'"
+            }
+            else {
+                registryRef.setUsername(registryOpts.userName)
+                        .setPassword(registryOpts.password)
+                log.debug "[AZURE BATCH] Connecting Azure Batch pool to Container Registry '$registryOpts.server'"
+            }
+            final containerRegistries = new ArrayList<ContainerRegistryReference>(1)
+            containerRegistries << registryRef
             containerConfig.setContainerRegistries(containerRegistries)
-            log.debug "[AZURE BATCH] Connecting Azure Batch pool to Container Registry '$registryOpts.server'"
         }
 
         final image = getImage(opts)

plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzRegistryOpts.groovy

@@ -48,6 +48,12 @@ class AzRegistryOpts implements ConfigScope {
     """)
     final String password
 
+    @ConfigOption
+    @Description("""
+        The ARM resource ID of a user-assigned managed identity used to authenticate to a private container registry (e.g. `/subscriptions/<sub>/resourceGroups/<rg>/providers/Microsoft.ManagedIdentity/userAssignedIdentities/<name>`). When set, it takes precedence over `userName` and `password`. Defaults to environment variable `AZURE_REGISTRY_MANAGED_RESOURCE_ID`. Note: the managed identity must already be attached to the pool nodes.
+    """)
+    final String managedIdentityResourceId
+
     AzRegistryOpts() {
         this(Collections.emptyMap())
     }
@@ -57,14 +63,21 @@ class AzRegistryOpts implements ConfigScope {
         this.server = config.server ?: 'docker.io'
         this.userName = config.userName ?: env.get('AZURE_REGISTRY_USER_NAME')
         this.password = config.password ?: env.get('AZURE_REGISTRY_PASSWORD')
+        this.managedIdentityResourceId = config.managedIdentityResourceId ?: env.get('AZURE_REGISTRY_MANAGED_RESOURCE_ID')
+    }
+
+    boolean usesManagedIdentity() {
+        return managedIdentityResourceId as boolean
     }
 
     boolean isConfigured() {
+        if( managedIdentityResourceId )
+            return true
         if( userName && password )
             return true
         if( !userName && !password )
             return false
-        throw new IllegalArgumentException("Invalid Container Registry configuration - Make sure userName and password are set for Container Registry")
+        throw new IllegalArgumentException("Invalid Container Registry configuration - Make sure managedIdentityResourceId or userName and password are set for Container Registry")
     }
 
 }

plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchServiceTest.groovy

@@ -26,6 +26,7 @@ import com.azure.compute.batch.models.AllocationState
 import com.azure.compute.batch.models.BatchPool
 import com.azure.compute.batch.models.BatchPoolState
 import com.azure.compute.batch.models.BatchJobCreateContent
+import com.azure.compute.batch.models.BatchSupportedImage
 import com.azure.compute.batch.models.ElevationLevel
 import com.azure.compute.batch.models.EnvironmentSetting
 import com.azure.compute.batch.models.ResizeError
@@ -466,6 +467,43 @@ class AzBatchServiceTest extends Specification {
 
     }
 
+    def 'should configure container registry with username and password' () {
+        given:
+        def exec = createExecutor(new AzConfig([registry: [server: 'reg.azurecr.io', userName: 'foo', password: 'bar']]))
+        AzBatchService svc = Spy(AzBatchService, constructorArgs: [exec])
+
+        when:
+        def vmConfig = svc.poolVmConfig(new AzPoolOpts())
+        then:
+        1 * svc.getImage(_) >> BatchSupportedImage.fromJson(com.azure.json.JsonProviders.createReader('{"nodeAgentSKUId":"sku","imageReference":{}}'))
+        and:
+        def registries = vmConfig.containerConfiguration.containerRegistries
+        registries.size() == 1
+        registries[0].registryServer == 'reg.azurecr.io'
+        registries[0].username == 'foo'
+        registries[0].password == 'bar'
+        registries[0].identityReference == null
+    }
+
+    def 'should configure container registry with managed identity' () {
+        given:
+        def RESOURCE_ID = '/subscriptions/sub/resourceGroups/rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/name'
+        def exec = createExecutor(new AzConfig([registry: [server: 'reg.azurecr.io', userName: 'foo', password: 'bar', managedIdentityResourceId: RESOURCE_ID]]))
+        AzBatchService svc = Spy(AzBatchService, constructorArgs: [exec])
+
+        when:
+        def vmConfig = svc.poolVmConfig(new AzPoolOpts())
+        then:
+        1 * svc.getImage(_) >> BatchSupportedImage.fromJson(com.azure.json.JsonProviders.createReader('{"nodeAgentSKUId":"sku","imageReference":{}}'))
+        and: 'managed identity takes precedence over username/password'
+        def registries = vmConfig.containerConfiguration.containerRegistries
+        registries.size() == 1
+        registries[0].registryServer == 'reg.azurecr.io'
+        registries[0].identityReference.resourceId == RESOURCE_ID
+        registries[0].username == null
+        registries[0].password == null
+    }
+
     def 'should check poolid' () {
         given:
         def exec = createExecutor()

plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzRegistryOptsTest.groovy

@@ -52,4 +52,38 @@ class AzRegistryOptsTest extends Specification {
         opts3.password == 'env-password'
     }
 
+    def 'should get managed identity resource id'() {
+        expect: 'config value is used'
+        new AzRegistryOpts([managedIdentityResourceId: 'res-id'], [:]).managedIdentityResourceId == 'res-id'
+
+        and: 'env fallback is used when config absent'
+        new AzRegistryOpts([:], [AZURE_REGISTRY_MANAGED_RESOURCE_ID: 'env-res-id']).managedIdentityResourceId == 'env-res-id'
+
+        and: 'config value wins over env'
+        new AzRegistryOpts([managedIdentityResourceId: 'res-id'], [AZURE_REGISTRY_MANAGED_RESOURCE_ID: 'env-res-id']).managedIdentityResourceId == 'res-id'
+
+        and: 'null when neither set'
+        new AzRegistryOpts([:], [:]).managedIdentityResourceId == null
+    }
+
+    def 'should validate isConfigured'() {
+        expect:
+        new AzRegistryOpts([:], [:]).isConfigured() == false
+        new AzRegistryOpts([userName: 'foo', password: 'bar'], [:]).isConfigured() == true
+        new AzRegistryOpts([managedIdentityResourceId: 'res-id'], [:]).isConfigured() == true
+        new AzRegistryOpts([managedIdentityResourceId: 'res-id'], [:]).usesManagedIdentity() == true
+        new AzRegistryOpts([userName: 'foo', password: 'bar'], [:]).usesManagedIdentity() == false
+
+        when: 'managed identity takes precedence over incomplete user/password'
+        def opts = new AzRegistryOpts([managedIdentityResourceId: 'res-id', userName: 'foo'], [:])
+        then:
+        opts.isConfigured() == true
+        opts.usesManagedIdentity() == true
+
+        when: 'partial user/password throws'
+        new AzRegistryOpts([userName: 'foo'], [:]).isConfigured()
+        then:
+        thrown(IllegalArgumentException)
+    }
+
 }