Support accelerator directive for local executor

Signed-off-by: Ben Sherman <[email protected]>

Author: bentsherman

docs/executor.md

@@ -227,6 +227,7 @@ The `local` executor is useful for developing and testing a pipeline script on y
 
 Resource requests and other job characteristics can be controlled via the following process directives:
 
+- {ref}`process-accelerator`
 - {ref}`process-cpus`
 - {ref}`process-memory`
 - {ref}`process-time`
@@ -241,6 +242,25 @@ The local executor supports two types of tasks:
 - Script tasks (processes with a `script` or `shell` block) - executed via a Bash wrapper
 - Native tasks (processes with an `exec` block) - executed directly in the JVM.
 
+(local-accelerators)=
+
+### Accelerators
+
+:::{versionadded} 25.10.0
+:::
+
+The local executor can use the `accelerator` directive to allocate accelerators such as GPUs. To use accelerators, set the corresponding environment variable:
+
+- `CUDA_VISIBLE_DEVICES` for [NVIDIA CUDA](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-environment-variables) applications
+
+- `HIP_VISIBLE_DEVICES` for [HIP](https://rocm.docs.amd.com/projects/HIP/en/docs-develop/reference/env_variables.html) applications
+
+- `ROCR_VISIBLE_DEVICES` for [AMD ROCm](https://rocm.docs.amd.com/en/latest/conceptual/gpu-isolation.html) applications
+
+Set the environment variable to a comma-separated list of device IDs that you want Nextflow to access. Nextflow uses the same environment variable to allocate accelerators for each task that requests them.
+
+For example, to use all GPUs on a node with four NVIDIA GPUs, set `CUDA_VISIBLE_DEVICES` to `0,1,2,3`. If four tasks each request one GPU, they will be executed with `CUDA_VISIBLE_DEVICES` set to `0`, `1`, `2`, and `3`, respectively.
+
 (lsf-executor)=
 
 ## LSF

docs/migrations/25-10.md

@@ -29,6 +29,12 @@ workflow {
 
 This syntax is simpler and easier to use with the {ref}`strict syntax <strict-syntax-page>`. See {ref}`workflow-handlers` for details.
 
+<h3>GPU scheduling for local executor</h3>
+
+The local executor can now schedule GPUs using the `accelerator` directive. This feature is useful when running Nextflow on a single machine with multiple GPUs.
+
+See {ref}`local-accelerators` for details.
+
 ## Breaking changes
 
 - The AWS Java SDK used by Nextflow was upgraded from v1 to v2, which introduced some breaking changes to the `aws.client` config options. See {ref}`the guide <aws-java-sdk-v2-page>` for details.

modules/nextflow/src/main/groovy/nextflow/executor/local/LocalTaskHandler.groovy

@@ -80,6 +80,10 @@ class LocalTaskHandler extends TaskHandler implements FusionAwareTask {
 
     private volatile TaskResult result
 
+    String acceleratorEnv
+
+    List<String> acceleratorIds
+
     LocalTaskHandler(TaskRun task, LocalExecutor executor) {
         super(task)
         // create the task handler
@@ -142,11 +146,13 @@ class LocalTaskHandler extends TaskHandler implements FusionAwareTask {
         final workDir = task.workDir.toFile()
         final logFile = new File(workDir, TaskRun.CMD_LOG)
 
-        return new ProcessBuilder()
+        final pb = new ProcessBuilder()
                 .redirectErrorStream(true)
                 .redirectOutput(logFile)
                 .directory(workDir)
                 .command(cmd)
+        prepareAccelerators(pb)
+        return pb
     }
 
     protected ProcessBuilder fusionProcessBuilder() {
@@ -162,10 +168,18 @@ class LocalTaskHandler extends TaskHandler implements FusionAwareTask {
 
         final logPath = Files.createTempFile('nf-task','.log')
 
-        return new ProcessBuilder()
+        final pb = new ProcessBuilder()
                 .redirectErrorStream(true)
                 .redirectOutput(logPath.toFile())
                 .command(List.of('sh','-c', cmd))
+        prepareAccelerators(pb)
+        return pb
+    }
+
+    protected void prepareAccelerators(ProcessBuilder pb) {
+        if( !acceleratorEnv )
+            return
+        pb.environment().put(acceleratorEnv, acceleratorIds.join(','))
     }
 
     protected ProcessBuilder createLaunchProcessBuilder() {

modules/nextflow/src/main/groovy/nextflow/processor/AcceleratorTracker.groovy

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2013-2024, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nextflow.processor
+
+import groovy.transform.CompileStatic
+import nextflow.SysEnv
+import nextflow.util.TrackingSemaphore
+
+/**
+ * Specialized semaphore that keeps track of accelerators by
+ * id. The id can be an integer or a UUID.
+ *
+ * @author Ben Sherman <[email protected]>
+ */
+@CompileStatic
+class AcceleratorTracker {
+
+    private static final List<String> DEVICE_ENV_NAMES = [
+        'CUDA_VISIBLE_DEVICES',
+        'HIP_VISIBLE_DEVICES',
+        'ROCR_VISIBLE_DEVICES'
+    ]
+
+    static AcceleratorTracker create() {
+        return DEVICE_ENV_NAMES.stream()
+            .filter(name -> SysEnv.containsKey(name))
+            .map((name) -> {
+                final ids = SysEnv.get(name).tokenize(',')
+                return new AcceleratorTracker(name, ids)
+            })
+            .findFirst().orElse(new AcceleratorTracker())
+    }
+
+    private final String name
+    private final TrackingSemaphore semaphore
+
+    AcceleratorTracker(String name, List<String> ids) {
+        this.name = name
+        this.semaphore = new TrackingSemaphore(ids)
+    }
+
+    AcceleratorTracker() {
+        this(null, [])
+    }
+
+    String name() {
+        return name
+    }
+
+    int total() {
+        return semaphore.totalPermits()
+    }
+
+    int available() {
+        return semaphore.availablePermits()
+    }
+
+    List<String> acquire(int permits) {
+        return semaphore.acquire(permits)
+    }
+
+    void release(List<String> ids) {
+        semaphore.release(ids)
+    }
+
+}

modules/nextflow/src/main/groovy/nextflow/processor/LocalPollingMonitor.groovy

@@ -24,6 +24,7 @@ import groovy.util.logging.Slf4j
 import nextflow.Session
 import nextflow.executor.ExecutorConfig
 import nextflow.exception.ProcessUnrecoverableException
+import nextflow.executor.local.LocalTaskHandler
 import nextflow.util.Duration
 import nextflow.util.MemoryUnit
 
@@ -59,6 +60,11 @@ class LocalPollingMonitor extends TaskPollingMonitor {
      */
     private final long maxMemory
 
+    /**
+     * Tracks the total and available accelerators in the system
+     */
+    private AcceleratorTracker acceleratorTracker
+
     /**
      * Create the task polling monitor with the provided named parameters object.
      * <p>
@@ -76,6 +82,7 @@ class LocalPollingMonitor extends TaskPollingMonitor {
         super(params)
         this.availCpus = maxCpus = params.cpus as int
         this.availMemory = maxMemory = params.memory as long
+        this.acceleratorTracker = AcceleratorTracker.create()
         assert availCpus>0, "Local avail `cpus` attribute cannot be zero"
         assert availMemory>0, "Local avail `memory` attribute cannot zero"
     }
@@ -154,6 +161,16 @@ class LocalPollingMonitor extends TaskPollingMonitor {
         handler.task.getConfig()?.getMemory()?.toBytes() ?: 1L
     }
 
+    /**
+     * @param handler
+     *      A {@link TaskHandler} instance
+     * @return
+     *      The number of accelerators requested to execute the specified task
+     */
+    private static int accelerators(TaskHandler handler) {
+        handler.task.getConfig()?.getAccelerator()?.getRequest() ?: 0
+    }
+
     /**
      * Determines if a task can be submitted for execution checking if the resources required
      * (cpus and memory) match the amount of avail resource
@@ -179,9 +196,13 @@ class LocalPollingMonitor extends TaskPollingMonitor {
         if( taskMemory>maxMemory)
             throw new ProcessUnrecoverableException("Process requirement exceeds available memory -- req: ${new MemoryUnit(taskMemory)}; avail: ${new MemoryUnit(maxMemory)}")
 
-        final result = super.canSubmit(handler) && taskCpus <= availCpus && taskMemory <= availMemory
+        final taskAccelerators = accelerators(handler)
+        if( taskAccelerators > acceleratorTracker.total() )
+            throw new ProcessUnrecoverableException("Process requirement exceeds available accelerators -- req: $taskAccelerators; avail: ${acceleratorTracker.total()}")
+
+        final result = super.canSubmit(handler) && taskCpus <= availCpus && taskMemory <= availMemory && taskAccelerators <= acceleratorTracker.available()
         if( !result && log.isTraceEnabled( ) ) {
-            log.trace "Task `${handler.task.name}` cannot be scheduled -- taskCpus: $taskCpus <= availCpus: $availCpus && taskMemory: ${new MemoryUnit(taskMemory)} <= availMemory: ${new MemoryUnit(availMemory)}"
+            log.trace "Task `${handler.task.name}` cannot be scheduled -- taskCpus: $taskCpus <= availCpus: $availCpus && taskMemory: ${new MemoryUnit(taskMemory)} <= availMemory: ${new MemoryUnit(availMemory)} && taskAccelerators: $taskAccelerators <= availAccelerators: ${acceleratorTracker.available()}"
         }
         return result
     }
@@ -194,9 +215,16 @@ class LocalPollingMonitor extends TaskPollingMonitor {
      */
     @Override
     protected void submit(TaskHandler handler) {
-        super.submit(handler)
         availCpus -= cpus(handler)
         availMemory -= mem(handler)
+
+        final taskAccelerators = accelerators(handler)
+        if ( taskAccelerators > 0 ) {
+            ((LocalTaskHandler) handler).acceleratorEnv = acceleratorTracker.name()
+            ((LocalTaskHandler) handler).acceleratorIds = acceleratorTracker.acquire(taskAccelerators)
+        }
+
+        super.submit(handler)
     }
 
     /**
@@ -209,11 +237,13 @@ class LocalPollingMonitor extends TaskPollingMonitor {
      *      {@code true} when the task is successfully removed from polling queue,
      *      {@code false} otherwise
      */
+    @Override
     protected boolean remove(TaskHandler handler) {
         final result = super.remove(handler)
         if( result ) {
             availCpus += cpus(handler)
             availMemory += mem(handler)
+            acceleratorTracker.release(((LocalTaskHandler) handler).acceleratorIds ?: Collections.<String>emptyList())
         }
         return result
     }

modules/nextflow/src/main/groovy/nextflow/util/TrackingSemaphore.groovy

@@ -0,0 +1,69 @@
+/*
+ * Copyright 2013-2024, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nextflow.util
+
+import java.util.concurrent.Semaphore
+
+import groovy.transform.CompileStatic
+
+/**
+ * Specialized semaphore that keeps track of which permits
+ * are being used.
+ *
+ * @author Ben Sherman <[email protected]>
+ */
+@CompileStatic
+class TrackingSemaphore {
+    private final Semaphore semaphore
+    private final Map<String,Boolean> availIds
+
+    TrackingSemaphore(List<String> ids) {
+        semaphore = new Semaphore(ids.size())
+        availIds = new HashMap<>(ids.size())
+        for( final id : ids )
+            availIds.put(id, true)
+    }
+
+    int totalPermits() {
+        return availIds.size()
+    }
+
+    int availablePermits() {
+        return semaphore.availablePermits()
+    }
+
+    List<String> acquire(int permits) {
+        semaphore.acquire(permits)
+        final result = new ArrayList<String>(permits)
+        for( final entry : availIds.entrySet() ) {
+            if( entry.getValue() ) {
+                entry.setValue(false)
+                result.add(entry.getKey())
+            }
+            if( result.size() == permits )
+                break
+        }
+        return result
+    }
+
+    void release(List<String> ids) {
+        semaphore.release(ids.size())
+        for( final id : ids )
+            availIds.put(id, true)
+    }
+
+}