kevin/ma
1
0
Fork 0
Code Issues 1 Packages Activity

feat(wip): sketch out the scheduling strategy

Browse Source
This commit is contained in:
Kevin
2026-06-10 22:51:28 +02:00
parent 1ff4446081
commit e9ff5ecf06
12 changed files with 271 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
agent/main.go
View File

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

63
nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sDVFSClient.groovy
View File

@@ -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) {
HttpRequest request = HttpRequest.newBuilder()
.uri(new URI("http://${nodeName}/cpu/frequency/${endpoint}"))
private OptionalLong getNodeFrequency(String node, String endpoint) {
return getLong(HttpRequest.newBuilder()
.uri(new URI("http://${agentAddress(node)}/cpu/frequency/${endpoint}"))
.GET()
.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()
}
.build())
}
boolean setNodeFrequency(String node, int frequency) {
String body = String.valueOf(frequency)
HttpRequest request = HttpRequest.newBuilder()
.uri(new URI("http://${nodeName}/cpu/frequency/current"))
.PUT(HttpRequ3est.BodyPublishers.ofString(body))
.uri(new URI("http://${agentAddress(node)}/cpu/frequency/current"))
.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)
}
}

2
nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sDriverLauncher.groovy
View File

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

2
nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sExecutor.groovy
View File

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

10
nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sNodeInitDeployer.groovy
View File

@@ -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() {
@@ -110,16 +108,16 @@ exit 0
private List<String> getNodes() {
final resp = client.nodeList()
ArrayList<String> nodes = new ArrayList<String>()
for ( Map item : resp.items ) {
for ( Map item: resp.items ) {
nodes.add(item.metadata.name as String)
}
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

4
nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sSchedulingRequest.groovy
View File

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

6
nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sSchedulingStrategy.groovy
View File

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

7
nextflow/plugins/nf-k8s/src/main/nextflow/k8s/K8sTaskScheduler.groovy
View File

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

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

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

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

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

2
nextflow/plugins/nf-k8s/src/test/nextflow/k8s/K8sDriverLauncherTest.groovy
View File

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

31
nextflow/plugins/nf-k8s/src/test/nextflow/k8s/K8sNodeInitDeployerTest.groovy
View File

@@ -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-run-foo-node-b')
1 * client.podDelete('nf-init-node-a')
1 * client.podDelete('nf-init-node-b')
0 * client._
}
}