From 008b8c7cd9ffbe841d1cfda88172bce7407e0e76 Mon Sep 17 00:00:00 2001 From: Kevin Trogant Date: Fri, 26 Jun 2026 21:52:06 +0200 Subject: [PATCH] feat: frequency assignment --- .../main/nextflow/k8s/K8sTaskScheduler.groovy | 31 ---- .../K8sDVFSSchedulingStrategy.groovy | 146 ++++++++++++++++-- .../K8sHashSchedulingStrategy.groovy | 7 +- 3 files changed, 139 insertions(+), 45 deletions(-) diff --git a/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sTaskScheduler.groovy b/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sTaskScheduler.groovy index 4524c79..bb0ded6 100644 --- a/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sTaskScheduler.groovy +++ b/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sTaskScheduler.groovy @@ -16,18 +16,10 @@ class K8sTaskScheduler implements Runnable { private synchronized boolean shouldStop - private HashMap> occupancy; - private HashMap 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()) - } - 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 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 getFreeNodes() { - def unoccupied = occupancy.findAll {it.value == null || it.value.size() == 0 } - unoccupied.keySet().toList() - } - List getNodes() { return nodes.toList() } diff --git a/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sDVFSSchedulingStrategy.groovy b/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sDVFSSchedulingStrategy.groovy index 2f53358..4a213fc 100644 --- a/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sDVFSSchedulingStrategy.groovy +++ b/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 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 nodes + private ArrayList nodes + private HashMap 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. - * If the task is on the critical path, just use the maximum available frequency. + /* Step 2.1: Determine if the task is on the critical path. + * 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 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 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 filterNodes(TaskRun task) { + final long reqBytes = task.config.getMemory().bytes + final int reqCPUs = task.config.hasCpus() ? task.config.getCpus() : 1 + ArrayList 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 + } } diff --git a/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sHashSchedulingStrategy.groovy b/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sHashSchedulingStrategy.groovy index b29a536..e1ff39a 100644 --- a/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sHashSchedulingStrategy.groovy +++ b/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 queue) { - if (!queue) - return false - final freeNodes = scheduler.freeNodes - return freeNodes.size() > 0 + return false } @Override