feat(wip): sketch out the scheduling strategy

agent/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"io"
 	"log/slog"
 	"math"
@@ -18,6 +19,11 @@ type cpuHandler struct {
 	logger    *slog.Logger
 }
 
+type memHandler struct {
+	amount int64
+	logger *slog.Logger
+}
+
 // readIntFromFile reads the given file until the end
 // and converts the contained string into an integer
 func readIntFromFile(path string) (int, error) {
@@ -28,6 +34,25 @@ func readIntFromFile(path string) (int, error) {
 	return strconv.Atoi(strings.TrimSpace(string(b)))
 }
 
+func readTotalMemoryKb() (int, error) {
+	b, err := os.ReadFile("/proc/meminfo")
+	if err != nil {
+		return 0, err
+	}
+	s := string(b)
+	for line := range strings.Lines(s) {
+		if !strings.Contains(line, "MemTotal:") {
+			continue
+		}
+		fields := strings.Fields(line)
+		if len(fields) < 2 {
+			return 0, fmt.Errorf("expected line '%s' to have at least 2 fields", line)
+		}
+		return strconv.Atoi(fields[1])
+	}
+	return 0, fmt.Errorf("failed to find MemTotal in /proc/meminfo")
+}
+
 func enumerateCPUCores(logger *slog.Logger) ([]string, int, int, error) {
 	files, err := os.ReadDir("/sys/devices/system/cpu/")
 	if err != nil {
@@ -135,6 +160,17 @@ func (s *cpuHandler) handleCurrent(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func (s *cpuHandler) handleCount(w http.ResponseWriter, r *http.Request) {
+	if r.Method != "GET" {
+		s.logger.Error("unexepcted method for cpu count", "method", r.Method)
+		writeResponse("invalid method", http.StatusBadRequest, w, s.logger)
+		return
+	}
+
+	count := len(s.corePaths)
+	writeResponse(strconv.Itoa(count), http.StatusOK, w, s.logger)
+}
+
 func (s *cpuHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	requestedPath := getStrippedRequestPath(r.URL, r.Pattern)
 	switch requestedPath {
@@ -147,12 +183,25 @@ func (s *cpuHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	case "frequency/current":
 		s.handleCurrent(w, r)
 		return
+	case "count":
+		s.handleCount(w, r)
+		return
 	default:
 		s.logger.Error("unknown resource", "path", requestedPath)
 		writeResponse("not found", http.StatusNotFound, w, s.logger)
 	}
 }
 
+func (m *memHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	requestedPath := getStrippedRequestPath(r.URL, r.Pattern)
+	if requestedPath == "amount" {
+		writeResponse(strconv.FormatInt(m.amount, 10), http.StatusOK, w, m.logger)
+	} else {
+		m.logger.Error("unknown resource", "path", requestedPath)
+		writeResponse("not found", http.StatusNotFound, w, m.logger)
+	}
+}
+
 func main() {
 	appEnv := os.Getenv("AGENT_ENV")
 	loggerOpts := &slog.HandlerOptions{
@@ -176,6 +225,18 @@ func main() {
 		minFreq:   minF,
 		corePaths: freqPaths,
 	}
+
+	memKb, err := readTotalMemoryKb()
+	if err != nil {
+		logger.Error("failed to read memory size", "err", err.Error())
+		os.Exit(1)
+	}
+	memHandler := &memHandler{
+		amount: int64(memKb) * 1024,
+		logger: logger,
+	}
+
 	http.Handle("/cpu/", logRequests(logger, cpuHandler))
+	http.Handle("/mem/", logRequests(logger, memHandler))
 	logger.Info("exit", "result", http.ListenAndServe(":8080", nil))
 }

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

@@ -17,46 +17,34 @@ class K8sDVFSClient {
         this.httpClient = HttpClient.newBuilder().build()
     }
 
-    OptionalInt getNodeCurrentFrequency(String node) {
+    OptionalLong getNodeCurrentFrequency(String node) {
         log.debug("Getting current frequency of node ${node}")
         return getNodeFrequency(node, "current")
     }
 
-    OptionalInt getNodeMaxFrequency(String node) {
+    OptionalLong getNodeMaxFrequency(String node) {
         log.debug("Getting max frequency of node ${node}")
         return getNodeFrequency(node, "max")
     }
 
-    OptionalInt getNodeMinFrequency(String node) {
+    OptionalLong getNodeMinFrequency(String node) {
         log.debug("Getting min frequency of node ${node}")
         return getNodeFrequency(node, "min")
     }
 
-    private OptionalInt getNodeFrequency(String node, String endpoint) {
+    private OptionalLong getNodeFrequency(String node, String endpoint) {
-        HttpRequest request = HttpRequest.newBuilder()
+        return getLong(HttpRequest.newBuilder()
-            .uri(new URI("http://${nodeName}/cpu/frequency/${endpoint}"))
+            .uri(new URI("http://${agentAddress(node)}/cpu/frequency/${endpoint}"))
             .GET()
-            .build()
+            .build())
-        HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString())
-        if (response.statusCode() != 200) {
-            log.error("Request GET ${request.uri().toString()} returned ${response.statusCode()}")
-            return OptionalInt.empty()
-        }
-        try {
-            double parsed = Double.parseDouble(response.body())
-            return OptionalInt.of(parsed)
-        } catch (NumberFormatException ex) {
-            log.error("Unexpected response ${response.body()} - ${ex.message}")
-            return OptionalInt.empty()
-        }
     }
 
     boolean setNodeFrequency(String node, int frequency) {
         String body = String.valueOf(frequency)
 
         HttpRequest request = HttpRequest.newBuilder()
-            .uri(new URI("http://${nodeName}/cpu/frequency/current"))
+            .uri(new URI("http://${agentAddress(node)}/cpu/frequency/current"))
-            .PUT(HttpRequ3est.BodyPublishers.ofString(body))
+            .PUT(HttpRequest.BodyPublishers.ofString(body))
             .build()
         HttpResponse<Void> response = httpClient.send(request, HttpResponse.BodyHandlers.discarding())
         if (response.statusCode() != 200) {
@@ -65,4 +53,37 @@ class K8sDVFSClient {
         }
         return true
     }
+
+    OptionalLong getCPUCount(String node) {
+        return getLong(HttpRequest.newBuilder()
+            .uri(new URI("http://${agentAddress(node)}/cpu/count"))
+            .GET()
+            .build())
+    }
+
+    OptionalLong getMemoryAmount(String node) {
+        return getInt(HttpRequest.newBuilder()
+            .uri(new URI("http://${agentAddress(node)}/mem/amount"))
+            .GET()
+            .build())
+    }
+
+    private OptionalLong getLong(HttpRequest request) {
+        HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString())
+        if (response.statusCode() != 200) {
+            log.error("Request GET ${request.uri().toString()} returned ${response.statusCode()}")
+            return OptionalLong.empty()
+        }
+        try {
+            long parsed = Long.parseLong(response.body())
+            return OptionalLong.of(parsed)
+        } catch (NumberFormatException ex) {
+            log.error("Unexpected response ${response.body()} - ${ex.message}")
+            return OptionalLong.empty()
+        }
+    }
+
+    private String agentAddress(String node) {
+        return K8sNodeInitDeployer.buildPodName(node)
+    }
 }

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

@@ -151,7 +151,7 @@ class K8sDriverLauncher {
         this.k8sConfig = makeK8sConfig(config.toMap())
         this.k8sClient = makeK8sClient(k8sConfig)
         this.k8sConfig.checkStorageAndPaths(k8sClient)
-        this.initDeployer = new K8sNodeInitDeployer(k8sClient, k8sConfig, runName)
+        this.initDeployer = new K8sNodeInitDeployer(k8sClient, k8sConfig)
 
         createK8sConfigMap()
 

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

@@ -103,7 +103,7 @@ class K8sExecutor extends Executor implements ExtensionPoint {
         if (k8sConfig.schedulingStrategy == "Hash") {
             strategy = new K8sHashSchedulingStrategy()
         } else if (k8sConfig.schedulingStrategy == "DVFS") {
-            strategy = new K8sDVFSSchedulingStrategy()
+            strategy = new K8sDVFSSchedulingStrategy(this.runtimeEstimator, new K8sDVFSClient())
         } else {
             log.error "[K8s] invalid scheduling strategy $k8sConfig.schedulingStrategy, falling back on \"Hash\""
             strategy = new K8sHashSchedulingStrategy()

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

@@ -9,12 +9,10 @@ import nextflow.k8s.model.PodSpecBuilder
 class K8sNodeInitDeployer {
     private K8sClient client
     private K8sConfig config
-    private String runName
 
-    K8sNodeInitDeployer(K8sClient client, K8sConfig config, String runName) {
+    K8sNodeInitDeployer(K8sClient client, K8sConfig config) {
         this.client = client
         this.config = config
-        this.runName = runName
     }
 
     void deploy() {
@@ -116,10 +114,10 @@ exit 0
         return nodes
     }
 
-    private String buildPodName(String nodeName) {
+    public static String buildPodName(String nodeName) {
         // TODO: Remove anything that is not lowercase alpha-numeric or dash
         String sanitizedNodeName = nodeName.toLowerCase()
-        String name = "nf-init-${runName}-${sanitizedNodeName}"
+        String name = "nf-init-${sanitizedNodeName}"
         if ( name.length() > 63 )
             name = name.substring(0, 63)
         return name

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

@@ -14,4 +14,8 @@ class K8sSchedulingRequest {
         this.task = handler.task
         this.submitTimeMillis = System.currentTimeMillis()
     }
+
+    String getTaskName() {
+        return task.processor.name
+    }
 }

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

@@ -18,4 +18,10 @@ interface K8sSchedulingStrategy {
      * @return {@code true} if @ref schedule should immediately be called
      */
     boolean scheduleImmediately(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue)
+
+    /**
+     * Called when a task has finished execution.
+     * @param task the task
+     */
+    void taskFinished(K8sTaskHandler task)
 }

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

@@ -60,6 +60,13 @@ class K8sTaskScheduler implements Runnable {
         }
         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.
+         * Because resources were freed right now */
+        if (queue.size() > 0)
+            drain()
     }
 
     protected synchronized void drain() {

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

@@ -1,22 +1,146 @@
 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
 
 /**
  * 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 class WorkerNode {
+        long maxFrequency
+        long minFrequency
+        long currentFrequency
+        long cpuCount
+        long memoryAmount
+
+        WorkerNode(long maxF, long minF, long curF, long cpus, long mem) {
+            this.maxFrequency = maxF
+            this.minFrequency = minF
+            this.currentFrequency = curF
+            this.cpuCount = cpus
+            this.memoryAmount = mem
+        }
+    }
+
+    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 HashMap<String, WorkerNode> nodes
+
+    private double averageRuntime
+    private long finishedTaskCount
+
+    K8sDVFSSchedulingStrategy(K8sRuntimeEstimator runtimeEstimator, K8sDVFSClient dvfsClient) {
+        this.runtimeEstimator = runtimeEstimator
+        this.dvfsClient = dvfsClient
+        this.nodes = 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.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.
+             */
+            final double taskEstimation = runtimeEstimator.estimate(req.handler)
+            final boolean isCriticalPath = taskEstimation > averageRuntime
+
+
+            /* Step 2.2: Filter nodes based on task requirements */
+
+            /* Step 2.3: Assign to node based on "best fit" - current node frequency is closest to determined frequency */
+        }
+
         return null
     }
 
     @Override
     boolean scheduleImmediately(K8sTaskScheduler scheduler, List<K8sSchedulingRequest> queue) {
-        return false
+        return queue != null && queue.size() > 0 && scheduler.freeNodes.size() > 0
+    }
+
+    @Override
+    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
+    }
+
+    private boolean initNodes(K8sTaskScheduler scheduler) {
+        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
+            }
+            this.nodes.put(node, new WorkerNode(max.asLong, min.asLong, cur.asLong, cpus.asLong, mem.asLong))
+        }
+        return !this.nodes.isEmpty()
     }
 }

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

@@ -4,6 +4,7 @@ import groovy.transform.CompileStatic
 import nextflow.k8s.K8sSchedulingDecision
 import nextflow.k8s.K8sSchedulingRequest
 import nextflow.k8s.K8sSchedulingStrategy
+import nextflow.k8s.K8sTaskHandler
 import nextflow.k8s.K8sTaskScheduler
 
 @CompileStatic
@@ -30,4 +31,7 @@ class K8sHashSchedulingStrategy implements K8sSchedulingStrategy {
         final freeNodes = scheduler.freeNodes
         return freeNodes.size() > 0
     }
+
+    @Override
+    void taskFinished(K8sTaskHandler task) { /* nop */ }
 }

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

@@ -644,7 +644,7 @@ class K8sDriverLauncherTest extends Specification {
         def driver = Spy(K8sDriverLauncher)
         driver.@k8sConfig = config
         driver.@runName = POD_NAME
-        driver.@initDeployer = new K8sNodeInitDeployer(driver.k8sClient, config, driver.runName)
+        driver.@initDeployer = new K8sNodeInitDeployer(driver.k8sClient, config)
 
         when:
         driver.shutdown()

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

@@ -14,7 +14,7 @@ class K8sNodeInitDeployerTest extends Specification {
         given:
         def client = Mock(K8sClient)
         def config = new K8sConfig(nodeInit: [enabled: false])
-        def deployer = new K8sNodeInitDeployer(client, config, 'run-foo')
+        def deployer = new K8sNodeInitDeployer(client, config)
 
         when:
         deployer.deploy()
@@ -32,7 +32,7 @@ class K8sNodeInitDeployerTest extends Specification {
             image: 'ubuntu:latest',
             command: ['/bin/bash', '-c', 'echo init']
         ])
-        def deployer = new K8sNodeInitDeployer(client, config, 'run-foo')
+        def deployer = new K8sNodeInitDeployer(client, config)
 
         when:
         deployer.deploy()
@@ -48,13 +48,13 @@ class K8sNodeInitDeployerTest extends Specification {
         then:
         1 * client.podCreate({ Map spec ->
             spec.kind == 'Pod'
-            spec.metadata.name == 'nf-init-run-foo-node-a'
+            spec.metadata.name == 'nf-init-node-a'
             spec.metadata.namespace == 'default'
             spec.spec.nodeName == 'node-a'
             spec.spec.restartPolicy == 'Never'
 
             def container = spec.spec.containers[0]
-            container.name == 'nf-init-run-foo-node-a'
+            container.name == 'nf-init-node-a'
             container.image == 'ubuntu:latest'
             container.command == ['/bin/bash', '-c', 'echo init']
             container.securityContext.privileged == true
@@ -66,11 +66,11 @@ class K8sNodeInitDeployerTest extends Specification {
         then:
         1 * client.podCreate({ Map spec ->
             spec.kind == 'Pod'
-            spec.metadata.name == 'nf-init-run-foo-node-b'
+            spec.metadata.name == 'nf-init-node-b'
             spec.spec.nodeName == 'node-b'
 
             def container = spec.spec.containers[0]
-            container.name == 'nf-init-run-foo-node-b'
+            container.name == 'nf-init-node-b'
             container.image == 'ubuntu:latest'
             container.command == ['/bin/bash', '-c', 'echo init']
             container.securityContext.privileged == true
@@ -82,13 +82,12 @@ class K8sNodeInitDeployerTest extends Specification {
     def 'should lowercase and truncate generated pod names' () {
         given:
         def client = Mock(K8sClient)
-        def runName = 'run-name-with-a-very-long-identifier'
         def config = new K8sConfig(nodeInit: [
             enabled: true,
             image: 'ubuntu:latest',
             command: ['true']
         ])
-        def deployer = new K8sNodeInitDeployer(client, config, runName)
+        def deployer = new K8sNodeInitDeployer(client, config)
 
         when:
         deployer.deploy()
@@ -96,16 +95,16 @@ class K8sNodeInitDeployerTest extends Specification {
         then:
         1 * client.nodeList() >> [
             items: [
-                [metadata: [name: 'NODE-WITH-A-VERY-LONG-NAME-ABCDEFGHIJKLMNOPQRSTUVWXYZ']]
+                [metadata: [name: 'NODE-WITH-A-VERY-LONG-NAME-ABCDEFGHIJKLMNOPQRSTUVWXYZ-0123']]
             ]
         ]
 
         then:
         1 * client.podCreate({ Map spec ->
-            spec.metadata.name == 'nf-init-run-name-with-a-very-long-identifier-node-with-a-very-l'
+            spec.metadata.name == 'nf-init-node-with-a-very-long-name-abcdefghijklmnopqrstuvwxyz-0'
             spec.metadata.name.size() == 63
             spec.metadata.name == spec.metadata.name.toLowerCase()
-            spec.spec.nodeName == 'NODE-WITH-A-VERY-LONG-NAME-ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+            spec.spec.nodeName == 'NODE-WITH-A-VERY-LONG-NAME-ABCDEFGHIJKLMNOPQRSTUVWXYZ-0123'
         })
 
         then:
@@ -116,7 +115,7 @@ class K8sNodeInitDeployerTest extends Specification {
         given:
         def client = Mock(K8sClient)
         def config = new K8sConfig(nodeInit: [enabled: false, cleanup: true])
-        def deployer = new K8sNodeInitDeployer(client, config, 'run-foo')
+        def deployer = new K8sNodeInitDeployer(client, config)
 
         when:
         deployer.cleanup()
@@ -130,7 +129,7 @@ class K8sNodeInitDeployerTest extends Specification {
         given:
         def client = Mock(K8sClient)
         def config = new K8sConfig(nodeInit: [enabled: true, cleanup: false])
-        def deployer = new K8sNodeInitDeployer(client, config, 'run-foo')
+        def deployer = new K8sNodeInitDeployer(client, config)
 
         when:
         deployer.cleanup()
@@ -144,7 +143,7 @@ class K8sNodeInitDeployerTest extends Specification {
         given:
         def client = Mock(K8sClient)
         def config = new K8sConfig(nodeInit: [enabled: true, cleanup: true])
-        def deployer = new K8sNodeInitDeployer(client, config, 'run-foo')
+        def deployer = new K8sNodeInitDeployer(client, config)
 
         when:
         deployer.cleanup()
@@ -158,8 +157,8 @@ class K8sNodeInitDeployerTest extends Specification {
         ]
 
         then:
-        1 * client.podDelete('nf-init-run-foo-node-a')
+        1 * client.podDelete('nf-init-node-a')
-        1 * client.podDelete('nf-init-run-foo-node-b')
+        1 * client.podDelete('nf-init-node-b')
         0 * client._
     }
 }