feat: linear regression for runtime estimation

This commit is contained in:
2026-06-08 22:05:28 +02:00
parent 933611af55
commit b92feeab5b
8 changed files with 222 additions and 240 deletions

View File

@@ -274,11 +274,11 @@ class K8sConfig implements ConfigScope {
nextflowImage = opts.nextflowImage ?: "nextflow/nextflow:${BuildInfo.version}"
schedulerInterval = opts.schedulerInterval as Duration ?: new Duration(10, TimeUnit.SECONDS)
recordTaskRuntimes = opts.recordTaskRuntimes as boolean ?: false
runtimeRecordPath = opts.runtimeRecordPath as String ?: "${launchDir}/work/runtimes.csv"
launchDir = opts.launchDir ?: "${storageMountPath}/${getUserName()}"
projectDir = opts.projectDir ?: "${storageMountPath}/projects"
workDir = opts.workDir ?: "${launchDir}/work"
runtimeRecordPath = opts.runtimeRecordPath as String ?: "${workDir}/runtimes.csv"
// -- shortcut to pod image pull-policy
if( imagePullPolicy )

View File

@@ -56,6 +56,7 @@ class K8sExecutor extends Executor implements ExtensionPoint {
private Thread schedulerThread
K8sRuntimeRecorder runtimeRecorder
K8sRuntimeEstimator runtimeEstimator
/**
* @return The Kubernetes HTTP client. Delegates to a Guava cache that refreshes
@@ -94,6 +95,7 @@ class K8sExecutor extends Executor implements ExtensionPoint {
log.debug "[K8s] config=$k8sConfig; API client config=$client.config"
this.runtimeRecorder = new K8sRuntimeRecorder(k8sConfig.recordTaskRuntimes, k8sConfig.runtimeRecordPath)
this.runtimeEstimator = new K8sRuntimeEstimator(k8sConfig.runtimeRecordPath)
this.taskScheduler = new K8sTaskScheduler(this, new K8sHashSchedulingStrategy())
this.schedulerThread = new Thread(this.taskScheduler)

View File

@@ -0,0 +1,157 @@
package nextflow.k8s
import groovy.util.logging.Slf4j
import java.nio.file.Files
import java.nio.file.Path
/**
* Provides runtime estimates for tasks based on task-name and input file size
*/
@Slf4j
class K8sRuntimeEstimator {
class Function {
private double m
private double n
Function(double m, double n) {
this.m = m
this.n = n
}
double estimate(long x) {
return m * (double)x + n
}
}
HashMap<String, Function> estimators;
/**
* Initialize the estimator with data recorded by K8sRuntimeRecorder
* @param dataFilePath
*/
K8sRuntimeEstimator(String dataFilePath) {
HashMap<String, ArrayList<Tuple2<Long, Long>>> data = new HashMap<>();
try {
BufferedReader reader = new BufferedReader(dataFilePath)
String line = reader.readLine()
while (line != null) {
// <task-name>,<input-size>,<runtime-ms>
String[] elems = line.split(",")
if (elems.length != 3) {
log.warn "[K8s] ${dataFilePath}: Unexpected line ${line}"
continue
}
ArrayList<Tuple2<Long, Long>> taskData = data.get(elems[0])
if (taskData == null) {
taskData = new ArrayList<>()
data.put(elems[0], taskData)
}
try {
taskData.add(new Tuple2(Long.parseLong(elems[1]), Long.parseLong(elems[2])))
} catch (NumberFormatException ex) {
log.error "[K8s] ${dataFilePath} invalid data: ${ex.message}"
}
line = reader.readLine()
}
reader.close()
} catch (IOException ex) {
log.error "[K8s] Failed to load ${dataFilePath}: ${ex.message}"
}
fit(data)
}
/**
* Initializes the runtime estimator with statically known data.
* @param data Map from task name to list of recordings, where each recording is a tuple (input-size, runtime-in-ms)
*/
K8sRuntimeEstimator(HashMap<String, ArrayList<Tuple2<Long, Long>>> data) {
fit(data)
}
/**
* Returns an estimation of the task runtime in milliseconds
* @param handler the task handler
* @return estimated task runtime in milliseconds OR infinity if the task is unknown.
*/
double estimate(K8sTaskHandler handler) {
long x = getTaskHandlerInputSize(handler)
return estimate(handler.task.processor.name, x)
}
/**
* Returns an estimation of the task runtime in milliseconds
* @param taskName the name of the task
* @param inputSize the total input size in bytes
* @return estimated task runtime in milliseconds OR infinity if the task is unknown.
*/
double estimate(String taskName, long inputSize) {
Function estimator = estimators.get(taskName)
if (estimator == null) {
log.warn "[K8s] Unable to estimate take ${taskName}: unknown task"
return Double.POSITIVE_INFINITY
}
return estimator.estimate(inputSize)
}
private void fit(HashMap<String, ArrayList<Tuple2<Long, Long>>> data) {
estimators = new HashMap<>()
for (Map.Entry<String, ArrayList<Tuple2<Long, Long>>> t : data) {
Function f = fit(t.value)
estimators.put(t.key, f)
}
}
/**
* Uses linear regression to fit a linear function (y = m * x + n) to the provided observations
* @param observations list of tuples (input size, runtime in ms)
* @return linear function fitted to the input
*/
private Function fit(ArrayList<Tuple2<Long, Long>> observations) throws IllegalArgumentException {
int n = observations.size()
if (n > 1) {
double sumX = 0, sumY = 0, sumXY = 0, sumX2 = 0
for (Tuple2<Long, Long> o : observations) {
double x = (double) o.get(0)
double y = (double) o.get(1)
sumX += x
sumY += y
sumXY += x * y
sumX2 += x * x
}
double m = (n * sumXY - sumX * sumY) / (n * sumX2 - sumX * sumX)
return new Function(
m,
(sumY - m * sumX) / n
)
} else if (n == 1) {
// Special case: We only have 1 measurement. We will just assume that the runtime is constant,
// because in our observed data, it is.
return new Function(0.0, (double) observations[0].get(1))
}
throw new IllegalArgumentException("requires at least 1 observation")
}
private long getTaskHandlerInputSize(K8sTaskHandler handler) {
long inputSizeSum = 0
// File input
def inputFiles = handler.task.getInputFilesMap()
for (Map.Entry<String, Path> f : inputFiles) {
try {
inputSizeSum += Files.size(f.value)
} catch (IOException ex) {
log.error "[K8s] failed to get size of input file ${f.value} of task ${task.task.name}: ${ex.message}"
}
}
// Non file input
def inputs = handler.task.getInputs()
for (Map.Entry i : inputs) {
inputSizeSum += i.value.toString().length()
}
return inputSizeSum
}
}

View File

@@ -29,6 +29,8 @@ class K8sRuntimeRecorder {
return;
long runtimeMilis = task.completeTimeMillis - task.startTimeMillis
long inputSizeSum = 0
// Input files
def inputFiles = task.task.getInputFilesMap()
for (Map.Entry<String, Path> f : inputFiles) {
try {
@@ -38,12 +40,27 @@ class K8sRuntimeRecorder {
}
}
// Non file input
def inputs = task.task.getInputs()
for (Map.Entry i : inputs) {
inputSizeSum += i.value.toString().length()
}
records.add(new K8sRuntimeRecord(task.task.processor.name, inputSizeSum, runtimeMilis))
}
void write() {
for (K8sRuntimeRecord record : records) {
log.info("[K8s] ${record.taskName} - ${record.inputSize} Bytes - ${record.runtimeMillis} ms")
if (!enabled)
return;
try {
FileWriter out = new FileWriter(recordPath)
for (K8sRuntimeRecord record : records) {
out.write("${record.taskName},${record.inputSize},${record.runtimeMillis}\n")
}
out.close()
log.info "[K8s] written runtime recording to ${recordPath}"
} catch (IOException err) {
log.error "[K8s] failed to write runtime recording ${recordPath}: ${err.message}"
}
}
}

View File

@@ -0,0 +1,42 @@
package nextflow.k8s
import spock.lang.Specification
class K8sRuntimeEstimatorTest extends Specification {
def 'does extract a function' () {
given:
def observations = [a: [new Tuple2(0, 0), new Tuple2(100, 100)]]
def estimator = new K8sRuntimeEstimator(observations)
when:
def estimation = estimator.estimate("a", 1234)
then:
estimation != Double.POSITIVE_INFINITY
}
def 'can extract a function from single measurement' () {
given:
def observations = [a: [new Tuple2(100, 100)]]
def estimator = new K8sRuntimeEstimator(observations)
when:
def estimation = estimator.estimate("a", 1234)
then:
estimation == 100.0 /* Single measurement implies constant runtime */
}
def 'unknown function returns positive infinity' () {
given:
def observations = [] as HashMap<String, ArrayList>
def estimator = new K8sRuntimeEstimator(observations)
when:
def estimation = estimator.estimate("b", 1234)
then:
estimation == Double.POSITIVE_INFINITY
}
}

View File

@@ -1,153 +0,0 @@
package nextflow.k8s
import nextflow.k8s.client.K8sClient
import nextflow.processor.TaskRun
import spock.lang.Specification
class K8sTaskSchedulerTest extends Specification {
def 'should enqueue task handler without immediately submitting it' () {
given:
def executor = Mock(K8sExecutor)
def strategy = Mock(K8sSchedulingStrategy)
def scheduler = new K8sTaskScheduler(executor, strategy)
def task = Mock(TaskRun)
def handler = Spy(K8sTaskHandler)
handler.task = task
when:
scheduler.submit(handler)
then:
1 * scheduler.() >> ['A', 'B', 'C']
0 * strategy._
0 * executor._
0 * handler.submitNow(_)
and:
scheduler.@queue.size() == 1
}
def 'should drain queued task and submit it to selected node' () {
given:
def client = Mock(K8sClient)
def executor = Mock(K8sExecutor)
def strategy = Mock(K8sSchedulingStrategy)
def scheduler = new K8sTaskScheduler(executor, strategy)
def task = Mock(TaskRun)
def handler = Spy(K8sTaskHandler)
handler.task = task
when:
scheduler.submit(handler)
scheduler.drain()
then:
1 * executor.getClient() >> client
1 * client.nodeList() >> [
items: [
[metadata: [name: 'node-a']],
[metadata: [name: 'node-b']]
]
]
then:
1 * strategy.schedule({ List<K8sSchedulingRequest> queue ->
queue.size() == 1 && queue[0].handler == handler && queue[0].task == task
}, ['node-a', 'node-b']) >> { List<K8sSchedulingRequest> queue, List<String> nodes ->
new K8sSchedulingDecision(queue[0], 'node-b')
}
then:
1 * handler.submitNow('node-b') >> {}
then:
1 * strategy.schedule([], ['node-a', 'node-b']) >> null
and:
scheduler.@queue.size() == 0
}
def 'should keep task queued when strategy returns no decision' () {
given:
def client = Mock(K8sClient)
def executor = Mock(K8sExecutor)
def strategy = Mock(K8sSchedulingStrategy)
def scheduler = new K8sTaskScheduler(executor, strategy)
def task = Mock(TaskRun)
def handler = Spy(K8sTaskHandler)
handler.task = task
when:
scheduler.submit(handler)
scheduler.drain()
then:
1 * executor.getClient() >> client
1 * client.nodeList() >> [
items: [
[metadata: [name: 'node-a']]
]
]
then:
1 * strategy.schedule({ List<K8sSchedulingRequest> queue ->
queue.size() == 1 && queue[0].handler == handler
}, ['node-a']) >> null
and:
0 * handler.submitNow(_)
scheduler.@queue.size() == 1
}
def 'should continue draining when selected request was already removed' () {
given:
def client = Mock(K8sClient)
def executor = Mock(K8sExecutor)
def strategy = Mock(K8sSchedulingStrategy)
def scheduler = new K8sTaskScheduler(executor, strategy)
def task1 = Mock(TaskRun)
def handler1 = Spy(K8sTaskHandler)
handler1.task = task1
def staleRequest = new K8sSchedulingRequest(handler1)
def task2 = Mock(TaskRun)
def handler2 = Spy(K8sTaskHandler)
handler2.task = task2
when:
scheduler.submit(handler2)
scheduler.drain()
then:
1 * executor.getClient() >> client
1 * client.nodeList() >> [
items: [
[metadata: [name: 'node-a']]
]
]
then:
1 * strategy.schedule(_, ['node-a']) >> new K8sSchedulingDecision(staleRequest, 'node-a')
then:
1 * strategy.schedule({ List<K8sSchedulingRequest> queue ->
queue.size() == 1 && queue[0].handler == handler2
}, ['node-a']) >> { List<K8sSchedulingRequest> queue, List<String> nodes ->
new K8sSchedulingDecision(queue[0], 'node-a')
}
then:
1 * handler2.submitNow('node-a') >> {}
then:
1 * strategy.schedule([], ['node-a']) >> null
and:
scheduler.@queue.size() == 0
}
}

View File

@@ -1,83 +0,0 @@
package nextflow.k8s.strategies
import nextflow.k8s.K8sSchedulingRequest
import nextflow.k8s.K8sTaskHandler
import nextflow.processor.TaskRun
import nextflow.util.HashBuilder
import spock.lang.Specification
class K8sHashSchedulingStrategyTest extends Specification {
def 'should return null when queue is empty' () {
given:
def strategy = new K8sHashSchedulingStrategy()
expect:
strategy.schedule([], ['node-a']) == null
}
def 'should return null when nodes are empty' () {
given:
def strategy = new K8sHashSchedulingStrategy()
def request = Mock(K8sSchedulingRequest)
expect:
strategy.schedule([request], []) == null
}
def 'should select first queued request and assign a node using task hash' () {
given:
def strategy = new K8sHashSchedulingStrategy()
def task = Mock(TaskRun)
def handler = Spy(K8sTaskHandler)
handler.task = task
def request = new K8sSchedulingRequest(handler)
def nodes = ['node-a', 'node-b', 'node-c']
when:
def decision = strategy.schedule([request], nodes)
then:
1 * task.getHash() >> taskHash
and:
decision.request == request
decision.nodeName == expectedNode
where:
taskHash | expectedNode
new HashBuilder().with(0).build() | 'node-a'
new HashBuilder().with(1).build() | 'node-b'
new HashBuilder().with(2).build() | 'node-c'
new HashBuilder().with(3).build() | 'node-b'
new HashBuilder().with(-1).build() | 'node-b'
}
def 'should always schedule the head of the queue' () {
given:
def strategy = new K8sHashSchedulingStrategy()
def task1 = Mock(TaskRun)
def task2 = Mock(TaskRun)
def handler1 = Spy(K8sTaskHandler)
handler1.task = task1
def handler2 = Spy(K8sTaskHandler)
handler2.task = task2
def request1 = new K8sSchedulingRequest(handler1)
def request2 = new K8sSchedulingRequest(handler2)
when:
def decision = strategy.schedule([request1, request2], ['node-a', 'node-b'])
then:
1 * task1.getHash() >> new HashBuilder().with(0).build()
0 * task2.getHash()
and:
decision.request == request1
decision.nodeName == 'node-a'
}
}

View File

@@ -22,7 +22,7 @@ k8s {
projectDir = '/workspace/projects'
cleanup = false
nextflowImage = 'gitea.kleine.eulenhexe.de/kevin/ma/nextflow-dvfs:0.5.3'
nextflowImage = 'gitea.kleine.eulenhexe.de/kevin/ma/nextflow-dvfs:0.5.5'
imagePullPolicy = 'IfNotPresent'
schedulerInterval = '10s'
recordTaskRuntimes = true