feat: linear regression for runtime estimation

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

@@ -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 )

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

@@ -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)

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

@@ -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
+    }
+}

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

@@ -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() {
+        if (!enabled)
+            return;
+        try {
+            FileWriter out = new FileWriter(recordPath)
             for (K8sRuntimeRecord record : records) {
-            log.info("[K8s] ${record.taskName} - ${record.inputSize} Bytes - ${record.runtimeMillis} ms")
+                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}"
         }
     }
 }

nextflow/plugins/nf-k8s/src/test/nextflow/k8s/K8sRuntimeEstimatorTest.groovy

@@ -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
+    }
+}

nextflow/plugins/nf-k8s/src/test/nextflow/k8s/K8sTaskSchedulerTest.groovy

@@ -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
-    }
-}

nextflow/plugins/nf-k8s/src/test/nextflow/k8s/strategies/K8sHashSchedulingStrategyTest.groovy

@@ -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'
-    }
-}

test/nextflow.config

@@ -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