diff --git a/docs/reference/config.md b/docs/reference/config.md index 0efb1ce6ae..55876dde9c 100644 --- a/docs/reference/config.md +++ b/docs/reference/config.md @@ -344,6 +344,11 @@ The following settings are available: `azure.batch.pools..lowPriority` : Enable the use of low-priority VMs (default: `false`). +`azure.batch.pools..managedIdentityId` +: :::{versionadded} 25.01.0-edge + ::: +: Specify the pool has a managed identity attached. This will be passed to the task as the environment variable `NXF_AZURE_MI_CLIENT_ID`. + `azure.batch.pools..maxVmCount` : Specify the max of virtual machine when using auto scale option. 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 cd3e4b37f5..a12b15baf3 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 @@ -47,6 +47,7 @@ import com.azure.compute.batch.models.ContainerConfiguration import com.azure.compute.batch.models.ContainerRegistryReference import com.azure.compute.batch.models.ContainerType import com.azure.compute.batch.models.ElevationLevel +import com.azure.compute.batch.models.EnvironmentSetting import com.azure.compute.batch.models.MetadataItem import com.azure.compute.batch.models.MountConfiguration import com.azure.compute.batch.models.NetworkConfiguration @@ -434,6 +435,13 @@ class AzBatchService implements Closeable { log.trace "[AZURE BATCH] Submitting task: $taskId, cpus=${task.config.getCpus()}, mem=${task.config.getMemory()?:'-'}, slots: $slots" + // Add environment variables for managed identity if configured + final env = [:] as Map + if( pool?.opts?.managedIdentityId ) { + env.put('AZCOPY_AUTO_LOGIN_TYPE', 'MSI') + env.put('AZCOPY_MSI_CLIENT_ID', pool.opts.managedIdentityId) + } + return new BatchTaskCreateContent(taskId, cmd) .setUserIdentity(userIdentity(pool.opts.privileged, pool.opts.runAs, AutoUserScope.TASK)) .setContainerSettings(containerOpts) @@ -441,8 +449,9 @@ class AzBatchService implements Closeable { .setOutputFiles(outputFileUrls(task, sas)) .setRequiredSlots(slots) .setConstraints(constraints) - - + .setEnvironmentSettings(env.collect { name, value -> + new EnvironmentSetting(name).setValue(value) + }) } AzTaskKey runTask(String poolId, String jobId, TaskRun task) { @@ -503,6 +512,13 @@ class AzBatchService implements Closeable { List result = new ArrayList<>(20) result << destFile(TaskRun.CMD_EXIT, task.workDir, sas) result << destFile(TaskRun.CMD_LOG, task.workDir, sas) + result << destFile(TaskRun.CMD_OUTFILE, task.workDir, sas) + result << destFile(TaskRun.CMD_ERRFILE, task.workDir, sas) + result << destFile(TaskRun.CMD_SCRIPT, task.workDir, sas) + result << destFile(TaskRun.CMD_RUN, task.workDir, sas) + result << destFile(TaskRun.CMD_STAGE, task.workDir, sas) + result << destFile(TaskRun.CMD_TRACE, task.workDir, sas) + result << destFile(TaskRun.CMD_ENV, task.workDir, sas) return result } diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchTaskHandler.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchTaskHandler.groovy index 0071e99bfa..2493b784e3 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchTaskHandler.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzBatchTaskHandler.groovy @@ -202,4 +202,5 @@ class AzBatchTaskHandler extends TaskHandler implements FusionAwareTask { } return machineInfo } + } diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzFileCopyStrategy.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzFileCopyStrategy.groovy index 4ac9e072fe..50626934bf 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzFileCopyStrategy.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/batch/AzFileCopyStrategy.groovy @@ -62,8 +62,9 @@ class AzFileCopyStrategy extends SimpleFileCopyStrategy { final result = new StringBuilder() final copy = environment ? new LinkedHashMap(environment) : new LinkedHashMap() copy.remove('PATH') - copy.put('PATH', '$PWD/.nextflow-bin:$AZ_BATCH_NODE_SHARED_DIR/bin/:$PATH') - copy.put('AZCOPY_LOG_LOCATION', '$PWD/.azcopy_log') + copy.put('PATH', '$AZ_BATCH_TASK_DIR/.nextflow-bin:$AZ_BATCH_NODE_SHARED_DIR/bin/:$PATH') + copy.put('AZCOPY_LOG_LOCATION', '$AZ_BATCH_TASK_DIR/.azcopy_log') + copy.put('AZCOPY_JOB_PLAN_LOCATION', '$AZ_BATCH_TASK_DIR/.azcopy_log') copy.put('AZ_SAS', sasToken) // finally render the environment diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy index 01a3570c4b..30866babff 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/config/AzPoolOpts.groovy @@ -68,6 +68,8 @@ class AzPoolOpts implements CacheFunnel { boolean lowPriority AzStartTaskOpts startTask + String managedIdentityId + AzPoolOpts() { this(Collections.emptyMap()) } @@ -92,6 +94,7 @@ class AzPoolOpts implements CacheFunnel { this.password = opts.password this.virtualNetwork = opts.virtualNetwork this.lowPriority = opts.lowPriority as boolean + this.managedIdentityId = opts.managedIdentityId } @Override @@ -114,6 +117,7 @@ class AzPoolOpts implements CacheFunnel { hasher.putBoolean(lowPriority) hasher.putUnencodedChars(startTask.script ?: '') hasher.putBoolean(startTask.privileged) + hasher.putUnencodedChars(managedIdentityId ?: '') return hasher } diff --git a/plugins/nf-azure/src/main/nextflow/cloud/azure/file/AzBashLib.groovy b/plugins/nf-azure/src/main/nextflow/cloud/azure/file/AzBashLib.groovy index 0653da019a..af7f6a4741 100644 --- a/plugins/nf-azure/src/main/nextflow/cloud/azure/file/AzBashLib.groovy +++ b/plugins/nf-azure/src/main/nextflow/cloud/azure/file/AzBashLib.groovy @@ -63,14 +63,23 @@ class AzBashLib extends BashFunLib { local base_name="$(basename "$name")" local dir_name="$(dirname "$name")" + local auth_args="" + if [[ ! -z "$AZCOPY_MSI_CLIENT_ID" ]]; then + # When using managed identity, no additional args needed + auth_args="" + else + # Use SAS token authentication + auth_args="?$AZ_SAS" + fi + if [[ -d $name ]]; then if [[ "$base_name" == "$name" ]]; then - azcopy cp "$name" "$target?$AZ_SAS" --recursive --block-blob-tier $AZCOPY_BLOCK_BLOB_TIER --block-size-mb $AZCOPY_BLOCK_SIZE_MB + azcopy cp "$name" "$target$auth_args" --recursive --block-blob-tier $AZCOPY_BLOCK_BLOB_TIER --block-size-mb $AZCOPY_BLOCK_SIZE_MB else - azcopy cp "$name" "$target/$dir_name?$AZ_SAS" --recursive --block-blob-tier $AZCOPY_BLOCK_BLOB_TIER --block-size-mb $AZCOPY_BLOCK_SIZE_MB + azcopy cp "$name" "$target/$dir_name$auth_args" --recursive --block-blob-tier $AZCOPY_BLOCK_BLOB_TIER --block-size-mb $AZCOPY_BLOCK_SIZE_MB fi else - azcopy cp "$name" "$target/$name?$AZ_SAS" --block-blob-tier $AZCOPY_BLOCK_BLOB_TIER --block-size-mb $AZCOPY_BLOCK_SIZE_MB + azcopy cp "$name" "$target/$name$auth_args" --block-blob-tier $AZCOPY_BLOCK_BLOB_TIER --block-size-mb $AZCOPY_BLOCK_SIZE_MB fi } @@ -79,12 +88,22 @@ class AzBashLib extends BashFunLib { local target=$2 local basedir=$(dirname $2) local ret + + local auth_args="" + if [[ ! -z "$AZCOPY_MSI_CLIENT_ID" ]]; then + # When using managed identity, no additional args needed + auth_args="" + else + # Use SAS token authentication + auth_args="?$AZ_SAS" + fi + mkdir -p "$basedir" - ret=$(azcopy cp "$source?$AZ_SAS" "$target" 2>&1) || { + ret=$(azcopy cp "$source$auth_args" "$target" 2>&1) || { ## if fails check if it was trying to download a directory mkdir -p $target - azcopy cp "$source/*?$AZ_SAS" "$target" --recursive >/dev/null || { + azcopy cp "$source/*$auth_args" "$target" --recursive >/dev/null || { rm -rf $target >&2 echo "Unable to download path: $source" exit 1 diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchTaskHandlerTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchTaskHandlerTest.groovy index 6d0b40d4c4..cc57d932f2 100644 --- a/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchTaskHandlerTest.groovy +++ b/plugins/nf-azure/src/test/nextflow/cloud/azure/batch/AzBatchTaskHandlerTest.groovy @@ -1,10 +1,13 @@ package nextflow.cloud.azure.batch +import nextflow.cloud.azure.config.AzPoolOpts import nextflow.cloud.types.CloudMachineInfo import nextflow.cloud.types.PriceModel +import nextflow.cloud.azure.batch.AzVmPoolSpec import nextflow.exception.ProcessUnrecoverableException import nextflow.executor.BashWrapperBuilder import nextflow.executor.Executor +import nextflow.processor.TaskBean import nextflow.processor.TaskConfig import nextflow.processor.TaskProcessor import nextflow.processor.TaskRun @@ -84,5 +87,4 @@ class AzBatchTaskHandlerTest extends Specification { trace.machineInfo.zone == 'west-eu' trace.machineInfo.priceModel == PriceModel.standard } - } diff --git a/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy b/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy index 094868ccc2..78d86f8e07 100644 --- a/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy +++ b/plugins/nf-azure/src/test/nextflow/cloud/azure/config/AzPoolOptsTest.groovy @@ -16,7 +16,7 @@ */ package nextflow.cloud.azure.config - +import nextflow.cloud.azure.config.AzPoolOpts import nextflow.util.Duration import spock.lang.Specification /**