feat: frequency assignment

nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sTaskScheduler.groovy

@@ -16,18 +16,10 @@ class K8sTaskScheduler implements Runnable {
 
     private synchronized boolean shouldStop
 
-    private HashMap<String, ArrayList<K8sTaskHandler>> occupancy;
-    private HashMap<String, String> taskToNode
-
     K8sTaskScheduler(String[] nodes, K8sSchedulingStrategy strategy) {
         this.executor = executor
         this.strategy = strategy
         this.nodes = nodes
-        this.occupancy = new HashMap<>()
-        for (String node : nodes) {
-            this.occupancy.put(node, new ArrayList<K8sTaskHandler>())
-        }
-        this.taskToNode = new HashMap<>()
     }
 
     /**
@@ -47,20 +39,6 @@ class K8sTaskScheduler implements Runnable {
      * @param handler
      */
     void taskFinished(K8sTaskHandler handler) {
-        /* Remove from occupancy list */
-        String nodeName = taskToNode.get(handler.task.name)
-        if (nodeName == null) {
-            log.error "[K8s] no node saved for task ${handler.task.name}"
-            return
-        }
-        ArrayList<K8sTaskHandler> nodeTasks = occupancy.get(nodeName)
-        if (nodeTasks == null) {
-            log.error "[K8s] tried to remove task ${handler.task.name} from node ${nodeName} but no tasks are saved for that node"
-            return
-        }
-        nodeTasks.remove(handler)
-        log.info "[K8s] removed task ${handler.task.name} from node ${nodeName}"
-
         strategy.taskFinished(handler)
 
         /* If we have pending tasks, now would be a good time to schedule a new one.
@@ -82,19 +60,10 @@ class K8sTaskScheduler implements Runnable {
 
             log.info "[K8s] launching queued task ${decision.request.task.name} on node: ${decision.nodeName}"
             decision.request.handler.submitNow(decision.nodeName)
-            ArrayList nodeTasks = occupancy.get(decision.nodeName)
-            assert nodeTasks != null
-            nodeTasks.add(decision.request.handler)
-            taskToNode.put(decision.request.task.name, decision.nodeName)
         }
     }
 
     /* Scheduling Strategy Interface */
-    List<String> getFreeNodes() {
-        def unoccupied = occupancy.findAll {it.value == null || it.value.size() == 0 }
-        unoccupied.keySet().toList()
-    }
-
     List<String> getNodes() {
         return nodes.toList()
     }

nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sDVFSSchedulingStrategy.groovy

@@ -9,6 +9,8 @@ import nextflow.k8s.K8sSchedulingRequest
 import nextflow.k8s.K8sSchedulingStrategy
 import nextflow.k8s.K8sTaskHandler
 import nextflow.k8s.K8sTaskScheduler
+import nextflow.processor.TaskRun
+import nextflow.util.ArrayTuple
 
 /**
  * Implements a scheduling strategy utilizing dvfs to reduce the energy consumption
@@ -19,18 +21,78 @@ import nextflow.k8s.K8sTaskScheduler
 class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {
 
     private class WorkerNode {
+        private class AssignedTask {
+            TaskRun task
+            long frequency
+
+            AssignedTask(TaskRun t, long f) {
+                this.task = t
+                this.frequency = f
+            }
+        }
+
+        String name
+
         long maxFrequency
         long minFrequency
         long currentFrequency
         long cpuCount
         long memoryAmount
 
-        WorkerNode(long maxF, long minF, long curF, long cpus, long mem) {
+        long allocatedMemory
+        long allocatedCPUs
+
+        ArrayList<AssignedTask> tasks
+
+        WorkerNode(String name, long maxF, long minF, long curF, long cpus, long mem) {
+            this.name = name
             this.maxFrequency = maxF
             this.minFrequency = minF
             this.currentFrequency = curF
             this.cpuCount = cpus
             this.memoryAmount = mem
+            this.allocatedCPUs = 0
+            this.allocatedMemory = 0
+
+            this.tasks = new ArrayList<>()
+        }
+
+        long getAvailableMemory() {
+            memoryAmount - allocatedMemory
+        }
+
+        long getAvailableCPUs() {
+            cpuCount - allocatedCPUs
+        }
+
+        // Sets the frequency to the max. requested frequency of all currently running tasks.
+        private void updateFrequency(K8sDVFSClient dvfsClient) {
+            if (tasks.size() == 0)
+                return
+            long max = Long.MIN_VALUE
+            for (AssignedTask t : tasks) {
+                max = Long.max(t.frequency, max)
+            }
+            dvfsClient.setNodeFrequency(name, (int)max)
+            this.currentFrequency = max
+        }
+
+        void assignTask(TaskRun task, long frequency, K8sDVFSClient dvfsClient) {
+            final long reqBytes = task.config.getMemory().bytes
+            final int reqCPUs = task.config.hasCpus() ? task.config.getCpus()  : 1
+            this.allocatedMemory += reqBytes
+            this.allocatedCPUs += reqCPUs
+            this.tasks.add(new AssignedTask(task, frequency))
+            updateFrequency(dvfsClient)
+        }
+
+        void taskFinished(TaskRun task, K8sDVFSClient dvfsClient) {
+            final long reqBytes = task.config.getMemory().bytes
+            final int reqCPUs = task.config.hasCpus() ? task.config.getCpus()  : 1
+            this.allocatedMemory -= reqBytes
+            this.allocatedCPUs -= reqCPUs
+            this.tasks.removeIf {it.task == task}
+            updateFrequency(dvfsClient)
         }
     }
 
@@ -63,15 +125,19 @@ class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {
     private K8sRuntimeEstimator runtimeEstimator
     private K8sDVFSClient dvfsClient
 
-    private HashMap<String, WorkerNode> nodes
+    private ArrayList<WorkerNode> nodes
+    private HashMap<String, WorkerNode> taskToNode
 
     private double averageRuntime
     private long finishedTaskCount
 
+    private long globalMaxFrequency
+
     K8sDVFSSchedulingStrategy(K8sRuntimeEstimator runtimeEstimator, K8sDVFSClient dvfsClient) {
         this.runtimeEstimator = runtimeEstimator
         this.dvfsClient = dvfsClient
-        this.nodes = new HashMap<>()
+        this.nodes = new ArrayList<>()
+        this.taskToNode = new HashMap<>();
     }
 
     @Override
@@ -88,21 +154,49 @@ class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {
          */
         SchedulingRequestComparator comparator = new SchedulingRequestComparator()
         comparator.runtimeEstimator = runtimeEstimator
+        comparator.avgRuntime = averageRuntime
         comparator.currentTime = System.currentTimeMillis()
         queue.sort(comparator)
 
         /* Step 2: For each task attempt to schedule it onto a node */
         for (K8sSchedulingRequest req : queue) {
-            /* Step 2.1: Determine the execution frequency.
+            /* Step 2.1: Determine if the task is on the critical path.
-             * If the task is on the critical path, just use the maximum available frequency.
+             * If yes, we just schedule it at max frequency on the node with the highest available
+             * frequency. If not, we determine a frequency (see below).
              */
             final double taskEstimation = runtimeEstimator.estimate(req.handler)
             final boolean isCriticalPath = taskEstimation > averageRuntime
-
+            long frequency = globalMaxFrequency
+            if (!isCriticalPath) {
+                /* Set frequency so that we expect the runtime to be close to the mean runtime. */
+                frequency = (long)Math.floor((taskEstimation * globalMaxFrequency) / averageRuntime)
+            }
 
             /* Step 2.2: Filter nodes based on task requirements */
+            ArrayList<WorkerNode> suitableNodes = filterNodes(req.task)
+            if (suitableNodes.size() == 0) {
+                if (!anyNode(req.task)) {
+                    log.error "[K8s] unable to schedule task ${req.task} - no node satisfies resource requirements"
+                    return null
+                }
+                /* No node can currently execute this task, but it should be possible in the future */
+                continue
+            }
 
             /* Step 2.3: Assign to node based on "best fit" - current node frequency is closest to determined frequency */
+            long minDist = Math.abs(suitableNodes[0].currentFrequency - frequency)
+            WorkerNode closest = suitableNodes[0]
+            for (WorkerNode node : suitableNodes) {
+                long dist = Math.abs(node.currentFrequency - frequency)
+                if (dist < minDist) {
+                    closest = node
+                    minDist = dist
+                }
+            }
+
+            closest.assignTask(req.task, frequency, dvfsClient)
+            taskToNode.put(req.task.hash.toString(), closest)
+            return new K8sSchedulingDecision(req, closest.name)
         }
 
         return null
@@ -110,11 +204,11 @@ class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {
 
     @Override
     boolean scheduleImmediately(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue) {
-        return queue != null && queue.size() > 0 && scheduler.freeNodes.size() > 0
+        return queue != null && queue.size() > 0 && taskToNode.size() < nodes.size()
     }
 
     @Override
-    void taskFinished(K8sTaskHandler task) {
+    synchronized void taskFinished(K8sTaskHandler task) {
         /* TODO: This just uses elapsed wall-clock time, regardless of the frequency used to execute the task.
          * This will skew the average towards longer runtimes, which is undesirable, because it will lead to more
          * tasks classified as "critical path".
@@ -124,9 +218,19 @@ class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {
         double runtime = (double)(task.getCompleteTimeMillis() - task.getStartTimeMillis())
         averageRuntime = (runtime + finishedTaskCount * averageRuntime) / (finishedTaskCount + 1)
         finishedTaskCount += 1
+
+        /* Free resources allocated by this task */
+        WorkerNode node = taskToNode.get(task.task.hash.toString())
+        if (node != null) {
+            node.taskFinished(task.task, dvfsClient)
+            taskToNode.remove(task.task.hash.toString())
+        } else {
+            log.warn "[K8s] no node recorded for task ${task.toString()}"
+        }
     }
 
     private boolean initNodes(K8sTaskScheduler scheduler) {
+        this.globalMaxFrequency = Long.MAX_VALUE
         final nodes = scheduler.getNodes()
         for (String node : nodes) {
             final cur = dvfsClient.getNodeCurrentFrequency(node)
@@ -139,8 +243,32 @@ class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {
                 log.error "[K8s] failed to query node $node information"
                 continue
             }
-            this.nodes.put(node, new WorkerNode(max.asLong, min.asLong, cur.asLong, cpus.asLong, mem.asLong))
+            globalMaxFrequency = Long.min(globalMaxFrequency, max.asLong)
+            this.nodes.add(new WorkerNode(node, max.asLong, min.asLong, cur.asLong, cpus.asLong, mem.asLong))
         }
         return !this.nodes.isEmpty()
     }
+
+    /* Returns a list of nodes that fulfill the tasks resource requirements
+     */
+    private ArrayList<WorkerNode> filterNodes(TaskRun task) {
+        final long reqBytes = task.config.getMemory().bytes
+        final int reqCPUs = task.config.hasCpus() ? task.config.getCpus()  : 1
+        ArrayList<WorkerNode> suitableNodes = new ArrayList<>()
+        for (WorkerNode n : nodes) {
+            if (n.availableMemory >= reqBytes && n.availableCPUs >= reqCPUs)
+                suitableNodes.add(n)
+        }
+        return suitableNodes
+    }
+
+    private boolean anyNode(TaskRun task) {
+        final long reqBytes = task.config.getMemory().bytes
+        final int reqCPUs = task.config.hasCpus() ? task.config.getCpus()  : 1
+        for (WorkerNode n : nodes) {
+            if (n.memoryAmount >= reqBytes && n.availableCPUs >= reqCPUs)
+                return true
+        }
+        return false
+    }
 }

nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sHashSchedulingStrategy.groovy

@@ -15,7 +15,7 @@ class K8sHashSchedulingStrategy implements K8sSchedulingStrategy {
         if (!queue)
             return null
 
-        final freeNodes = scheduler.freeNodes
+        final freeNodes = scheduler.nodes
         if (freeNodes.size() == 0)
             return null
 
@@ -26,10 +26,7 @@ class K8sHashSchedulingStrategy implements K8sSchedulingStrategy {
 
     @Override
     boolean scheduleImmediately(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue) {
-        if (!queue)
         return false
-        final freeNodes = scheduler.freeNodes
-        return freeNodes.size() > 0
     }
 
     @Override