From 898c411ca79b8c90705efe7a5441cbfb142f2823 Mon Sep 17 00:00:00 2001 From: adamrtalbot <12817534+adamrtalbot@users.noreply.github.com> Date: Thu, 10 Jul 2025 18:36:37 +0100 Subject: [PATCH 1/2] Fix Azure Low Priority VM deprecation for Batch Managed accounts (#6258) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Azure Low Priority VMs are deprecated in Batch Managed pool allocation mode but still supported in User Subscription mode. This adds validation to: - Block low priority VMs for Batch Managed accounts with clear error message - Allow low priority VMs for User Subscription accounts - Gracefully handle unknown allocation modes with warnings Adds subscriptionId to config to enable retrieval of this value. If not supplied Nextflow will raise a warning and continue. This does open us up to being able to do more with Azure Batch because we have access to the management API. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude Signed-off-by: adamrtalbot <12817534+adamrtalbot@users.noreply.github.com> --- plugins/nf-azure/build.gradle | 3 + .../cloud/azure/batch/AzBatchExecutor.groovy | 33 +++++ .../cloud/azure/batch/AzBatchService.groovy | 96 +++++++++++++ .../cloud/azure/config/AzBatchOpts.groovy | 2 + .../azure/batch/AzBatchExecutorTest.groovy | 134 ++++++++++++++++++ .../cloud/azure/config/AzBatchOptsTest.groovy | 22 +++ 6 files changed, 290 insertions(+) create mode 100644 plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy diff --git a/plugins/nf-azure/build.gradle b/plugins/nf-azure/build.gradle index 79b615cc88..313c44a144 100644 --- a/plugins/nf-azure/build.gradle +++ b/plugins/nf-azure/build.gradle @@ -46,6 +46,9 @@ dependencies { api('com.azure:azure-identity:1.15.1') { exclude group: 'org.slf4j', module: 'slf4j-api' } + api('com.azure.resourcemanager:azure-resourcemanager-batch:1.1.0-beta.4') { + exclude group: 'org.slf4j', module: 'slf4j-api' + } // address security vulnerabilities runtimeOnly 'io.netty:netty-handler:4.1.118.Final' diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy index af4ecd9c7e..6714fa5cf4 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy @@ -86,6 +86,38 @@ class AzBatchExecutor extends Executor implements ExtensionPoint { } } + protected void validateLowPriorityVMs() { + // Check if any pool has lowPriority enabled + def lowPriorityPools = config.batch().pools.findAll { poolName, poolOpts -> + poolOpts.lowPriority + } + + if( lowPriorityPools ) { + def poolNames = lowPriorityPools.keySet().join(', ') + + // Get the pool allocation mode to determine if low priority VMs are allowed + def poolAllocationMode = batchService.getPoolAllocationMode() + log.debug "[AZURE BATCH] Pool allocation mode determined as: ${poolAllocationMode}" + + if( poolAllocationMode == 'BATCH_SERVICE' || poolAllocationMode == 'BatchService' ) { + throw new AbortOperationException("Azure Low Priority VMs are deprecated and no longer supported for Batch Managed pool allocation mode. " + + "Please update your configuration to use standard VMs instead, or migrate to User Subscription pool allocation mode. " + + "Affected pools: ${poolNames}. " + + "Remove 'lowPriority: true' from your pool configuration or set 'lowPriority: false'.") + } else if( poolAllocationMode == 'USER_SUBSCRIPTION' || poolAllocationMode == 'UserSubscription' ) { + // Low Priority VMs are still supported in User Subscription mode, proceed without warning + log.debug "[AZURE BATCH] User Subscription mode detected, allowing low priority VMs in pools: ${poolNames}" + } else { + // If we can't determine the pool allocation mode, show a warning but allow execution + log.warn "[AZURE BATCH] Unable to determine pool allocation mode (got: ${poolAllocationMode}). " + + "Low Priority VMs are configured in pools: ${poolNames}. " + + "Please note that Low Priority VMs are deprecated in Batch Managed accounts. " + + "If you're using a Batch Managed account, please update your configuration to use standard VMs. " + + "To enable automatic detection, set azure.batch.subscriptionId in your config or AZURE_SUBSCRIPTION_ID environment variable." + } + } + } + protected void uploadBinDir() { /* * upload local binaries @@ -120,6 +152,7 @@ class AzBatchExecutor extends Executor implements ExtensionPoint { initBatchService() validateWorkDir() validatePathDir() + validateLowPriorityVMs() uploadBinDir() } diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy index 2c83d8e8b9..c7e1be5c41 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy @@ -388,6 +388,102 @@ class AzBatchService implements Closeable { return builder.buildClient() } + /** + * Determines the pool allocation mode of the Azure Batch account + * @return The pool allocation mode ('BatchService' or 'UserSubscription'), or null if it cannot be determined + */ + @Memoized + protected String getPoolAllocationMode() { + try { + // Get batch account name from endpoint + final accountName = extractAccountName(config.batch().endpoint) + if (!accountName) { + log.debug "[AZURE BATCH] Cannot extract account name from endpoint" + return null + } + + // Get subscription ID + final subscriptionId = config.batch().subscriptionId + if (!subscriptionId) { + log.debug "[AZURE BATCH] No subscription ID configured. Set azure.batch.subscriptionId or AZURE_SUBSCRIPTION_ID" + return null + } + + // Get Azure credentials + final credential = getAzureCredential() + if (!credential) { + log.debug "[AZURE BATCH] No valid credentials for Azure Resource Manager" + return null + } + + // Create BatchManager with proper configuration + final batchManager = createBatchManager(credential, subscriptionId) + + // Find the batch account + return findBatchAccountPoolMode(batchManager, accountName) + + } catch (Exception e) { + log.warn "[AZURE BATCH] Failed to determine pool allocation mode: ${e.message}", e + return null + } + } + + /** + * Extract account name from batch endpoint URL + */ + private String extractAccountName(String endpoint) { + if (!endpoint) return null + // Format: https://accountname.region.batch.azure.com + return endpoint.split('\\.')[0].replace('https://', '') + } + + /** + * Get Azure credentials based on configuration + */ + private TokenCredential getAzureCredential() { + if (config.managedIdentity().isConfigured()) { + return createBatchCredentialsWithManagedIdentity() + } else if (config.activeDirectory().isConfigured()) { + return createBatchCredentialsWithServicePrincipal() + } + return null + } + + /** + * Create and configure BatchManager + */ + private com.azure.resourcemanager.batch.BatchManager createBatchManager(TokenCredential credential, String subscriptionId) { + // AzureProfile constructor: (tenantId, subscriptionId, environment) + final profile = new com.azure.core.management.profile.AzureProfile( + null, // tenantId - null to use default + subscriptionId, + com.azure.core.management.AzureEnvironment.AZURE + ) + + // Use configure().authenticate() pattern to ensure proper initialization + return com.azure.resourcemanager.batch.BatchManager + .configure() + .authenticate(credential, profile) + } + + /** + * Find batch account and return its pool allocation mode + */ + private String findBatchAccountPoolMode(com.azure.resourcemanager.batch.BatchManager batchManager, String accountName) { + log.debug "[AZURE BATCH] Searching for account '${accountName}'" + + for (batchAccount in batchManager.batchAccounts().list()) { + if (batchAccount.name() == accountName) { + final poolMode = batchAccount.poolAllocationMode() + log.debug "[AZURE BATCH] Found account with pool allocation mode: ${poolMode}" + return poolMode?.toString() + } + } + + log.debug "[AZURE BATCH] Account '${accountName}' not found in subscription" + return null + } + AzTaskKey submitTask(TaskRun task) { final poolId = getOrCreatePool(task) final jobId = getOrCreateJob(poolId, task) diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzBatchOpts.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzBatchOpts.groovy index 323d682203..a209e030b7 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzBatchOpts.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzBatchOpts.groovy @@ -48,6 +48,7 @@ class AzBatchOpts implements CloudTransferOptions { String accountKey String endpoint String location + String subscriptionId Boolean autoPoolMode Boolean allowPoolCreation Boolean terminateJobsOnCompletion @@ -67,6 +68,7 @@ class AzBatchOpts implements CloudTransferOptions { accountKey = config.accountKey ?: sysEnv.get('AZURE_BATCH_ACCOUNT_KEY') endpoint = config.endpoint location = config.location + subscriptionId = config.subscriptionId ?: sysEnv.get('AZURE_SUBSCRIPTION_ID') autoPoolMode = config.autoPoolMode allowPoolCreation = config.allowPoolCreation terminateJobsOnCompletion = config.terminateJobsOnCompletion != Boolean.FALSE diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy new file mode 100644 index 0000000000..f32603963b --- /dev/null +++ b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy @@ -0,0 +1,134 @@ +package nextflow.cloud.azure.batch + +import nextflow.Session +import nextflow.cloud.azure.config.AzConfig +import nextflow.cloud.azure.config.AzBatchOpts +import nextflow.cloud.azure.config.AzPoolOpts +import nextflow.exception.AbortOperationException +import spock.lang.Specification + +/** + * Test for AzBatchExecutor validation logic + */ +class AzBatchExecutorTest extends Specification { + + def 'should validate low priority VMs for BatchService allocation mode'() { + given: + def CONFIG = [ + batch: [ + endpoint: 'https://testaccount.eastus.batch.azure.com', + pools: [ + 'pool1': [vmType: 'Standard_D2_v2', lowPriority: true], + 'pool2': [vmType: 'Standard_D2_v2', lowPriority: false] + ] + ] + ] + + and: + def config = new AzConfig(CONFIG) + def batchService = Mock(AzBatchService) { + getPoolAllocationMode() >> 'BATCH_SERVICE' + } + + and: + def executor = new AzBatchExecutor() + executor.config = config + executor.batchService = batchService + + when: + executor.validateLowPriorityVMs() + + then: + def e = thrown(AbortOperationException) + e.message.contains('Azure Low Priority VMs are deprecated and no longer supported for Batch Managed pool allocation mode') + e.message.contains('pool1') + } + + def 'should allow low priority VMs for UserSubscription allocation mode'() { + given: + def CONFIG = [ + batch: [ + endpoint: 'https://testaccount.eastus.batch.azure.com', + pools: [ + 'pool1': [vmType: 'Standard_D2_v2', lowPriority: true] + ] + ] + ] + + and: + def config = new AzConfig(CONFIG) + def batchService = Mock(AzBatchService) { + getPoolAllocationMode() >> 'USER_SUBSCRIPTION' + } + + and: + def executor = new AzBatchExecutor() + executor.config = config + executor.batchService = batchService + + when: + executor.validateLowPriorityVMs() + + then: + noExceptionThrown() + } + + def 'should handle unknown allocation mode gracefully'() { + given: + def CONFIG = [ + batch: [ + endpoint: 'https://testaccount.eastus.batch.azure.com', + pools: [ + 'pool1': [vmType: 'Standard_D2_v2', lowPriority: true] + ] + ] + ] + + and: + def config = new AzConfig(CONFIG) + def batchService = Mock(AzBatchService) { + getPoolAllocationMode() >> null + } + + and: + def executor = new AzBatchExecutor() + executor.config = config + executor.batchService = batchService + + when: + executor.validateLowPriorityVMs() + + then: + noExceptionThrown() + } + + def 'should not validate when no low priority VMs configured'() { + given: + def CONFIG = [ + batch: [ + endpoint: 'https://testaccount.eastus.batch.azure.com', + pools: [ + 'pool1': [vmType: 'Standard_D2_v2', lowPriority: false], + 'pool2': [vmType: 'Standard_D2_v2'] + ] + ] + ] + + and: + def config = new AzConfig(CONFIG) + def batchService = Mock(AzBatchService) { + getPoolAllocationMode() >> 'BATCH_SERVICE' + } + + and: + def executor = new AzBatchExecutor() + executor.config = config + executor.batchService = batchService + + when: + executor.validateLowPriorityVMs() + + then: + noExceptionThrown() + } +} \ No newline at end of file diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzBatchOptsTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzBatchOptsTest.groovy index 271c7d004e..3ee986bcc0 100644 --- a/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzBatchOptsTest.groovy +++ b/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzBatchOptsTest.groovy @@ -105,4 +105,26 @@ class AzBatchOptsTest extends Specification { then: opts3.jobMaxWallClockTime.toString() == '12h' } + + def 'should set subscription ID from config or environment' () { + when: + def opts1 = new AzBatchOpts([:], [:]) + then: + opts1.subscriptionId == null + + when: + def opts2 = new AzBatchOpts([subscriptionId: 'config-sub-id'], [:]) + then: + opts2.subscriptionId == 'config-sub-id' + + when: + def opts3 = new AzBatchOpts([:], [AZURE_SUBSCRIPTION_ID: 'env-sub-id']) + then: + opts3.subscriptionId == 'env-sub-id' + + when: + def opts4 = new AzBatchOpts([subscriptionId: 'config-sub-id'], [AZURE_SUBSCRIPTION_ID: 'env-sub-id']) + then: + opts4.subscriptionId == 'config-sub-id' // config takes precedence over environment + } } From 9f9e8d53cc0c4c0d8f545654e82c935a25c9af51 Mon Sep 17 00:00:00 2001 From: adamrtalbot <12817534+adamrtalbot@users.noreply.github.com> Date: Thu, 10 Jul 2025 19:26:12 +0100 Subject: [PATCH 2/2] Simplifies and resolves some issues. Tidies up code. Signed-off-by: adamrtalbot <12817534+adamrtalbot@users.noreply.github.com> --- .../cloud/azure/batch/AzBatchExecutor.groovy | 16 +++++++--------- .../cloud/azure/batch/AzBatchService.groovy | 5 +++-- .../cloud/azure/batch/AzBatchExecutorTest.groovy | 9 +++++---- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy index 6714fa5cf4..ade28d0918 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchExecutor.groovy @@ -100,20 +100,18 @@ class AzBatchExecutor extends Executor implements ExtensionPoint { log.debug "[AZURE BATCH] Pool allocation mode determined as: ${poolAllocationMode}" if( poolAllocationMode == 'BATCH_SERVICE' || poolAllocationMode == 'BatchService' ) { - throw new AbortOperationException("Azure Low Priority VMs are deprecated and no longer supported for Batch Managed pool allocation mode. " + - "Please update your configuration to use standard VMs instead, or migrate to User Subscription pool allocation mode. " + - "Affected pools: ${poolNames}. " + - "Remove 'lowPriority: true' from your pool configuration or set 'lowPriority: false'.") + throw new AbortOperationException( + "Low Priority VMs are not supported with Batch Managed pool allocation mode. " + + "Update your configuration to use standard VMs or switch to User Subscription mode. " + + "Pools: ${poolNames}." + ) } else if( poolAllocationMode == 'USER_SUBSCRIPTION' || poolAllocationMode == 'UserSubscription' ) { // Low Priority VMs are still supported in User Subscription mode, proceed without warning log.debug "[AZURE BATCH] User Subscription mode detected, allowing low priority VMs in pools: ${poolNames}" } else { // If we can't determine the pool allocation mode, show a warning but allow execution - log.warn "[AZURE BATCH] Unable to determine pool allocation mode (got: ${poolAllocationMode}). " + - "Low Priority VMs are configured in pools: ${poolNames}. " + - "Please note that Low Priority VMs are deprecated in Batch Managed accounts. " + - "If you're using a Batch Managed account, please update your configuration to use standard VMs. " + - "To enable automatic detection, set azure.batch.subscriptionId in your config or AZURE_SUBSCRIPTION_ID environment variable." + log.warn "[AZURE BATCH] Unable to determine pool allocation mode. Low Priority VMs are configured in pools: ${poolNames}. " + + "Low Priority VMs may not be supported. Set 'azure.batch.subscriptionId' in your config or 'AZURE_SUBSCRIPTION_ID' environment variable for automatic detection." } } } diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy index c7e1be5c41..63e9665118 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchService.groovy @@ -453,9 +453,10 @@ class AzBatchService implements Closeable { * Create and configure BatchManager */ private com.azure.resourcemanager.batch.BatchManager createBatchManager(TokenCredential credential, String subscriptionId) { - // AzureProfile constructor: (tenantId, subscriptionId, environment) + // AzureProfile requires: (tenantId, subscriptionId, environment) + // We pass null for tenantId to use the default from the credential final profile = new com.azure.core.management.profile.AzureProfile( - null, // tenantId - null to use default + null, subscriptionId, com.azure.core.management.AzureEnvironment.AZURE ) diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy index f32603963b..f74bb57084 100644 --- a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy +++ b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchExecutorTest.groovy @@ -27,7 +27,7 @@ class AzBatchExecutorTest extends Specification { and: def config = new AzConfig(CONFIG) def batchService = Mock(AzBatchService) { - getPoolAllocationMode() >> 'BATCH_SERVICE' + getPoolAllocationMode() >> 'BatchService' } and: @@ -40,7 +40,8 @@ class AzBatchExecutorTest extends Specification { then: def e = thrown(AbortOperationException) - e.message.contains('Azure Low Priority VMs are deprecated and no longer supported for Batch Managed pool allocation mode') + e.message.contains('Low Priority VMs are not supported with Batch Managed pool allocation mode') + e.message.contains('Update your configuration to use standard VMs or switch to User Subscription mode') e.message.contains('pool1') } @@ -58,7 +59,7 @@ class AzBatchExecutorTest extends Specification { and: def config = new AzConfig(CONFIG) def batchService = Mock(AzBatchService) { - getPoolAllocationMode() >> 'USER_SUBSCRIPTION' + getPoolAllocationMode() >> 'UserSubscription' } and: @@ -117,7 +118,7 @@ class AzBatchExecutorTest extends Specification { and: def config = new AzConfig(CONFIG) def batchService = Mock(AzBatchService) { - getPoolAllocationMode() >> 'BATCH_SERVICE' + getPoolAllocationMode() >> 'BatchService' } and: