Files
ma/nextflow/plugins/nf-k8s/src/main/nextflow/k8s/strategies/K8sDVFSSchedulingStrategy.groovy
2026-07-10 21:27:36 +02:00

292 lines
11 KiB
Groovy

package nextflow.k8s.strategies
import groovy.transform.CompileStatic
import groovy.util.logging.Slf4j
import nextflow.k8s.K8sDVFSClient
import nextflow.k8s.K8sRuntimeEstimator
import nextflow.k8s.K8sSchedulingDecision
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
* of workflow execution, while attempting to maintain the same makespan.
*/
@Slf4j
@CompileStatic
class K8sDVFSSchedulingStrategy implements K8sSchedulingStrategy {
private static long getTaskMemoryRequirment(TaskRun task) {
return task.config.getMemory() ? task.config.getMemory().bytes : 64 * 1024 * 1024
}
private static int getTaskCPURequirement(TaskRun task) {
return task.config.hasCpus() ? task.config.getCpus() : 1
}
@Slf4j
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
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 = getTaskMemoryRequirment(task)
final int reqCPUs = getTaskCPURequirement(task)
this.allocatedMemory += reqBytes
this.allocatedCPUs += reqCPUs
log.info "[K8s] node ${name}: task ${task.name} assigned, now have ${availableMemory} bytes and ${availableCPUs} cpus [${reqBytes}, ${reqCPUs}]"
this.tasks.add(new AssignedTask(task, frequency))
updateFrequency(dvfsClient)
}
void taskFinished(TaskRun task, K8sDVFSClient dvfsClient) {
final long reqBytes = getTaskMemoryRequirment(task)
final int reqCPUs = getTaskCPURequirement(task)
this.allocatedMemory -= reqBytes
this.allocatedCPUs -= reqCPUs
log.info "[K8s] node ${name}: task ${task.name} finished, now have ${availableMemory} bytes and ${availableCPUs} cpus [${reqBytes}, ${reqCPUs}]"
this.tasks.removeIf {it.task == task}
updateFrequency(dvfsClient)
}
}
class SchedulingRequestComparator implements Comparator<K8sSchedulingRequest> {
K8sRuntimeEstimator runtimeEstimator
long currentTime
double avgRuntime
@Override
int compare(K8sSchedulingRequest o1, K8sSchedulingRequest o2) {
// First, check if one of the tasks is (estimated to be) on the critical path
double t1 = runtimeEstimator.estimate(o1.handler)
double t2 = runtimeEstimator.estimate(o2.handler)
if (t1 > avgRuntime && t2 <= avgRuntime)
return -1
else if (t1 < avgRuntime && t2 > avgRuntime)
return 1
// Both are not on the critical path. Sort based on the time they spent in the queue
long w1 = currentTime - o1.submitTimeMillis
long w2 = currentTime - o2.submitTimeMillis
if (w1 > w2)
return -1
else if (w2 > w1)
return 1
return 0
}
}
private K8sRuntimeEstimator runtimeEstimator
private K8sDVFSClient dvfsClient
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 ArrayList<>()
this.taskToNode = new HashMap<>();
}
@Override
K8sSchedulingDecision schedule(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue) {
if (nodes.isEmpty()) {
if (!initNodes(scheduler))
return null
}
/* Step 1: Sort by task priority. We will attempt to schedule tasks "in order", so that the
* highest priority tasks are assigned to nodes as soon as possible.
*
* Priority is based on a) an estimation if the task is on the critical path and b) the wait time of the task.
*/
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 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<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
}
@Override
boolean scheduleImmediately(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue) {
return queue != null && queue.size() > 0 && taskToNode.size() < nodes.size()
}
@Override
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".
* A simple (rough) solution could be to keep track of the tasks "relative" frequency and just scale the
* elapsed time based on that.
*/
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)
final min = dvfsClient.getNodeMinFrequency(node)
final max = dvfsClient.getNodeMaxFrequency(node)
final cpus = dvfsClient.getCPUCount(node)
final mem = dvfsClient.getMemoryAmount(node)
if (cur.empty || min.empty || max.empty || cpus.empty || mem.empty) {
log.error "[K8s] failed to query node $node information"
continue
}
globalMaxFrequency = Long.min(globalMaxFrequency, max.asLong)
log.info "[K8s] node ${node}: ${cpus.asLong} CPUs, ${mem.asLong} bytes RAM ${min.asLong} Hz - ${max.asLong} Hz current ${cur.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 = getTaskMemoryRequirment(task)
final int reqCPUs = getTaskCPURequirement(task)
ArrayList<WorkerNode> suitableNodes = new ArrayList<>()
for (WorkerNode n : nodes) {
if (n.availableMemory >= reqBytes && n.availableCPUs >= reqCPUs) {
log.info "[K8s] task ${task.name}: ${reqBytes} bytes ${reqCPUs} cpus: node ${n.name} has ${n.availableMemory} bytes, ${n.availableCPUs} cpus"
suitableNodes.add(n)
}
}
return suitableNodes
}
private boolean anyNode(TaskRun task) {
final long reqBytes = getTaskMemoryRequirment(task)
final int reqCPUs = getTaskCPURequirement(task)
for (WorkerNode n : nodes) {
if (n.memoryAmount >= reqBytes && n.availableCPUs >= reqCPUs)
return true
}
return false
}
}